@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7a9406690e

[KV Connector] Fix DecodeBench fp8 fill values and add a startup fill mode (#58472)

Signed-off-by: Zijing Liu <liuzijing2014@gmail.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Zijing Liu committed Sep 24, 2026 at 16:55 UTC 7a9406690ebe2fc47236a81cf7ec33188c6565e2
3 files changed +549 -38
docs/serving/expert_parallel_deployment.md
+1 -1
@@ -315,6 +315,6 @@ except Exception as e:
315
316 ### Benchmarking
317
318 -- To simulate the decode deployment of disaggregated serving, pass `--kv-transfer-config '{"kv_connector":"DecodeBenchConnector","kv_role":"kv_both"}'` to the `vllm serve` invocation. The connector populates KV cache with random values so decode can be profiled in isolation.
318 +- To simulate the decode deployment of disaggregated serving, pass `--kv-transfer-config '{"kv_connector":"DecodeBenchConnector","kv_role":"kv_both"}'` to the `vllm serve` invocation. The connector populates KV cache with random values so decode can be profiled in isolation. Add `"kv_connector_extra_config":{"startup_fill":true}` to fill the whole KV cache once at startup instead of in the decode steps that admit new requests (see the connector docstring for its limitations).
319
320 - **CUDAGraph capture**: Use `--compilation_config '{"cudagraph_mode": "FULL_DECODE_ONLY"}'` to enable CUDA graph capture for decode only and save KV cache.
tests/v1/kv_connector/unit/test_decode_bench_connector.py
+325 -1
@@ -18,6 +18,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector import
18 DecodeBenchConnectorMetadata,
19 )
20 from vllm.forward_context import ForwardContext
21 +from vllm.platforms import current_platform
22 from vllm.utils.hashing import sha256
23 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
24 from vllm.v1.core.kv_cache_utils import (
@@ -34,6 +35,8 @@ from vllm.v1.kv_cache_interface import (
35 FullAttentionSpec,
36 KVCacheConfig,
37 KVCacheGroupSpec,
38 + KVQuantMode,
39 + MambaSpec,
40 SlidingWindowSpec,
41 )
42 from vllm.v1.request import Request
@@ -50,7 +53,12 @@ from .utils import (
53 class DecodeBenchTestRunner:
54 """Test runner for DecodeBenchConnector."""
55
53 - def __init__(self, block_size: int, num_gpu_blocks: int):
56 + def __init__(
57 + self,
58 + block_size: int,
59 + num_gpu_blocks: int,
60 + kv_connector_extra_config: dict | None = None,
61 + ):
62 self.block_size = block_size
63 self.num_gpu_blocks = num_gpu_blocks
64
@@ -61,6 +69,7 @@ class DecodeBenchTestRunner:
69 block_size=block_size,
70 max_num_batched_tokens=1000,
71 kv_connector="DecodeBenchConnector",
72 + kv_connector_extra_config=kv_connector_extra_config,
73 )
74
75 self.vllm_config = vllm_config
@@ -780,5 +789,320 @@ def test_decode_bench_connector_concurrent_requests():
789 assert len(metadata2.reqs_to_fill) == 0
790
791
792 +def _fill_single_block(cache_dtype, spec_dtype, kv_quant_mode, cache, fill_std=0.0):
793 + """Fill block 1 of a one-layer cache through the worker connector."""
794 + vllm_config = create_vllm_config(
795 + block_size=16,
796 + kv_connector="DecodeBenchConnector",
797 + kv_connector_extra_config={"fill_std": fill_std},
798 + cache_dtype=cache_dtype,
799 + )
800 + kv_cache_spec = FullAttentionSpec(
801 + block_size=16,
802 + num_kv_heads=1,
803 + head_size=8,
804 + dtype=spec_dtype,
805 + kv_quant_mode=kv_quant_mode,
806 + )
807 + connector = DecodeBenchConnector(
808 + vllm_config,
809 + KVConnectorRole.WORKER,
810 + KVCacheConfig(
811 + num_blocks=cache.shape[0],
812 + kv_cache_tensors=[],
813 + kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
814 + ),
815 + )
816 + connector.register_kv_caches({"layer": cache})
817 + connector.bind_connector_metadata(
818 + DecodeBenchConnectorMetadata(reqs_to_fill={"request": (([1],), 16)})
819 + )
820 + connector.start_load_kv(
821 + ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
822 + )
823 +
824 +
825 +@pytest.mark.parametrize(
826 + "cache_dtype,spec_dtype,kv_quant_mode,fp8_dtype",
827 + [
828 + (
829 + "fp8",
830 + torch.uint8,
831 + KVQuantMode.FP8_PER_TENSOR,
832 + current_platform.fp8_dtype(),
833 + ),
834 + ("fp8_e5m2", torch.uint8, KVQuantMode.FP8_PER_TENSOR, torch.float8_e5m2),
835 + ("auto", torch.float8_e4m3fn, KVQuantMode.NONE, torch.float8_e4m3fn),
836 + ],
837 +)
838 +@pytest.mark.parametrize("fill_std", [0.0, 0.1])
839 +def test_decode_bench_connector_fills_fp8_caches(
840 + cache_dtype, spec_dtype, kv_quant_mode, fp8_dtype, fill_std
841 +):
842 + """fp8 caches, including ones stored as uint8, get encoded fill values."""
843 + cache = torch.zeros(4, 16, 8, dtype=spec_dtype)
844 + _fill_single_block(cache_dtype, spec_dtype, kv_quant_mode, cache, fill_std)
845 +
846 + values = cache.view(fp8_dtype).float()
847 + assert torch.count_nonzero(values[[0, 2, 3]]) == 0
848 + filled = values[1]
849 + if fill_std == 0:
850 + expected = torch.tensor(0.015).to(fp8_dtype).float()
851 + torch.testing.assert_close(filled, expected.expand_as(filled))
852 + else:
853 + assert torch.isfinite(filled).all()
854 + assert filled.std() > 0
855 +
856 +
857 +def test_decode_bench_connector_zero_fills_packed_uint8_caches():
858 + """uint8 layouts that do not hold plain fp8 values are zeroed."""
859 + cache = torch.full((4, 16, 8), 7, dtype=torch.uint8)
860 + _fill_single_block("nvfp4", torch.uint8, KVQuantMode.NVFP4, cache)
861 +
862 + assert torch.count_nonzero(cache[1]) == 0
863 + assert torch.all(cache[[0, 2, 3]] == 7)
864 +
865 +
866 +@pytest.mark.parametrize("fill_std", [0.0, 0.1])
867 +def test_decode_bench_connector_startup_fill_fills_registered_caches(fill_std):
868 + """Startup fill writes every cache view in place, including strided views."""
869 + vllm_config = create_vllm_config(
870 + block_size=16,
871 + kv_connector="DecodeBenchConnector",
872 + kv_connector_extra_config={"startup_fill": True, "fill_std": fill_std},
873 + )
874 + kv_cache_config = KVCacheConfig(
875 + num_blocks=4,
876 + kv_cache_tensors=[],
877 + kv_cache_groups=[
878 + KVCacheGroupSpec(
879 + ["attention", "transposed"],
880 + FullAttentionSpec(
881 + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32
882 + ),
883 + ),
884 + ],
885 + )
886 + connector = DecodeBenchConnector(
887 + vllm_config, KVConnectorRole.WORKER, kv_cache_config
888 + )
889 + attention = torch.full((4, 2, 16), -1.0)
890 + transposed = torch.full((2, 4, 16), -1.0).transpose(0, 1)
891 + connector.register_kv_caches({"attention": attention, "transposed": transposed})
892 +
893 + for filled in (attention, transposed):
894 + if fill_std == 0:
895 + torch.testing.assert_close(filled, torch.full_like(filled, 0.015))
896 + else:
897 + assert torch.isfinite(filled).all()
898 + assert filled.std() > 0
899 +
900 +
901 +@pytest.mark.parametrize("fill_std", [0.0, 0.1])
902 +def test_decode_bench_connector_startup_fill_encodes_fp8_caches(fill_std):
903 + """Startup fill encodes values in the fp8 dtype of a uint8 fp8 cache."""
904 + vllm_config = create_vllm_config(
905 + block_size=16,
906 + kv_connector="DecodeBenchConnector",
907 + kv_connector_extra_config={"startup_fill": True, "fill_std": fill_std},
908 + cache_dtype="fp8",
909 + )
910 + kv_cache_spec = FullAttentionSpec(
911 + block_size=16,
912 + num_kv_heads=1,
913 + head_size=8,
914 + dtype=torch.uint8,
915 + kv_quant_mode=KVQuantMode.FP8_PER_TENSOR,
916 + )
917 + connector = DecodeBenchConnector(
918 + vllm_config,
919 + KVConnectorRole.WORKER,
920 + KVCacheConfig(
921 + num_blocks=4,
922 + kv_cache_tensors=[],
923 + kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
924 + ),
925 + )
926 + cache = torch.zeros(4, 16, 8, dtype=torch.uint8)
927 + connector.register_kv_caches({"layer": cache})
928 +
929 + values = cache.view(current_platform.fp8_dtype()).float()
930 + if fill_std == 0:
931 + expected = torch.tensor(0.015).to(current_platform.fp8_dtype()).float()
932 + torch.testing.assert_close(values, expected.expand_as(values))
933 + else:
934 + assert torch.isfinite(values).all()
935 + assert values.std() > 0
936 +
937 +
938 +def test_decode_bench_connector_startup_fill_keeps_zeroed_groups_per_request():
939 + """Attention blocks the engine zeroes on allocation are filled per request."""
940 + vllm_config = create_vllm_config(
941 + block_size=16,
942 + max_num_batched_tokens=1000,
943 + kv_connector="DecodeBenchConnector",
944 + kv_connector_extra_config={"startup_fill": True},
945 + )
946 + # Mamba layers make the engine zero newly allocated attention blocks.
947 + kv_cache_config = KVCacheConfig(
948 + num_blocks=4,
949 + kv_cache_tensors=[],
950 + kv_cache_groups=[
951 + KVCacheGroupSpec(
952 + ["attention"],
953 + FullAttentionSpec(
954 + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32
955 + ),
956 + ),
957 + KVCacheGroupSpec(
958 + ["mamba"],
959 + MambaSpec(
960 + block_size=16,
961 + shapes=((4,), (4,)),
962 + dtypes=(torch.float32, torch.float32),
963 + ),
964 + ),
965 + ],
966 + )
967 + assert kv_cache_config.needs_kv_cache_zeroing
968 + scheduler_connector = DecodeBenchConnector(
969 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
970 + )
971 + num_external_tokens = 17
972 + request = create_request(
973 + request_id=1, num_tokens=num_external_tokens + 1, block_size=16
974 + )
975 + blocks = KVCacheBlocks(
976 + (
977 + [KVCacheBlock(block_id=1), KVCacheBlock(block_id=2)],
978 + [KVCacheBlock(block_id=3)],
979 + )
980 + )
981 + scheduler_connector.update_state_after_alloc(request, blocks, num_external_tokens)
982 + metadata = scheduler_connector.build_connector_meta(SchedulerOutput.make_empty())
983 + assert isinstance(metadata, DecodeBenchConnectorMetadata)
984 + assert metadata.reqs_to_fill == {
985 + request.request_id: (([1, 2], []), num_external_tokens)
986 + }
987 +
988 + worker_connector = DecodeBenchConnector(
989 + vllm_config, KVConnectorRole.WORKER, kv_cache_config
990 + )
991 + attention = torch.full((4, 2), -1.0)
992 + mamba_states = (torch.zeros(4, 4), torch.zeros(4, 4))
993 + worker_connector.register_kv_caches({"attention": attention, "mamba": mamba_states})
994 + assert torch.all(attention == -1.0)
995 +
996 + worker_connector.bind_connector_metadata(metadata)
997 + worker_connector.start_load_kv(
998 + ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
999 + )
1000 + torch.testing.assert_close(attention[1:3], torch.full_like(attention[1:3], 0.015))
1001 + assert torch.all(attention[[0, 3]] == -1.0)
1002 +
1003 +
1004 +def test_decode_bench_connector_startup_fill_skips_per_step_fill():
1005 + """Admitted requests need no per-step fill and leave existing KV intact."""
1006 + block_size = 16
1007 + runner = DecodeBenchTestRunner(
1008 + block_size=block_size,
1009 + num_gpu_blocks=100,
1010 + kv_connector_extra_config={"startup_fill": True},
1011 + )
1012 + for kv_cache in runner.kv_caches.values():
1013 + torch.testing.assert_close(kv_cache, torch.full_like(kv_cache, 0.015))
1014 + kv_cache.fill_(7.0)
1015 +
1016 + req1 = runner.new_request([1] * (block_size * 2))
1017 + scheduler_output, metadata = runner.run_single_step()
1018 + assert scheduler_output.num_scheduled_tokens == {req1.request_id: 1}
1019 + assert scheduler_output.scheduled_new_reqs[0].num_computed_tokens == (
1020 + block_size * 2 - 1
1021 + )
1022 + assert metadata.reqs_to_fill == {}
1023 +
1024 + req2 = runner.new_request([2] * (block_size * 3))
1025 + scheduler_output, metadata = runner.run_single_step()
1026 + assert scheduler_output.num_scheduled_tokens == {
1027 + req1.request_id: 1,
1028 + req2.request_id: 1,
1029 + }
1030 + assert metadata.reqs_to_fill == {}
1031 +
1032 + # A re-admitted (e.g. preempted) request is not reported again.
1033 + assert runner.scheduler_connector.get_num_new_matched_tokens(req1, 0) == (
1034 + 0,
1035 + False,
1036 + )
1037 + for kv_cache in runner.kv_caches.values():
1038 + assert torch.all(kv_cache == 7.0)
1039 +
1040 +
1041 +def test_decode_bench_connector_startup_fill_keeps_circular_buffer_zero_fill():
1042 + """Circular buffers are still zeroed per request with startup fill."""
1043 + vllm_config = create_vllm_config(
1044 + block_size=16,
1045 + max_num_batched_tokens=1000,
1046 + kv_connector="DecodeBenchConnector",
1047 + kv_connector_extra_config={"startup_fill": True},
1048 + )
1049 + kv_cache_config = KVCacheConfig(
1050 + num_blocks=8,
1051 + kv_cache_tensors=[],
1052 + kv_cache_groups=[
1053 + KVCacheGroupSpec(
1054 + ["full_attention"],
1055 + FullAttentionSpec(
1056 + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32
1057 + ),
1058 + ),
1059 + KVCacheGroupSpec(
1060 + ["compressor_ring"],
1061 + CircularBufferSpec(
1062 + block_size=32, num_kv_heads=1, head_size=1, dtype=torch.float32
1063 + ),
1064 + ),
1065 + ],
1066 + )
1067 + scheduler_connector = DecodeBenchConnector(
1068 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
1069 + )
1070 + num_external_tokens = 17
1071 + request = create_request(
1072 + request_id=1, num_tokens=num_external_tokens + 1, block_size=16
1073 + )
1074 + blocks = KVCacheBlocks(
1075 + (
1076 + [KVCacheBlock(block_id=1), KVCacheBlock(block_id=2)],
1077 + [KVCacheBlock(block_id=4)],
1078 + )
1079 + )
1080 + scheduler_connector.update_state_after_alloc(request, blocks, num_external_tokens)
1081 + metadata = scheduler_connector.build_connector_meta(SchedulerOutput.make_empty())
1082 + assert isinstance(metadata, DecodeBenchConnectorMetadata)
1083 + assert metadata.reqs_to_fill == {
1084 + request.request_id: (([], [4]), num_external_tokens)
1085 + }
1086 +
1087 + worker_connector = DecodeBenchConnector(
1088 + vllm_config, KVConnectorRole.WORKER, kv_cache_config
1089 + )
1090 + attention_cache = torch.zeros(8, 2)
1091 + ring_cache = torch.ones(8, 2)
1092 + worker_connector.register_kv_caches(
1093 + {"full_attention": attention_cache, "compressor_ring": ring_cache}
1094 + )
1095 + assert torch.all(ring_cache == 1.0)
1096 + worker_connector.bind_connector_metadata(metadata)
1097 + worker_connector.start_load_kv(
1098 + ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
1099 + )
1100 +
1101 + torch.testing.assert_close(attention_cache, torch.full_like(attention_cache, 0.015))
1102 + expected_ring = torch.ones_like(ring_cache)
1103 + expected_ring[4] = 0
1104 + assert torch.equal(ring_cache, expected_ring)
1105 +
1106 +
1107 if __name__ == "__main__":
1108 pytest.main([__file__, "-v"])
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+223 -36
@@ -28,6 +28,16 @@ Example:
28 - fill_mean (float): Mean value for random normal fill (default: 0.015)
29 - fill_std (float): Standard deviation for random fill (default: 0.0)
30 Set to 0 for constant values, >0 for random sampling
31 + - startup_fill (bool): Fill the whole KV cache once at startup instead
32 + of filling each new request's blocks in the step that schedules it
33 + (default: False). Circular-buffer caches are still zeroed per
34 + request, and attention caches whose new blocks the engine zeroes
35 + before use (hybrid Mamba or mixed-precision KV caches) are still
36 + filled per request. The fill is not redone after sleep mode
37 + discards the KV cache.
38 +
39 + The dummy KV cache is only meant for performance measurement; outputs are
40 + not meaningful for accuracy evaluation.
41
42 """
43
@@ -45,13 +55,21 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
55 SupportsHMA,
56 )
57 from vllm.logger import init_logger
58 +from vllm.platforms import current_platform
59 from vllm.utils.math_utils import cdiv
60 from vllm.v1.attention.backend import AttentionMetadata
61 from vllm.v1.core.kv_cache_utils import (
62 dcp_world_size_for_kv_cache_spec,
63 resolve_dcp_kv_block_size,
64 )
54 -from vllm.v1.kv_cache_interface import CircularBufferSpec, iter_layer_specs
65 +from vllm.v1.kv_cache_interface import (
66 + AttentionSpec,
67 + CircularBufferSpec,
68 + KVCacheSpec,
69 + KVQuantMode,
70 + UniformTypeKVCacheSpecs,
71 + iter_layer_specs,
72 +)
73
74 if TYPE_CHECKING:
75 from vllm.config import VllmConfig
@@ -64,6 +82,45 @@ if TYPE_CHECKING:
82 logger = init_logger(__name__)
83
84
85 +def _get_fp8_dtype(spec: KVCacheSpec, cache_dtype: str) -> torch.dtype | None:
86 + """The fp8 dtype of a KV cache that stores plain fp8 values as uint8.
87 +
88 + Returns None for other layouts, including packed ones that embed scales
89 + or other metadata in the uint8 page.
90 + """
91 + if (
92 + not isinstance(spec, AttentionSpec)
93 + or spec.dtype != torch.uint8
94 + or spec.kv_quant_mode != KVQuantMode.FP8_PER_TENSOR
95 + or spec.state_content_bytes is not None
96 + ):
97 + return None
98 + cache_dtype = getattr(spec, "cache_dtype_str", None) or cache_dtype
99 + if cache_dtype in ("fp8", "fp8_e4m3"):
100 + return current_platform.fp8_dtype()
101 + if cache_dtype == "fp8_e5m2":
102 + return torch.float8_e5m2
103 + return None
104 +
105 +
106 +def _get_startup_fill_group_ids(kv_cache_config: "KVCacheConfig") -> set[int]:
107 + """KV cache groups that startup fill covers.
108 +
109 + Circular buffers are zero-filled per request. Attention groups whose new
110 + blocks the engine zeroes before use are filled per request, since the
111 + zeroing would wipe the startup fill.
112 + """
113 + group_ids = set()
114 + for group_idx, group in enumerate(kv_cache_config.kv_cache_groups):
115 + spec = group.kv_cache_spec
116 + if any(isinstance(s, CircularBufferSpec) for s in iter_layer_specs(spec)):
117 + continue
118 + if kv_cache_config.needs_kv_cache_zeroing and isinstance(spec, AttentionSpec):
119 + continue
120 + group_ids.add(group_idx)
121 + return group_ids
122 +
123 +
124 @dataclass
125 class DecodeBenchConnectorMetadata(KVConnectorMetadata):
126 """Metadata for DecodeBenchConnector.
@@ -191,6 +248,14 @@ class DecodeBenchConnectorScheduler:
248 """Scheduler-side implementation for DecodeBenchConnector."""
249
250 def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
251 + kv_transfer_config = vllm_config.kv_transfer_config
252 + assert kv_transfer_config is not None
253 + self.startup_fill = kv_transfer_config.get_from_extra_config(
254 + "startup_fill", False
255 + )
256 + self._startup_fill_group_ids = (
257 + _get_startup_fill_group_ids(kv_cache_config) if self.startup_fill else set()
258 + )
259 dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
260 self.kv_cache_groups = kv_cache_config.kv_cache_groups
261 self.group_block_sizes = tuple(
@@ -275,6 +340,9 @@ class DecodeBenchConnectorScheduler:
340 strict=True,
341 )
342 ):
343 + if group_idx in self._startup_fill_group_ids:
344 + block_ids_per_group_list.append([])
345 + continue
346 is_circular_buffer = all(
347 isinstance(spec, CircularBufferSpec)
348 for spec in iter_layer_specs(group.kv_cache_spec)
@@ -310,6 +378,9 @@ class DecodeBenchConnectorScheduler:
378 )
379 block_ids_per_group_list.append(block_ids)
380 block_ids_per_group = tuple(block_ids_per_group_list)
381 + self._filled_requests.add(req_id)
382 + if self.startup_fill and not any(block_ids_per_group):
383 + return
384
385 # Store the blocks to fill for all group. _pending_fills doesn't need cleanup
386 # as it's cleared after build_connector_meta
@@ -317,7 +388,6 @@ class DecodeBenchConnectorScheduler:
388 block_ids_per_group,
389 num_external_tokens,
390 )
320 - self._filled_requests.add(req_id)
391
392 block_counts = tuple(len(group) for group in block_ids_per_group)
393 logger.debug(
@@ -356,10 +426,27 @@ class DecodeBenchConnectorWorker:
426 assert kv_transfer_config is not None
427 self.fill_mean = kv_transfer_config.get_from_extra_config("fill_mean", 0.015)
428 self.fill_std = kv_transfer_config.get_from_extra_config("fill_std", 0.0)
429 + self.startup_fill = kv_transfer_config.get_from_extra_config(
430 + "startup_fill", False
431 + )
432 + self._startup_fill_group_ids = (
433 + _get_startup_fill_group_ids(kv_cache_config) if self.startup_fill else set()
434 + )
435
436 # Will be populated via register_kv_caches
437 self.kv_caches: dict[str, torch.Tensor] | None = None
438
439 + # Layers whose uint8 KV cache stores plain fp8 values
440 + self._fp8_dtypes: dict[str, torch.dtype] = {}
441 + for group in kv_cache_config.kv_cache_groups:
442 + for layer_name in group.layer_names:
443 + spec = group.kv_cache_spec
444 + if isinstance(spec, UniformTypeKVCacheSpecs):
445 + spec = spec.kv_cache_specs[layer_name]
446 + fp8_dtype = _get_fp8_dtype(spec, vllm_config.cache_config.cache_dtype)
447 + if fp8_dtype is not None:
448 + self._fp8_dtypes[layer_name] = fp8_dtype
449 +
450 # Mapping from KV cache group index to list of layer names in that group
451 self.group_to_layers = {
452 group_idx: list(group.layer_names)
@@ -373,6 +460,17 @@ class DecodeBenchConnectorWorker:
460 for spec in iter_layer_specs(group.kv_cache_spec)
461 )
462 }
463 + num_zeroed_groups = len(
464 + self.group_to_layers.keys()
465 + - self._startup_fill_group_ids
466 + - self._zero_fill_group_ids
467 + )
468 + if self.startup_fill and num_zeroed_groups:
469 + logger.info(
470 + "DecodeBenchConnector: the engine zeroes new blocks of %d KV "
471 + "cache groups before use, so they are still filled per request.",
472 + num_zeroed_groups,
473 + )
474
475 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
476 """Store references to the KV cache tensors."""
@@ -383,6 +481,48 @@ class DecodeBenchConnectorWorker:
481 len(kv_caches),
482 )
483
484 + if self.startup_fill:
485 + self._fill_all_kv_caches()
486 +
487 + def _fill_all_kv_caches(self):
488 + """Fill the KV cache tensors of the startup-filled groups once, in
489 + place, with dummy values, so that no per-request fill is needed.
490 + """
491 + assert self.kv_caches is not None
492 +
493 + num_layers = 0
494 + for group_idx in sorted(self._startup_fill_group_ids):
495 + for layer_name in self.group_to_layers[group_idx]:
496 + kv_cache = self.kv_caches.get(layer_name)
497 + if isinstance(kv_cache, torch.Tensor):
498 + tensors = [kv_cache]
499 + elif isinstance(kv_cache, (list, tuple)) and all(
500 + isinstance(t, torch.Tensor) for t in kv_cache
501 + ):
502 + tensors = list(kv_cache)
503 + else:
504 + logger.warning_once(
505 + "DecodeBenchConnector: skipping startup fill for layer "
506 + "%s whose KV cache is %s, not a tensor or a list/tuple "
507 + "of tensors.",
508 + layer_name,
509 + type(kv_cache).__name__,
510 + )
511 + continue
512 + for tensor in tensors:
513 + fill_dtype = self._fp8_dtypes.get(layer_name, tensor.dtype)
514 + self._fill_tensor(tensor, self.fill_mean, self.fill_std, fill_dtype)
515 + num_layers += 1
516 +
517 + logger.info(
518 + "DecodeBenchConnector: Filled %d KV cache layers at startup with %s "
519 + "values (mean=%.3f, std=%.3f)",
520 + num_layers,
521 + "random" if self.fill_std > 0 else "constant",
522 + self.fill_mean,
523 + self.fill_std,
524 + )
525 +
526 def start_fill_kv(self, metadata: DecodeBenchConnectorMetadata):
527 """Fill the allocated KV cache blocks with dummy values.
528
@@ -455,12 +595,17 @@ class DecodeBenchConnectorWorker:
595 # dimension — so fill each tensor in its entirety with the same
596 # dummy values.
597 if isinstance(kv_cache, torch.Tensor):
458 - self._fill_block_tensor(kv_cache, block_ids, fill_mean, fill_std)
598 + fill_dtype = self._fp8_dtypes.get(layer_name, kv_cache.dtype)
599 + self._fill_block_tensor(
600 + kv_cache, block_ids, fill_mean, fill_std, fill_dtype
601 + )
602 elif isinstance(kv_cache, (list, tuple)) and all(
603 isinstance(t, torch.Tensor) for t in kv_cache
604 ):
605 for state_tensor in kv_cache:
463 - self._fill_state_tensor(state_tensor, fill_mean, fill_std)
606 + self._fill_tensor(
607 + state_tensor, fill_mean, fill_std, state_tensor.dtype
608 + )
609 else:
610 logger.warning_once(
611 "DecodeBenchConnector: skipping fill for layer %s whose KV "
@@ -486,6 +631,7 @@ class DecodeBenchConnectorWorker:
631 block_ids: list[int],
632 fill_mean: float,
633 fill_std: float,
634 + fill_dtype: torch.dtype,
635 ):
636 """Fill the requested block rows of a block-indexed KV cache tensor.
637
@@ -495,6 +641,8 @@ class DecodeBenchConnectorWorker:
641 tensor's first dim are ignored.
642 fill_mean: Mean value for the fill.
643 fill_std: Standard deviation for the fill.
644 + fill_dtype: Dtype the values are encoded as, e.g. the fp8 dtype
645 + of a uint8 fp8 cache.
646
647 """
648 # Convert block_ids to tensor on device
@@ -509,46 +657,85 @@ class DecodeBenchConnectorWorker:
657 if len(valid_block_ids) == 0:
658 return
659
512 - # Create fill values - either constant or random
513 - block_shape = kv_cache.shape[1:]
514 - if fill_std > 0:
515 - # Random normal sampling
516 - fill_values = torch.normal(
517 - mean=fill_mean,
518 - std=fill_std,
519 - size=(len(valid_block_ids),) + block_shape,
520 - dtype=kv_cache.dtype,
521 - device=kv_cache.device,
522 - )
523 - else:
524 - # Constant fill value
525 - fill_values = torch.full(
526 - (len(valid_block_ids),) + block_shape,
527 - fill_mean,
528 - dtype=kv_cache.dtype,
529 - device=kv_cache.device,
530 - )
660 + fill_values = self._make_fill_values(
661 + (len(valid_block_ids),) + kv_cache.shape[1:],
662 + kv_cache.device,
663 + fill_dtype,
664 + fill_mean,
665 + fill_std,
666 + )
667
668 # Batch fill operation
533 - kv_cache[valid_block_ids] = fill_values
669 + kv_cache[valid_block_ids] = fill_values.view(kv_cache.dtype)
670 +
671 + def _make_fill_values(
672 + self,
673 + size: tuple[int, ...],
674 + device: torch.device,
675 + dtype: torch.dtype,
676 + fill_mean: float,
677 + fill_std: float,
678 + ) -> torch.Tensor:
679 + """Create constant or random fill values of ``dtype``, clamped to its
680 + finite range. Non-floating dtypes (packed layouts) are filled with
681 + zeros.
682 + """
683 + if not dtype.is_floating_point:
684 + if fill_mean or fill_std:
685 + logger.warning_once(
686 + "DecodeBenchConnector: %s KV caches do not hold plain "
687 + "floating-point values; filling them with zeros.",
688 + dtype,
689 + )
690 + return torch.zeros(size, dtype=dtype, device=device)
691 +
692 + finfo = torch.finfo(dtype)
693 + if fill_std > 0:
694 + sample_dtype = torch.float32
695 + if dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
696 + sample_dtype = dtype
697 + values = torch.normal(
698 + fill_mean, fill_std, size=size, device=device, dtype=sample_dtype
699 + )
700 + return values.clamp_(finfo.min, finfo.max).to(dtype)
701 + fill_mean = min(max(fill_mean, finfo.min), finfo.max)
702 + return torch.full(size, fill_mean, dtype=dtype, device=device)
703
535 - def _fill_state_tensor(
536 - self, kv_cache: torch.Tensor, fill_mean: float, fill_std: float
704 + def _fill_tensor(
705 + self,
706 + kv_cache: torch.Tensor,
707 + fill_mean: float,
708 + fill_std: float,
709 + fill_dtype: torch.dtype,
710 ):
538 - """Fill an entire non-block-indexed state tensor with dummy values.
711 + """Fill an entire tensor in place with dummy values.
712
540 - Hybrid / linear-attention layers (e.g. Mamba, Kimi Delta Attention)
541 - store their per-layer state as tensors with no num_blocks dimension,
542 - so the whole tensor is filled with the same constant or random values
543 - used for block fills, rather than selected block rows.
713 + Used for startup fills, and for hybrid / linear-attention layers (e.g.
714 + Mamba, Kimi Delta Attention) whose per-layer state tensors are filled
715 + in their entirety with the same constant or random values used for
716 + block fills, rather than selected block rows.
717
718 Args:
546 - kv_cache: A state tensor to fill in its entirety.
719 + kv_cache: A tensor to fill in its entirety.
720 fill_mean: Mean value for the fill.
721 fill_std: Standard deviation for the fill.
722 + fill_dtype: Dtype the values are encoded as, e.g. the fp8 dtype
723 + of a uint8 fp8 cache.
724
725 """
551 - if fill_std > 0:
552 - kv_cache.normal_(mean=fill_mean, std=fill_std)
553 - else:
554 - kv_cache.fill_(fill_mean)
726 + view = kv_cache.view(fill_dtype)
727 + if fill_std == 0 or not fill_dtype.is_floating_point:
728 + view.fill_(
729 + self._make_fill_values(
730 + (), torch.device("cpu"), fill_dtype, fill_mean, fill_std
731 + )
732 + )
733 + return
734 + # Sample in chunks to bound the temporary memory.
735 + rows = max(1, (1 << 24) // max(1, view[0].numel()))
736 + for chunk in view.split(rows):
737 + chunk.copy_(
738 + self._make_fill_values(
739 + chunk.shape, chunk.device, fill_dtype, fill_mean, fill_std
740 + )
741 + )