[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>

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.

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"])

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
+                )
+            )