From 7a9406690ebe2fc47236a81cf7ec33188c6565e2 Mon Sep 17 00:00:00 2001 From: Zijing Liu Date: Thu, 24 Sep 2026 16:55:47 -0700 Subject: [PATCH] [KV Connector] Fix DecodeBench fp8 fill values and add a startup fill mode (#58472) Signed-off-by: Zijing Liu Signed-off-by: Yifan Qiao Co-authored-by: Yifan Qiao Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- docs/serving/expert_parallel_deployment.md | 2 +- .../unit/test_decode_bench_connector.py | 326 +++++++++++++++++- .../kv_connector/v1/decode_bench_connector.py | 259 ++++++++++++-- 3 files changed, 549 insertions(+), 38 deletions(-) diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index d1dff32fc9..ea60233645 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -315,6 +315,6 @@ except Exception as e: ### Benchmarking -- 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. +- 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). - **CUDAGraph capture**: Use `--compilation_config '{"cudagraph_mode": "FULL_DECODE_ONLY"}'` to enable CUDA graph capture for decode only and save KV cache. diff --git a/tests/v1/kv_connector/unit/test_decode_bench_connector.py b/tests/v1/kv_connector/unit/test_decode_bench_connector.py index e1054d9ef8..693c10d40a 100644 --- a/tests/v1/kv_connector/unit/test_decode_bench_connector.py +++ b/tests/v1/kv_connector/unit/test_decode_bench_connector.py @@ -18,6 +18,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector import DecodeBenchConnectorMetadata, ) from vllm.forward_context import ForwardContext +from vllm.platforms import current_platform from vllm.utils.hashing import sha256 from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.kv_cache_utils import ( @@ -34,6 +35,8 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVQuantMode, + MambaSpec, SlidingWindowSpec, ) from vllm.v1.request import Request @@ -50,7 +53,12 @@ from .utils import ( class DecodeBenchTestRunner: """Test runner for DecodeBenchConnector.""" - def __init__(self, block_size: int, num_gpu_blocks: int): + def __init__( + self, + block_size: int, + num_gpu_blocks: int, + kv_connector_extra_config: dict | None = None, + ): self.block_size = block_size self.num_gpu_blocks = num_gpu_blocks @@ -61,6 +69,7 @@ class DecodeBenchTestRunner: block_size=block_size, max_num_batched_tokens=1000, kv_connector="DecodeBenchConnector", + kv_connector_extra_config=kv_connector_extra_config, ) self.vllm_config = vllm_config @@ -780,5 +789,320 @@ def test_decode_bench_connector_concurrent_requests(): assert len(metadata2.reqs_to_fill) == 0 +def _fill_single_block(cache_dtype, spec_dtype, kv_quant_mode, cache, fill_std=0.0): + """Fill block 1 of a one-layer cache through the worker connector.""" + vllm_config = create_vllm_config( + block_size=16, + kv_connector="DecodeBenchConnector", + kv_connector_extra_config={"fill_std": fill_std}, + cache_dtype=cache_dtype, + ) + kv_cache_spec = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=8, + dtype=spec_dtype, + kv_quant_mode=kv_quant_mode, + ) + connector = DecodeBenchConnector( + vllm_config, + KVConnectorRole.WORKER, + KVCacheConfig( + num_blocks=cache.shape[0], + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)], + ), + ) + connector.register_kv_caches({"layer": cache}) + connector.bind_connector_metadata( + DecodeBenchConnectorMetadata(reqs_to_fill={"request": (([1],), 16)}) + ) + connector.start_load_kv( + ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) + ) + + +@pytest.mark.parametrize( + "cache_dtype,spec_dtype,kv_quant_mode,fp8_dtype", + [ + ( + "fp8", + torch.uint8, + KVQuantMode.FP8_PER_TENSOR, + current_platform.fp8_dtype(), + ), + ("fp8_e5m2", torch.uint8, KVQuantMode.FP8_PER_TENSOR, torch.float8_e5m2), + ("auto", torch.float8_e4m3fn, KVQuantMode.NONE, torch.float8_e4m3fn), + ], +) +@pytest.mark.parametrize("fill_std", [0.0, 0.1]) +def test_decode_bench_connector_fills_fp8_caches( + cache_dtype, spec_dtype, kv_quant_mode, fp8_dtype, fill_std +): + """fp8 caches, including ones stored as uint8, get encoded fill values.""" + cache = torch.zeros(4, 16, 8, dtype=spec_dtype) + _fill_single_block(cache_dtype, spec_dtype, kv_quant_mode, cache, fill_std) + + values = cache.view(fp8_dtype).float() + assert torch.count_nonzero(values[[0, 2, 3]]) == 0 + filled = values[1] + if fill_std == 0: + expected = torch.tensor(0.015).to(fp8_dtype).float() + torch.testing.assert_close(filled, expected.expand_as(filled)) + else: + assert torch.isfinite(filled).all() + assert filled.std() > 0 + + +def test_decode_bench_connector_zero_fills_packed_uint8_caches(): + """uint8 layouts that do not hold plain fp8 values are zeroed.""" + cache = torch.full((4, 16, 8), 7, dtype=torch.uint8) + _fill_single_block("nvfp4", torch.uint8, KVQuantMode.NVFP4, cache) + + assert torch.count_nonzero(cache[1]) == 0 + assert torch.all(cache[[0, 2, 3]] == 7) + + +@pytest.mark.parametrize("fill_std", [0.0, 0.1]) +def test_decode_bench_connector_startup_fill_fills_registered_caches(fill_std): + """Startup fill writes every cache view in place, including strided views.""" + vllm_config = create_vllm_config( + block_size=16, + kv_connector="DecodeBenchConnector", + kv_connector_extra_config={"startup_fill": True, "fill_std": fill_std}, + ) + kv_cache_config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["attention", "transposed"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + ], + ) + connector = DecodeBenchConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + attention = torch.full((4, 2, 16), -1.0) + transposed = torch.full((2, 4, 16), -1.0).transpose(0, 1) + connector.register_kv_caches({"attention": attention, "transposed": transposed}) + + for filled in (attention, transposed): + if fill_std == 0: + torch.testing.assert_close(filled, torch.full_like(filled, 0.015)) + else: + assert torch.isfinite(filled).all() + assert filled.std() > 0 + + +@pytest.mark.parametrize("fill_std", [0.0, 0.1]) +def test_decode_bench_connector_startup_fill_encodes_fp8_caches(fill_std): + """Startup fill encodes values in the fp8 dtype of a uint8 fp8 cache.""" + vllm_config = create_vllm_config( + block_size=16, + kv_connector="DecodeBenchConnector", + kv_connector_extra_config={"startup_fill": True, "fill_std": fill_std}, + cache_dtype="fp8", + ) + kv_cache_spec = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=8, + dtype=torch.uint8, + kv_quant_mode=KVQuantMode.FP8_PER_TENSOR, + ) + connector = DecodeBenchConnector( + vllm_config, + KVConnectorRole.WORKER, + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)], + ), + ) + cache = torch.zeros(4, 16, 8, dtype=torch.uint8) + connector.register_kv_caches({"layer": cache}) + + values = cache.view(current_platform.fp8_dtype()).float() + if fill_std == 0: + expected = torch.tensor(0.015).to(current_platform.fp8_dtype()).float() + torch.testing.assert_close(values, expected.expand_as(values)) + else: + assert torch.isfinite(values).all() + assert values.std() > 0 + + +def test_decode_bench_connector_startup_fill_keeps_zeroed_groups_per_request(): + """Attention blocks the engine zeroes on allocation are filled per request.""" + vllm_config = create_vllm_config( + block_size=16, + max_num_batched_tokens=1000, + kv_connector="DecodeBenchConnector", + kv_connector_extra_config={"startup_fill": True}, + ) + # Mamba layers make the engine zero newly allocated attention blocks. + kv_cache_config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=16, + shapes=((4,), (4,)), + dtypes=(torch.float32, torch.float32), + ), + ), + ], + ) + assert kv_cache_config.needs_kv_cache_zeroing + scheduler_connector = DecodeBenchConnector( + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config + ) + num_external_tokens = 17 + request = create_request( + request_id=1, num_tokens=num_external_tokens + 1, block_size=16 + ) + blocks = KVCacheBlocks( + ( + [KVCacheBlock(block_id=1), KVCacheBlock(block_id=2)], + [KVCacheBlock(block_id=3)], + ) + ) + scheduler_connector.update_state_after_alloc(request, blocks, num_external_tokens) + metadata = scheduler_connector.build_connector_meta(SchedulerOutput.make_empty()) + assert isinstance(metadata, DecodeBenchConnectorMetadata) + assert metadata.reqs_to_fill == { + request.request_id: (([1, 2], []), num_external_tokens) + } + + worker_connector = DecodeBenchConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + attention = torch.full((4, 2), -1.0) + mamba_states = (torch.zeros(4, 4), torch.zeros(4, 4)) + worker_connector.register_kv_caches({"attention": attention, "mamba": mamba_states}) + assert torch.all(attention == -1.0) + + worker_connector.bind_connector_metadata(metadata) + worker_connector.start_load_kv( + ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) + ) + torch.testing.assert_close(attention[1:3], torch.full_like(attention[1:3], 0.015)) + assert torch.all(attention[[0, 3]] == -1.0) + + +def test_decode_bench_connector_startup_fill_skips_per_step_fill(): + """Admitted requests need no per-step fill and leave existing KV intact.""" + block_size = 16 + runner = DecodeBenchTestRunner( + block_size=block_size, + num_gpu_blocks=100, + kv_connector_extra_config={"startup_fill": True}, + ) + for kv_cache in runner.kv_caches.values(): + torch.testing.assert_close(kv_cache, torch.full_like(kv_cache, 0.015)) + kv_cache.fill_(7.0) + + req1 = runner.new_request([1] * (block_size * 2)) + scheduler_output, metadata = runner.run_single_step() + assert scheduler_output.num_scheduled_tokens == {req1.request_id: 1} + assert scheduler_output.scheduled_new_reqs[0].num_computed_tokens == ( + block_size * 2 - 1 + ) + assert metadata.reqs_to_fill == {} + + req2 = runner.new_request([2] * (block_size * 3)) + scheduler_output, metadata = runner.run_single_step() + assert scheduler_output.num_scheduled_tokens == { + req1.request_id: 1, + req2.request_id: 1, + } + assert metadata.reqs_to_fill == {} + + # A re-admitted (e.g. preempted) request is not reported again. + assert runner.scheduler_connector.get_num_new_matched_tokens(req1, 0) == ( + 0, + False, + ) + for kv_cache in runner.kv_caches.values(): + assert torch.all(kv_cache == 7.0) + + +def test_decode_bench_connector_startup_fill_keeps_circular_buffer_zero_fill(): + """Circular buffers are still zeroed per request with startup fill.""" + vllm_config = create_vllm_config( + block_size=16, + max_num_batched_tokens=1000, + kv_connector="DecodeBenchConnector", + kv_connector_extra_config={"startup_fill": True}, + ) + kv_cache_config = KVCacheConfig( + num_blocks=8, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full_attention"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["compressor_ring"], + CircularBufferSpec( + block_size=32, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + ], + ) + scheduler_connector = DecodeBenchConnector( + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config + ) + num_external_tokens = 17 + request = create_request( + request_id=1, num_tokens=num_external_tokens + 1, block_size=16 + ) + blocks = KVCacheBlocks( + ( + [KVCacheBlock(block_id=1), KVCacheBlock(block_id=2)], + [KVCacheBlock(block_id=4)], + ) + ) + scheduler_connector.update_state_after_alloc(request, blocks, num_external_tokens) + metadata = scheduler_connector.build_connector_meta(SchedulerOutput.make_empty()) + assert isinstance(metadata, DecodeBenchConnectorMetadata) + assert metadata.reqs_to_fill == { + request.request_id: (([], [4]), num_external_tokens) + } + + worker_connector = DecodeBenchConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + attention_cache = torch.zeros(8, 2) + ring_cache = torch.ones(8, 2) + worker_connector.register_kv_caches( + {"full_attention": attention_cache, "compressor_ring": ring_cache} + ) + assert torch.all(ring_cache == 1.0) + worker_connector.bind_connector_metadata(metadata) + worker_connector.start_load_kv( + ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) + ) + + torch.testing.assert_close(attention_cache, torch.full_like(attention_cache, 0.015)) + expected_ring = torch.ones_like(ring_cache) + expected_ring[4] = 0 + assert torch.equal(ring_cache, expected_ring) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py index cff2367966..a9255f7cda 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py @@ -28,6 +28,16 @@ Example: - fill_mean (float): Mean value for random normal fill (default: 0.015) - fill_std (float): Standard deviation for random fill (default: 0.0) Set to 0 for constant values, >0 for random sampling + - startup_fill (bool): Fill the whole KV cache once at startup instead + of filling each new request's blocks in the step that schedules it + (default: False). Circular-buffer caches are still zeroed per + request, and attention caches whose new blocks the engine zeroes + before use (hybrid Mamba or mixed-precision KV caches) are still + filled per request. The fill is not redone after sleep mode + discards the KV cache. + + The dummy KV cache is only meant for performance measurement; outputs are + not meaningful for accuracy evaluation. """ @@ -45,13 +55,21 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( SupportsHMA, ) from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.kv_cache_utils import ( dcp_world_size_for_kv_cache_spec, resolve_dcp_kv_block_size, ) -from vllm.v1.kv_cache_interface import CircularBufferSpec, iter_layer_specs +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + CircularBufferSpec, + KVCacheSpec, + KVQuantMode, + UniformTypeKVCacheSpecs, + iter_layer_specs, +) if TYPE_CHECKING: from vllm.config import VllmConfig @@ -64,6 +82,45 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _get_fp8_dtype(spec: KVCacheSpec, cache_dtype: str) -> torch.dtype | None: + """The fp8 dtype of a KV cache that stores plain fp8 values as uint8. + + Returns None for other layouts, including packed ones that embed scales + or other metadata in the uint8 page. + """ + if ( + not isinstance(spec, AttentionSpec) + or spec.dtype != torch.uint8 + or spec.kv_quant_mode != KVQuantMode.FP8_PER_TENSOR + or spec.state_content_bytes is not None + ): + return None + cache_dtype = getattr(spec, "cache_dtype_str", None) or cache_dtype + if cache_dtype in ("fp8", "fp8_e4m3"): + return current_platform.fp8_dtype() + if cache_dtype == "fp8_e5m2": + return torch.float8_e5m2 + return None + + +def _get_startup_fill_group_ids(kv_cache_config: "KVCacheConfig") -> set[int]: + """KV cache groups that startup fill covers. + + Circular buffers are zero-filled per request. Attention groups whose new + blocks the engine zeroes before use are filled per request, since the + zeroing would wipe the startup fill. + """ + group_ids = set() + for group_idx, group in enumerate(kv_cache_config.kv_cache_groups): + spec = group.kv_cache_spec + if any(isinstance(s, CircularBufferSpec) for s in iter_layer_specs(spec)): + continue + if kv_cache_config.needs_kv_cache_zeroing and isinstance(spec, AttentionSpec): + continue + group_ids.add(group_idx) + return group_ids + + @dataclass class DecodeBenchConnectorMetadata(KVConnectorMetadata): """Metadata for DecodeBenchConnector. @@ -191,6 +248,14 @@ class DecodeBenchConnectorScheduler: """Scheduler-side implementation for DecodeBenchConnector.""" def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"): + kv_transfer_config = vllm_config.kv_transfer_config + assert kv_transfer_config is not None + self.startup_fill = kv_transfer_config.get_from_extra_config( + "startup_fill", False + ) + self._startup_fill_group_ids = ( + _get_startup_fill_group_ids(kv_cache_config) if self.startup_fill else set() + ) dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size self.kv_cache_groups = kv_cache_config.kv_cache_groups self.group_block_sizes = tuple( @@ -275,6 +340,9 @@ class DecodeBenchConnectorScheduler: strict=True, ) ): + if group_idx in self._startup_fill_group_ids: + block_ids_per_group_list.append([]) + continue is_circular_buffer = all( isinstance(spec, CircularBufferSpec) for spec in iter_layer_specs(group.kv_cache_spec) @@ -310,6 +378,9 @@ class DecodeBenchConnectorScheduler: ) block_ids_per_group_list.append(block_ids) block_ids_per_group = tuple(block_ids_per_group_list) + self._filled_requests.add(req_id) + if self.startup_fill and not any(block_ids_per_group): + return # Store the blocks to fill for all group. _pending_fills doesn't need cleanup # as it's cleared after build_connector_meta @@ -317,7 +388,6 @@ class DecodeBenchConnectorScheduler: block_ids_per_group, num_external_tokens, ) - self._filled_requests.add(req_id) block_counts = tuple(len(group) for group in block_ids_per_group) logger.debug( @@ -356,10 +426,27 @@ class DecodeBenchConnectorWorker: assert kv_transfer_config is not None self.fill_mean = kv_transfer_config.get_from_extra_config("fill_mean", 0.015) self.fill_std = kv_transfer_config.get_from_extra_config("fill_std", 0.0) + self.startup_fill = kv_transfer_config.get_from_extra_config( + "startup_fill", False + ) + self._startup_fill_group_ids = ( + _get_startup_fill_group_ids(kv_cache_config) if self.startup_fill else set() + ) # Will be populated via register_kv_caches self.kv_caches: dict[str, torch.Tensor] | None = None + # Layers whose uint8 KV cache stores plain fp8 values + self._fp8_dtypes: dict[str, torch.dtype] = {} + for group in kv_cache_config.kv_cache_groups: + for layer_name in group.layer_names: + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + spec = spec.kv_cache_specs[layer_name] + fp8_dtype = _get_fp8_dtype(spec, vllm_config.cache_config.cache_dtype) + if fp8_dtype is not None: + self._fp8_dtypes[layer_name] = fp8_dtype + # Mapping from KV cache group index to list of layer names in that group self.group_to_layers = { group_idx: list(group.layer_names) @@ -373,6 +460,17 @@ class DecodeBenchConnectorWorker: for spec in iter_layer_specs(group.kv_cache_spec) ) } + num_zeroed_groups = len( + self.group_to_layers.keys() + - self._startup_fill_group_ids + - self._zero_fill_group_ids + ) + if self.startup_fill and num_zeroed_groups: + logger.info( + "DecodeBenchConnector: the engine zeroes new blocks of %d KV " + "cache groups before use, so they are still filled per request.", + num_zeroed_groups, + ) def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Store references to the KV cache tensors.""" @@ -383,6 +481,48 @@ class DecodeBenchConnectorWorker: len(kv_caches), ) + if self.startup_fill: + self._fill_all_kv_caches() + + def _fill_all_kv_caches(self): + """Fill the KV cache tensors of the startup-filled groups once, in + place, with dummy values, so that no per-request fill is needed. + """ + assert self.kv_caches is not None + + num_layers = 0 + for group_idx in sorted(self._startup_fill_group_ids): + for layer_name in self.group_to_layers[group_idx]: + kv_cache = self.kv_caches.get(layer_name) + if isinstance(kv_cache, torch.Tensor): + tensors = [kv_cache] + elif isinstance(kv_cache, (list, tuple)) and all( + isinstance(t, torch.Tensor) for t in kv_cache + ): + tensors = list(kv_cache) + else: + logger.warning_once( + "DecodeBenchConnector: skipping startup fill for layer " + "%s whose KV cache is %s, not a tensor or a list/tuple " + "of tensors.", + layer_name, + type(kv_cache).__name__, + ) + continue + for tensor in tensors: + fill_dtype = self._fp8_dtypes.get(layer_name, tensor.dtype) + self._fill_tensor(tensor, self.fill_mean, self.fill_std, fill_dtype) + num_layers += 1 + + logger.info( + "DecodeBenchConnector: Filled %d KV cache layers at startup with %s " + "values (mean=%.3f, std=%.3f)", + num_layers, + "random" if self.fill_std > 0 else "constant", + self.fill_mean, + self.fill_std, + ) + def start_fill_kv(self, metadata: DecodeBenchConnectorMetadata): """Fill the allocated KV cache blocks with dummy values. @@ -455,12 +595,17 @@ class DecodeBenchConnectorWorker: # dimension — so fill each tensor in its entirety with the same # dummy values. if isinstance(kv_cache, torch.Tensor): - self._fill_block_tensor(kv_cache, block_ids, fill_mean, fill_std) + fill_dtype = self._fp8_dtypes.get(layer_name, kv_cache.dtype) + self._fill_block_tensor( + kv_cache, block_ids, fill_mean, fill_std, fill_dtype + ) elif isinstance(kv_cache, (list, tuple)) and all( isinstance(t, torch.Tensor) for t in kv_cache ): for state_tensor in kv_cache: - self._fill_state_tensor(state_tensor, fill_mean, fill_std) + self._fill_tensor( + state_tensor, fill_mean, fill_std, state_tensor.dtype + ) else: logger.warning_once( "DecodeBenchConnector: skipping fill for layer %s whose KV " @@ -486,6 +631,7 @@ class DecodeBenchConnectorWorker: block_ids: list[int], fill_mean: float, fill_std: float, + fill_dtype: torch.dtype, ): """Fill the requested block rows of a block-indexed KV cache tensor. @@ -495,6 +641,8 @@ class DecodeBenchConnectorWorker: tensor's first dim are ignored. fill_mean: Mean value for the fill. fill_std: Standard deviation for the fill. + fill_dtype: Dtype the values are encoded as, e.g. the fp8 dtype + of a uint8 fp8 cache. """ # Convert block_ids to tensor on device @@ -509,46 +657,85 @@ class DecodeBenchConnectorWorker: if len(valid_block_ids) == 0: return - # Create fill values - either constant or random - block_shape = kv_cache.shape[1:] - if fill_std > 0: - # Random normal sampling - fill_values = torch.normal( - mean=fill_mean, - std=fill_std, - size=(len(valid_block_ids),) + block_shape, - dtype=kv_cache.dtype, - device=kv_cache.device, - ) - else: - # Constant fill value - fill_values = torch.full( - (len(valid_block_ids),) + block_shape, - fill_mean, - dtype=kv_cache.dtype, - device=kv_cache.device, - ) + fill_values = self._make_fill_values( + (len(valid_block_ids),) + kv_cache.shape[1:], + kv_cache.device, + fill_dtype, + fill_mean, + fill_std, + ) # Batch fill operation - kv_cache[valid_block_ids] = fill_values + kv_cache[valid_block_ids] = fill_values.view(kv_cache.dtype) - def _fill_state_tensor( - self, kv_cache: torch.Tensor, fill_mean: float, fill_std: float + def _make_fill_values( + self, + size: tuple[int, ...], + device: torch.device, + dtype: torch.dtype, + fill_mean: float, + fill_std: float, + ) -> torch.Tensor: + """Create constant or random fill values of ``dtype``, clamped to its + finite range. Non-floating dtypes (packed layouts) are filled with + zeros. + """ + if not dtype.is_floating_point: + if fill_mean or fill_std: + logger.warning_once( + "DecodeBenchConnector: %s KV caches do not hold plain " + "floating-point values; filling them with zeros.", + dtype, + ) + return torch.zeros(size, dtype=dtype, device=device) + + finfo = torch.finfo(dtype) + if fill_std > 0: + sample_dtype = torch.float32 + if dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64): + sample_dtype = dtype + values = torch.normal( + fill_mean, fill_std, size=size, device=device, dtype=sample_dtype + ) + return values.clamp_(finfo.min, finfo.max).to(dtype) + fill_mean = min(max(fill_mean, finfo.min), finfo.max) + return torch.full(size, fill_mean, dtype=dtype, device=device) + + def _fill_tensor( + self, + kv_cache: torch.Tensor, + fill_mean: float, + fill_std: float, + fill_dtype: torch.dtype, ): - """Fill an entire non-block-indexed state tensor with dummy values. + """Fill an entire tensor in place with dummy values. - Hybrid / linear-attention layers (e.g. Mamba, Kimi Delta Attention) - store their per-layer state as tensors with no num_blocks dimension, - so the whole tensor is filled with the same constant or random values - used for block fills, rather than selected block rows. + Used for startup fills, and for hybrid / linear-attention layers (e.g. + Mamba, Kimi Delta Attention) whose per-layer state tensors are filled + in their entirety with the same constant or random values used for + block fills, rather than selected block rows. Args: - kv_cache: A state tensor to fill in its entirety. + kv_cache: A tensor to fill in its entirety. fill_mean: Mean value for the fill. fill_std: Standard deviation for the fill. + fill_dtype: Dtype the values are encoded as, e.g. the fp8 dtype + of a uint8 fp8 cache. """ - if fill_std > 0: - kv_cache.normal_(mean=fill_mean, std=fill_std) - else: - kv_cache.fill_(fill_mean) + view = kv_cache.view(fill_dtype) + if fill_std == 0 or not fill_dtype.is_floating_point: + view.fill_( + self._make_fill_values( + (), torch.device("cpu"), fill_dtype, fill_mean, fill_std + ) + ) + return + # Sample in chunks to bound the temporary memory. + rows = max(1, (1 << 24) // max(1, view[0].numel())) + for chunk in view.split(rows): + chunk.copy_( + self._make_fill_values( + chunk.shape, chunk.device, fill_dtype, fill_mean, fill_std + ) + )