From 5093126bbd277c65cd817e4a16cd6e3703fd3008 Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Mon, 20 Jul 2026 12:31:21 -0500 Subject: [PATCH] feat(dcp): add depth-N native CKV prefetch --- .../test_b12x_ckv_prefetch_policy.py | 310 +++++++++++++ vllm/envs.py | 6 + .../attention/backends/mla/b12x_mla_sparse.py | 419 +++++++++++++----- 3 files changed, 623 insertions(+), 112 deletions(-) create mode 100644 tests/v1/attention/test_b12x_ckv_prefetch_policy.py diff --git a/tests/v1/attention/test_b12x_ckv_prefetch_policy.py b/tests/v1/attention/test_b12x_ckv_prefetch_policy.py new file mode 100644 index 000000000000..d9e22918c1c2 --- /dev/null +++ b/tests/v1/attention/test_b12x_ckv_prefetch_policy.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import gc +from types import SimpleNamespace + +import pytest +import torch + +import vllm.v1.worker.workspace as workspace +from vllm.v1.attention.backends.mla.b12x_mla_sparse import ( + B12xMLASparseImpl, + _ckv_prefetch_ring_slots, + _ckv_prefetch_supports_format, + _ckv_prefetch_target_indices, + _ckv_workspace_identity, + _CKVPrefetchStateRegistry, +) + + +@pytest.mark.parametrize( + ("depth", "expected_slots", "expected_targets"), + [ + (0, 1, []), + (1, 2, [2]), + (3, 4, [2, 3, 4]), + ], +) +def test_ckv_prefetch_depth_controls_ring_and_targets( + depth, expected_slots, expected_targets +): + caches = [torch.empty(0) for _ in range(6)] + + assert _ckv_prefetch_ring_slots(depth) == expected_slots + assert _ckv_prefetch_target_indices(1, depth, caches, {}) == expected_targets + + +def test_ckv_prefetch_supports_native_full_record_formats(): + assert _ckv_prefetch_supports_format("nvfp4_ds_mla") + assert _ckv_prefetch_supports_format("fp8_ds_mla") + assert not _ckv_prefetch_supports_format("auto") + + +def test_ckv_prefetch_targets_stop_at_first_unregistered_layer(): + caches = [torch.empty(0), torch.empty(0), torch.empty(0), None, torch.empty(0)] + pending = {2: (object(), 0)} + + assert _ckv_prefetch_target_indices(1, 3, caches, pending) == [] + + +def test_ckv_workspace_reuses_local_staging_across_ring_slots(): + impl = object.__new__(B12xMLASparseImpl) + impl._ckv_gather_enabled = True + impl._ckv_workspace_slots = 4 + impl._ckv_local_capacity = 8 + impl._kv_record_bytes = 432 + impl.dcp_world_size = 4 + impl.device = torch.device("cpu") + impl._ckv_workspace_nbytes = ( + (1 + impl._ckv_workspace_slots * impl.dcp_world_size) + * impl._ckv_local_capacity + * impl._kv_record_bytes + ) + workspace = torch.empty(impl._ckv_workspace_nbytes, dtype=torch.uint8) + + local_0, gathered_0 = impl._ckv_workspace_views(workspace, 0) + local_3, gathered_3 = impl._ckv_workspace_views(workspace, 3) + + assert local_0.data_ptr() == local_3.data_ptr() + assert gathered_0.shape == gathered_3.shape == (32, 432) + assert gathered_0.data_ptr() != gathered_3.data_ptr() + + +def test_ckv_workspace_rejects_ring_slot_outside_depth(): + impl = object.__new__(B12xMLASparseImpl) + impl._ckv_gather_enabled = True + impl._ckv_workspace_slots = 2 + impl._ckv_local_capacity = 1 + impl._kv_record_bytes = 432 + impl.dcp_world_size = 2 + impl.device = torch.device("cpu") + impl._ckv_workspace_nbytes = 5 * impl._kv_record_bytes + workspace = torch.empty(impl._ckv_workspace_nbytes, dtype=torch.uint8) + + with pytest.raises(ValueError, match="outside"): + impl._ckv_workspace_views(workspace, 2) + + +class _FakeEvent: + def __init__(self): + self.wait_calls = 0 + + def wait(self): + self.wait_calls += 1 + + +def test_ckv_prefetch_incomplete_step_recovers_with_stream_wait(): + registry = _CKVPrefetchStateRegistry() + state = registry.for_workspace(torch.empty(16, dtype=torch.uint8)) + cache = torch.empty(0) + event = _FakeEvent() + state.register_cache(3, cache) + state.pending_layers[3] = (event, 0) + state.last_layer_idx = 2 + + state.enter_layer(0) + + assert event.wait_calls == 1 + assert state.pending_layers == {} + assert state.layer_caches[3] is cache + assert state.last_layer_idx == 0 + + +def test_ckv_prefetch_first_request_discovers_caches_without_lookahead(): + registry = _CKVPrefetchStateRegistry() + state = registry.for_workspace(torch.empty(16, dtype=torch.uint8)) + caches = [torch.empty(0) for _ in range(4)] + + for layer_idx, cache in enumerate(caches): + state.enter_layer(layer_idx) + state.register_cache(layer_idx, cache) + + assert ( + _ckv_prefetch_target_indices( + layer_idx, 3, state.layer_caches, state.pending_layers + ) + == [] + ) + assert state.gather_stream is None + + state.enter_layer(0) + + assert _ckv_prefetch_target_indices(0, 3, state.layer_caches, {}) == [1, 2, 3] + + +def test_ckv_prefetch_target_and_draft_lifecycles_are_isolated(monkeypatch): + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2) + (target_workspace,) = manager.get_simultaneous(((16,), torch.uint8)) + with workspace.use_workspace_lane(1): + (draft_workspace,) = manager.get_simultaneous(((16,), torch.uint8)) + + registry = _CKVPrefetchStateRegistry() + target_state = registry.for_workspace(target_workspace) + draft_state = registry.for_workspace(draft_workspace) + target_cache = torch.empty(0) + draft_cache = torch.empty(0) + target_event = _FakeEvent() + target_ring = target_state.get_ckv_workspace(64) + draft_ring = draft_state.get_ckv_workspace(64) + + target_state.register_cache(1, target_cache) + target_state.pending_layers[1] = (target_event, 1) + draft_state.register_cache(1, draft_cache) + + assert target_state is not draft_state + assert target_ring.untyped_storage().data_ptr() != ( + draft_ring.untyped_storage().data_ptr() + ) + assert target_state.layer_caches[1] is target_cache + assert target_state.pending_layers[1] == (target_event, 1) + assert draft_state.layer_caches[1] is draft_cache + assert draft_state.pending_layers == {} + + +def test_ckv_prefetch_lazily_owns_one_stream_per_workspace_lane(monkeypatch): + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2) + (target_workspace,) = manager.get_simultaneous(((16,), torch.uint8)) + (target_workspace_reused,) = manager.get_simultaneous(((16,), torch.uint8)) + with workspace.use_workspace_lane(1): + (draft_workspace,) = manager.get_simultaneous(((16,), torch.uint8)) + + created_streams = [] + + def create_stream(*, device): + stream = SimpleNamespace(device=device) + created_streams.append(stream) + return stream + + monkeypatch.setattr(torch.cuda, "Stream", create_stream) + registry = _CKVPrefetchStateRegistry() + target_state = registry.for_workspace(target_workspace) + target_state_reused = registry.for_workspace(target_workspace_reused) + draft_state = registry.for_workspace(draft_workspace) + + assert created_streams == [] + assert target_state_reused is target_state + assert target_state.get_gather_stream() is target_state.get_gather_stream() + assert draft_state.get_gather_stream() is draft_state.get_gather_stream() + assert target_state.gather_stream is not draft_state.gather_stream + assert len(created_streams) == 2 + + +def test_ckv_prefetch_ring_survives_intervening_workspace_borrow(monkeypatch): + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=1) + (lane_workspace,) = manager.get_simultaneous(((256,), torch.uint8)) + registry = _CKVPrefetchStateRegistry() + state = registry.for_workspace(lane_workspace) + + assert state.ckv_workspace is None + ring = state.get_ckv_workspace(64) + ring.fill_(0xA5) + + # WorkspaceManager callers all borrow from offset zero. An intervening + # indexer/MoE scratch allocation must not alias cross-layer CKV state. + (intervening_workspace,) = manager.get_simultaneous(((128,), torch.uint8)) + intervening_workspace.zero_() + + assert ring.untyped_storage().data_ptr() != ( + intervening_workspace.untyped_storage().data_ptr() + ) + assert torch.all(ring == 0xA5) + + +def test_ckv_prefetch_ring_resize_drains_pending_generation(): + registry = _CKVPrefetchStateRegistry() + state = registry.for_workspace(torch.empty(16, dtype=torch.uint8)) + + first_ring = state.get_ckv_workspace(64) + assert state.get_ckv_workspace(64) is first_ring + assert state.ckv_workspace_generation == 1 + + event = _FakeEvent() + state.pending_layers[1] = (event, 0) + resized_ring = state.get_ckv_workspace(128) + + assert resized_ring is not first_ring + assert resized_ring.numel() == 128 + assert state.ckv_workspace_generation == 2 + assert event.wait_calls == 1 + assert state.pending_layers == {} + + +def test_ckv_prefetch_workspace_identity_invalidates_changed_geometry(): + registry = _CKVPrefetchStateRegistry() + workspace_buffer = torch.empty(16, dtype=torch.uint8) + same_geometry = workspace_buffer.view_as(workspace_buffer) + changed_geometry = workspace_buffer[:8] + old_state = registry.for_workspace(workspace_buffer) + event = _FakeEvent() + old_state.pending_layers[1] = (event, 0) + + assert registry.for_workspace(same_geometry) is old_state + + resized_state = registry.for_workspace(changed_geometry) + + assert resized_state is not old_state + assert event.wait_calls == 1 + assert len(registry.states) == 1 + + +def test_ckv_prefetch_workspace_identity_tracks_manager_resize(monkeypatch): + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2) + (first_workspace,) = manager.get_simultaneous(((16,), torch.uint8)) + with workspace.use_workspace_lane(1): + (draft_workspace,) = manager.get_simultaneous(((16,), torch.uint8)) + registry = _CKVPrefetchStateRegistry() + cache = torch.empty(0) + first_state = registry.for_workspace(first_workspace, 0, cache) + first_state.register_cache(0, cache) + draft_cache = torch.empty(0) + draft_state = registry.for_workspace(draft_workspace, 0, draft_cache) + draft_state.register_cache(0, draft_cache) + event = _FakeEvent() + first_state.pending_layers[1] = (event, 0) + + (resized_workspace,) = manager.get_simultaneous(((257,), torch.uint8)) + resized_state = registry.for_workspace(resized_workspace, 0, cache) + + assert ( + _ckv_workspace_identity(first_workspace).storage_generation + != _ckv_workspace_identity(resized_workspace).storage_generation + ) + assert resized_state is not first_state + assert resized_state.layer_caches == [] + assert event.wait_calls == 1 + assert registry.for_workspace(draft_workspace) is draft_state + assert len(registry.states) == 2 + + +def test_ckv_prefetch_registry_retires_released_profile_workspace(): + registry = _CKVPrefetchStateRegistry() + profile_workspace = torch.empty(16, dtype=torch.uint8) + state = registry.for_workspace(profile_workspace) + event = _FakeEvent() + state.pending_layers[1] = (event, 0) + + del profile_workspace + gc.collect() + registry.begin_step() + + assert event.wait_calls == 1 + assert registry.states == {} + + +def test_ckv_gather_uses_capture_fallback_without_reading_prefetch_state( + monkeypatch, +): + impl = object.__new__(B12xMLASparseImpl) + impl._ckv_gather_enabled = True + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + assert not impl.dcp_prefill_ckv_gather_eligible(SimpleNamespace(), 128) diff --git a/vllm/envs.py b/vllm/envs.py index e3be1405db8b..b37216ff1888 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -79,6 +79,7 @@ VLLM_B12X_MLA_CKV_GATHER: bool = False VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS: int = 16 VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS: int = 524288 + VLLM_B12X_MLA_CKV_PREFETCH_DEPTH: int = 1 VLLM_MINIMAX_M3_ENABLE_TORCH_COMPILE: bool = False VLLM_B12X_CUDAGRAPH_PIECEWISE_PREWARM: bool = False VLLM_B12X_MOE_FORCE_MODELOPT_PREP: bool = False @@ -1180,6 +1181,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS": lambda: int( os.getenv("VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS", "524288") ), + # Number of future full-CKV layer gathers to queue. Zero keeps the + # synchronous gather path without allocating lookahead ring slots. + "VLLM_B12X_MLA_CKV_PREFETCH_DEPTH": lambda: int( + os.getenv("VLLM_B12X_MLA_CKV_PREFETCH_DEPTH", "1") + ), # Diagnostic flag retained for local experiments. MiniMax M3 compile is # fail-closed in the model until the no-break path is validated. "VLLM_MINIMAX_M3_ENABLE_TORCH_COMPILE": lambda: bool( diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 4df3117dc4ef..9b2478eda476 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -27,6 +27,7 @@ import inspect import os +import weakref from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, cast @@ -77,7 +78,6 @@ _EXTEND_PREWARM_DONE: set[ tuple[int | None, int, int, int, int, int, bool, str, bool] ] = set() -_CKV_GATHER_WORKSPACES: dict[tuple[str, int | None], torch.Tensor] = {} _KV_FP8_ROPE_REQUESTED = os.getenv("KV_FP8_ROPE", "0") == "1" @@ -146,18 +146,173 @@ def _env_int(name: str, default: int) -> int: return parsed -def _get_ckv_gather_workspace(device: torch.device, nbytes: int) -> torch.Tensor: - key = (device.type, device.index) - workspace = _CKV_GATHER_WORKSPACES.get(key) - if workspace is None: - workspace = torch.empty((nbytes,), dtype=torch.uint8, device=device) - _CKV_GATHER_WORKSPACES[key] = workspace - elif workspace.numel() < nbytes: - raise RuntimeError( - "CKV gather workspace cannot grow after attention layers retain " - f"aliases: existing={workspace.numel()} requested={nbytes}" +def _ckv_prefetch_supports_format(kv_cache_dtype: str) -> bool: + return kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla") + + +def _ckv_prefetch_ring_slots(depth: int) -> int: + return max(0, int(depth)) + 1 + + +def _ckv_prefetch_target_indices( + layer_idx: int, + depth: int, + layer_caches: list[torch.Tensor | None], + pending_layers: dict[int, tuple[Any, int]], +) -> list[int]: + targets: list[int] = [] + for distance in range(1, max(0, int(depth)) + 1): + target_idx = layer_idx + distance + if target_idx in pending_layers: + continue + if target_idx >= len(layer_caches) or layer_caches[target_idx] is None: + break + targets.append(target_idx) + return targets + + +@dataclass(frozen=True) +class _CKVWorkspaceIdentity: + device: torch.device + storage_generation: int + data_ptr: int + storage_offset: int + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: torch.dtype + + +def _ckv_workspace_identity(workspace: torch.Tensor) -> _CKVWorkspaceIdentity: + storage = workspace.untyped_storage() + return _CKVWorkspaceIdentity( + device=workspace.device, + # WorkspaceManager does not expose its allocation generation. PyTorch's + # storage handle changes when MRV2 replaces the backing allocation. + storage_generation=int(storage._cdata), + data_ptr=workspace.data_ptr(), + storage_offset=workspace.storage_offset(), + shape=tuple(workspace.shape), + stride=tuple(workspace.stride()), + dtype=workspace.dtype, + ) + + +class _CKVPrefetchState: + """Cross-layer state for one workspace allocation and execution lane.""" + + def __init__( + self, + workspace_identity: _CKVWorkspaceIdentity, + workspace: torch.Tensor, + ) -> None: + self.workspace_identity = workspace_identity + self.workspace_storage_ref = weakref.ref(workspace.untyped_storage()) + self.layer_caches: list[torch.Tensor | None] = [] + self.pending_layers: dict[int, tuple[Any, int]] = {} + self.gather_stream: torch.cuda.Stream | None = None + self.ckv_workspace: torch.Tensor | None = None + self.ckv_workspace_generation = 0 + self.last_layer_idx: int | None = None + + def begin_step(self) -> None: + for event, _ in self.pending_layers.values(): + # Preserve ring ordering without blocking the host indefinitely. + # The next main-stream gather is enqueued after these dependencies. + event.wait() + self.pending_layers.clear() + self.last_layer_idx = None + + def enter_layer(self, layer_idx: int) -> None: + if self.last_layer_idx is not None and layer_idx <= self.last_layer_idx: + self.begin_step() + self.last_layer_idx = layer_idx + + def register_cache(self, layer_idx: int, kv_cache: torch.Tensor) -> None: + while len(self.layer_caches) <= layer_idx: + self.layer_caches.append(None) + self.layer_caches[layer_idx] = kv_cache + + def get_gather_stream(self) -> torch.cuda.Stream: + if self.gather_stream is None: + self.gather_stream = torch.cuda.Stream( + device=self.workspace_identity.device + ) + return self.gather_stream + + def get_ckv_workspace(self, nbytes: int) -> torch.Tensor: + if nbytes <= 0: + raise ValueError(f"CKV workspace size must be positive, got {nbytes}") + workspace = self.ckv_workspace + if workspace is None or workspace.numel() != nbytes: + # Pending side-stream writes must be ordered before replacing their + # backing allocation. record_stream() in _dcp_gather_ckv keeps the + # old storage alive until those writes complete. + self.begin_step() + workspace = torch.empty( + (nbytes,), + dtype=torch.uint8, + device=self.workspace_identity.device, + ) + self.ckv_workspace = workspace + self.ckv_workspace_generation += 1 + return workspace + + +class _CKVPrefetchStateRegistry: + """Builder-owned states partitioned by lane-scoped CKV workspace.""" + + def __init__(self) -> None: + self.states: dict[_CKVWorkspaceIdentity, _CKVPrefetchState] = {} + + def _retire(self, identities: list[_CKVWorkspaceIdentity]) -> None: + for identity in identities: + self.states.pop(identity).begin_step() + + def _prune_released_workspaces(self) -> None: + self._retire( + [ + identity + for identity, state in self.states.items() + if state.workspace_storage_ref() is None + ] ) - return workspace[:nbytes] + + def begin_step(self) -> None: + self._prune_released_workspaces() + for state in self.states.values(): + state.begin_step() + + def for_workspace( + self, + workspace: torch.Tensor, + layer_idx: int | None = None, + kv_cache: torch.Tensor | None = None, + ) -> _CKVPrefetchState: + self._prune_released_workspaces() + identity = _ckv_workspace_identity(workspace) + state = self.states.get(identity) + if state is None: + # A resized view may retain its address, while an allocation + # replacement changes it. A known layer cache identifies the same + # execution lane across the latter without merging target/draft. + stale_identities = [ + existing + for existing, existing_state in self.states.items() + if ( + existing.device == identity.device + and existing.data_ptr == identity.data_ptr + ) + or ( + layer_idx is not None + and kv_cache is not None + and layer_idx < len(existing_state.layer_caches) + and existing_state.layer_caches[layer_idx] is kv_cache + ) + ] + self._retire(stale_identities) + state = _CKVPrefetchState(identity, workspace) + self.states[identity] = state + return state def _dcp_all_gather_current_stream( @@ -531,6 +686,7 @@ class B12xMLASparseMetadata(AttentionMetadata): dcp_local_total_tokens: int = 0 dcp_padded_total_tokens: int = 0 dcp_ckv_gather_eligible: bool = False + ckv_prefetch_registry: _CKVPrefetchStateRegistry | None = None block_size: int = 64 topk_tokens: int = 2048 @@ -571,6 +727,9 @@ def __init__( from vllm import envs as envs_mod ckv_gather_requested = envs_mod.VLLM_B12X_MLA_CKV_GATHER + self.ckv_prefetch_registry = ( + _CKVPrefetchStateRegistry() if ckv_gather_requested else None + ) # Max-batched-token scratch buffers so cudagraph capture sees stable # allocations (sliced per build()). self.cache_seq_lens_per_token_buffer = torch.empty( @@ -672,16 +831,8 @@ def build( from vllm import envs as envs_mod - if envs_mod.VLLM_B12X_MLA_CKV_GATHER: - # Reset the cross-layer prefetch pipeline once per step so the - # first layer always sync-gathers. The event/buf-idx are class - # state that would otherwise leak across chunks (the last layer - # of a chunk consumes but never re-arms), scheduling layer 0 - # of subsequent chunks onto a stale gathered buffer. The - # layer->cache registry is intentionally left intact (stable - # cache pointers across chunks). - B12xMLASparseImpl._shared_gather_event = None - B12xMLASparseImpl._shared_gather_buf_idx = 0 + if self.ckv_prefetch_registry is not None: + self.ckv_prefetch_registry.begin_step() if ( use_dcp @@ -892,6 +1043,7 @@ def build( dcp_local_total_tokens=dcp_local_total_tokens, dcp_padded_total_tokens=dcp_padded_total_tokens, dcp_ckv_gather_eligible=dcp_ckv_gather_eligible, + ckv_prefetch_registry=self.ckv_prefetch_registry, block_size=self.kv_cache_spec.block_size, topk_tokens=self.topk_tokens, ) @@ -1157,30 +1309,6 @@ def _make_plan( int(self._extend_plan.layout.nbytes), ) - # Pre-touch q-concat + the attention scratch TOGETHER so the workspace - # manager grows during warmup (before lock_workspace() runs - # post-cudagraph-capture) and so the two always come from ONE - # get_simultaneous call -> distinct, non-overlapping offsets. The manager - # packs every call from offset 0, so borrowing q and the scratch the kernel - # writes in separate calls would alias them. - workspace_specs: list[tuple[tuple[int, ...], torch.dtype]] = [ - ( - (max_batched, self._kernel_num_heads, self.q_head_dim), - torch.bfloat16, - ) - ] - if self._pad_heads: - workspace_specs.append( - ( - (max_batched, self._input_num_heads, self.kv_lora_rank), - torch.bfloat16, - ) - ) - workspace_specs.append(((self._scratch_nbytes,), torch.uint8)) - self._workspace_specs = tuple(workspace_specs) - self._borrow_workspaces() - self._prewarm_extend_kernels_once(max_batched) - # CKV gather setup (Fix B). from vllm import envs as envs_mod @@ -1201,6 +1329,20 @@ def _make_plan( self._ckv_kernel_num_heads = self.num_heads self._ckv_gather_max_tokens = envs_mod.VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS self._ckv_gather_min_tokens = envs_mod.VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS + configured_prefetch_depth = int(envs_mod.VLLM_B12X_MLA_CKV_PREFETCH_DEPTH) + if configured_prefetch_depth < 0: + logger.warning_once( + "Ignoring negative VLLM_B12X_MLA_CKV_PREFETCH_DEPTH=%d; using 0", + configured_prefetch_depth, + ) + self._ckv_prefetch_supported = ( + self._ckv_gather_enabled + and _ckv_prefetch_supports_format(self.kv_cache_dtype) + ) + self._ckv_prefetch_depth = ( + max(0, configured_prefetch_depth) if self._ckv_prefetch_supported else 0 + ) + self._ckv_workspace_slots = _ckv_prefetch_ring_slots(self._ckv_prefetch_depth) self._ckv_local_capacity = ( _cdiv( _cdiv(self._ckv_gather_max_tokens, max(1, self.dcp_world_size)) @@ -1210,18 +1352,12 @@ def _make_plan( * self.block_size ) self._ckv_workspace_nbytes = ( - 2 - * (self.dcp_world_size + 1) + (1 + self._ckv_workspace_slots * self.dcp_world_size) * self._ckv_local_capacity * self._kv_record_bytes if self._ckv_gather_enabled else 0 ) - self._ckv_workspace = ( - _get_ckv_gather_workspace(self.device, self._ckv_workspace_nbytes) - if self._ckv_gather_enabled - else None - ) # Separate extend plan for the gathered-cache path: full local heads # (no head all-gather), global seq lens. @@ -1236,19 +1372,34 @@ def _make_plan( else: self._ckv_extend_plan = None - # Layer prefetch (side stream + events + ping-pong). - # _shared_* are class-level: layer L kicks off the prefetch for - # layer L+1, and layer L+1 (a different impl instance) consumes it. - self._ckv_prefetch_supported = self._ckv_gather_enabled and ( - self.kv_cache_dtype == "fp8_ds_mla" or self._kv_fp8_rope - ) + # Pre-touch q-concat and attention scratch together. Cross-layer CKV + # data cannot live in WorkspaceManager: every caller borrows from + # offset zero, so intervening indexer/MoE scratch would alias it. The + # builder-owned state allocates a dedicated ring per workspace lane. + workspace_specs: list[tuple[tuple[int, ...], torch.dtype]] = [ + ( + (max_batched, self._kernel_num_heads, self.q_head_dim), + torch.bfloat16, + ) + ] + if self._pad_heads: + workspace_specs.append( + ( + (max_batched, self._input_num_heads, self.kv_lora_rank), + torch.bfloat16, + ) + ) + workspace_specs.append(((self._scratch_nbytes,), torch.uint8)) + self._workspace_specs = tuple(workspace_specs) + self._borrow_workspaces() + self._prewarm_extend_kernels_once(max_batched) + + # The builder-owned registry lazily creates one stream per workspace + # lane after cache discovery finds the first lookahead target. CUDA + # capture keeps the existing non-CKV fallback. if self._ckv_gather_enabled: - self._ckv_gather_stream = torch.cuda.Stream(device=self.device) self._ckv_current_chunk_kv_c: torch.Tensor | None = None self._ckv_current_chunk_kpe: torch.Tensor | None = None - B12xMLASparseImpl._all_layer_kv_caches: list[torch.Tensor | None] = [] - B12xMLASparseImpl._shared_gather_event: torch.cuda.Event | None = None - B12xMLASparseImpl._shared_gather_buf_idx = 0 if not self._ckv_prefetch_supported: logger.warning_once( "CKV gather prefetch disabled for kv_cache_dtype=%s " @@ -1256,8 +1407,14 @@ def _make_plan( self.kv_cache_dtype, int(self._kv_fp8_rope), ) + elif self._ckv_prefetch_depth > 0: + logger.info_once( + "Using native CKV layer prefetch with depth=%d and " + "%d workspace slots.", + self._ckv_prefetch_depth, + self._ckv_workspace_slots, + ) else: - self._ckv_gather_stream = None self._ckv_current_chunk_kv_c = None self._ckv_current_chunk_kpe = None @@ -1593,6 +1750,27 @@ def _validate_ckv_workspace(self, ckv_workspace: torch.Tensor) -> None: ): raise RuntimeError("B12X CKV gather borrowed an invalid workspace") + def _ckv_workspace_views( + self, ckv_workspace: torch.Tensor, buf_idx: int + ) -> tuple[torch.Tensor, torch.Tensor]: + self._validate_ckv_workspace(ckv_workspace) + if not 0 <= int(buf_idx) < self._ckv_workspace_slots: + raise ValueError( + f"CKV gather buffer index {buf_idx} is outside " + f"[0, {self._ckv_workspace_slots})" + ) + records = ckv_workspace.view(-1, self._kv_record_bytes) + local_buffer = records[: self._ckv_local_capacity] + gathered_base = ( + self._ckv_local_capacity + + buf_idx * self.dcp_world_size * self._ckv_local_capacity + ) + gathered_buffer = records[ + gathered_base : gathered_base + + self.dcp_world_size * self._ckv_local_capacity + ] + return local_buffer, gathered_buffer + def dcp_prefill_ckv_gather_eligible( self, attn_metadata: B12xMLASparseMetadata, @@ -1650,18 +1828,11 @@ def _dcp_gather_ckv( assert attn_metadata.dcp_local_cu_seq_lens is not None padded_tokens = attn_metadata.dcp_padded_total_tokens local_tokens = attn_metadata.dcp_local_total_tokens - self._validate_ckv_workspace(ckv_workspace) - half_nbytes = ( - (self.dcp_world_size + 1) * self._ckv_local_capacity * self._kv_record_bytes + local_buffer, gathered_buffer = self._ckv_workspace_views( + ckv_workspace, buf_idx ) - ws_half = ckv_workspace.view(-1, self._kv_record_bytes) - base = buf_idx * (half_nbytes // self._kv_record_bytes) - local_buffer = ws_half[base : base + self._ckv_local_capacity] - gathered_buffer = ws_half[ - base + self._ckv_local_capacity : base - + self._ckv_local_capacity * (self.dcp_world_size + 1) - ] if stream is not None: + ckv_workspace.record_stream(stream) # The side stream must observe the default stream's prior writes # to the paged KV cache (this and earlier steps' do_kv_cache_update) # before gathering history off it; there is otherwise no ordering @@ -1810,7 +1981,7 @@ def _append_current_chunk_to_gathered( slots, k_scale, ) - elif self.kv_cache_dtype == "fp8_ds_mla": + elif self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla"): ops.concat_and_cache_mla( kv_c, k_pe_flat, @@ -1823,7 +1994,7 @@ def _append_current_chunk_to_gathered( raise RuntimeError( "CKV gather prefetch append is not yet supported for " f"kv_cache_dtype={self.kv_cache_dtype!r}; disable prefetch " - "or use fp8_ds_mla / KV_FP8_ROPE." + "or use fp8_ds_mla / nvfp4_ds_mla." ) def _sync_warmup(self) -> None: @@ -1953,7 +2124,6 @@ def forward_mqa( workspace_tensors = self._borrow_workspaces() q_workspace = workspace_tensors[0] dense_out_workspace = workspace_tensors[1] if self._pad_heads else None - ckv_workspace = self._ckv_workspace scratch_storage = workspace_tensors[-1] expected_input_heads = ( self.num_heads if use_ckv_gather else self._input_num_heads @@ -2113,28 +2283,28 @@ def forward_mqa( f"got stride={tuple(kv_cache.stride())}" ) if use_ckv_gather: - if ckv_workspace is None: - raise RuntimeError("CKV gather workspace was not borrowed") layer_idx = self._resolve_layer_index(layer) - if layer_idx is not None: - while len(B12xMLASparseImpl._all_layer_kv_caches) <= layer_idx: - B12xMLASparseImpl._all_layer_kv_caches.append(None) - B12xMLASparseImpl._all_layer_kv_caches[layer_idx] = kv_cache - if B12xMLASparseImpl._shared_gather_event is not None: - B12xMLASparseImpl._shared_gather_event.wait() - half_nbytes = ( - (self.dcp_world_size + 1) - * self._ckv_local_capacity - * self._kv_record_bytes - ) - ws_half = ckv_workspace.view(-1, self._kv_record_bytes) - base = B12xMLASparseImpl._shared_gather_buf_idx * ( - half_nbytes // self._kv_record_bytes + prefetch_registry = attn_metadata.ckv_prefetch_registry + if prefetch_registry is None: + raise RuntimeError("CKV gather requires a prefetch state registry") + prefetch_state = prefetch_registry.for_workspace( + q_workspace, layer_idx, kv_cache + ) + ckv_workspace = prefetch_state.get_ckv_workspace(self._ckv_workspace_nbytes) + if layer_idx is not None and self._ckv_prefetch_depth > 0: + prefetch_state.enter_layer(layer_idx) + prefetch_state.register_cache(layer_idx, kv_cache) + pending = ( + prefetch_state.pending_layers.pop(layer_idx, None) + if layer_idx is not None and self._ckv_prefetch_depth > 0 + else None + ) + if pending is not None: + gather_event, current_buf_idx = pending + gather_event.wait() + _, gathered_buffer = self._ckv_workspace_views( + ckv_workspace, current_buf_idx ) - gathered_buffer = ws_half[ - base + self._ckv_local_capacity : base - + self._ckv_local_capacity * (self.dcp_world_size + 1) - ] kv_cache = gathered_buffer[ : self.dcp_world_size * self._ckv_local_capacity ].view(-1, self.block_size, self._kv_record_bytes) @@ -2142,7 +2312,17 @@ def forward_mqa( kv_cache, attn_metadata, layer, num_actual_toks ) else: - kv_cache = self._dcp_gather_ckv(kv_cache, attn_metadata, ckv_workspace) + current_buf_idx = ( + layer_idx % self._ckv_workspace_slots + if layer_idx is not None + else 0 + ) + kv_cache = self._dcp_gather_ckv( + kv_cache, + attn_metadata, + ckv_workspace, + buf_idx=current_buf_idx, + ) logger.info_once( "Using transient full-CKV gather for B12X sparse MLA prefill " "(capacity=%d logical tokens)", @@ -2150,24 +2330,39 @@ def forward_mqa( ) if ( self._ckv_prefetch_supported + and self._ckv_prefetch_depth > 0 and layer_idx is not None - and layer_idx + 1 < len(B12xMLASparseImpl._all_layer_kv_caches) - and B12xMLASparseImpl._all_layer_kv_caches[layer_idx + 1] is not None ): - next_kv = B12xMLASparseImpl._all_layer_kv_caches[layer_idx + 1] - next_buf_idx = 1 - B12xMLASparseImpl._shared_gather_buf_idx - self._dcp_gather_ckv( - next_kv, - attn_metadata, - ckv_workspace, - buf_idx=next_buf_idx, - stream=self._ckv_gather_stream, + # The first eligible request only discovers one layer cache at + # a time. It intentionally has no lookahead and primes the + # cache registry for subsequent requests. + targets = _ckv_prefetch_target_indices( + layer_idx, + self._ckv_prefetch_depth, + prefetch_state.layer_caches, + prefetch_state.pending_layers, ) - B12xMLASparseImpl._shared_gather_event = torch.cuda.Event( - blocking=False + prefetch_stream = ( + prefetch_state.get_gather_stream() if targets else None ) - B12xMLASparseImpl._shared_gather_event.record(self._ckv_gather_stream) - B12xMLASparseImpl._shared_gather_buf_idx = next_buf_idx + for target_idx in targets: + assert prefetch_stream is not None + target_kv = prefetch_state.layer_caches[target_idx] + assert target_kv is not None + target_buf_idx = target_idx % self._ckv_workspace_slots + self._dcp_gather_ckv( + target_kv, + attn_metadata, + ckv_workspace, + buf_idx=target_buf_idx, + stream=prefetch_stream, + ) + target_event = torch.cuda.Event(blocking=False) + target_event.record(prefetch_stream) + prefetch_state.pending_layers[target_idx] = ( + target_event, + target_buf_idx, + ) use_decode_kernel = attn_metadata.max_query_len <= 1 or ( self.spec_extend_as_decode