diff --git a/tests/v1/attention/test_b12x_sparse_ckv_decode_policy.py b/tests/v1/attention/test_b12x_sparse_ckv_decode_policy.py new file mode 100644 index 000000000000..465cf777fb9b --- /dev/null +++ b/tests/v1/attention/test_b12x_sparse_ckv_decode_policy.py @@ -0,0 +1,371 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.distributed import parallel_state +from vllm.v1.attention.backends.mla import b12x_sparse_ckv_decode +from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + build_dense_union_remap, + dense_union_remap_reference, + owner_and_local_ordinal, + plan_sparse_ckv_decode, + sparse_decode_batch_eligible, + sparse_decode_prefetch_targets, +) + + +def _layout(*, dcp: int = 4, requests: int = 8, pool: int = 0): + return plan_sparse_ckv_decode( + dcp_world_size=dcp, + topk=2048, + rows_per_request=4, + max_requests=requests, + pool_records=pool, + record_bytes=432, + prefetch_depth=3, + ) + + +@pytest.mark.parametrize("dcp", [2, 3, 4, 5, 6, 7, 8]) +def test_sparse_decode_layout_supports_dcp_two_through_eight(dcp): + layout = _layout(dcp=dcp) + + assert layout.dcp_world_size == dcp + assert layout.per_request_capacity == 8192 + assert layout.pool_records == 65536 + assert layout.max_fast_requests == 8 + assert layout.record_bytes == 432 + assert layout.workspace_slots == 4 + + +@pytest.mark.parametrize("dcp", [1, 9]) +def test_sparse_decode_layout_rejects_unsupported_dcp(dcp): + with pytest.raises(ValueError, match="DCP world sizes 2-8"): + _layout(dcp=dcp) + + +@pytest.mark.parametrize("concurrency", [1, 2, 4, 8]) +def test_sparse_decode_batch_policy_accepts_c1_through_c8(concurrency): + layout = _layout() + rows = concurrency * layout.rows_per_request + + assert sparse_decode_batch_eligible( + layout, + num_requests=concurrency, + num_rows=rows, + max_query_len=layout.rows_per_request, + num_actual_tokens=rows, + has_required_metadata=True, + ) + assert layout.active_records(concurrency) == concurrency * 8192 + + +def test_sparse_decode_pool_overflow_falls_back_before_transport(): + layout = _layout(pool=3 * 8192) + + assert layout.max_fast_requests == 3 + assert sparse_decode_batch_eligible( + layout, + num_requests=3, + num_rows=12, + max_query_len=4, + num_actual_tokens=12, + has_required_metadata=True, + ) + assert not sparse_decode_batch_eligible( + layout, + num_requests=4, + num_rows=16, + max_query_len=4, + num_actual_tokens=16, + has_required_metadata=True, + ) + with pytest.raises(ValueError, match="outside sparse capacity"): + layout.active_records(4) + + +def test_mtp3_union_is_per_sequence_dense_stable_and_exact(): + indices = torch.tensor( + [ + [ + [10, 11, 10, -1], + [12, 11, 13, 10], + [13, 14, 12, -1], + [10, 15, 14, 15], + ], + [ + [10, 20, 10, -1], + [21, 20, 22, 10], + [22, 23, 21, -1], + [10, 24, 23, 24], + ], + ], + dtype=torch.int32, + ) + + union, remap, counts = dense_union_remap_reference(indices) + + assert union[0, :6].tolist() == [10, 11, 12, 13, 14, 15] + assert union[1, :6].tolist() == [10, 20, 21, 22, 23, 24] + assert counts.tolist() == [6, 6] + for request in range(indices.shape[0]): + for row in range(indices.shape[1]): + for column in range(indices.shape[2]): + token = int(indices[request, row, column]) + slot = int(remap[request, row, column]) + if token < 0: + assert slot == -1 + else: + assert int(union[request, slot]) == token + + +def test_destination_unions_and_remaps_are_independent_and_exact(): + indices = torch.tensor( + [ + [[[1, 2], [2, 3]], [[10, 11], [11, 12]]], + [[[4, 5], [5, 6]], [[10, 13], [13, 14]]], + ], + dtype=torch.int32, + ) + + union, remap, counts = dense_union_remap_reference(indices) + + assert counts.tolist() == [[3, 3], [3, 3]] + assert union[0, 0, :3].tolist() == [1, 2, 3] + assert union[1, 0, :3].tolist() == [4, 5, 6] + assert union[0, 1, :3].tolist() == [10, 11, 12] + assert union[1, 1, :3].tolist() == [10, 13, 14] + reconstructed = torch.gather( + union.unsqueeze(-2).expand(*indices.shape[:-2], indices.shape[-2], -1), + -1, + remap.to(torch.int64), + ) + assert torch.equal(reconstructed, indices) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cuda_union_accepts_runtime_rows_below_reserved_mtp_rows(): + device = torch.device("cuda") + indices = torch.tensor([[[[1, 2]]], [[[3, 4]]]], dtype=torch.int32, device=device) + union = torch.empty((2, 2, 8), dtype=torch.int32, device=device) + remap = torch.empty((2, 2, 4, 2), dtype=torch.int32, device=device) + counts = torch.empty((2, 2), dtype=torch.int32, device=device) + hash_keys = torch.empty((2, 2, 16), dtype=torch.int32, device=device) + hash_first = torch.empty_like(hash_keys) + first_to_dense = torch.empty_like(union) + + build_dense_union_remap( + indices, + union, + remap, + counts, + hash_keys, + hash_first, + first_to_dense, + num_requests=1, + ) + torch.cuda.synchronize() + + assert counts[:, 0].cpu().tolist() == [2, 2] + assert union[0, 0, :2].cpu().tolist() == [1, 2] + assert union[1, 0, :2].cpu().tolist() == [3, 4] + assert remap[:, 0, 0, :].cpu().tolist() == [[0, 1], [0, 1]] + + +def test_union_preflight_caches_rank_consistent_success(monkeypatch): + process_group = object() + loads = 0 + reductions = 0 + monkeypatch.setattr(b12x_sparse_ckv_decode, "_UNION_PREFLIGHT_RESULTS", {}) + + def load_extension(): + nonlocal loads + loads += 1 + + def all_reduce(status, *, op, group): + nonlocal reductions + reductions += 1 + assert int(status.item()) == 1 + assert op is torch.distributed.ReduceOp.MIN + assert group is process_group + + monkeypatch.setattr( + b12x_sparse_ckv_decode, "preload_dense_union_extension", load_extension + ) + monkeypatch.setattr(b12x_sparse_ckv_decode.dist, "all_reduce", all_reduce) + + for _ in range(2): + b12x_sparse_ckv_decode.preload_dense_union_extension_consistently( + process_group, torch.device("cpu") + ) + + assert loads == 1 + assert reductions == 1 + + +def test_union_preflight_propagates_remote_rank_failure(monkeypatch): + process_group = object() + reductions = 0 + monkeypatch.setattr(b12x_sparse_ckv_decode, "_UNION_PREFLIGHT_RESULTS", {}) + + def all_reduce(status, *, op, group): + nonlocal reductions + reductions += 1 + status.zero_() + + monkeypatch.setattr( + b12x_sparse_ckv_decode, "preload_dense_union_extension", lambda: None + ) + monkeypatch.setattr(b12x_sparse_ckv_decode.dist, "all_reduce", all_reduce) + + for _ in range(2): + with pytest.raises(RuntimeError, match="another DCP rank"): + b12x_sparse_ckv_decode.preload_dense_union_extension_consistently( + process_group, torch.device("cpu") + ) + + assert reductions == 1 + + +def test_selected_record_cleanup_closes_once_and_clears_registries(monkeypatch): + class Exchange: + def __init__(self): + self.close_calls = 0 + + def close(self): + self.close_calls += 1 + + first = Exchange() + second = Exchange() + monkeypatch.setattr( + b12x_sparse_ckv_decode, + "_SELECTED_RECORD_EXCHANGES", + {("target",): first, ("target-alias",): first, ("draft",): second}, + ) + monkeypatch.setattr( + b12x_sparse_ckv_decode, + "_SELECTED_RECORD_STREAMS", + {id(first): object(), id(second): object()}, + ) + monkeypatch.setattr( + b12x_sparse_ckv_decode, + "_UNION_PREFLIGHT_RESULTS", + {(1, "cuda", 0): (True, "")}, + ) + + b12x_sparse_ckv_decode.close_selected_record_exchanges() + + assert first.close_calls == 1 + assert second.close_calls == 1 + assert b12x_sparse_ckv_decode._SELECTED_RECORD_EXCHANGES == {} + assert b12x_sparse_ckv_decode._SELECTED_RECORD_STREAMS == {} + assert b12x_sparse_ckv_decode._UNION_PREFLIGHT_RESULTS == {} + + +def test_model_parallel_cleanup_runs_before_group_destroy(monkeypatch): + events = [] + + class Group: + def destroy(self): + events.append("group") + + monkeypatch.setattr( + parallel_state, + "_MODEL_PARALLEL_CLEANUP_HOOKS", + [lambda: events.append("cleanup")], + ) + monkeypatch.setattr(parallel_state, "_TP", Group()) + for name in ( + "_DCP", + "_QUERY_SPLIT", + "_DCP_CKV_PREFETCH", + "_PCP", + "_PP", + "_DP", + "_EP", + "_EPLB", + ): + monkeypatch.setattr(parallel_state, name, None) + + parallel_state.destroy_model_parallel() + + assert events == ["cleanup", "group"] + + +@pytest.mark.parametrize("concurrency", [2, 4, 8]) +def test_multi_sequence_union_never_aliases_request_local_positions(concurrency): + rows = [] + for request in range(concurrency): + base = request * 1000 + rows.append( + [ + [0, base + 1, base + 2], + [base + 2, base + 3, 0], + [base + 3, base + 4, base + 1], + [base + 4, base + 5, 0], + ] + ) + indices = torch.tensor(rows, dtype=torch.int32) + + union, remap, counts = dense_union_remap_reference(indices) + + assert counts.tolist() == [6] * concurrency + pooled_union = union.reshape(-1) + for request in range(concurrency): + reconstructed = torch.gather( + union[request].expand(indices.shape[1], -1), + 1, + remap[request].to(torch.int64), + ) + assert torch.equal(reconstructed, indices[request]) + pooled_slots = remap[request].to(torch.int64) + request * union.shape[-1] + assert torch.equal(pooled_union[pooled_slots], indices[request]) + + +def test_shared_layer_prefetch_is_exactly_full_to_s1_s2_s3(): + # GLM: Full 0-2, Shared 3-5, then Full 6. + emits_topk = [True, True, True, False, False, False, True, False] + + assert sparse_decode_prefetch_targets(2, 3, emits_topk) == [3, 4, 5] + assert sparse_decode_prefetch_targets(1, 3, emits_topk) == [] + assert sparse_decode_prefetch_targets(6, 3, emits_topk) == [7] + + +@pytest.mark.parametrize("dcp", [2, 3, 4, 5, 6, 7, 8]) +@pytest.mark.parametrize("interleave", [1, 2, 4]) +def test_owner_arithmetic_round_trips_global_token(dcp, interleave): + for logical_token in range(1024): + owner, local = owner_and_local_ordinal( + logical_token, + dcp_world_size=dcp, + interleave=interleave, + ) + group = local // interleave + lane = local % interleave + reconstructed = (group * dcp + owner) * interleave + lane + + assert 0 <= owner < dcp + assert reconstructed == logical_token + + +def test_sparse_decode_requires_uniform_rows_and_complete_metadata(): + layout = _layout() + + assert not sparse_decode_batch_eligible( + layout, + num_requests=2, + num_rows=7, + max_query_len=4, + num_actual_tokens=7, + has_required_metadata=True, + ) + assert not sparse_decode_batch_eligible( + layout, + num_requests=2, + num_rows=8, + max_query_len=4, + num_actual_tokens=8, + has_required_metadata=False, + ) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 922265d1c8d9..69e74e7ca8a4 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1439,16 +1439,12 @@ def get_pp_group() -> GroupCoordinator: return _PP -def checkpoint_b12x_graph_channels() -> tuple[ - tuple[Callable[[Any], None], Any], ... -]: +def checkpoint_b12x_graph_channels() -> tuple[tuple[Callable[[Any], None], Any], ...]: """Snapshot B12X channels used by disposable CUDA graph captures.""" checkpoints: list[tuple[Callable[[Any], None], Any]] = [] seen_communicators: set[int] = set() for group in (_TP, _DCP, _PP): - device_communicator = ( - None if group is None else group.device_communicator - ) + device_communicator = None if group is None else group.device_communicator communicator = getattr(device_communicator, "ca_comm", None) if communicator is None or id(communicator) in seen_communicators: continue @@ -1975,14 +1971,15 @@ def initialize_model_parallel( group_name="query_split", ) - # A dedicated communicator over the DCP ranks for the transient ckv - # prefetch gather (Fix B). The prefetch runs on a side stream and would - # otherwise share the DCP communicator with the indexer's DCP top-k - # merge on the default stream; concurrent collectives on one NCCL - # communicator from two streams is unsupported. Same ranks as ``_DCP``. + # A dedicated communicator over the DCP ranks for full-CKV prefetch and + # selected-record decode exchange. These paths run on a side stream and + # must not share the DCP communicator with the indexer's top-k merge on + # the default stream. Same ranks as ``_DCP``. global _DCP_CKV_PREFETCH assert _DCP_CKV_PREFETCH is None, "DCP ckv prefetch group is already initialized" - if decode_context_model_parallel_size > 1 and envs.VLLM_B12X_MLA_CKV_GATHER: + if decode_context_model_parallel_size > 1 and ( + envs.VLLM_B12X_MLA_CKV_GATHER or envs.VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER + ): _DCP_CKV_PREFETCH = init_model_parallel_group( group_ranks, get_world_group().local_rank, @@ -2186,6 +2183,15 @@ def model_parallel_is_initialized(): _TP_STATE_PATCHED = False +_MODEL_PARALLEL_CLEANUP_HOOKS: list[Callable[[], None]] = [] + + +def register_model_parallel_cleanup_hook(callback: Callable[[], None]) -> None: + """Register idempotent cleanup that must run before groups are destroyed.""" + if callback not in _MODEL_PARALLEL_CLEANUP_HOOKS: + _MODEL_PARALLEL_CLEANUP_HOOKS.append(callback) + + def get_tensor_model_parallel_world_size() -> int: """Return world size for the tensor model parallel group.""" return get_tp_group().world_size @@ -2204,6 +2210,12 @@ def get_node_count() -> int: def destroy_model_parallel(): """Set the groups to none and destroy them.""" + for cleanup in _MODEL_PARALLEL_CLEANUP_HOOKS: + try: + cleanup() + except Exception: + logger.exception("Model-parallel cleanup hook failed") + global _TP if _TP: diff --git a/vllm/envs.py b/vllm/envs.py index b37216ff1888..1522ff3bb169 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -77,6 +77,9 @@ VLLM_DCP_GLOBAL_TOPK: bool = True VLLM_DCP_QUERY_SPLIT: bool = False VLLM_B12X_MLA_CKV_GATHER: bool = False + VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER: bool = False + VLLM_B12X_MLA_SPARSE_DECODE_MAX_SEQS: int = 8 + VLLM_B12X_MLA_SPARSE_DECODE_POOL_RECORDS: int = 0 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 @@ -1175,6 +1178,18 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_B12X_MLA_CKV_GATHER": lambda: ( os.getenv("VLLM_B12X_MLA_CKV_GATHER", "0").lower() in ("1", "true", "yes", "on") ), + # Exchange only sparse-MLA-selected native CKV records during DCP decode. + "VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER": lambda: ( + os.getenv("VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER", "0").lower() + in ("1", "true", "yes", "on") + ), + "VLLM_B12X_MLA_SPARSE_DECODE_MAX_SEQS": lambda: int( + os.getenv("VLLM_B12X_MLA_SPARSE_DECODE_MAX_SEQS", "8") + ), + # Zero sizes the pool for the full configured request count. + "VLLM_B12X_MLA_SPARSE_DECODE_POOL_RECORDS": lambda: int( + os.getenv("VLLM_B12X_MLA_SPARSE_DECODE_POOL_RECORDS", "0") + ), "VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS": lambda: int( os.getenv("VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS", "16") ), diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index d87241456b91..5134c8764f89 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1114,12 +1114,22 @@ def forward_impl( workspace_gather_used = False ckv_gather_used = False if self.impl.dcp_world_size > 1: - ckv_gather_selector = getattr( + prefill_ckv_gather_selector = getattr( self.impl, "dcp_prefill_ckv_gather_eligible", None ) - ckv_gather_used = bool( - callable(ckv_gather_selector) - and ckv_gather_selector(attn_metadata, num_mqa_tokens) + sparse_decode_ckv_gather_selector = getattr( + self.impl, "dcp_sparse_decode_ckv_gather_eligible", None + ) + prefill_ckv_gather_used = bool( + callable(prefill_ckv_gather_selector) + and prefill_ckv_gather_selector(attn_metadata, num_mqa_tokens) + ) + sparse_decode_ckv_gather_used = bool( + callable(sparse_decode_ckv_gather_selector) + and sparse_decode_ckv_gather_selector(attn_metadata, num_mqa_tokens) + ) + ckv_gather_used = ( + prefill_ckv_gather_used or sparse_decode_ckv_gather_used ) if not ckv_gather_used: if not self.impl.can_return_lse_for_decode: @@ -1176,8 +1186,7 @@ def forward_impl( mqa_q = torch.cat(mqa_q, dim=-1) if ckv_gather_used: logger.info_once( - "Keeping local query heads for transient full-CKV " - "B12X sparse MLA prefill" + "Keeping local query heads for B12X sparse MLA CKV gather" ) elif dcp_use_b12x: mqa_q = dcp_b12x_all_gather_heads( diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 9b2478eda476..0af5756b914e 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -208,7 +208,9 @@ def __init__( self.workspace_identity = workspace_identity self.workspace_storage_ref = weakref.ref(workspace.untyped_storage()) self.layer_caches: list[torch.Tensor | None] = [] + self.layer_impls: list[B12xMLASparseImpl | None] = [] self.pending_layers: dict[int, tuple[Any, int]] = {} + self.sparse_decode_state: Any | None = None self.gather_stream: torch.cuda.Stream | None = None self.ckv_workspace: torch.Tensor | None = None self.ckv_workspace_generation = 0 @@ -232,6 +234,28 @@ def register_cache(self, layer_idx: int, kv_cache: torch.Tensor) -> None: self.layer_caches.append(None) self.layer_caches[layer_idx] = kv_cache + def register_impl(self, layer_idx: int, impl: "B12xMLASparseImpl") -> None: + while len(self.layer_impls) <= layer_idx: + self.layer_impls.append(None) + self.layer_impls[layer_idx] = impl + + def get_sparse_decode_state(self, layout, exchange): + if self.sparse_decode_state is None: + from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + SparseCKVDecodeState, + ) + + self.sparse_decode_state = SparseCKVDecodeState( + layout=layout, + device=self.workspace_identity.device, + exchange=exchange, + ) + elif self.sparse_decode_state.layout != layout: + raise RuntimeError("sparse CKV layout changed within one execution lane") + elif self.sparse_decode_state.exchange is not exchange: + raise RuntimeError("sparse CKV exchange changed within one execution lane") + return self.sparse_decode_state + def get_gather_stream(self) -> torch.cuda.Stream: if self.gather_stream is None: self.gather_stream = torch.cuda.Stream( @@ -350,6 +374,94 @@ def _dcp_all_gather_current_stream( output_tensor.copy_(gathered) +@triton.jit +def _find_sparse_decode_current_token_slot_kernel( + union_indices_ptr, + req_id_ptr, + global_seq_len_ptr, + output_slot_ptr, + union_stride0, + capacity, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK_N) + req = tl.load(req_id_ptr + row) + selected = tl.load( + union_indices_ptr + req * union_stride0 + offsets, + mask=offsets < capacity, + other=-1, + ) + current_token = tl.load(global_seq_len_ptr + row) - 1 + matched = tl.max(tl.where(selected == current_token, offsets, -1), axis=0) + flattened_slot = tl.where(matched >= 0, req * capacity + matched, -1) + tl.store(output_slot_ptr + row, flattened_slot) + + +def _find_sparse_decode_current_token_slot( + union_indices: torch.Tensor, + req_id_per_token: torch.Tensor, + global_seq_len: torch.Tensor, + output_slot: torch.Tensor, +) -> None: + if union_indices.ndim != 2 or union_indices.dtype != torch.int32: + raise TypeError("sparse CKV unions must be rank-2 int32") + if req_id_per_token.shape != global_seq_len.shape: + raise ValueError("sparse CKV request ids must match causal lengths") + if output_slot.shape != global_seq_len.shape or output_slot.dtype != torch.int64: + raise TypeError("sparse CKV patch slots must be matching int64 rows") + capacity = int(union_indices.shape[1]) + _find_sparse_decode_current_token_slot_kernel[(global_seq_len.numel(),)]( + union_indices, + req_id_per_token, + global_seq_len, + output_slot, + union_indices.stride(0), + capacity, + BLOCK_N=triton.next_power_of_2(capacity), + ) + + +@triton.jit +def _offset_sparse_decode_remap_kernel( + remap_ptr, + numel, + request_stride, + rows_per_request, + topk, + BLOCK_SIZE: tl.constexpr, +): + positions = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = positions < numel + values = tl.load(remap_ptr + positions, mask=mask, other=-1) + request = positions // (rows_per_request * topk) + values = tl.where(values >= 0, values + request * request_stride, values) + tl.store(remap_ptr + positions, values, mask=mask) + + +def _offset_sparse_decode_remap( + remap: torch.Tensor, + *, + request_stride: int, + rows_per_request: int, + topk: int, +) -> None: + if remap.dtype != torch.int32 or not remap.is_contiguous(): + raise TypeError("sparse CKV remap must be contiguous int32") + numel = int(remap.numel()) + if numel == 0: + return + block_size = 256 + _offset_sparse_decode_remap_kernel[(_cdiv(numel, block_size),)]( + remap, + numel, + int(request_stride), + int(rows_per_request), + int(topk), + BLOCK_SIZE=block_size, + ) + + @triton.jit def _mask_page_table_after_nsa_len_kernel( page_table_ptr, @@ -727,8 +839,10 @@ def __init__( from vllm import envs as envs_mod ckv_gather_requested = envs_mod.VLLM_B12X_MLA_CKV_GATHER + sparse_decode_requested = envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER + ckv_metadata_requested = ckv_gather_requested or sparse_decode_requested self.ckv_prefetch_registry = ( - _CKVPrefetchStateRegistry() if ckv_gather_requested else None + _CKVPrefetchStateRegistry() if ckv_metadata_requested else None ) # Max-batched-token scratch buffers so cudagraph capture sees stable # allocations (sliced per build()). @@ -1037,7 +1151,11 @@ def build( dcp_local_cu_seq_lens=dcp_local_cu_seq_lens, global_cache_seq_lens_per_req=( cm.seq_lens[: cm.num_reqs] - if use_dcp and envs_mod.VLLM_B12X_MLA_CKV_GATHER + if use_dcp + and ( + envs_mod.VLLM_B12X_MLA_CKV_GATHER + or envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER + ) else None ), dcp_local_total_tokens=dcp_local_total_tokens, @@ -1222,6 +1340,8 @@ def __init__( spec = getattr(vllm_config, "speculative_config", None) if spec is not None and getattr(spec, "num_speculative_tokens", None): q_per_req = 1 + int(spec.num_speculative_tokens) + self._sparse_decode_rows_per_request = q_per_req + self._sparse_decode_emits_topk = indexer is not None if self.spec_extend_as_decode: q_per_req = max(q_per_req, self.spec_decode_max_q) self._decode_max_rows = min(max_num_seqs * q_per_req, max_batched) @@ -1302,16 +1422,103 @@ def _make_plan( self._extend_plan = _make_plan( "extend", max_batched, self._kernel_num_heads, max_num_seqs ) + + from vllm import envs as envs_mod + from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + get_selected_record_exchange, + plan_sparse_ckv_decode, + preload_dense_union_extension_consistently, + ) + + sparse_decode_requested = envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER + self._sparse_decode_enabled = sparse_decode_requested and ( + 2 <= self.dcp_world_size <= 8 + and self.num_heads % _HEAD_ALIGNMENT == 0 + and self.dcp_workspace_non_dbo + and self.kv_cache_dtype == "nvfp4_ds_mla" + and self._kv_record_bytes == 432 + and not self._kv_fp8_rope + ) + sparse_max_requests = min( + max_num_seqs, + max(1, int(envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_MAX_SEQS)), + ) + self._sparse_decode_layout = None + self._sparse_decode_exchange = None + if self._sparse_decode_enabled: + self._sparse_decode_layout = plan_sparse_ckv_decode( + dcp_world_size=self.dcp_world_size, + topk=self.topk_tokens, + rows_per_request=self._sparse_decode_rows_per_request, + max_requests=sparse_max_requests, + pool_records=int(envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_POOL_RECORDS), + record_bytes=self._kv_record_bytes, + prefetch_depth=max(0, int(envs_mod.VLLM_B12X_MLA_CKV_PREFETCH_DEPTH)), + ) + graph_rows = int( + vllm_config.compilation_config.max_cudagraph_capture_size or 0 + ) + required_rows = ( + self._sparse_decode_layout.max_fast_requests + * self._sparse_decode_layout.rows_per_request + ) + if graph_rows and graph_rows < required_rows: + raise ValueError( + "selected-record CKV decode requires max cudagraph " + f"capture size >= {required_rows}; configured {graph_rows}" + ) + + from vllm.distributed.parallel_state import get_dcp_ckv_prefetch_group + + process_group = get_dcp_ckv_prefetch_group().device_group + if process_group is None: + raise RuntimeError("DCP CKV group has no device process group") + preload_dense_union_extension_consistently(process_group, self.device) + model_type = str( + getattr(vllm_config.model_config.hf_config, "model_type", "target") + ) + self._sparse_decode_exchange = get_selected_record_exchange( + process_group=process_group, + device=self.device, + layout=self._sparse_decode_layout, + lane_key=(model_type, id(vllm_config)), + ) + elif sparse_decode_requested: + logger.warning_once( + "Ignoring selected-record CKV decode for unsupported " + "topology/format: DCP=%d local_heads=%d dtype=%s bytes=%d " + "KV_FP8_ROPE=%d DBO=%d", + self.dcp_world_size, + self.num_heads, + self.kv_cache_dtype, + self._kv_record_bytes, + int(self._kv_fp8_rope), + int(not self.dcp_workspace_non_dbo), + ) + + self._sparse_decode_plan = ( + _make_plan( + "decode", + self._sparse_decode_layout.max_fast_requests + * self._sparse_decode_layout.rows_per_request, + self.num_heads, + self._sparse_decode_layout.max_fast_requests + * self._sparse_decode_layout.rows_per_request, + ) + if self._sparse_decode_layout is not None + else None + ) # One caller-owned uint8 scratch tensor covers either path (the larger # layout); the per-mode materializer carves its views from the prefix. self._scratch_nbytes = max( int(self._decode_plan.layout.nbytes), int(self._extend_plan.layout.nbytes), + int(self._sparse_decode_plan.layout.nbytes) + if self._sparse_decode_plan is not None + else 0, ) # CKV gather setup (Fix B). - from vllm import envs as envs_mod - ckv_gather_requested = envs_mod.VLLM_B12X_MLA_CKV_GATHER self._ckv_gather_enabled = ckv_gather_requested and ( self.dcp_world_size > 1 @@ -1394,20 +1601,20 @@ def _make_plan( 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: + # The registry isolates target/draft and graph workspaces. Full-CKV + # prefetch creates its lane stream lazily; the stream-affine selected- + # record exchange owns one dedicated stream created during init. + if self._ckv_gather_enabled or self._sparse_decode_enabled: self._ckv_current_chunk_kv_c: torch.Tensor | None = None self._ckv_current_chunk_kpe: torch.Tensor | None = None - if not self._ckv_prefetch_supported: + if self._ckv_gather_enabled and not self._ckv_prefetch_supported: logger.warning_once( "CKV gather prefetch disabled for kv_cache_dtype=%s " "(KV_FP8_ROPE=%s); falling back to synchronous gather.", self.kv_cache_dtype, int(self._kv_fp8_rope), ) - elif self._ckv_prefetch_depth > 0: + elif self._ckv_gather_enabled and self._ckv_prefetch_depth > 0: logger.info_once( "Using native CKV layer prefetch with depth=%d and " "%d workspace slots.", @@ -1803,6 +2010,234 @@ def dcp_prefill_ckv_gather_eligible( ) ) + def dcp_sparse_decode_ckv_gather_eligible( + self, + attn_metadata: B12xMLASparseMetadata, + num_tokens: int, + ) -> bool: + """Select bounded, uniform-row decode batches for sparse CKV exchange.""" + if self._sparse_decode_layout is None: + return False + from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + sparse_decode_batch_eligible, + ) + + has_required_metadata = all( + tensor is not None + for tensor in ( + attn_metadata.req_id_per_token, + attn_metadata.page_table_1, + attn_metadata.nsa_cache_seqlens, + attn_metadata.global_cache_seq_lens_per_req, + attn_metadata.ckv_prefetch_registry, + ) + ) + eligible = sparse_decode_batch_eligible( + self._sparse_decode_layout, + num_requests=int(attn_metadata.num_reqs), + num_rows=int(num_tokens), + max_query_len=int(attn_metadata.max_query_len), + num_actual_tokens=int(attn_metadata.num_actual_tokens), + has_required_metadata=has_required_metadata, + ) + if ( + self._sparse_decode_enabled + and int(attn_metadata.num_reqs) + > self._sparse_decode_layout.max_fast_requests + ): + logger.info_once( + "Selected-record CKV decode falling back to stock DCP: " + "active requests=%d exceeds pooled maximum=%d", + int(attn_metadata.num_reqs), + self._sparse_decode_layout.max_fast_requests, + ) + return bool(self._sparse_decode_enabled and eligible) + + def _dcp_gather_sparse_decode_ckv( + self, + kv_cache: torch.Tensor, + attn_metadata: B12xMLASparseMetadata, + topk_indices: torch.Tensor, + prefetch_state: _CKVPrefetchState, + *, + buf_idx: int, + build_union: bool, + wait_for_completion: bool, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.cuda.Event, + ]: + """Exchange one dense per-sequence MTP union through generic B12X P2P.""" + layout = self._sparse_decode_layout + exchange = self._sparse_decode_exchange + if layout is None or exchange is None: + raise RuntimeError("selected-record CKV decode was not initialized") + num_rows = int(topk_indices.shape[0]) + num_requests = int(attn_metadata.num_reqs) + rows_per_request = num_rows // num_requests + active_records = layout.active_records(num_requests) + if not self.dcp_sparse_decode_ckv_gather_eligible(attn_metadata, num_rows): + raise RuntimeError( + "selected-record CKV gather received an ineligible batch" + ) + if not 0 <= int(buf_idx) < layout.workspace_slots: + raise ValueError( + f"selected-record workspace slot {buf_idx} is outside " + f"[0, {layout.workspace_slots})" + ) + if ( + kv_cache.dtype != torch.uint8 + or kv_cache.ndim != 3 + or tuple(kv_cache.shape[1:]) != (self.block_size, layout.record_bytes) + or not kv_cache.is_contiguous() + ): + raise ValueError( + "selected-record CKV decode requires contiguous native paged cache" + ) + + assert attn_metadata.req_id_per_token is not None + assert attn_metadata.page_table_1 is not None + assert attn_metadata.nsa_cache_seqlens is not None + assert attn_metadata.global_cache_seq_lens_per_req is not None + state = prefetch_state.get_sparse_decode_state(layout, exchange) + output = state.payload_workspace[buf_idx, :active_records] + caller_stream = torch.cuda.current_stream() + + from vllm.distributed.parallel_state import get_dcp_ckv_prefetch_group + from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + build_dense_union_remap, + get_selected_record_stream, + ) + + gather_stream = get_selected_record_stream(exchange) + gather_stream.wait_stream(caller_stream) + output.record_stream(gather_stream) + + with torch.cuda.stream(gather_stream): + if build_union: + local_topk = topk_indices[:num_rows, : layout.topk] + all_rank_topk = state.all_rank_topk[: layout.dcp_world_size * num_rows] + _dcp_all_gather_current_stream( + get_dcp_ckv_prefetch_group(), + local_topk, + all_rank_topk, + ) + destination_topk = all_rank_topk.view( + layout.dcp_world_size, + num_requests, + rows_per_request, + layout.topk, + ) + build_dense_union_remap( + destination_topk, + state.union_indices, + state.remap, + state.union_counts, + state.hash_keys, + state.hash_first_positions, + state.first_to_dense, + num_requests=num_requests, + ) + for request in range(num_requests): + row_start = request * rows_per_request + state.selected_indices[ + row_start : row_start + rows_per_request + ].copy_( + state.remap[ + self.dcp_rank, + request, + :rows_per_request, + ], + non_blocking=True, + ) + selected = state.selected_indices[:num_rows, : layout.topk] + _offset_sparse_decode_remap( + selected, + request_stride=layout.per_request_capacity, + rows_per_request=rows_per_request, + topk=layout.topk, + ) + + transport_indices = state.transport_local_slots[num_requests] + for destination in range(layout.dcp_world_size): + destination_slots = state.local_slots[destination, :num_requests] + triton_filter_and_convert_dcp_index( + state.request_ids[:num_requests], + attn_metadata.block_table, + state.union_indices[destination, :num_requests], + dcp_size=layout.dcp_world_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=layout.per_request_capacity, + compact_valid_to_front=False, + out=destination_slots, + ) + transport_indices[destination].copy_( + destination_slots.reshape(-1), non_blocking=True + ) + exchange.exchange( + kv_cache.reshape(-1, layout.record_bytes), + transport_indices, + output, + ) + complete = state.complete_events[buf_idx] + complete.record(gather_stream) + + if wait_for_completion: + caller_stream.wait_event(complete) + selected_indices = state.selected_indices[:num_rows, : layout.topk] + nsa_cache_seqlens = attn_metadata.nsa_cache_seqlens[:num_rows] + if wait_for_completion: + global_causal_lens = _global_causal_lens_for_ckv_gather( + attn_metadata.global_cache_seq_lens_per_req, + attn_metadata.query_start_loc, + attn_metadata.req_id_per_token, + num_rows, + ) + nsa_cache_seqlens.copy_(global_causal_lens, non_blocking=True) + nsa_cache_seqlens.clamp_(max=layout.topk) + _mask_page_table_after_nsa_len(selected_indices, nsa_cache_seqlens) + gathered_cache = output.view(-1, self.block_size, layout.record_bytes) + return gathered_cache, selected_indices, nsa_cache_seqlens, complete + + def _append_current_token_to_sparse_decode_gathered( + self, + gathered_cache: torch.Tensor, + attn_metadata: B12xMLASparseMetadata, + union_indices: torch.Tensor, + global_causal_lens: torch.Tensor, + patch_slots: torch.Tensor, + layer: AttentionLayer, + ) -> None: + """Patch records written after a Shared layer's history prefetch.""" + if self._ckv_current_chunk_kv_c is None: + raise RuntimeError("selected-record prefetch is missing current CKV") + if self._ckv_current_chunk_kpe is None: + raise RuntimeError("selected-record prefetch is missing current RoPE") + if attn_metadata.req_id_per_token is None: + raise RuntimeError("selected-record prefetch is missing request ids") + num_rows = int(global_causal_lens.numel()) + _find_sparse_decode_current_token_slot( + union_indices, + attn_metadata.req_id_per_token[:num_rows], + global_causal_lens, + patch_slots[:num_rows], + ) + k_pe = self._ckv_current_chunk_kpe[:num_rows] + if k_pe.ndim == 3: + k_pe = k_pe.squeeze(1) + ops.concat_and_cache_mla( + self._ckv_current_chunk_kv_c[:num_rows], + k_pe, + gathered_cache, + patch_slots[:num_rows], + self.kv_cache_dtype, + getattr(layer, "_k_scale", None), + ) + def _dcp_gather_ckv( self, kv_cache: torch.Tensor, @@ -2121,14 +2556,18 @@ def forward_mqa( use_ckv_gather = self.dcp_prefill_ckv_gather_eligible( attn_metadata, int(query_rows) ) + use_sparse_decode_ckv_gather = self.dcp_sparse_decode_ckv_gather_eligible( + attn_metadata, int(query_rows) + ) + use_local_query_heads = use_ckv_gather or use_sparse_decode_ckv_gather workspace_tensors = self._borrow_workspaces() q_workspace = workspace_tensors[0] dense_out_workspace = workspace_tensors[1] if self._pad_heads else None scratch_storage = workspace_tensors[-1] expected_input_heads = ( - self.num_heads if use_ckv_gather else self._input_num_heads + self.num_heads if use_local_query_heads else self._input_num_heads ) - if use_ckv_gather: + if use_local_query_heads: local_q_numel = ( self._max_batched * self._ckv_kernel_num_heads * self.q_head_dim ) @@ -2175,6 +2614,16 @@ def forward_mqa( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] per_token_cache = attn_metadata.cache_seq_lens_per_token[:num_actual_toks] + sparse_decode_global_causal_lens = None + if use_sparse_decode_ckv_gather: + assert attn_metadata.global_cache_seq_lens_per_req is not None + assert attn_metadata.req_id_per_token is not None + sparse_decode_global_causal_lens = _global_causal_lens_for_ckv_gather( + attn_metadata.global_cache_seq_lens_per_req, + attn_metadata.query_start_loc, + attn_metadata.req_id_per_token, + num_actual_toks, + ) if use_ckv_gather: assert attn_metadata.req_id_per_token is not None assert attn_metadata.ckv_page_table_1 is not None @@ -2217,6 +2666,12 @@ def forward_mqa( out=nsa_cache_seqlens, ) _mask_page_table_after_nsa_len(selected_indices, nsa_cache_seqlens) + elif use_sparse_decode_ckv_gather: + # The dense remap is materialized after the current layer's native + # cache is available below. Keep placeholders here so the stock + # rank-local conversion is not entered. + selected_indices = topk_indices + nsa_cache_seqlens = per_token_cache elif self.dcp_world_size > 1: # The indexer globally merges logical top-k ids across DCP ranks. # Compact just this rank's winners into local physical cache slots; @@ -2282,6 +2737,127 @@ def forward_mqa( "B12X_MLA_SPARSE requires a contiguous native paged KV cache; " f"got stride={tuple(kv_cache.stride())}" ) + if use_sparse_decode_ckv_gather: + layout = self._sparse_decode_layout + exchange = self._sparse_decode_exchange + if layout is None or exchange is None: + raise RuntimeError("selected-record CKV decode was not initialized") + if sparse_decode_global_causal_lens is None: + raise RuntimeError( + "selected-record CKV decode is missing global causal lengths" + ) + layer_idx = self._resolve_layer_index(layer) + registry = attn_metadata.ckv_prefetch_registry + if registry is None: + raise RuntimeError( + "selected-record CKV decode is missing its lane registry" + ) + prefetch_state = registry.for_workspace(q_workspace, layer_idx, kv_cache) + if layer_idx is not None: + prefetch_state.enter_layer(layer_idx) + prefetch_state.register_cache(layer_idx, kv_cache) + prefetch_state.register_impl(layer_idx, self) + sparse_state = prefetch_state.get_sparse_decode_state(layout, exchange) + pending = ( + prefetch_state.pending_layers.pop(layer_idx, None) + if layer_idx is not None + else None + ) + if pending is not None: + gather_event, current_buf_idx = pending + torch.cuda.current_stream().wait_event(gather_event) + active_records = layout.active_records(int(attn_metadata.num_reqs)) + kv_cache = sparse_state.payload_workspace[ + current_buf_idx, :active_records + ].view(-1, self.block_size, layout.record_bytes) + self._append_current_token_to_sparse_decode_gathered( + kv_cache, + attn_metadata, + sparse_state.union_indices[ + self.dcp_rank, : int(attn_metadata.num_reqs) + ], + sparse_decode_global_causal_lens, + sparse_state.patch_slots, + layer, + ) + selected_indices = sparse_state.selected_indices[ + :num_actual_toks, : layout.topk + ] + nsa_cache_seqlens = attn_metadata.nsa_cache_seqlens[:num_actual_toks] + nsa_cache_seqlens.copy_( + sparse_decode_global_causal_lens, non_blocking=True + ) + nsa_cache_seqlens.clamp_(max=layout.topk) + _mask_page_table_after_nsa_len(selected_indices, nsa_cache_seqlens) + else: + current_buf_idx = ( + layer_idx % layout.workspace_slots if layer_idx is not None else 0 + ) + ( + kv_cache, + selected_indices, + nsa_cache_seqlens, + _, + ) = self._dcp_gather_sparse_decode_ckv( + kv_cache, + attn_metadata, + topk_indices, + prefetch_state, + buf_idx=current_buf_idx, + build_union=True, + wait_for_completion=True, + ) + + logger.info_once( + "Using selected-record CKV decode for C<=%d/MTP%d " + "(DCP%d topk=%d pool=%d record_bytes=%d depth=%d)", + layout.max_fast_requests, + layout.rows_per_request - 1, + layout.dcp_world_size, + layout.topk, + layout.pool_records, + layout.record_bytes, + layout.prefetch_depth, + ) + if ( + self._sparse_decode_emits_topk + and layer_idx is not None + and layout.prefetch_depth > 0 + ): + from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + sparse_decode_prefetch_targets, + ) + + emits_topk_by_layer = [ + None if impl is None else impl._sparse_decode_emits_topk + for impl in prefetch_state.layer_impls + ] + targets = sparse_decode_prefetch_targets( + layer_idx, + layout.prefetch_depth, + emits_topk_by_layer, + ) + for target_idx in targets: + if target_idx in prefetch_state.pending_layers: + continue + target_impl = prefetch_state.layer_impls[target_idx] + target_kv = prefetch_state.layer_caches[target_idx] + if target_impl is None or target_kv is None: + break + target_buf_idx = target_idx % layout.workspace_slots + _, _, _, target_event = target_impl._dcp_gather_sparse_decode_ckv( + target_kv, + attn_metadata, + topk_indices, + prefetch_state, + buf_idx=target_buf_idx, + build_union=False, + wait_for_completion=False, + ) + prefetch_state.pending_layers[target_idx] = ( + target_event, + target_buf_idx, + ) if use_ckv_gather: layer_idx = self._resolve_layer_index(layer) prefetch_registry = attn_metadata.ckv_prefetch_registry @@ -2364,33 +2940,52 @@ def forward_mqa( target_buf_idx, ) - use_decode_kernel = attn_metadata.max_query_len <= 1 or ( - self.spec_extend_as_decode - and attn_metadata.max_query_len <= self.spec_decode_max_q - and num_actual_toks <= attn_metadata.num_reqs * self.spec_decode_max_q - and num_actual_toks <= self._decode_max_rows + use_decode_kernel = ( + use_sparse_decode_ckv_gather + or (attn_metadata.max_query_len <= 1) + or ( + self.spec_extend_as_decode + and attn_metadata.max_query_len <= self.spec_decode_max_q + and num_actual_toks <= attn_metadata.num_reqs * self.spec_decode_max_q + and num_actual_toks <= self._decode_max_rows + ) ) if use_decode_kernel: cache_seqlens = ( - attn_metadata.cache_seq_lens_per_req + sparse_decode_global_causal_lens + if use_sparse_decode_ckv_gather + else attn_metadata.cache_seq_lens_per_req if attn_metadata.max_query_len <= 1 else attn_metadata.cache_seq_lens_per_token[:num_actual_toks] ) + if cache_seqlens is None: + raise RuntimeError( + "selected-record CKV decode is missing global causal lengths" + ) decode_q = q_all - if self._pad_heads: + if self._pad_heads and not use_sparse_decode_ckv_gather: decode_q = q_buffer[:, : self._kernel_num_heads] decode_q[:, self._input_num_heads :, :].zero_() # Eager bind maps caller-owned scratch into views. forced_num_splits # pins the planner choice for this captured graph; the merge kernel is # specialized on that count and needs no device-side control fill. - binding = self._decode_plan.bind( + decode_plan = ( + self._sparse_decode_plan + if use_sparse_decode_ckv_gather + else self._decode_plan + ) + if decode_plan is None: + raise RuntimeError( + "selected-record CKV decode plan was not initialized" + ) + binding = decode_plan.bind( scratch=scratch_storage, q=decode_q, selected_indices=selected_indices, cache_seqlens_int32=cache_seqlens, nsa_cache_seqlens_int32=nsa_cache_seqlens, ) - if self.need_to_return_lse_for_decode: + if self.need_to_return_lse_for_decode and not use_sparse_decode_ckv_gather: out, lse = cast( tuple[torch.Tensor, torch.Tensor], self._sparse_mla_decode_forward( @@ -2422,7 +3017,7 @@ def forward_mqa( **kernel_format_kwargs, ), ) - if self._pad_heads: + if self._pad_heads and not use_sparse_decode_ckv_gather: assert dense_out_workspace is not None dense_out = dense_out_workspace[:num_actual_toks] dense_out.copy_(out[:, : self._input_num_heads, :]) diff --git a/vllm/v1/attention/backends/mla/b12x_sparse_ckv_decode.py b/vllm/v1/attention/backends/mla/b12x_sparse_ckv_decode.py new file mode 100644 index 000000000000..4234b0a05f73 --- /dev/null +++ b/vllm/v1/attention/backends/mla/b12x_sparse_ckv_decode.py @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Planning and deterministic union construction for sparse DCP CKV decode. + +The B12X transport only exchanges fixed-width records. This module owns the +vLLM policy which decides which records each request needs and how MTP rows map +into one per-request, densely packed sparse workspace. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup +from torch.utils.cpp_extension import load + +from vllm.distributed.parallel_state import register_model_parallel_cleanup_hook + + +@dataclass(frozen=True) +class SparseCKVDecodeLayout: + dcp_world_size: int + topk: int + rows_per_request: int + max_requests: int + per_request_capacity: int + pool_records: int + max_fast_requests: int + record_bytes: int + prefetch_depth: int + workspace_slots: int + + @property + def workspace_bytes(self) -> int: + return self.workspace_slots * self.pool_records * self.record_bytes + + def active_records(self, num_requests: int) -> int: + if not 1 <= num_requests <= self.max_fast_requests: + raise ValueError( + f"request count {num_requests} is outside sparse capacity " + f"[1, {self.max_fast_requests}]" + ) + return num_requests * self.per_request_capacity + + +def plan_sparse_ckv_decode( + *, + dcp_world_size: int, + topk: int, + rows_per_request: int, + max_requests: int, + pool_records: int, + record_bytes: int, + prefetch_depth: int, +) -> SparseCKVDecodeLayout: + """Return the static pooled-workspace layout for selected-record decode.""" + if not 2 <= dcp_world_size <= 8: + raise ValueError("selected-record CKV decode supports DCP world sizes 2-8") + if min(topk, rows_per_request, max_requests, record_bytes) <= 0: + raise ValueError("topk, rows, requests, and record width must be positive") + if prefetch_depth < 0: + raise ValueError("prefetch depth cannot be negative") + + per_request_capacity = topk * rows_per_request + requested_pool = int(pool_records) + if requested_pool <= 0: + requested_pool = max_requests * per_request_capacity + max_fast_requests = min(max_requests, requested_pool // per_request_capacity) + if max_fast_requests <= 0: + raise ValueError( + "selected-record pool cannot hold one request's worst-case MTP union" + ) + # Capacity beyond the largest schedulable batch only wastes VRAM. + effective_pool = min( + requested_pool, + max_requests * per_request_capacity, + ) + return SparseCKVDecodeLayout( + dcp_world_size=dcp_world_size, + topk=topk, + rows_per_request=rows_per_request, + max_requests=max_requests, + per_request_capacity=per_request_capacity, + pool_records=effective_pool, + max_fast_requests=max_fast_requests, + record_bytes=record_bytes, + prefetch_depth=prefetch_depth, + workspace_slots=prefetch_depth + 1, + ) + + +def sparse_decode_batch_eligible( + layout: SparseCKVDecodeLayout, + *, + num_requests: int, + num_rows: int, + max_query_len: int, + num_actual_tokens: int, + has_required_metadata: bool, +) -> bool: + """Pure policy gate used before entering the selected-record fast path.""" + if not has_required_metadata or not 1 <= num_requests <= layout.max_fast_requests: + return False + if not 1 <= num_rows <= num_requests * layout.rows_per_request: + return False + if num_rows % num_requests: + return False + rows_per_request = num_rows // num_requests + return ( + rows_per_request <= layout.rows_per_request + and max_query_len == rows_per_request + and num_actual_tokens == num_rows + ) + + +def sparse_decode_prefetch_targets( + layer_index: int, + depth: int, + emits_topk_by_layer: list[bool | None], +) -> list[int]: + """Return consecutive Shared layers following one Full/indexer layer.""" + targets: list[int] = [] + for distance in range(1, max(0, depth) + 1): + target = layer_index + distance + if target >= len(emits_topk_by_layer): + break + emits_topk = emits_topk_by_layer[target] + if emits_topk is None or emits_topk: + break + targets.append(target) + return targets + + +def owner_and_local_ordinal( + logical_token: int, + *, + dcp_world_size: int, + interleave: int, +) -> tuple[int, int]: + """Map one global logical token to its DCP owner and owner-local ordinal.""" + if logical_token < 0: + return -1, -1 + if dcp_world_size <= 0 or interleave <= 0: + raise ValueError("DCP world size and interleave must be positive") + owner = (logical_token // interleave) % dcp_world_size + local_ordinal = ( + logical_token // (dcp_world_size * interleave) * interleave + + logical_token % interleave + ) + return owner, local_ordinal + + +def dense_union_remap_reference( + indices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """CPU oracle: stable dense union and exact remap for every group. + + ``indices`` has shape ``[..., rows, topk]``. Every leading group receives + an independent stable union in flattened row order. + """ + if indices.device.type != "cpu" or indices.dtype != torch.int32: + raise ValueError("reference union requires a CPU int32 tensor") + if indices.ndim < 2: + raise ValueError("indices must include row and top-k dimensions") + rows, topk = indices.shape[-2:] + capacity = rows * topk + groups = indices.reshape(-1, capacity) + union = torch.full_like(groups, -1) + remap = torch.full_like(groups, -1) + counts = torch.zeros(groups.shape[0], dtype=torch.int32) + for group_index, row in enumerate(groups.tolist()): + slots: dict[int, int] = {} + for input_index, token in enumerate(row): + if token < 0: + continue + slot = slots.get(token) + if slot is None: + slot = len(slots) + slots[token] = slot + union[group_index, slot] = token + remap[group_index, input_index] = slot + counts[group_index] = len(slots) + leading = indices.shape[:-2] + return ( + union.reshape(*leading, capacity), + remap.reshape(indices.shape), + counts.reshape(leading), + ) + + +class SparseCKVDecodeState: + """Persistent graph-safe union/remap state for one execution lane.""" + + def __init__( + self, + *, + layout: SparseCKVDecodeLayout, + device: torch.device, + exchange, + ) -> None: + self.layout = layout + self.exchange = exchange + capacity = layout.per_request_capacity + hash_capacity = 1 << (2 * capacity - 1).bit_length() + max_rows = layout.max_requests * layout.rows_per_request + state_shape = (layout.dcp_world_size, layout.max_requests) + self.all_rank_topk = torch.empty( + (layout.dcp_world_size * max_rows, layout.topk), + dtype=torch.int32, + device=device, + ) + self.union_indices = torch.empty( + (*state_shape, capacity), dtype=torch.int32, device=device + ) + self.remap = torch.empty( + ( + *state_shape, + layout.rows_per_request, + layout.topk, + ), + dtype=torch.int32, + device=device, + ) + self.union_counts = torch.empty(state_shape, dtype=torch.int32, device=device) + self.hash_keys = torch.empty( + (*state_shape, hash_capacity), dtype=torch.int32, device=device + ) + self.hash_first_positions = torch.empty_like(self.hash_keys) + self.first_to_dense = torch.empty_like(self.union_indices) + self.local_slots = torch.empty_like(self.union_indices) + # PCIeSelectedRecordExchange requires one contiguous destination-major + # index plane. A [:num_requests] view of local_slots retains the + # max-request row stride and is therefore non-contiguous. Preallocate + # the eight small exact-shape planes once so decode never allocates or + # copies through a temporary contiguous() tensor during graph replay. + self.transport_local_slots = { + request_count: torch.empty( + ( + layout.dcp_world_size, + request_count * layout.per_request_capacity, + ), + dtype=torch.int32, + device=device, + ) + for request_count in range(1, layout.max_fast_requests + 1) + } + self.selected_indices = torch.empty( + (max_rows, layout.topk), dtype=torch.int32, device=device + ) + # Cross-layer prefetch outlives a WorkspaceManager borrow. Keep the + # payload in lane-owned storage so unrelated scratch users cannot + # overwrite S1/S2/S3 records before attention consumes them. + self.payload_workspace = torch.empty( + ( + layout.workspace_slots, + layout.pool_records, + layout.record_bytes, + ), + dtype=torch.uint8, + device=device, + ) + self.request_ids = torch.arange( + layout.max_requests, dtype=torch.int32, device=device + ) + self.patch_slots = torch.empty(max_rows, dtype=torch.int64, device=device) + self.complete_events = [ + torch.cuda.Event(blocking=False) for _ in range(layout.workspace_slots) + ] + + +_SELECTED_RECORD_EXCHANGES: dict[tuple[object, ...], object] = {} +_SELECTED_RECORD_STREAMS: dict[int, torch.cuda.Stream] = {} +_UNION_PREFLIGHT_RESULTS: dict[tuple[int, str, int | None], tuple[bool, str]] = {} + + +def close_selected_record_exchanges() -> None: + """Release IPC slabs before their DCP process groups are destroyed.""" + exchanges = list( + {id(value): value for value in _SELECTED_RECORD_EXCHANGES.values()}.values() + ) + _SELECTED_RECORD_EXCHANGES.clear() + _SELECTED_RECORD_STREAMS.clear() + _UNION_PREFLIGHT_RESULTS.clear() + + first_error: Exception | None = None + for exchange in exchanges: + try: + exchange.close() + except Exception as exc: # pragma: no cover - defensive shutdown path + if first_error is None: + first_error = exc + if first_error is not None: + raise RuntimeError( + "failed to close a selected-record exchange" + ) from first_error + + +register_model_parallel_cleanup_hook(close_selected_record_exchanges) + + +def get_selected_record_exchange( + *, + process_group: ProcessGroup, + device: torch.device, + layout: SparseCKVDecodeLayout, + lane_key: tuple[str, int], +): + """Return one generic B12X exchange per target/draft execution lane.""" + key = ( + id(process_group), + device.type, + device.index, + layout.dcp_world_size, + layout.pool_records, + layout.record_bytes, + lane_key, + ) + exchange = _SELECTED_RECORD_EXCHANGES.get(key) + if exchange is None: + from sparkinfer.comm.pcie import SelectedRecordExchange + + exchange = SelectedRecordExchange.from_process_group( + process_group=process_group, + device=device, + max_records=layout.pool_records, + record_bytes=layout.record_bytes, + ) + _SELECTED_RECORD_EXCHANGES[key] = exchange + if id(exchange) not in _SELECTED_RECORD_STREAMS: + _SELECTED_RECORD_STREAMS[id(exchange)] = torch.cuda.Stream(device=device) + return exchange + + +def get_selected_record_stream(exchange) -> torch.cuda.Stream: + """Return the dedicated stream created with one stream-affine exchange.""" + stream = _SELECTED_RECORD_STREAMS.get(id(exchange)) + if stream is None: + raise RuntimeError("selected-record exchange has no execution stream") + return stream + + +@lru_cache(maxsize=1) +def _load_union_extension(): + source = Path(__file__).with_name("b12x_sparse_ckv_union.cu") + return load( + name="vllm_b12x_sparse_ckv_union_ext", + sources=[str(source)], + extra_cuda_cflags=["-O3"], + verbose=False, + ) + + +def preload_dense_union_extension() -> None: + """Compile the union helper before CUDA graph warmup or capture.""" + _load_union_extension() + + +def preload_dense_union_extension_consistently( + process_group: ProcessGroup, + device: torch.device, +) -> None: + """Preload on every DCP rank and fail or proceed as one group.""" + key = (id(process_group), device.type, device.index) + cached = _UNION_PREFLIGHT_RESULTS.get(key) + if cached is not None: + ok, message = cached + if not ok: + raise RuntimeError(message) + return + + local_error: Exception | None = None + try: + preload_dense_union_extension() + except Exception as exc: # pragma: no cover - exercised by fault injection + local_error = exc + + status = torch.tensor( + [0 if local_error is not None else 1], + dtype=torch.int32, + device=device, + ) + dist.all_reduce(status, op=dist.ReduceOp.MIN, group=process_group) + globally_ok = bool(status.item()) + if globally_ok: + _UNION_PREFLIGHT_RESULTS[key] = (True, "") + return + + message = ( + f"selected-record CKV union extension failed on this DCP rank: {local_error}" + if local_error is not None + else "selected-record CKV union extension failed on another DCP rank" + ) + _UNION_PREFLIGHT_RESULTS[key] = (False, message) + raise RuntimeError(message) + + +def build_dense_union_remap( + indices: torch.Tensor, + union_indices: torch.Tensor, + remap: torch.Tensor, + union_counts: torch.Tensor, + hash_keys: torch.Tensor, + hash_first_positions: torch.Tensor, + first_to_dense: torch.Tensor, + *, + num_requests: int, +) -> None: + """Build stable dense per-destination/per-request unions on CUDA.""" + _load_union_extension().build_dense_union_remap( + indices, + union_indices, + remap, + union_counts, + hash_keys, + hash_first_positions, + first_to_dense, + int(num_requests), + ) diff --git a/vllm/v1/attention/backends/mla/b12x_sparse_ckv_union.cu b/vllm/v1/attention/backends/mla/b12x_sparse_ckv_union.cu new file mode 100644 index 000000000000..d1be542dc45e --- /dev/null +++ b/vllm/v1/attention/backends/mla/b12x_sparse_ckv_union.cu @@ -0,0 +1,261 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +__device__ __forceinline__ uint32_t hash_token_index(uint32_t value) { + value ^= value >> 16; + value *= 0x7feb352dU; + value ^= value >> 15; + value *= 0x846ca68bU; + value ^= value >> 16; + return value; +} + +__device__ __forceinline__ int32_t find_first_position( + int32_t key, + const int32_t* hash_keys, + const int32_t* hash_first_positions, + int64_t hash_capacity) { + const uint32_t mask = static_cast(hash_capacity - 1); + uint32_t slot = hash_token_index(static_cast(key)) & mask; + for (int64_t probe = 0; probe < hash_capacity; ++probe) { + const int32_t stored = hash_keys[slot]; + if (stored == key) { + return hash_first_positions[slot]; + } + if (stored == -1) { + return -1; + } + slot = (slot + 1) & mask; + } + return -1; +} + +__global__ void build_dense_union_remap_kernel( + const int32_t* __restrict__ indices, + int32_t* __restrict__ union_indices, + int32_t* __restrict__ remap, + int32_t* __restrict__ union_counts, + int32_t* __restrict__ hash_keys, + int32_t* __restrict__ hash_first_positions, + int32_t* __restrict__ first_to_dense, + int destinations, + int num_requests, + int max_requests, + int64_t input_entries, + int64_t capacity, + int64_t hash_capacity) { + const int group = static_cast(blockIdx.x); + const int destination = group / num_requests; + const int request = group % num_requests; + if (destination >= destinations) { + return; + } + const int state_group = destination * max_requests + request; + const int64_t input_base = static_cast(group) * input_entries; + const int64_t state_base = static_cast(state_group) * capacity; + const int64_t hash_base = static_cast(state_group) * hash_capacity; + + for (int64_t i = threadIdx.x; i < hash_capacity; i += blockDim.x) { + hash_keys[hash_base + i] = -1; + hash_first_positions[hash_base + i] = std::numeric_limits::max(); + } + for (int64_t i = threadIdx.x; i < capacity; i += blockDim.x) { + union_indices[state_base + i] = -1; + first_to_dense[state_base + i] = -1; + } + for (int64_t i = threadIdx.x; i < input_entries; i += blockDim.x) { + remap[state_base + i] = -1; + } + if (threadIdx.x == 0) { + union_counts[state_group] = 0; + } + __syncthreads(); + + const uint32_t mask = static_cast(hash_capacity - 1); + for (int64_t i = threadIdx.x; i < input_entries; i += blockDim.x) { + const int32_t key = indices[input_base + i]; + if (key < 0) { + continue; + } + uint32_t slot = hash_token_index(static_cast(key)) & mask; + for (int64_t probe = 0; probe < hash_capacity; ++probe) { + const int32_t previous = atomicCAS(hash_keys + hash_base + slot, -1, key); + if (previous == -1 || previous == key) { + atomicMin( + hash_first_positions + hash_base + slot, + static_cast(i)); + break; + } + slot = (slot + 1) & mask; + } + } + __syncthreads(); + + __shared__ int32_t first_flags[256]; + __shared__ int32_t running_count; + if (threadIdx.x == 0) { + running_count = 0; + } + __syncthreads(); + + for (int64_t tile = 0; tile < input_entries; tile += blockDim.x) { + const int64_t i = tile + threadIdx.x; + int32_t is_first = 0; + int32_t key = -1; + if (i < input_entries) { + key = indices[input_base + i]; + if (key >= 0) { + is_first = find_first_position( + key, + hash_keys + hash_base, + hash_first_positions + hash_base, + hash_capacity) == i; + } + } + first_flags[threadIdx.x] = is_first; + __syncthreads(); + if (threadIdx.x == 0) { + int32_t prefix = running_count; + const int64_t remaining = input_entries - tile; + const int32_t valid = static_cast( + remaining < static_cast(blockDim.x) + ? remaining + : static_cast(blockDim.x)); + for (int32_t lane = 0; lane < valid; ++lane) { + const int32_t flag = first_flags[lane]; + first_flags[lane] = prefix; + prefix += flag; + } + running_count = prefix; + } + __syncthreads(); + if (i < input_entries && is_first) { + const int32_t dense_slot = first_flags[threadIdx.x]; + first_to_dense[state_base + i] = dense_slot; + union_indices[state_base + dense_slot] = key; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + union_counts[state_group] = running_count; + } + __syncthreads(); + + for (int64_t i = threadIdx.x; i < input_entries; i += blockDim.x) { + const int32_t key = indices[input_base + i]; + if (key < 0) { + continue; + } + const int32_t first = find_first_position( + key, + hash_keys + hash_base, + hash_first_positions + hash_base, + hash_capacity); + if (first >= 0) { + remap[state_base + i] = first_to_dense[state_base + first]; + } + } +} + +void validate_cuda_int32_contiguous(const torch::Tensor& value, const char* name) { + TORCH_CHECK(value.is_cuda(), name, " must be CUDA"); + TORCH_CHECK(value.scalar_type() == torch::kInt32, name, " must be int32"); + TORCH_CHECK(value.is_contiguous(), name, " must be contiguous"); +} + +void build_dense_union_remap( + torch::Tensor indices, + torch::Tensor union_indices, + torch::Tensor remap, + torch::Tensor union_counts, + torch::Tensor hash_keys, + torch::Tensor hash_first_positions, + torch::Tensor first_to_dense, + int64_t num_requests) { + validate_cuda_int32_contiguous(indices, "indices"); + validate_cuda_int32_contiguous(union_indices, "union_indices"); + validate_cuda_int32_contiguous(remap, "remap"); + validate_cuda_int32_contiguous(union_counts, "union_counts"); + validate_cuda_int32_contiguous(hash_keys, "hash_keys"); + validate_cuda_int32_contiguous(hash_first_positions, "hash_first_positions"); + validate_cuda_int32_contiguous(first_to_dense, "first_to_dense"); + TORCH_CHECK(indices.dim() == 4, "indices must be [destination, request, row, topk]"); + TORCH_CHECK(num_requests > 0 && num_requests == indices.size(1), "active request count must match indices"); + TORCH_CHECK( + union_indices.dim() == 3, + "union_indices must be [destination, request, capacity]"); + TORCH_CHECK( + hash_keys.dim() == 3, + "hash_keys must be [destination, request, hash capacity]"); + + const auto device = indices.device(); + TORCH_CHECK(union_indices.device() == device, "union_indices device mismatch"); + TORCH_CHECK(remap.device() == device, "remap device mismatch"); + TORCH_CHECK(union_counts.device() == device, "union_counts device mismatch"); + TORCH_CHECK(hash_keys.device() == device, "hash_keys device mismatch"); + TORCH_CHECK(hash_first_positions.device() == device, "hash_first_positions device mismatch"); + TORCH_CHECK(first_to_dense.device() == device, "first_to_dense device mismatch"); + + const int destinations = static_cast(indices.size(0)); + const int max_requests = static_cast(union_indices.size(1)); + const int64_t input_entries = indices.size(2) * indices.size(3); + const int64_t capacity = union_indices.size(2); + const int64_t hash_capacity = hash_keys.size(2); + TORCH_CHECK(capacity >= input_entries, "union capacity is too small"); + TORCH_CHECK((hash_capacity & (hash_capacity - 1)) == 0, "hash capacity must be a power of two"); + TORCH_CHECK(hash_capacity >= 2 * input_entries, "hash capacity must be at least 2x input entries"); + TORCH_CHECK(union_indices.size(0) == destinations, "union destination mismatch"); + TORCH_CHECK(num_requests <= max_requests, "active requests exceed union state capacity"); + TORCH_CHECK( + remap.dim() == 4 && remap.size(0) == destinations && + remap.size(1) == max_requests && remap.size(2) >= indices.size(2) && + remap.size(3) == indices.size(3) && + remap.size(2) * remap.size(3) == capacity, + "remap shape mismatch"); + TORCH_CHECK( + union_counts.dim() == 2 && union_counts.size(0) == destinations && + union_counts.size(1) == max_requests, + "count shape mismatch"); + TORCH_CHECK(hash_keys.sizes() == hash_first_positions.sizes(), "hash shape mismatch"); + TORCH_CHECK( + hash_keys.size(0) == destinations && hash_keys.size(1) == max_requests, + "hash state shape mismatch"); + TORCH_CHECK(first_to_dense.sizes() == union_indices.sizes(), "dense-slot shape mismatch"); + + constexpr int threads = 256; + const int blocks = destinations * static_cast(num_requests); + const auto stream = c10::cuda::getCurrentCUDAStream(indices.get_device()); + build_dense_union_remap_kernel<<>>( + indices.data_ptr(), + union_indices.data_ptr(), + remap.data_ptr(), + union_counts.data_ptr(), + hash_keys.data_ptr(), + hash_first_positions.data_ptr(), + first_to_dense.data_ptr(), + destinations, + static_cast(num_requests), + max_requests, + input_entries, + capacity, + hash_capacity); + AT_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("build_dense_union_remap", &build_dense_union_remap); +}