diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 1227669e39ab..198d30615635 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -35,7 +35,18 @@ B12xMLASparseImpl, B12xMLASparseMetadata, B12xMLASparseMetadataBuilder, + _ckv_prefetch_depth_within_budget, + _ckv_prefetch_execution_lanes, + _ckv_prefetch_ring_slots, + _ckv_prefetch_target_indices, + _ckv_prefetch_workspace_nbytes, + _CKVPrefetchStateRegistry, + _CKVPrefetchWorkspacePool, + _global_causal_lens_for_ckv_gather, + _is_glm_next_ckv_source_layout, _selected_index_block_stride_rows, + _use_b12x_full_ckv_gather, + _use_b12x_sparse_decode_plan, ) from vllm.v1.attention.backends.mla.sparse_utils import _remap_tiling from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -170,10 +181,12 @@ def test_b12x_glm5_next_cache_spec_and_layout(monkeypatch) -> None: assert invalid_reasons == [] assert unidentified == probe assert packed_by_glm_backend.state_content_bytes == 528 - assert packed_by_glm_backend.page_size_padded == 64 * 528 + 16 * 128 * 2 + assert packed_by_glm_backend.page_size_padded is None + assert packed_by_glm_backend.page_size_bytes == 64 * (528 + 33) assert packed_by_glm_backend.model_version == "glm5_next" assert packed.state_content_bytes == 528 - assert packed.page_size_padded == 64 * 528 + 16 * 128 * 2 + assert packed.page_size_padded is None + assert packed.page_size_bytes == 64 * (528 + 33) assert packed.model_version == "glm5_next" assert packed_without_config_context == packed assert layouts == (KVCacheLayout.BLHNC,) @@ -256,6 +269,170 @@ def test_b12x_glm5_next_accepts_dcp_with_speculation(monkeypatch) -> None: assert invalid_reasons == [] +@pytest.mark.parametrize( + ("max_query_len", "is_spec_decode", "force", "expected"), + [ + (1, False, False, True), + (6, True, False, True), + (6, False, False, False), + (6, False, True, True), + (9, True, False, False), + ], +) +def test_b12x_sparse_routes_only_verifier_extends_to_decode( + max_query_len: int, + is_spec_decode: bool, + force: bool, + expected: bool, +) -> None: + assert ( + _use_b12x_sparse_decode_plan( + max_query_len=max_query_len, + num_tokens=max_query_len * 4, + num_reqs=4, + is_spec_decode=is_spec_decode, + spec_extend_as_decode=True, + spec_extend_as_decode_force=force, + spec_decode_max_q=8, + max_tokens=4096, + ) + is expected + ) + + +@pytest.mark.parametrize( + ("max_query_len", "is_spec_decode", "num_tokens", "expected"), + [ + (1, False, 32, False), + (6, True, 192, False), + (6, False, 192, True), + (128, False, 8192, True), + (128, False, 600000, False), + ], +) +def test_b12x_full_ckv_gather_excludes_decode_and_mtp_batches( + max_query_len: int, + is_spec_decode: bool, + num_tokens: int, + expected: bool, +) -> None: + assert ( + _use_b12x_full_ckv_gather( + enabled=True, + is_glm_next=True, + dcp_world_size=4, + max_query_len=max_query_len, + num_tokens=num_tokens, + is_spec_decode=is_spec_decode, + min_tokens=16, + max_tokens=524288, + ) + is expected + ) + + +def test_b12x_full_ckv_gather_uses_global_causal_lengths() -> None: + global_seq_lens = torch.tensor([5, 12], dtype=torch.int32) + query_start_loc = torch.tensor([0, 2, 5], dtype=torch.int32) + req_id_per_token = torch.tensor([0, 0, 1, 1, 1], dtype=torch.int32) + + actual = _global_causal_lens_for_ckv_gather( + global_seq_lens, + query_start_loc, + req_id_per_token, + num_actual_tokens=5, + ) + + assert actual.tolist() == [4, 5, 10, 11, 12] + + +@pytest.mark.parametrize( + ("depth", "expected_slots", "expected_targets"), + [ + (0, 1, []), + (1, 2, [2]), + (3, 4, [2, 3, 4]), + ], +) +def test_b12x_ckv_prefetch_depth_controls_ring_and_targets( + depth: int, + expected_slots: int, + expected_targets: list[int], +) -> None: + caches = [torch.empty(0) for _ in range(5)] + assert _ckv_prefetch_ring_slots(depth) == expected_slots + assert _ckv_prefetch_target_indices(1, depth, caches, {}) == expected_targets + + +def test_b12x_ckv_prefetch_budget_caps_depth_but_keeps_sync_slot() -> None: + args = dict(dcp_world_size=4, local_capacity=1024, record_bytes=256) + assert _ckv_prefetch_workspace_nbytes(0, **args) == 5 * 1024 * 256 + assert _ckv_prefetch_workspace_nbytes(2, **args) == 13 * 1024 * 256 + assert _ckv_prefetch_depth_within_budget(3, 13 * 1024 * 256, **args) == 2 + assert _ckv_prefetch_depth_within_budget(3, 4 * 1024 * 256, **args) == 0 + assert _ckv_prefetch_depth_within_budget(3, 0, **args) == 3 + + +@pytest.mark.parametrize( + ("num_ubatches", "speculative", "expected"), + [(1, False, 1), (2, False, 2), (1, True, 2), (2, True, 4)], +) +def test_b12x_ckv_prefetch_reserves_execution_lanes( + num_ubatches: int, + speculative: bool, + expected: int, +) -> None: + assert _ckv_prefetch_execution_lanes(num_ubatches, speculative) == expected + + +def test_b12x_ckv_workspace_supports_unresolved_layer_index() -> None: + pool = _CKVPrefetchWorkspacePool(torch.device("cpu"), 64, 1) + registry = _CKVPrefetchStateRegistry() + query_workspace = torch.empty(8, dtype=torch.uint8) + + state = registry.for_workspace( + query_workspace, + layer_idx=None, + kv_cache=None, + workspace_pool=pool, + ) + + assert state.get_ckv_workspace(64).numel() == 64 + registry.clear() + + +def test_b12x_ckv_prefetch_targets_stop_at_unknown_layer() -> None: + caches = [torch.empty(0), torch.empty(0), None, torch.empty(0)] + assert _ckv_prefetch_target_indices(0, 3, caches, {}) == [1] + + +def test_b12x_ckv_prefetch_appends_current_chunk_to_rank_ordered_slots() -> None: + calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + impl = object.__new__(B12xMLASparseImpl) + impl.dcp_world_size = 4 + impl._ckv_current_chunk_kv_c = torch.arange(5 * 512).view(5, 512) + impl._concat_and_cache_glm_next_mla = lambda *args: calls.append(args) + cache = torch.empty((1, 64, 528), dtype=torch.uint8) + metadata = SimpleNamespace( + num_reqs=2, + req_id_per_token=torch.tensor([0, 0, 1, 1, 1], dtype=torch.int32), + global_cache_seq_lens_per_req=torch.tensor([5, 12], dtype=torch.int32), + query_start_loc=torch.tensor([0, 2, 5], dtype=torch.int32), + cp_kv_cache_interleave_size=4, + dcp_rank_req_starts=torch.tensor( + [[0, 4], [0, 1], [0, 0], [0, 0]], dtype=torch.int32 + ), + dcp_padded_total_tokens=8, + ) + + impl._append_current_chunk_to_gathered(cache, metadata, 5) + + assert len(calls) == 1 + assert torch.equal(calls[0][0], impl._ckv_current_chunk_kv_c) + assert calls[0][1] is cache + assert calls[0][2].tolist() == [3, 8, 17, 18, 19] + + def test_b12x_glm5_next_accepts_dcp_with_prefix_caching(monkeypatch) -> None: monkeypatch.setattr(b12x_mla_sparse, "get_b12x_sparse_mla", lambda: object()) with set_current_vllm_config( @@ -339,6 +516,8 @@ def test_b12x_glm5_next_selected_indices_use_physical_slots() -> None: ) == 64 ) + assert _is_glm_next_ckv_source_layout(cache, page_size=64) + assert not _is_glm_next_ckv_source_layout(cache[:, :, ::2], page_size=64) def test_sparse_index_remap_tiling_covers_glm5_next_width() -> None: @@ -397,6 +576,8 @@ def plan(caps): impl._model_type = 1 impl._decode_plan = SimpleNamespace() impl._extend_plan = SimpleNamespace() + impl._ckv_gather_enabled = False + impl._ckv_capacity_tokens = 0 owner = SimpleNamespace(impl=impl, indexer=None) cache = torch.empty((2, 1, 2304, 528), dtype=torch.uint8) @@ -414,10 +595,49 @@ def plan(caps): ] +def test_b12x_glm5_next_pretouches_largest_attention_workspace(monkeypatch) -> None: + calls: list[tuple[tuple[tuple[int, ...], torch.dtype], ...]] = [] + + class RecordingWorkspace: + def get_simultaneous(self, *specs): + calls.append(specs) + return [torch.empty(shape, dtype=dtype) for shape, dtype in specs] + + monkeypatch.setattr( + b12x_mla_sparse, + "current_workspace_manager", + lambda: RecordingWorkspace(), + ) + impl = object.__new__(B12xMLASparseImpl) + impl._max_tokens = 4 + impl._input_num_heads = 8 + impl.num_heads = 2 + impl._q_head_dim = 16 + impl._decode_plan = SimpleNamespace( + shapes_and_dtypes=lambda: (((64,), torch.uint8),) + ) + impl._extend_plan = SimpleNamespace( + shapes_and_dtypes=lambda: (((128,), torch.uint8),) + ) + impl._ckv_extend_plan = SimpleNamespace( + shapes_and_dtypes=lambda: (((4096,), torch.uint8),) + ) + + impl._pretouch_attention_workspace() + + assert calls == [ + ( + ((4, 2, 16), torch.bfloat16), + ((4096,), torch.uint8), + ) + ] + + def _bare_glm_selector_metadata_builder() -> B12xMLASparseMetadataBuilder: builder = B12xMLASparseMetadataBuilder.__new__(B12xMLASparseMetadataBuilder) builder.requires_glm_next_selector_metadata = True builder.supports_draft_decode_metadata_update = True + builder.ckv_prefetch_registry = None builder.dcp_world_size = 1 builder._capture_default_state_slot_ids = torch.arange(4, dtype=torch.int32) builder._capture_state_slot_ids = torch.empty(4, dtype=torch.int32) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index cbccbfc9a2ac..c1bbded4c35a 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -27,6 +27,7 @@ MLAAttention, QueryLenSupport, _DecodeConcatQuantFP8, + _select_mqa_query, build_mla_chunked_context_metadata, ) from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape @@ -68,6 +69,32 @@ DEVICE_TYPE = current_platform.device_type +@pytest.mark.cpu_test +def test_full_ckv_dcp_prefers_local_query_geometry() -> None: + q = torch.zeros((2, 2, 6)) + q_dcp_replicated = torch.ones((2, 8, 6)) + + selected, replicated = _select_mqa_query( + q, + q_dcp_replicated, + num_mqa_tokens=1, + full_ckv_dcp=True, + ) + assert selected.shape == (1, 2, 6) + assert not replicated + assert torch.equal(selected, q[:1]) + + selected, replicated = _select_mqa_query( + q, + q_dcp_replicated, + num_mqa_tokens=1, + full_ckv_dcp=False, + ) + assert selected.shape == (1, 8, 6) + assert replicated + assert torch.equal(selected, q_dcp_replicated[:1]) + + @pytest.mark.parametrize( ("cache_dtype", "expected_quant_mode"), [ diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 39f18c6d4f7c..190c16d97109 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -2079,6 +2079,30 @@ def test_get_kv_cache_configs_attention_free(): ] +def test_get_kv_cache_configs_preserves_model_sliding_window_retention(): + """Generic spec-decode planning must not erase model-specific retention.""" + model_config = ModelConfig(max_model_len=4096) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + spec = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + sliding_window=2048, + extra_retained_tokens=2048, + ) + + configs = get_kv_cache_configs( + vllm_config, + [{"draft": spec}], + [spec.page_size_bytes * 1024], + ) + + assert configs[0].kv_cache_groups[0].kv_cache_spec.extra_retained_tokens == 2048 + + def test_generate_uniform_type_kv_cache_specs(): # All layers are full attention, can be merged kv_cache_specs = { @@ -2232,6 +2256,62 @@ def test_group_and_unify_kv_cache_specs_mixed_page_size_groups(): assert layer_names == {"mla.0", "mla.1", "swa.0"} +def test_group_dcp_replicated_dflash_draft(): + target = new_mla_spec() + draft = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + sliding_window=2048, + dcp_replicated=True, + ) + assert target.page_size_bytes != draft.page_size_bytes + + specs = {"model.layers.0": target, "draft.layers.0": draft} + # DeepSeek-V4's UniformType tuple planner is not needed for DFlash. + assert group_and_unify_kv_cache_specs(specs) is None + + groups = get_kv_cache_groups(_grouping_config(), specs) + draft_group = next( + group for group in groups if isinstance(group.kv_cache_spec, SlidingWindowSpec) + ) + assert all(group.kv_cache_spec.block_size == 16 for group in groups) + assert draft_group.kv_cache_spec.dcp_replicated is True + + +def test_group_dcp_replicated_dflash_with_hybrid_mla_target(): + target_full = new_mla_spec(block_size=16) + target_swa = SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + sliding_window=2048, + ) + draft = SlidingWindowSpec( + block_size=256, + num_kv_heads=4, + head_size=128, + dtype=torch.bfloat16, + sliding_window=2048, + dcp_replicated=True, + ) + config = _grouping_config() + groups = get_kv_cache_groups( + config, + {"target.full": target_full, "target.swa": target_swa, "draft": draft}, + ) + + assert len(groups) == 3 + assert [group.layer_names for group in groups] == [ + ["target.full"], + ["target.swa"], + ["draft"], + ] + assert groups[-1].kv_cache_spec.dcp_replicated is True + + def new_indexer_mla_spec(block_size=16): # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller # page size than the main MLA attention (uint8, small head), so their pages diff --git a/tests/v1/core/test_mamba_align_chunk_split.py b/tests/v1/core/test_mamba_align_chunk_split.py index bc809fc3b530..8acd19c5efc1 100644 --- a/tests/v1/core/test_mamba_align_chunk_split.py +++ b/tests/v1/core/test_mamba_align_chunk_split.py @@ -82,6 +82,7 @@ def _split( use_eagle: bool = True, partial_hit: bool = False, num_prefill_checkpoint_blocks: int = 0, + has_sliding_eagle_group: bool = False, ) -> int: """Call the real `Scheduler._mamba_block_aligned_split` on a stub self.""" stub = SimpleNamespace( @@ -95,10 +96,25 @@ def _split( mamba_has_prefill_checkpoint_blocks=( num_prefill_checkpoint_blocks > 0 and not use_eagle ), + mamba_has_sliding_eagle_group=has_sliding_eagle_group, ) return Scheduler._mamba_block_aligned_split(stub, request, num_new_tokens) +def test_dflash_keeps_final_mamba_cache_boundary() -> None: + """DFlash drops its group-local draft block, not the target Mamba state.""" + (request,) = create_requests(1, num_tokens=10355, block_size=ATTN_BLOCK_SIZE) + request.num_computed_tokens = 5 * MAMBA_BLOCK_SIZE + assert ( + _split( + request, + request.num_tokens - request.num_computed_tokens, + has_sliding_eagle_group=True, + ) + == MAMBA_BLOCK_SIZE + ) + + @pytest.mark.parametrize( ("prompt_len", "num_new_tokens", "use_eagle", "expected"), [ diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 6fc2131b6e87..4004e70cbe81 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -3492,6 +3492,93 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(): assert [len(blocks) for blocks in computed_blocks.blocks] == [3, 12] +def test_dcp_hybrid_dflash_reuses_chunked_prompt_boundary(): + """DCP-sharded target + Mamba + replicated DFlash keeps one common hit.""" + hash_block_size = 2304 + scheduler_block_size = hash_block_size * 4 + kv_cache_config = KVCacheConfig( + num_blocks=2000, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["target"], + MLAAttentionSpec( + # DCP4 expands this to the 9,216-token effective target + # page used by GLM-5.3 Flash at runtime. + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float16, + tokens_per_state=4, + ), + is_eagle_group=True, + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=hash_block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + is_eagle_group=True, + ), + KVCacheGroupSpec( + ["draft"], + SlidingWindowSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float16, + sliding_window=2048, + dcp_replicated=True, + extra_retained_tokens=2048, + ), + is_eagle_group=True, + ), + ], + prefix_cache_retention_interval=0, + ) + manager = make_kv_cache_manager( + kv_cache_config, + max_model_len=524288, + max_in_flight_tokens=4096, + enable_caching=True, + # Replicated DFlash storage requires dense checkpoints independently + # of the aggregate coordinator flag. + use_eagle=False, + hash_block_size=hash_block_size, + scheduler_block_size=scheduler_block_size, + dcp_world_size=4, + ) + + token_ids = list(range(10355)) + fill = make_request("fill", token_ids, hash_block_size, sha256) + for chunk in ( + hash_block_size, + hash_block_size, + hash_block_size, + hash_block_size, + len(token_ids) - 4 * hash_block_size, + ): + blocks = manager.allocate_slots(fill, chunk) + assert blocks is not None + fill.num_computed_tokens += chunk + manager.free(fill) + + replay = make_request("replay", token_ids, hash_block_size, sha256) + _, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group( + replay.block_hashes, replay.num_tokens - 1 + ) + assert per_group_hits == ( + hash_block_size * 3, + hash_block_size * 4, + hash_block_size * 3, + ) + _, num_computed_tokens, _ = manager.get_computed_blocks(replay) + assert num_computed_tokens == hash_block_size * 3 + + def test_block_lookup_cache_single_block_per_key(): cache = BlockHashToBlockMap() key0 = BlockHashWithGroupId(b"hash0") diff --git a/tests/v1/spec_decode/test_dflash_dcp.py b/tests/v1/spec_decode/test_dflash_dcp.py new file mode 100644 index 000000000000..eac56eb4acc1 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_dcp.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import torch + +from vllm.model_executor.models.qwen3_dflash import DFlashAttention +from vllm.v1.attention.backend import AttentionType +from vllm.v1.kv_cache_interface import SlidingWindowSpec + + +def test_dflash_sliding_window_cache_is_replicated_under_dcp(): + attention = SimpleNamespace( + sliding_window=2048, + attn_type=AttentionType.DECODER, + num_kv_heads=1, + head_size=128, + head_size_v=128, + kv_cache_torch_dtype=torch.float8_e4m3fn, + kv_cache_dtype="fp8", + ) + config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=16), + parallel_config=SimpleNamespace(decode_context_parallel_size=4), + ) + + spec = DFlashAttention.get_kv_cache_spec(attention, config) + + assert isinstance(spec, SlidingWindowSpec) + assert spec.sliding_window == 2048 + assert spec.extra_retained_tokens == 2048 + assert spec.dcp_replicated is True diff --git a/tests/v1/worker/test_cp_utils.py b/tests/v1/worker/test_cp_utils.py index b38ae4d0636c..e57fdb46abd8 100644 --- a/tests/v1/worker/test_cp_utils.py +++ b/tests/v1/worker/test_cp_utils.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens +from vllm.v1.worker import cp_utils from vllm.v1.worker.cp_utils import should_skip_dcp_context_attention @@ -14,6 +17,67 @@ def test_skip_gate_only_for_zero_context(): ) +def test_non_mtp_speculator_does_not_require_mtp_interleave_support(monkeypatch): + layer_impl = SimpleNamespace( + supports_mtp_with_cp_non_trivial_interleave_size=False, + need_to_return_lse_for_decode=True, + ) + monkeypatch.setattr( + cp_utils, + "get_layers_from_vllm_config", + lambda *_args, **_kwargs: {"draft": SimpleNamespace(impl=layer_impl)}, + ) + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + decode_context_parallel_size=4, + cp_kv_cache_interleave_size=4, + ), + speculative_config=SimpleNamespace(method="dflash"), + ) + + cp_utils.check_attention_cp_compatibility(config) + + config.speculative_config.method = "mtp" + with pytest.raises(AssertionError, match="MTP with cp_kv_cache_interleave_size"): + cp_utils.check_attention_cp_compatibility(config) + + +def test_replicated_draft_attention_executes_as_local_dcp(monkeypatch): + layer_impl = SimpleNamespace( + supports_mtp_with_cp_non_trivial_interleave_size=False, + need_to_return_lse_for_decode=False, + dcp_world_size=4, + dcp_rank=2, + total_cp_world_size=4, + total_cp_rank=2, + ) + layer = SimpleNamespace( + impl=layer_impl, + get_kv_cache_spec=lambda _config: SimpleNamespace(dcp_replicated=True), + ) + monkeypatch.setattr( + cp_utils, + "get_layers_from_vllm_config", + lambda *_args, **_kwargs: {"draft": layer}, + ) + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + decode_context_parallel_size=4, + cp_kv_cache_interleave_size=4, + ), + speculative_config=SimpleNamespace(method="dflash"), + ) + + cp_utils.check_attention_cp_compatibility(config) + + assert layer_impl.dcp_world_size == 1 + assert layer_impl.dcp_rank == 0 + assert layer_impl.total_cp_world_size == 1 + assert layer_impl.total_cp_rank == 0 + + @pytest.mark.parametrize( "dcp_world_size,interleave_size,context_len", [(2, 16, 10), (4, 16, 10), (8, 16, 10), (4, 1, 2)], diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index ee44ff24d581..9ac2b2783326 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -176,6 +176,45 @@ def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int): assert torch.equal(actual, expected) +def test_dcp_slot_mapping_with_replicated_draft_group(): + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[16, 16], + max_num_reqs=1, + max_num_batched_tokens=32, + max_num_blocks_per_group=[2, 2], + device=device, + kernel_block_sizes=[16, 16], + cp_size=4, + cp_rank=1, + cp_interleave=4, + group_cp_sizes=[4, 1], + ) + block_tables.append_block_ids( + req_index=0, + new_block_ids=([5], [7, 8]), + overwrite=True, + ) + block_tables.apply_staged_writes() + + idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 32], dtype=torch.int32, device=device) + positions = torch.arange(32, dtype=torch.int64, device=device) + actual = block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + num_tokens_padded=32, + ) + + expected_target = torch.full((32,), -1, dtype=torch.int64, device=device) + expected_target[4:8] = torch.arange(80, 84, dtype=torch.int64, device=device) + expected_target[20:24] = torch.arange(84, 88, dtype=torch.int64, device=device) + expected_draft = torch.arange(112, 144, dtype=torch.int64, device=device) + assert torch.equal(actual[0], expected_target) + assert torch.equal(actual[1], expected_draft) + + def test_v1_block_table_move_row_clears_vacated_row(): """condense() moves the last row into a freed slot; the vacated row must not keep stale block ids. Padded dummy-run batches dereference stale rows diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 9353dbfe3609..2a4ef3ccc160 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1408,6 +1408,22 @@ def get_dcp_group() -> GroupCoordinator: return _DCP +_DCP_CKV_PREFETCH: GroupCoordinator | None = None + + +def get_dcp_ckv_prefetch_group() -> GroupCoordinator: + """Return the process group used to prefetch CKV across DCP ranks. + + Returns: + The initialized DCP CKV prefetch group coordinator. + + Raises: + AssertionError: If the DCP CKV prefetch group is not initialized. + """ + assert _DCP_CKV_PREFETCH is not None, "DCP CKV prefetch group is not initialized" + return _DCP_CKV_PREFETCH + + _PP: GroupCoordinator | None = None @@ -1875,6 +1891,19 @@ def initialize_model_parallel( group_name="dcp", ) + # CKV lookahead gathers run on a side stream. Give them a distinct NCCL + # communicator so they cannot overlap the indexer's default-stream DCP + # collectives on the same communicator. + global _DCP_CKV_PREFETCH + assert _DCP_CKV_PREFETCH is None, "DCP CKV prefetch group is already initialized" + if dcp_size > 1 and envs.VLLM_B12X_MLA_CKV_GATHER: + _DCP_CKV_PREFETCH = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="dcp_ckv_prefetch", + ) + global _PCP assert _PCP is None, "prefill context parallel group is already initialized" group_ranks = ( @@ -2103,6 +2132,11 @@ def destroy_model_parallel(): _DCP.destroy() _DCP = None + global _DCP_CKV_PREFETCH + if _DCP_CKV_PREFETCH: + _DCP_CKV_PREFETCH.destroy() + _DCP_CKV_PREFETCH = None + global _PCP if _PCP: _PCP.destroy() diff --git a/vllm/envs.py b/vllm/envs.py index 657f484407bd..18813337066d 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -190,6 +190,11 @@ VLLM_HUMMING_USE_F16_ACCUM: bool = False VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None VLLM_B12X_MOE_FP4_FORCE_A16: bool = False + 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_B12X_MLA_CKV_PREFETCH_WORKSPACE_MIB: int = 1024 VLLM_PLE_CPU_OFFLOAD: bool = False VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False VLLM_V1_USE_OUTLINES_CACHE: bool = False @@ -1627,6 +1632,24 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_B12X_MOE_FP4_FORCE_A16": lambda: bool( int(os.getenv("VLLM_B12X_MOE_FP4_FORCE_A16", "0")) ), + # Gather DCP-sharded C4 records before B12X sparse-MLA prefill. This avoids + # query replication plus the per-rank LSE combine and is opt-in while the + # path is being qualified on GLM5Next. + "VLLM_B12X_MLA_CKV_GATHER": lambda: ( + os.getenv("VLLM_B12X_MLA_CKV_GATHER", "0").lower() in ("1", "true", "yes", "on") + ), + "VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS": lambda: int( + os.getenv("VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS", "16") + ), + "VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS": lambda: int( + os.getenv("VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS", "524288") + ), + "VLLM_B12X_MLA_CKV_PREFETCH_DEPTH": lambda: int( + os.getenv("VLLM_B12X_MLA_CKV_PREFETCH_DEPTH", "1") + ), + "VLLM_B12X_MLA_CKV_PREFETCH_WORKSPACE_MIB": lambda: int( + os.getenv("VLLM_B12X_MLA_CKV_PREFETCH_WORKSPACE_MIB", "1024") + ), # Qwen3.8-Flash-Next only. Store PLE table payloads in CUDA-mapped host # memory unless additional_config.ple_table_memory is explicitly set. "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "0"))), diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 0a9a332becda..0421e790e592 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -346,6 +346,29 @@ def _detect_output_quant_key( return kFp8StaticTensorSym +def _select_mqa_query( + q: torch.Tensor, + q_dcp_replicated: torch.Tensor | None, + *, + num_mqa_tokens: int, + full_ckv_dcp: bool, +) -> tuple[torch.Tensor, bool]: + """Select local or replicated query geometry for MLA decode/prefill. + + Args: + q: Query tensor in local DCP geometry. + q_dcp_replicated: Query tensor replicated across the DCP group, if built. + num_mqa_tokens: Number of MQA query rows required by the backend. + full_ckv_dcp: Whether the backend attends a globally gathered CKV cache. + + Returns: + The selected query tensor and whether replicated geometry was selected. + """ + if q_dcp_replicated is not None and not full_ckv_dcp: + return q_dcp_replicated[:num_mqa_tokens], True + return q[:num_mqa_tokens], False + + def _canonicalize_sparse_mla_kv_cache_dtype( attn_backend: type[AttentionBackend], kv_cache_dtype: CacheDType, @@ -903,12 +926,17 @@ def forward_impl( ) if num_mqa_tokens > 0: - if q_dcp_replicated is not None: - mqa_q = q_dcp_replicated[:num_mqa_tokens] - qrep_decode = True - else: - mqa_q = q[:num_mqa_tokens] - qrep_decode = False + full_ckv_dcp = self.impl.uses_full_ckv_dcp(attn_metadata, num_mqa_tokens) + # Full-CKV prefill already makes every rank's cache visible to + # its local query heads. Prefer the local projection even when + # dcp_q_replicate retained a global query for ordinary DCP decode; + # the replicated query does not fit the local-head CKV plan. + mqa_q, qrep_decode = _select_mqa_query( + q, + q_dcp_replicated, + num_mqa_tokens=num_mqa_tokens, + full_ckv_dcp=full_ckv_dcp, + ) mqa_output_slice = output[:num_mqa_tokens] mqa_q_nope, mqa_q_pe = mqa_q.split( @@ -998,17 +1026,21 @@ def forward_impl( if isinstance(mqa_q, tuple): # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) mqa_q = torch.cat(mqa_q, dim=-1) - if not qrep_decode: + if not qrep_decode and not full_ckv_dcp: assert self.dcp_manager.query_gather is not None mqa_q = self.dcp_manager.query_gather(mqa_q) # call decode attn if not self.impl.is_sparse: assert attn_metadata.decode is not None + if full_ckv_dcp: + ckv_setter = getattr(self.impl, "set_ckv_current_chunk_kv", None) + if callable(ckv_setter): + ckv_setter(k_c_normed, k_pe) attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined] # correct dcp attn_out with lse. - if self.impl.dcp_world_size > 1: + if self.impl.dcp_world_size > 1 and not full_ckv_dcp: assert lse is not None assert self.dcp_manager is not None decode_metadata = getattr(attn_metadata, "decode", None) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index f5a7bff8fc6e..ba117479e345 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import dataclasses import io from collections.abc import Iterable @@ -35,6 +36,12 @@ from vllm.transformers_utils.config import set_default_rope_theta from vllm.transformers_utils.repo_utils import get_hf_file_bytes from vllm.v1.attention.backend import AttentionType +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheSpec, + SlidingWindowSpec, + get_kv_quant_mode, +) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( get_eagle3_aux_layers_from_config, ) @@ -169,6 +176,33 @@ def _resolve_layer_attention( return sliding_window, _dflash_layer_causal(config, layer_idx) +class DFlashAttention(Attention): + """Attention whose small draft KV is replicated across DCP ranks.""" + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + dcp_replicated = vllm_config.parallel_config.decode_context_parallel_size > 1 + if self.sliding_window is not None: + assert self.attn_type == AttentionType.DECODER + return SlidingWindowSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + head_size_v=self.head_size_v, + dtype=self.kv_cache_torch_dtype, + sliding_window=self.sliding_window, + # Prefix lookup verifies one lookahead block and then drops it. + # Keep one additional local window alive during chunked prefill + # so the proof block is not recycled before it can be hashed. + extra_retained_tokens=self.sliding_window, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + dcp_replicated=dcp_replicated, + ) + spec = super().get_kv_cache_spec(vllm_config) + if dcp_replicated and isinstance(spec, FullAttentionSpec): + spec = dataclasses.replace(spec, dcp_replicated=True) + return spec + + class DFlashQwen3Attention(nn.Module): """Attention for DFlash speculative decoding. @@ -244,7 +278,7 @@ def __init__( ) self.sliding_window = sliding_window - self.attn = Attention( + self.attn = DFlashAttention( self.num_heads, self.head_dim, self.scaling, diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 4fdfafaaea9a..0bfaa4240e67 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -985,6 +985,18 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]): supports_pcp: bool = True + def uses_full_ckv_dcp(self, attn_metadata: T, num_tokens: int) -> bool: + """Report whether this call attends a globally gathered DCP cache. + + Args: + attn_metadata: Backend-specific attention metadata for this call. + num_tokens: Number of scheduled input tokens. + + Returns: + True when the implementation will use globally gathered CKV. + """ + return False + @abstractmethod def __init__( self, diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 3d0244520c99..6a057ee36511 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -409,15 +409,17 @@ def __init__( self.max_num_splits = 0 # No upper bound on the number of splits. self.aot_schedule = get_flash_attn_version() == 3 - try: - from vllm.distributed.parallel_state import get_dcp_group - - self.dcp_world_size = get_dcp_group().world_size - self.dcp_rank = get_dcp_group().rank_in_group - except AssertionError: - # DCP might not be initialized in testing + if getattr(kv_cache_spec, "dcp_replicated", False): self.dcp_world_size = 1 self.dcp_rank = 0 + else: + try: + self.dcp_world_size = get_dcp_group().world_size + self.dcp_rank = get_dcp_group().rank_in_group + except AssertionError: + # DCP might not be initialized in testing + self.dcp_world_size = 1 + self.dcp_rank = 0 # Fused draft decode reuses the captured metadata object across draft # steps. For DCP, build-time host-side decisions such as diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index b9aaf1ea25ce..2b0d6c1cb05b 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -2,16 +2,23 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """B12x sparse MLA attention backend.""" +import os +import weakref from dataclasses import dataclass, replace -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np import torch +import torch.distributed as dist +import triton +import triton.language as tl from vllm import _custom_ops as ops +from vllm import envs from vllm.config import VllmConfig, get_current_vllm_config_or_none from vllm.config.cache import CacheDType from vllm.distributed import get_dcp_group +from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, @@ -46,6 +53,8 @@ _GLM_NEXT_CACHE_RECORD_BYTES = 528 _GLM_NEXT_INDEX_TAIL_BYTES_PER_TOKEN = 132 // 4 +logger = init_logger(__name__) + def _is_glm_next_config(hf_config: object | None) -> bool: return getattr(hf_config, "model_type", None) in _GLM_NEXT_MODEL_TYPES @@ -113,6 +122,569 @@ def _selected_index_block_stride_rows( return int(kv_cache.stride(0)) // record_width +def _is_glm_next_ckv_source_layout( + kv_cache: torch.Tensor, + *, + page_size: int, +) -> bool: + return ( + kv_cache.dtype == torch.uint8 + and kv_cache.ndim == 3 + and tuple(kv_cache.shape[1:]) == (page_size, _GLM_NEXT_CACHE_RECORD_BYTES) + and kv_cache.stride(1) == _GLM_NEXT_CACHE_RECORD_BYTES + and kv_cache.stride(2) == 1 + ) + + +def _use_b12x_sparse_decode_plan( + *, + max_query_len: int, + num_tokens: int, + num_reqs: int, + is_spec_decode: bool, + spec_extend_as_decode: bool, + spec_extend_as_decode_force: bool, + spec_decode_max_q: int, + max_tokens: int, +) -> bool: + if max_query_len <= 1: + return True + use_spec_decode = spec_extend_as_decode and ( + spec_extend_as_decode_force or is_spec_decode + ) + return ( + use_spec_decode + and max_query_len <= spec_decode_max_q + and num_tokens <= num_reqs * spec_decode_max_q + and num_tokens <= max_tokens + ) + + +def _use_b12x_full_ckv_gather( + *, + enabled: bool, + is_glm_next: bool, + dcp_world_size: int, + max_query_len: int, + num_tokens: int, + is_spec_decode: bool, + min_tokens: int, + max_tokens: int, +) -> bool: + return ( + enabled + and is_glm_next + and dcp_world_size > 1 + and max_query_len > 1 + and not is_spec_decode + and num_tokens > min_tokens + and num_tokens <= max_tokens + ) + + +def _ckv_prefetch_ring_slots(depth: int) -> int: + return max(0, int(depth)) + 1 + + +def _ckv_prefetch_workspace_nbytes( + depth: int, + dcp_world_size: int, + local_capacity: int, + record_bytes: int, +) -> int: + """Return one lane's local staging plus gathered-cache ring size.""" + return ( + (1 + _ckv_prefetch_ring_slots(depth) * int(dcp_world_size)) + * int(local_capacity) + * int(record_bytes) + ) + + +def _ckv_prefetch_execution_lanes(num_ubatches: int, speculative: bool) -> int: + return max(1, int(num_ubatches)) * (2 if speculative else 1) + + +def _ckv_prefetch_depth_within_budget( + requested_depth: int, + workspace_budget_bytes: int, + dcp_world_size: int, + local_capacity: int, + record_bytes: int, +) -> int: + """Cap lookahead depth without removing the synchronous gather slot.""" + requested_depth = max(0, int(requested_depth)) + workspace_budget_bytes = int(workspace_budget_bytes) + if workspace_budget_bytes <= 0: + return requested_depth + for depth in range(requested_depth, -1, -1): + if ( + _ckv_prefetch_workspace_nbytes( + depth, + dcp_world_size, + local_capacity, + record_bytes, + ) + <= workspace_budget_bytes + ): + return depth + return 0 + + +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 + + +class _CKVPrefetchWorkspacePool: + """Preallocated CKV rings shared by attention layers on one device.""" + + def __init__( + self, + device: torch.device, + slot_nbytes: int, + max_slots: int, + ) -> None: + if slot_nbytes <= 0 or max_slots <= 0: + raise ValueError( + "CKV workspace pool requires positive slot size and count, got " + f"slot_nbytes={slot_nbytes} max_slots={max_slots}" + ) + self.device = device + self.slot_nbytes = int(slot_nbytes) + self.max_slots = int(max_slots) + self.storage = torch.empty( + (self.slot_nbytes * self.max_slots,), + dtype=torch.uint8, + device=device, + ) + self._free_slots = list(reversed(range(self.max_slots))) + self._leased_slots: set[int] = set() + + def acquire(self) -> tuple[int, torch.Tensor]: + if not self._free_slots: + raise RuntimeError( + "CKV prefetch workspace pool exhausted. The runtime created more " + f"than {self.max_slots} execution lanes." + ) + slot = self._free_slots.pop() + self._leased_slots.add(slot) + start = slot * self.slot_nbytes + return slot, self.storage.narrow(0, start, self.slot_nbytes) + + def release(self, slot: int) -> None: + if slot not in self._leased_slots: + raise RuntimeError(f"CKV workspace slot {slot} is not leased") + self._leased_slots.remove(slot) + self._free_slots.append(slot) + + +_CKV_PREFETCH_WORKSPACE_POOLS: dict[ + tuple[str, int | None, int, int], _CKVPrefetchWorkspacePool +] = {} + + +def _get_ckv_prefetch_workspace_pool( + device: torch.device, + slot_nbytes: int, + max_slots: int, +) -> _CKVPrefetchWorkspacePool: + key = (device.type, device.index, int(slot_nbytes), int(max_slots)) + pool = _CKV_PREFETCH_WORKSPACE_POOLS.get(key) + if pool is None: + pool = _CKVPrefetchWorkspacePool(device, slot_nbytes, max_slots) + _CKV_PREFETCH_WORKSPACE_POOLS[key] = pool + return pool + + +@dataclass(frozen=True) +class _CKVWorkspaceIdentity: + device: torch.device + storage_data_ptr: int + storage_nbytes: 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, + storage_data_ptr=storage.data_ptr(), + storage_nbytes=storage.nbytes(), + 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, + workspace_pool: _CKVPrefetchWorkspacePool, + ) -> None: + self.workspace_identity = workspace_identity + self.workspace_storage_ref = weakref.ref(workspace.untyped_storage()) + self.workspace_pool = workspace_pool + 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_slot: int | None = None + self.last_layer_idx: int | None = None + + def begin_step(self) -> None: + self.wait_for_pending_writes() + self.pending_layers.clear() + self.last_layer_idx = None + + def wait_for_pending_writes(self) -> None: + for event, _ in self.pending_layers.values(): + event.wait() + + 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 != self.workspace_pool.slot_nbytes: + raise ValueError( + "CKV workspace size changed after persistent allocation: " + f"pool={self.workspace_pool.slot_nbytes} requested={nbytes}" + ) + if self.ckv_workspace is None: + slot, workspace = self.workspace_pool.acquire() + self.ckv_workspace_slot = slot + self.ckv_workspace = workspace + return self.ckv_workspace + + def close(self) -> None: + for event, _ in self.pending_layers.values(): + event.synchronize() + self.pending_layers.clear() + self.last_layer_idx = None + if self.ckv_workspace_slot is not None: + self.workspace_pool.release(self.ckv_workspace_slot) + self.ckv_workspace_slot = None + self.ckv_workspace = None + + +_CKV_PREFETCH_STATE_REGISTRIES: weakref.WeakSet = weakref.WeakSet() + + +class _CKVPrefetchStateRegistry: + """Builder-owned states partitioned by lane-scoped query workspace.""" + + def __init__(self) -> None: + self.states: dict[_CKVWorkspaceIdentity, _CKVPrefetchState] = {} + self.workspace_pool: _CKVPrefetchWorkspacePool | None = None + _CKV_PREFETCH_STATE_REGISTRIES.add(self) + + def _bind_workspace_pool(self, pool: _CKVPrefetchWorkspacePool) -> None: + if self.workspace_pool is None: + self.workspace_pool = pool + elif self.workspace_pool is not pool: + raise RuntimeError("CKV prefetch registry cannot switch workspace pools") + + def _retire(self, identities: list[_CKVWorkspaceIdentity]) -> None: + for identity in identities: + self.states.pop(identity).close() + + def _prune_released_workspaces(self) -> None: + self._retire( + [ + identity + for identity, state in self.states.items() + if state.workspace_storage_ref() is None + ] + ) + + def begin_step(self) -> None: + self._prune_released_workspaces() + for state in self.states.values(): + state.begin_step() + + def clear(self) -> None: + self._retire(list(self.states)) + + def for_workspace( + self, + workspace: torch.Tensor, + layer_idx: int | None, + kv_cache: torch.Tensor | None, + workspace_pool: _CKVPrefetchWorkspacePool, + ) -> _CKVPrefetchState: + self._prune_released_workspaces() + self._bind_workspace_pool(workspace_pool) + identity = _ckv_workspace_identity(workspace) + state = self.states.get(identity) + if state is None: + 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) + assert self.workspace_pool is not None + state = _CKVPrefetchState(identity, workspace, self.workspace_pool) + self.states[identity] = state + return state + + +def _dcp_all_gather_current_stream( + group, + input_tensor: torch.Tensor, + output_tensor: torch.Tensor, +) -> None: + if not input_tensor.is_contiguous() or not output_tensor.is_contiguous(): + raise ValueError("CKV all-gather tensors must be contiguous") + if output_tensor.numel() != input_tensor.numel() * group.world_size: + raise ValueError("CKV all-gather tensors have incompatible sizes") + + communicator = getattr(group, "device_communicator", None) + pynccl_comm = getattr(communicator, "pynccl_comm", None) + if pynccl_comm is not None and not getattr(pynccl_comm, "disabled", False): + pynccl_comm.all_gather(output_tensor, input_tensor) + return + + device_group = getattr(group, "device_group", None) + if device_group is None: + device_group = getattr(communicator, "device_group", None) + if device_group is not None: + dist.all_gather_into_tensor( + output_tensor, + input_tensor, + group=device_group, + async_op=False, + ) + return + + output_tensor.copy_(group.all_gather(input_tensor, dim=0)) + + +@triton.jit +def _mask_page_table_after_nsa_len_kernel( + page_table_ptr, + nsa_len_ptr, + page_stride0, + page_stride1, + width: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + offs = tile * BLOCK_N + tl.arange(0, BLOCK_N) + valid = offs < width + nsa_len = tl.load(nsa_len_ptr + row) + tl.store( + page_table_ptr + row * page_stride0 + offs * page_stride1, + -1, + mask=valid & (offs >= nsa_len), + ) + + +def _mask_page_table_after_nsa_len( + page_table: torch.Tensor, + nsa_cache_seqlens: torch.Tensor, +) -> None: + width = page_table.shape[1] + if width == 0 or page_table.shape[0] == 0: + return + block_n = 128 + _mask_page_table_after_nsa_len_kernel[ + (page_table.shape[0], triton.cdiv(width, block_n)) + ]( + page_table, + nsa_cache_seqlens, + page_table.stride(0), + page_table.stride(1), + width, + BLOCK_N=block_n, + ) + + +def _global_causal_lens_for_ckv_gather( + global_seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + req_id_per_token: torch.Tensor, + num_actual_tokens: int, +) -> torch.Tensor: + """Return each query token's causal length in the gathered global cache.""" + num_reqs = global_seq_lens.shape[0] + qsl = query_start_loc[: num_reqs + 1].to(torch.int32) + req_ids = req_id_per_token[:num_actual_tokens].to(torch.int64) + chunk_start = qsl[:-1][req_ids] + chunk_len = (qsl[1:] - qsl[:-1])[req_ids] + full_seq = global_seq_lens[req_ids].to(torch.int32) + token_idx = torch.arange( + num_actual_tokens, + device=global_seq_lens.device, + dtype=torch.int32, + ) + return full_seq - chunk_len + (token_idx - chunk_start) + 1 + + +@triton.jit +def _map_global_topk_to_gathered_ckv_kernel( + req_id_ptr, + token_indices_ptr, + rank_req_starts_ptr, + rank_req_lens_ptr, + out_ptr, + valid_count_ptr, + starts_stride0, + starts_stride1, + lens_stride0, + lens_stride1, + ti_stride0, + ti_stride1, + out_stride0, + out_stride1, + padded_rank_tokens, + DCP_SIZE: tl.constexpr, + DCP_INTERLEAVE: tl.constexpr, + NUM_TOPK_TOKENS: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + cols = tile * BLOCK_N + tl.arange(0, BLOCK_N) + col_mask = cols < NUM_TOPK_TOKENS + req = tl.load(req_id_ptr + row) + tok = tl.load( + token_indices_ptr + row * ti_stride0 + cols * ti_stride1, + mask=col_mask, + other=-1, + ) + owner = (tok // DCP_INTERLEAVE) % DCP_SIZE + local_idx = ( + tok // (DCP_SIZE * DCP_INTERLEAVE) + ) * DCP_INTERLEAVE + tok % DCP_INTERLEAVE + valid_tok = col_mask & (tok >= 0) + req_start = tl.load( + rank_req_starts_ptr + owner * starts_stride0 + req * starts_stride1, + mask=valid_tok, + other=0, + ) + req_len = tl.load( + rank_req_lens_ptr + owner * lens_stride0 + req * lens_stride1, + mask=valid_tok, + other=0, + ) + valid = valid_tok & (local_idx >= 0) & (local_idx < req_len) + gathered_slot = owner * padded_rank_tokens + req_start + local_idx + valid_i32 = valid.to(tl.int32) + local_offset = tl.cumsum(valid_i32) - valid_i32 + tile_valid_count = tl.sum(valid_i32) + output_base = tl.atomic_add(valid_count_ptr + row, tile_valid_count) + tl.store( + out_ptr + row * out_stride0 + (output_base + local_offset) * out_stride1, + gathered_slot, + mask=valid, + ) + + +def _map_global_topk_to_gathered_ckv( + req_ids: torch.Tensor, + token_indices: torch.Tensor, + rank_req_starts: torch.Tensor, + rank_req_lens: torch.Tensor, + out: torch.Tensor, + valid_counts: torch.Tensor, + *, + dcp_size: int, + cp_kv_cache_interleave_size: int, + padded_rank_tokens: int, +) -> None: + if token_indices.shape != out.shape: + raise ValueError("CKV gather index output shape does not match top-k input") + if rank_req_starts.shape != rank_req_lens.shape: + raise ValueError("CKV gather request starts/lens shapes do not match") + if rank_req_starts.shape[0] != dcp_size: + raise ValueError("CKV gather request metadata does not match DCP size") + if any( + tensor.dtype != torch.int32 + for tensor in ( + req_ids, + token_indices, + rank_req_starts, + rank_req_lens, + out, + valid_counts, + ) + ): + raise TypeError("CKV gather index metadata must be int32") + + block_n = 128 + out.fill_(-1) + valid_counts.zero_() + _map_global_topk_to_gathered_ckv_kernel[ + (token_indices.shape[0], triton.cdiv(token_indices.shape[1], block_n)) + ]( + req_ids, + token_indices, + rank_req_starts, + rank_req_lens, + out, + valid_counts, + rank_req_starts.stride(0), + rank_req_starts.stride(1), + rank_req_lens.stride(0), + rank_req_lens.stride(1), + token_indices.stride(0), + token_indices.stride(1), + out.stride(0), + out.stride(1), + padded_rank_tokens, + DCP_SIZE=dcp_size, + DCP_INTERLEAVE=cp_kv_cache_interleave_size, + NUM_TOPK_TOKENS=token_indices.shape[1], + BLOCK_N=block_n, + ) + + class B12xMLASparseBackend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ @@ -262,6 +834,7 @@ class B12xMLASparseMetadata(AttentionMetadata): num_decodes: int num_prefills: int num_decode_tokens: int + is_spec_decode: bool = False prefill_max_seq_len: int = 0 prefill: MLACommonPrefillMetadata | None = None prefill_query_lens_cpu: torch.Tensor | None = None @@ -273,6 +846,16 @@ class B12xMLASparseMetadata(AttentionMetadata): selector_state_is_fresh: torch.Tensor | None = None selector_num_accepted_tokens: torch.Tensor | None = None selector_is_prefilling: torch.Tensor | None = None + ckv_selected_indices: torch.Tensor | None = None + ckv_active_counts: torch.Tensor | None = None + dcp_rank_req_starts: torch.Tensor | None = None + dcp_rank_req_lens: torch.Tensor | None = None + dcp_local_cu_seq_lens: torch.Tensor | None = None + global_cache_seq_lens_per_req: torch.Tensor | None = None + dcp_local_total_tokens: int = 0 + dcp_padded_total_tokens: int = 0 + dcp_ckv_gather_eligible: bool = False + ckv_prefetch_registry: _CKVPrefetchStateRegistry | None = None class B12xMLASparseMetadataBuilder( @@ -301,12 +884,17 @@ def __init__( ) self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 scheduler_config = vllm_config.scheduler_config + speculative_config = vllm_config.speculative_config + self.num_speculative_tokens = int( + getattr(speculative_config, "num_speculative_tokens", 0) or 0 + ) max_tokens = scheduler_config.max_num_batched_tokens + max_reqs = int(scheduler_config.max_num_seqs) + self._ckv_max_reqs = max_reqs self.cache_seq_lens_per_token_buffer = torch.empty( (max_tokens,), dtype=torch.int32, device=device ) if self.requires_glm_next_selector_metadata: - max_reqs = int(scheduler_config.max_num_seqs) self._capture_default_state_slot_ids = torch.arange( max_reqs, dtype=torch.int32, device=device ) @@ -322,6 +910,41 @@ def __init__( self._capture_is_prefilling = torch.zeros( max_reqs, dtype=torch.bool, device=device ) + self._ckv_gather_requested = ( + self.requires_glm_next_selector_metadata + and self.dcp_world_size > 1 + and envs.VLLM_B12X_MLA_CKV_GATHER + ) + ckv_workspace_requested = ( + self.dcp_world_size > 1 and envs.VLLM_B12X_MLA_CKV_GATHER + ) + self.ckv_prefetch_registry = ( + _CKVPrefetchStateRegistry() if ckv_workspace_requested else None + ) + if self._ckv_gather_requested: + hf_config = vllm_config.model_config.hf_text_config + ckv_topk_tokens = int(hf_config.index_topk) + int(hf_config.index_kpool) - 1 + self.ckv_selected_indices_buffer = torch.empty( + (max_tokens, ckv_topk_tokens), dtype=torch.int32, device=device + ) + self.ckv_active_counts_buffer = torch.empty( + (max_tokens,), dtype=torch.int32, device=device + ) + self.dcp_rank_req_lens_buffer = torch.empty( + (self.dcp_world_size, max_reqs), dtype=torch.int32, device=device + ) + self.dcp_rank_req_starts_buffer = torch.empty( + (self.dcp_world_size, max_reqs), dtype=torch.int32, device=device + ) + self.dcp_local_cu_seq_lens_buffer = torch.empty( + (max_reqs + 1,), dtype=torch.int32, device=device + ) + else: + self.ckv_selected_indices_buffer = None + self.ckv_active_counts_buffer = None + self.dcp_rank_req_lens_buffer = None + self.dcp_rank_req_starts_buffer = None + self.dcp_local_cu_seq_lens_buffer = None num_q_heads = vllm_config.model_config.get_num_attention_heads( vllm_config.parallel_config ) @@ -431,6 +1054,8 @@ def _build( selector_num_accepted_tokens: torch.Tensor | None = None, selector_is_prefilling: torch.Tensor | None = None, ) -> B12xMLASparseMetadata: + if self.ckv_prefetch_registry is not None: + self.ckv_prefetch_registry.begin_step() metadata = super().build( common_prefix_len, common_attn_metadata, fast_build=fast_build ) @@ -443,6 +1068,7 @@ def _build( else common.seq_lens ) metadata.seq_lens = seq_lens + metadata.ckv_prefetch_registry = self.ckv_prefetch_registry if common.max_query_len <= 1 and num_tokens == common.num_reqs: per_token_lens = seq_lens[:num_tokens] @@ -484,12 +1110,94 @@ def _build( per_token_lens = self.cache_seq_lens_per_token_buffer[:num_tokens] metadata.cache_seq_lens_per_token = per_token_lens + metadata.is_spec_decode = False + if ( + self.num_speculative_tokens > 0 + and 1 < common.max_query_len <= self.num_speculative_tokens + 1 + and common.is_prefilling is not None + ): + metadata.is_spec_decode = not bool( + torch.any(common.is_prefilling[: common.num_reqs]) + ) if metadata.num_prefills: prefill_start = metadata.num_decodes prefill_end = prefill_start + metadata.num_prefills + 1 metadata.prefill_query_lens_cpu = torch.diff( common.query_start_loc_cpu[prefill_start:prefill_end] ) + if ( + _use_b12x_full_ckv_gather( + enabled=self._ckv_gather_requested, + is_glm_next=self.requires_glm_next_selector_metadata, + dcp_world_size=self.dcp_world_size, + max_query_len=common.max_query_len, + num_tokens=num_tokens, + is_spec_decode=metadata.is_spec_decode, + min_tokens=envs.VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS, + max_tokens=envs.VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS, + ) + and metadata.num_decode_tokens == 0 + ): + assert self.ckv_selected_indices_buffer is not None + assert self.ckv_active_counts_buffer is not None + assert self.dcp_rank_req_lens_buffer is not None + assert self.dcp_rank_req_starts_buffer is not None + assert self.dcp_local_cu_seq_lens_buffer is not None + global_seq_lens = common.seq_lens[: common.num_reqs] + all_rank_lens = get_dcp_local_seq_lens( + global_seq_lens, + self.dcp_world_size, + dcp_rank=None, + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + ).transpose(0, 1) + rank_req_lens = self.dcp_rank_req_lens_buffer[ + : self.dcp_world_size, : common.num_reqs + ] + rank_req_lens.copy_(all_rank_lens) + rank_req_starts = self.dcp_rank_req_starts_buffer[ + : self.dcp_world_size, : common.num_reqs + ] + rank_req_starts[:, 0].zero_() + if common.num_reqs > 1: + torch.cumsum(rank_req_lens[:, :-1], dim=1, out=rank_req_starts[:, 1:]) + local_cu_seq_lens = self.dcp_local_cu_seq_lens_buffer[: common.num_reqs + 1] + local_cu_seq_lens[0].zero_() + torch.cumsum( + rank_req_lens[self.dcp_rank], + dim=0, + out=local_cu_seq_lens[1:], + ) + rank_totals = rank_req_lens.sum(dim=1).tolist() + local_total_tokens = int(rank_totals[self.dcp_rank]) + page_size = int(self.kv_cache_spec.block_size) + padded_total_tokens = ( + (max(int(total) for total in rank_totals) + page_size - 1) + // page_size + * page_size + ) + max_local_capacity = ( + ( + (envs.VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS + self.dcp_world_size - 1) + // self.dcp_world_size + + self._ckv_max_reqs * self.cp_kv_cache_interleave_size + + page_size + - 1 + ) + // page_size + * page_size + ) + if 0 < padded_total_tokens <= max_local_capacity: + metadata.ckv_selected_indices = self.ckv_selected_indices_buffer[ + :num_tokens + ] + metadata.ckv_active_counts = self.ckv_active_counts_buffer[:num_tokens] + metadata.dcp_rank_req_lens = rank_req_lens + metadata.dcp_rank_req_starts = rank_req_starts + metadata.dcp_local_cu_seq_lens = local_cu_seq_lens + metadata.global_cache_seq_lens_per_req = global_seq_lens + metadata.dcp_local_total_tokens = local_total_tokens + metadata.dcp_padded_total_tokens = padded_total_tokens + metadata.dcp_ckv_gather_eligible = True ( metadata.selector_state_slot_ids, metadata.selector_state_is_fresh, @@ -560,6 +1268,17 @@ class B12xMLASparseImpl(SparseMLACommonImpl[B12xMLASparseMetadata]): supports_dense_mha_prefill = False supports_pcp = False + @classmethod + def reset_kv_cache_binding_state(cls) -> None: + """Release class-wide CKV state before a cache allocation is replaced. + + This hook resets shared binding registries for every instance of this + implementation class. The worker therefore invokes it once per concrete + implementation type during cache unbinding. + """ + for registry in tuple(_CKV_PREFETCH_STATE_REGISTRIES): + registry.clear() + def __init__( self, num_heads: int, @@ -659,16 +1378,147 @@ def __init__( ) self._max_tokens = max_tokens self._max_seqs = max_seqs + self._spec_decode_max_q = int(os.getenv("VLLM_B12X_MLA_SPEC_DECODE_MAX_Q", "8")) + spec_decode_mode = ( + os.getenv("VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE", "off").strip().lower() + ) + disabled_modes = {"0", "false", "off", "no"} + forced_modes = {"1", "true", "on", "yes"} + if spec_decode_mode not in {"auto", *disabled_modes, *forced_modes}: + raise ValueError( + "VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE must be auto, 0, or 1 " + f"(got {spec_decode_mode!r})" + ) + self._spec_extend_as_decode = spec_decode_mode not in disabled_modes + self._spec_extend_as_decode_force = spec_decode_mode in forced_modes self._kv_dtype = torch.uint8 kernel_page_size = ( int(vllm_config.cache_config.block_size) if self._is_glm_next else 64 ) + self._ckv_gather_enabled = ( + self._is_glm_next + and self.dcp_world_size > 1 + and envs.VLLM_B12X_MLA_CKV_GATHER + ) + max_ckv_tokens = envs.VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS + cp_kv_cache_interleave_size = int( + vllm_config.parallel_config.cp_kv_cache_interleave_size + ) + self._ckv_capacity_tokens = ( + max_ckv_tokens + self.dcp_world_size - 1 + ) // self.dcp_world_size + max_seqs * cp_kv_cache_interleave_size + self._ckv_local_capacity = 0 self._module = module self._kernel_page_size = 0 self._set_kernel_page_size(kernel_page_size) + configured_prefetch_depth = max(0, int(envs.VLLM_B12X_MLA_CKV_PREFETCH_DEPTH)) + configured_workspace_mib = max( + 0, int(envs.VLLM_B12X_MLA_CKV_PREFETCH_WORKSPACE_MIB) + ) + workspace_budget_bytes = configured_workspace_mib * 1024 * 1024 + requested_prefetch_depth = ( + configured_prefetch_depth if self._ckv_gather_enabled else 0 + ) + self._ckv_prefetch_depth = _ckv_prefetch_depth_within_budget( + requested_prefetch_depth, + workspace_budget_bytes, + self.dcp_world_size, + self._ckv_local_capacity, + _GLM_NEXT_CACHE_RECORD_BYTES, + ) + self._ckv_workspace_slots = _ckv_prefetch_ring_slots(self._ckv_prefetch_depth) + self._ckv_workspace_nbytes = ( + _ckv_prefetch_workspace_nbytes( + self._ckv_prefetch_depth, + self.dcp_world_size, + self._ckv_local_capacity, + _GLM_NEXT_CACHE_RECORD_BYTES, + ) + if self._ckv_gather_enabled + else 0 + ) + execution_lanes = _ckv_prefetch_execution_lanes( + vllm_config.parallel_config.num_ubatches, + vllm_config.speculative_config is not None, + ) + device = torch.device("cuda", torch.accelerator.current_device_index()) + self._ckv_workspace_pool = ( + _get_ckv_prefetch_workspace_pool( + device, + self._ckv_workspace_nbytes, + execution_lanes, + ) + if self._ckv_workspace_nbytes > 0 + else None + ) + self._ckv_current_chunk_kv_c: torch.Tensor | None = None + if self._ckv_workspace_pool is not None: + logger.info_once( + "Using CKV layer prefetch depth=%d with %.1f MiB for %d " + "execution lane(s)", + self._ckv_prefetch_depth, + self._ckv_workspace_pool.storage.numel() / (1024 * 1024), + execution_lanes, + ) + if self._ckv_prefetch_depth < requested_prefetch_depth: + logger.info_once( + "Capped CKV prefetch depth from %d to %d for the %d MiB " + "per-lane workspace budget", + requested_prefetch_depth, + self._ckv_prefetch_depth, + configured_workspace_mib, + ) + self._pretouch_attention_workspace() self.supports_quant_query_input = False + def _pretouch_attention_workspace(self) -> None: + """Reserve the largest planned attention scratch before KV profiling.""" + candidates = [ + ( + ( + (self._max_tokens, self._input_num_heads, self._q_head_dim), + torch.bfloat16, + ), + *self._decode_plan.shapes_and_dtypes(), + ), + ( + ( + (self._max_tokens, self._input_num_heads, self._q_head_dim), + torch.bfloat16, + ), + *self._extend_plan.shapes_and_dtypes(), + ), + ] + if self._ckv_extend_plan is not None: + candidates.append( + ( + ( + (self._max_tokens, self.num_heads, self._q_head_dim), + torch.bfloat16, + ), + *self._ckv_extend_plan.shapes_and_dtypes(), + ) + ) + + def workspace_bytes(specs: tuple[tuple[tuple[int, ...], torch.dtype], ...]): + total = 0 + for shape, dtype in specs: + numel = 1 + for dim in shape: + numel *= int(dim) + nbytes = numel * torch.empty((), dtype=dtype).element_size() + total += (nbytes + 255) // 256 * 256 + return total + + largest = max(candidates, key=workspace_bytes) + reserved_bytes = workspace_bytes(largest) + current_workspace_manager().get_simultaneous(*largest) + logger.info_once( + "Preallocated %.1f MiB of B12X sparse MLA scratch before KV profiling", + reserved_bytes / (1024 * 1024), + ) + def _set_kernel_page_size(self, kernel_page_size: int) -> None: if kernel_page_size <= 0 or kernel_page_size % 64: raise ValueError( @@ -678,10 +1528,10 @@ def _set_kernel_page_size(self, kernel_page_size: int) -> None: if kernel_page_size == self._kernel_page_size: return - def make_plan(mode: str): + def make_plan(mode: str, num_q_heads: int = self._input_num_heads): caps_kwargs = dict( device=torch.device("cuda", torch.accelerator.current_device_index()), - num_q_heads=self._input_num_heads, + num_q_heads=num_q_heads, max_q_rows=self._max_tokens, max_width=self._topk_tokens, dtype=torch.bfloat16, @@ -701,6 +1551,14 @@ def make_plan(mode: str): extend_plan = make_plan("extend") self._decode_plan = decode_plan self._extend_plan = extend_plan + self._ckv_extend_plan = ( + make_plan("extend", self.num_heads) if self._ckv_gather_enabled else None + ) + self._ckv_local_capacity = ( + (self._ckv_capacity_tokens + kernel_page_size - 1) + // kernel_page_size + * kernel_page_size + ) self._kernel_page_size = kernel_page_size def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: @@ -747,6 +1605,191 @@ def do_kv_cache_update( slot_mapping.flatten(), ) + def uses_full_ckv_dcp( + self, + attn_metadata: B12xMLASparseMetadata, + num_tokens: int, + ) -> bool: + if torch.cuda.is_current_stream_capturing(): + return False + return ( + self._ckv_gather_enabled + and attn_metadata.dcp_ckv_gather_eligible + and attn_metadata.num_decode_tokens == 0 + and num_tokens == attn_metadata.num_actual_tokens + and 0 < attn_metadata.dcp_padded_total_tokens <= self._ckv_local_capacity + and attn_metadata.dcp_local_total_tokens + <= attn_metadata.dcp_padded_total_tokens + and all( + value is not None + for value in ( + attn_metadata.ckv_selected_indices, + attn_metadata.ckv_active_counts, + attn_metadata.dcp_rank_req_starts, + attn_metadata.dcp_rank_req_lens, + attn_metadata.dcp_local_cu_seq_lens, + attn_metadata.global_cache_seq_lens_per_req, + ) + ) + ) + + def _gather_full_ckv( + self, + kv_cache: torch.Tensor, + attn_metadata: B12xMLASparseMetadata, + local_buffer: torch.Tensor, + gathered_buffer: torch.Tensor, + stream: torch.cuda.Stream | None = None, + ) -> torch.Tensor: + if not self.uses_full_ckv_dcp(attn_metadata, attn_metadata.num_actual_tokens): + raise RuntimeError("full CKV gather called for an ineligible batch") + if not _is_glm_next_ckv_source_layout( + kv_cache, page_size=self._kernel_page_size + ): + raise ValueError( + "GLM5Next CKV gather requires native 528-byte records; " + f"got shape={tuple(kv_cache.shape)}, stride={kv_cache.stride()}" + ) + expected_local_shape = ( + self._ckv_local_capacity, + _GLM_NEXT_CACHE_RECORD_BYTES, + ) + expected_gathered_shape = ( + self.dcp_world_size * self._ckv_local_capacity, + _GLM_NEXT_CACHE_RECORD_BYTES, + ) + if tuple(local_buffer.shape) != expected_local_shape: + raise RuntimeError("CKV local workspace has an invalid shape") + if tuple(gathered_buffer.shape) != expected_gathered_shape: + raise RuntimeError("CKV gathered workspace has an invalid shape") + + assert attn_metadata.dcp_local_cu_seq_lens is not None + local_tokens = attn_metadata.dcp_local_total_tokens + padded_tokens = attn_metadata.dcp_padded_total_tokens + if stream is not None: + local_buffer.record_stream(stream) + gathered_buffer.record_stream(stream) + stream.wait_stream(torch.cuda.current_stream()) + stream_context = torch.cuda.stream(stream) + else: + stream_context = torch.cuda.stream(torch.cuda.current_stream()) + with stream_context: + if local_tokens: + ops.cp_gather_cache( + src_cache=kv_cache, + dst=local_buffer[:local_tokens], + block_table=attn_metadata.block_table, + cu_seq_lens=attn_metadata.dcp_local_cu_seq_lens, + batch_size=attn_metadata.num_reqs, + ) + if local_tokens < padded_tokens: + local_buffer[local_tokens:padded_tokens].zero_() + if stream is None: + dcp_group = get_dcp_group() + else: + from vllm.distributed.parallel_state import ( + get_dcp_ckv_prefetch_group, + ) + + dcp_group = get_dcp_ckv_prefetch_group() + _dcp_all_gather_current_stream( + dcp_group, + local_buffer[:padded_tokens].view(-1), + gathered_buffer[: self.dcp_world_size * padded_tokens].view(-1), + ) + return gathered_buffer.view( + -1, self._kernel_page_size, _GLM_NEXT_CACHE_RECORD_BYTES + ) + + def _ckv_workspace_views( + self, + workspace: torch.Tensor, + buf_idx: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not 0 <= buf_idx < self._ckv_workspace_slots: + raise ValueError(f"CKV workspace slot out of range: {buf_idx}") + local_nbytes = self._ckv_local_capacity * _GLM_NEXT_CACHE_RECORD_BYTES + gathered_nbytes = self.dcp_world_size * local_nbytes + local_buffer = workspace.narrow(0, 0, local_nbytes).view( + self._ckv_local_capacity, + _GLM_NEXT_CACHE_RECORD_BYTES, + ) + gathered_offset = local_nbytes + buf_idx * gathered_nbytes + gathered_buffer = workspace.narrow(0, gathered_offset, gathered_nbytes).view( + self.dcp_world_size * self._ckv_local_capacity, + _GLM_NEXT_CACHE_RECORD_BYTES, + ) + return local_buffer, gathered_buffer + + def set_ckv_current_chunk_kv( + self, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + ) -> None: + del k_pe + self._ckv_current_chunk_kv_c = kv_c_normed + + @staticmethod + def _resolve_layer_index(layer: AttentionLayer) -> int | None: + layer_idx = getattr(layer, "layer_idx", None) + if layer_idx is not None: + try: + return int(layer_idx) + except (TypeError, ValueError): + return None + layer_name = getattr(layer, "layer_name", None) + if not layer_name: + return None + from vllm.model_executor.models.utils import extract_layer_index + + try: + return extract_layer_index(layer_name) + except (ValueError, AssertionError, IndexError): + return None + + def _append_current_chunk_to_gathered( + self, + gathered_cache: torch.Tensor, + attn_metadata: B12xMLASparseMetadata, + num_tokens: int, + ) -> None: + if self._ckv_current_chunk_kv_c is None or num_tokens == 0: + return + assert self._concat_and_cache_glm_next_mla is not None + assert attn_metadata.global_cache_seq_lens_per_req is not None + assert attn_metadata.dcp_rank_req_starts is not None + req_ids = attn_metadata.req_id_per_token[:num_tokens].to(torch.int64) + global_seq_lens = attn_metadata.global_cache_seq_lens_per_req[ + : attn_metadata.num_reqs + ] + seq_len_per_token = global_seq_lens[req_ids].to(torch.int32) + query_start_loc = attn_metadata.query_start_loc[ + : attn_metadata.num_reqs + 1 + ].to(torch.int32) + chunk_start = query_start_loc[:-1][req_ids] + chunk_len = (query_start_loc[1:] - query_start_loc[:-1])[req_ids] + token_idx = torch.arange( + num_tokens, + device=gathered_cache.device, + dtype=torch.int32, + ) + global_pos = seq_len_per_token - chunk_len + (token_idx - chunk_start) + interleave = attn_metadata.cp_kv_cache_interleave_size + owner = ((global_pos // interleave) % self.dcp_world_size).to(torch.int64) + local_pos = ( + global_pos // (self.dcp_world_size * interleave) * interleave + + global_pos % interleave + ).to(torch.int64) + rank_req_starts = attn_metadata.dcp_rank_req_starts + flat_idx = owner * attn_metadata.num_reqs + req_ids + rank_start = rank_req_starts.reshape(-1)[flat_idx].to(torch.int64) + slots = owner * attn_metadata.dcp_padded_total_tokens + rank_start + local_pos + self._concat_and_cache_glm_next_mla( + self._ckv_current_chunk_kv_c[:num_tokens], + gathered_cache, + slots, + ) + def forward_mqa( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], @@ -754,7 +1797,6 @@ def forward_mqa( attn_metadata: B12xMLASparseMetadata, layer: AttentionLayer, ) -> tuple[torch.Tensor, torch.Tensor | None]: - del layer cache_page_size = int(kv_c_and_k_pe_cache.shape[1]) metadata_page_size = int(attn_metadata.block_size) if self._is_glm_next and ( @@ -766,18 +1808,62 @@ def forward_mqa( f"cache={cache_page_size}, metadata={metadata_page_size}, " f"plan={self._kernel_page_size}" ) - plan = ( - self._decode_plan if attn_metadata.max_query_len <= 1 else self._extend_plan + num_tokens = int(q[0].shape[0] if isinstance(q, tuple) else q.shape[0]) + use_decode = _use_b12x_sparse_decode_plan( + max_query_len=attn_metadata.max_query_len, + num_tokens=num_tokens, + num_reqs=attn_metadata.num_reqs, + is_spec_decode=attn_metadata.is_spec_decode, + spec_extend_as_decode=self._spec_extend_as_decode, + spec_extend_as_decode_force=self._spec_extend_as_decode_force, + spec_decode_max_q=self._spec_decode_max_q, + max_tokens=self._max_tokens, ) + use_ckv_gather = self.uses_full_ckv_dcp(attn_metadata, num_tokens) + layer_idx = self._resolve_layer_index(layer) if use_ckv_gather else None + prefetch_registry = attn_metadata.ckv_prefetch_registry + if use_ckv_gather and self._ckv_workspace_pool is None: + raise RuntimeError("CKV gather requires a persistent workspace pool") + if use_ckv_gather and prefetch_registry is None: + raise RuntimeError("CKV gather requires a prefetch state registry") + use_persistent_ckv = use_ckv_gather + if use_ckv_gather: + assert self._ckv_extend_plan is not None + plan = self._ckv_extend_plan + logger.info_once("Using full-CKV gather for GLM5Next B12X DCP prefill") + else: + plan = self._decode_plan if use_decode else self._extend_plan + if use_decode and attn_metadata.max_query_len > 1: + logger.info_once("Using B12X decode plan for GLM5Next MTP verification") + input_num_heads = self.num_heads if use_ckv_gather else self._input_num_heads q_spec = ( - (self._max_tokens, self._input_num_heads, self._q_head_dim), + (self._max_tokens, input_num_heads, self._q_head_dim), torch.bfloat16, ) + plan_specs = plan.shapes_and_dtypes() + ckv_specs = ( + ( + ( + (self._ckv_local_capacity, _GLM_NEXT_CACHE_RECORD_BYTES), + torch.uint8, + ), + ( + ( + self.dcp_world_size * self._ckv_local_capacity, + _GLM_NEXT_CACHE_RECORD_BYTES, + ), + torch.uint8, + ), + ) + if use_ckv_gather and not use_persistent_ckv + else () + ) workspaces = current_workspace_manager().get_simultaneous( - q_spec, *plan.shapes_and_dtypes() + q_spec, *plan_specs, *ckv_specs ) q_buffer = workspaces[0] - scratch = workspaces[1:] + scratch_end = 1 + len(plan_specs) + scratch = workspaces[1:scratch_end] if isinstance(q, tuple): q_nope, q_pe = q @@ -792,20 +1878,148 @@ def forward_mqa( q_all = q_buffer[:num_tokens] q_all.copy_(q) - if int(q_all.shape[1]) != self._input_num_heads: + if int(q_all.shape[1]) != input_num_heads: raise ValueError( "B12X sparse MLA query heads do not match the planned head " - f"count: {q_all.shape[1]} != {self._input_num_heads}." + f"count: {q_all.shape[1]} != {input_num_heads}." ) assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_tokens] - block_stride_rows = _selected_index_block_stride_rows( - kv_c_and_k_pe_cache, - block_size=attn_metadata.block_size, - is_glm_next=self._is_glm_next, - ) - if self.dcp_world_size > 1: + kv_cache_for_run = kv_c_and_k_pe_cache + prefetch_state: _CKVPrefetchState | None = None + ckv_workspace: torch.Tensor | None = None + if use_ckv_gather: + if use_persistent_ckv: + assert prefetch_registry is not None + assert self._ckv_workspace_pool is not None + prefetch_state = prefetch_registry.for_workspace( + q_buffer, + layer_idx, + kv_c_and_k_pe_cache, + self._ckv_workspace_pool, + ) + if layer_idx is not None: + prefetch_state.enter_layer(layer_idx) + prefetch_state.register_cache(layer_idx, kv_c_and_k_pe_cache) + ckv_workspace = prefetch_state.get_ckv_workspace( + self._ckv_workspace_nbytes + ) + 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 + gather_event.wait() + _, gathered_buffer = self._ckv_workspace_views( + ckv_workspace, current_buf_idx + ) + kv_cache_for_run = gathered_buffer.view( + -1, + self._kernel_page_size, + _GLM_NEXT_CACHE_RECORD_BYTES, + ) + self._append_current_chunk_to_gathered( + kv_cache_for_run, + attn_metadata, + num_tokens, + ) + else: + prefetch_state.wait_for_pending_writes() + current_buf_idx = ( + layer_idx % self._ckv_workspace_slots + if layer_idx is not None + else 0 + ) + local_buffer, gathered_buffer = self._ckv_workspace_views( + ckv_workspace, current_buf_idx + ) + kv_cache_for_run = self._gather_full_ckv( + kv_c_and_k_pe_cache, + attn_metadata, + local_buffer, + gathered_buffer, + ) + else: + local_buffer, gathered_buffer = workspaces[scratch_end:] + kv_cache_for_run = self._gather_full_ckv( + kv_c_and_k_pe_cache, + attn_metadata, + local_buffer, + gathered_buffer, + ) + + if ( + prefetch_state is not None + and ckv_workspace is not None + and layer_idx is not None + ): + targets = _ckv_prefetch_target_indices( + layer_idx, + self._ckv_prefetch_depth, + prefetch_state.layer_caches, + prefetch_state.pending_layers, + ) + prefetch_stream = ( + prefetch_state.get_gather_stream() if targets else None + ) + for target_idx in targets: + assert prefetch_stream is not None + target_cache = prefetch_state.layer_caches[target_idx] + assert target_cache is not None + target_buf_idx = target_idx % self._ckv_workspace_slots + target_local, target_gathered = self._ckv_workspace_views( + ckv_workspace, target_buf_idx + ) + self._gather_full_ckv( + target_cache, + attn_metadata, + target_local, + target_gathered, + 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, + ) + assert attn_metadata.ckv_selected_indices is not None + assert attn_metadata.ckv_active_counts is not None + assert attn_metadata.dcp_rank_req_starts is not None + assert attn_metadata.dcp_rank_req_lens is not None + selected_indices = attn_metadata.ckv_selected_indices[ + :num_tokens, : topk_indices.shape[1] + ] + active_counts = attn_metadata.ckv_active_counts[:num_tokens] + _map_global_topk_to_gathered_ckv( + attn_metadata.req_id_per_token[:num_tokens], + topk_indices, + attn_metadata.dcp_rank_req_starts, + attn_metadata.dcp_rank_req_lens, + selected_indices, + active_counts, + dcp_size=self.dcp_world_size, + cp_kv_cache_interleave_size=(attn_metadata.cp_kv_cache_interleave_size), + padded_rank_tokens=attn_metadata.dcp_padded_total_tokens, + ) + assert attn_metadata.global_cache_seq_lens_per_req is not None + cache_seq_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_tokens, + ).contiguous() + torch.minimum(active_counts, cache_seq_lens, out=active_counts) + _mask_page_table_after_nsa_len(selected_indices, active_counts) + elif self.dcp_world_size > 1: + block_stride_rows = _selected_index_block_stride_rows( + kv_c_and_k_pe_cache, + block_size=attn_metadata.block_size, + is_glm_next=self._is_glm_next, + ) selected_indices, active_counts = triton_filter_and_convert_dcp_index( attn_metadata.req_id_per_token[:num_tokens], attn_metadata.block_table, @@ -819,6 +2033,11 @@ def forward_mqa( return_valid_counts=True, ) else: + block_stride_rows = _selected_index_block_stride_rows( + kv_c_and_k_pe_cache, + block_size=attn_metadata.block_size, + is_glm_next=self._is_glm_next, + ) selected_indices, active_counts = triton_convert_req_index_to_global_index( attn_metadata.req_id_per_token[:num_tokens], attn_metadata.block_table, @@ -829,9 +2048,10 @@ def forward_mqa( return_valid_counts=True, ) - cache_seq_lens = attn_metadata.cache_seq_lens_per_token - assert cache_seq_lens is not None - cache_seq_lens = cache_seq_lens[:num_tokens].contiguous() + if not use_ckv_gather: + cache_seq_lens = attn_metadata.cache_seq_lens_per_token + assert cache_seq_lens is not None + cache_seq_lens = cache_seq_lens[:num_tokens].contiguous() binding = plan.bind( scratch=scratch, q=q_all, @@ -842,7 +2062,7 @@ def forward_mqa( run = self._run_decode if plan is self._decode_plan else self._run_extend run_kwargs = dict( binding=binding, - kv_cache=kv_c_and_k_pe_cache, + kv_cache=kv_cache_for_run, sm_scale=self.scale, v_head_dim=self.kv_lora_rank, return_lse=self.need_to_return_lse_for_decode, diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 767ea0033483..b18ed0168752 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -598,6 +598,16 @@ def __init__( # can be a multiple of hash_block_size. self.hash_block_size = hash_block_size self.dcp_world_size = dcp_world_size + self.pcp_world_size = pcp_world_size + self.has_dcp_replicated_group = any( + getattr(group.kv_cache_spec, "dcp_replicated", False) + for group in kv_cache_config.kv_cache_groups + ) + self.has_replicated_sliding_group = any( + isinstance(group.kv_cache_spec, SlidingWindowSpec) + and group.kv_cache_spec.dcp_replicated + for group in kv_cache_config.kv_cache_groups + ) group_block_sizes = [ manager.block_size for manager in self.single_type_managers ] @@ -610,14 +620,16 @@ def __init__( ) assert pcp_world_size == 1, "PCP not support hybrid attn now." if dcp_world_size > 1: - # DCP shards full-attention KV across ranks and replicates Mamba - # state; other spec types (e.g. sliding window) have no DCP-aware - # handling yet, so reject them explicitly. + # Target attention remains DCP-sharded. Small speculative draft + # groups may instead replicate their cache and execute locally. for g in kv_cache_config.kv_cache_groups: - assert isinstance(g.kv_cache_spec, (FullAttentionSpec, MambaSpec)), ( + spec = g.kv_cache_spec + assert isinstance(spec, (FullAttentionSpec, MambaSpec)) or getattr( + spec, "dcp_replicated", False + ), ( "DCP with hybrid KV cache layouts only supports " - "full-attention and Mamba groups, got: " - f"{type(g.kv_cache_spec).__name__}." + "full-attention, Mamba, and replicated draft groups, got: " + f"{type(spec).__name__}." ) # Fine-grained hash hits require Mamba "align" and compatible cache # managers in every group. TP needs hashing finer than the Mamba block; @@ -649,6 +661,9 @@ def __init__( "cache managers require block-aligned lookups: %s.", ", ".join(sorted(unsupported_partial_hit_managers)), ) + prefix_cache_alignment_tokens = self._cache_hit_alignment_tokens + for manager in self.single_type_managers: + manager.prefix_cache_alignment_tokens = prefix_cache_alignment_tokens self.verify_and_split_kv_cache_groups() @property @@ -712,6 +727,18 @@ def verify_and_split_kv_cache_groups(self) -> None: for gid in group.group_ids: self.single_type_managers[gid].use_eagle = True + def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: + if ( + self.dcp_world_size > 1 + and self.pcp_world_size == 1 + and self.has_dcp_replicated_group + ): + # Avoid enabling cascade attention for only the sharded target side + # of a target+replicated-draft hybrid. Concrete prefix replay still + # happens through find_longest_cache_hit(). + return [0] * len(self.kv_cache_config.kv_cache_groups) + return super().get_num_common_prefix_blocks(running_request_id) + def _align_cacheable(self, num_tokens: int) -> int: """Largest prefix of ``num_tokens`` a future cache hit could match. @@ -748,10 +775,22 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: # (``scheduler_block_size``); retention is passed separately so it # can keep both the coarse segment tails and the fine replay # boundary (which needs the fine value). + retention_interval = self.retention_interval + if ( + retention_interval == 0 + and self.has_replicated_sliding_group + and isinstance(manager.kv_cache_spec, (MambaSpec, SlidingWindowSpec)) + ): + # A DFlash sliding draft drops its own lookahead block, so the + # common reusable boundary can be an earlier scheduler-page + # boundary rather than the prompt's final hash boundary. Keep + # the small recurrent/draft caches dense, matching the prefix + # behavior before sparse retention became the default. + retention_interval = None manager.cache_blocks( request, num_tokens_to_cache, - retention_interval=self.retention_interval, + retention_interval=retention_interval, ) def find_longest_cache_hit( @@ -912,6 +951,9 @@ def find_longest_cache_hit_per_group( kv_cache_spec=spec, drop_eagle_block=use_eagle, alignment_tokens=self._cache_hit_alignment_tokens, + dcp_world_size=( + self.dcp_world_size if isinstance(spec, FullAttentionSpec) else 1 + ), ) for gid, blks in zip(group_ids, blocks): hit_blocks[gid] = blks diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 337a55aa6cd0..f457573b23a1 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -670,12 +670,16 @@ def resolve_kv_cache_block_sizes( groups = kv_cache_config.kv_cache_groups if len(groups) <= 1: - bs = cache_config.block_size * dcp + dcp_replicated = len(groups) == 1 and getattr( + groups[0].kv_cache_spec, "dcp_replicated", False + ) + bs = cache_config.block_size * (1 if dcp_replicated else dcp) return bs, bs group_block_sizes = [ g.kv_cache_spec.block_size * dcp if isinstance(g.kv_cache_spec, AttentionSpec) + and not getattr(g.kv_cache_spec, "dcp_replicated", False) else g.kv_cache_spec.block_size for g in groups ] @@ -1595,6 +1599,44 @@ def group_and_unify_kv_cache_specs( return [mla_uniform_spec, *swa_uniform_specs] +def group_dcp_replicated_draft_kv_cache_specs( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Keep a replicated speculative draft separate from a sharded target. + + DFlash's small sliding-window cache has different allocation and DCP + semantics from the target cache. When both sides are independently + uniform, retain their concrete specs and native block sizes instead of + promoting or page-size-unifying the draft with the target. + """ + replicated = { + name: spec + for name, spec in kv_cache_spec.items() + if getattr(spec, "dcp_replicated", False) + } + if not replicated: + return None + sharded = { + name: spec + for name, spec in kv_cache_spec.items() + if not getattr(spec, "dcp_replicated", False) + } + if not sharded: + return None + if not is_kv_cache_spec_uniform(replicated): + return None + # The target need not itself be uniform. GLM-5.3, for example, mixes MLA + # cache layouts that the normal hybrid grouping path already understands. + # Re-enter grouping without the replicated draft so that path can preserve + # the target's native groups instead of page-unifying it with the draft. + sharded_groups = get_kv_cache_groups(vllm_config, dict(sharded)) + return [ + *sharded_groups, + *_get_kv_cache_groups_uniform_spec(replicated), + ] + + def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: """Pick a chunk size that minimizes total upward padding. @@ -1779,6 +1821,10 @@ def get_kv_cache_groups( # full attention, or all layers are sliding window attention with the # same window size). Put all layers into one group. return _get_kv_cache_groups_uniform_type(uniform_spec) + elif replicated_groups := group_dcp_replicated_draft_kv_cache_specs( + vllm_config, kv_cache_spec + ): + return replicated_groups elif grouped_specs := group_and_unify_kv_cache_specs(kv_cache_spec): # DeepseekV4 case: All layers need the same number of token slots, # yet some layers are full attention while others are sliding window @@ -2111,9 +2157,11 @@ def get_kv_cache_configs( # This is to prevent that some layers are initialized with unregistered specs. KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) - # When speculating with more than 1 speculative module (e.g. multi-layered MTP) + # When speculating with more than 1 speculative module (e.g. multi-layered MTP), # tag every SlidingWindowSpec with how many extra tokens to retain in the window. - extra_retained_tokens = ( + # A model-specific cache spec may require a larger retention window (DFlash's + # EAGLE prefix proof is one example), so never erase that value here. + mtp_extra_retained_tokens = ( vllm_config.speculative_config.num_speculative_tokens - 1 if vllm_config.speculative_config is not None and vllm_config.speculative_config.use_multi_module_mtp() @@ -2122,7 +2170,10 @@ def get_kv_cache_configs( for layer_name, layer_spec in merged_kv_cache_specs.items(): if isinstance(layer_spec, SlidingWindowSpec): merged_kv_cache_specs[layer_name] = replace( - layer_spec, extra_retained_tokens=extra_retained_tokens + layer_spec, + extra_retained_tokens=max( + layer_spec.extra_retained_tokens, mtp_extra_retained_tokens + ), ) # Get global KV cache groups. This also handles spec unification for diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index c373d5cf5c87..65e3cc0ca0b5 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -53,7 +53,7 @@ ) from vllm.v1.core.sched.utils import check_stop, remove_all from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec, SlidingWindowSpec from vllm.v1.metrics.perf import ModelMetrics, PerfStats from vllm.v1.metrics.stats import ( PrefixCacheStats, @@ -325,6 +325,11 @@ def __init__( self.need_mamba_block_aligned_split = ( self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" ) + self.mamba_has_sliding_eagle_group = self.use_eagle and any( + isinstance(group.kv_cache_spec, SlidingWindowSpec) + and group.kv_cache_spec.dcp_replicated + for group in kv_cache_config.kv_cache_groups + ) self.mamba_has_prefill_checkpoint_blocks = ( self.has_mamba_layers # TODO: support spec decoding @@ -413,7 +418,7 @@ def _mamba_block_aligned_split( # Eagle, FullAttn prunes the last matching block, so back off one # block to avoid a Mamba cache miss. last_cache_position = request.num_tokens - request.num_tokens % block_size - if self.use_eagle: + if self.use_eagle and not getattr(self, "mamba_has_sliding_eagle_group", False): last_cache_position = max(last_cache_position - block_size, 0) end = start + num_new_tokens diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index ceb56a59875f..d103c2baa558 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -71,11 +71,14 @@ def __init__( block until the request finishes. """ self.scheduler_block_size = scheduler_block_size + # The coordinator may lower this to the shared hash granularity when + # every group supports fine-grained prefix lookup. + self.prefix_cache_alignment_tokens = scheduler_block_size # The block size for this manager; used for actual block allocation. self.block_size = kv_cache_spec.block_size self.dcp_world_size = dcp_world_size self.pcp_world_size = pcp_world_size - if dcp_world_size > 1: + if dcp_world_size > 1 and not getattr(kv_cache_spec, "dcp_replicated", False): self.block_size *= dcp_world_size self.kv_cache_spec = kv_cache_spec self.block_pool = block_pool @@ -455,7 +458,7 @@ def cache_blocks( block_mask = self.reachable_block_mask( start_block=num_cached_blocks, end_block=num_full_blocks, - alignment_tokens=self.scheduler_block_size, + alignment_tokens=self.prefix_cache_alignment_tokens, kv_cache_spec=self.kv_cache_spec, use_eagle=self.use_eagle, retention_interval=retention_interval, @@ -700,7 +703,7 @@ def find_longest_cache_hit( "and chunked local attention groups" ) block_size = kv_cache_spec.block_size - if dcp_world_size > 1: + if dcp_world_size > 1 and not getattr(kv_cache_spec, "dcp_replicated", False): # DCP shards each block's KV across ranks; hashes must be viewed at # the sharded block size. block_size *= dcp_world_size @@ -915,8 +918,12 @@ def find_longest_cache_hit( assert isinstance(kv_cache_spec, SlidingWindowSpec), ( "SlidingWindowManager can only be used for sliding window groups" ) - assert dcp_world_size == 1, "DCP not support sliding window attn now." - assert pcp_world_size == 1, "PCP not support sliding window attn now." + assert dcp_world_size == 1 or kv_cache_spec.dcp_replicated, ( + "DCP only supports sliding-window KV when it is replicated." + ) + assert pcp_world_size == 1 or kv_cache_spec.dcp_replicated, ( + "PCP only supports sliding-window KV when it is replicated." + ) # Fine-grained partial hits are not supported for sliding window now assert alignment_tokens % kv_cache_spec.block_size == 0, ( "SlidingWindowManager does not support fine-grained (partial) cache hits" diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 3cecb4fbee8a..8f9483960136 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -228,8 +228,11 @@ def is_uniform_with_collection( f"Unsupported KV cache spec type: {type(self)}. " "Please register it using @register_kv_cache_spec decorator." ) + dcp_replicated = getattr(self, "dcp_replicated", False) return all( - isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values() + isinstance(spec, uniform_type_base_spec) + and getattr(spec, "dcp_replicated", False) == dcp_replicated + for spec in kv_cache_specs.values() ) @@ -428,7 +431,11 @@ def real_page_size_bytes(self) -> int: def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: parallel_config = vllm_config.parallel_config - kv_shard_count = parallel_config.decode_context_parallel_size + kv_shard_count = ( + 1 + if getattr(self, "dcp_replicated", False) + else parallel_config.decode_context_parallel_size + ) return cdiv(max_len, self.block_size * kv_shard_count) @@ -458,10 +465,12 @@ class FullAttentionSpec(AttentionSpec): cache layout itself. """ + dcp_replicated: bool = False + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_model_len = vllm_config.model_config.max_model_len dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size - if dcp_world_size > 1: + if dcp_world_size > 1 and not self.dcp_replicated: max_model_len = cdiv(max_model_len, dcp_world_size) return cdiv(max_model_len, self.block_size) * self.page_size_bytes @@ -498,6 +507,11 @@ def merge(cls, specs: list[Self]) -> Self: assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), ( "MLAAttentionSpec should be merged in MLAAttentionSpec.merge" ) + dcp_replicated = {spec.dcp_replicated for spec in specs} + assert len(dcp_replicated) == 1, ( + "All attention layers in one KV cache group must use the same " + "DCP replication mode." + ) merged_spec = cls( block_size=specs[0].block_size, num_kv_heads=specs[0].num_kv_heads, @@ -514,6 +528,7 @@ def merge(cls, specs: list[Self]) -> Self: # If any layer in the group is non-causal, treat the group as # non-causal so the engine core disables incompatible scheduling. non_causal=any(spec.non_causal for spec in specs), + dcp_replicated=dcp_replicated.pop(), ) for spec in specs: for f in fields(AttentionSpec): @@ -703,6 +718,7 @@ def is_uniform_with_collection( @dataclass(frozen=True, kw_only=True) class SlidingWindowSpec(AttentionSpec): sliding_window: int + dcp_replicated: bool = False # The trailing edge of the window is extended by ``extra_retained_tokens`` # so that those extra trailing tokens' blocks are retained (but not # attended). This is needed for multi-module spec decoding which can @@ -739,9 +755,10 @@ def max_admission_blocks_per_request( return cdiv(num_tokens, self.block_size) + 1 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: - assert vllm_config.parallel_config.decode_context_parallel_size == 1, ( - "DCP not support sliding window." - ) + assert ( + vllm_config.parallel_config.decode_context_parallel_size == 1 + or self.dcp_replicated + ), "DCP only supports sliding-window KV when it is replicated." max_blocks = self.max_admission_blocks_per_request( max_in_flight_tokens=vllm_config.max_in_flight_tokens, max_model_len=vllm_config.model_config.max_model_len, @@ -754,6 +771,7 @@ def is_uniform_with_collection( return all( isinstance(spec, SlidingWindowSpec) and spec.sliding_window == self.sliding_window + and spec.dcp_replicated == self.dcp_replicated for spec in kv_cache_specs.values() ) @@ -980,6 +998,13 @@ class UniformTypeKVCacheSpecs(KVCacheSpec): def page_size_bytes(self) -> int: return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values()) + @property + def dcp_replicated(self) -> bool: + return all( + getattr(spec, "dcp_replicated", False) + for spec in self.kv_cache_specs.values() + ) + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_num_pages = max( cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes) diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index 92d8383c1f12..1f89bf0487bb 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -37,7 +37,27 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: layer_impl = getattr(layer, "impl", None) if layer_impl is None: continue - if vllm_config.speculative_config is not None and interleave_size > 1: + get_spec = getattr(layer, "get_kv_cache_spec", None) + if get_spec is not None: + try: + spec = get_spec(vllm_config) + except Exception: + spec = None + if getattr(spec, "dcp_replicated", False): + # Replicated draft KV contains the complete sequence on + # every rank, so its attention executes as a local DCP1 op. + layer_impl.dcp_world_size = 1 + layer_impl.dcp_rank = 0 + layer_impl.total_cp_world_size = 1 + layer_impl.total_cp_rank = 0 + layer_impl.need_to_return_lse_for_decode = False + continue + speculative_config = vllm_config.speculative_config + if ( + speculative_config is not None + and speculative_config.method == "mtp" + and interleave_size > 1 + ): assert layer_impl.supports_mtp_with_cp_non_trivial_interleave_size, ( "MTP with cp_kv_cache_interleave_size > 1 is not " f"supported in {layer_impl.__class__.__name__}." diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index dae1630a2c19..18258ab2ae98 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -26,6 +26,7 @@ def __init__( cp_size: int = 1, cp_rank: int = 0, cp_interleave: int = 1, + group_cp_sizes: list[int] | None = None, ): self.block_sizes = block_sizes self.kernel_block_sizes = kernel_block_sizes @@ -36,6 +37,12 @@ def __init__( self.cp_size = cp_size self.cp_rank = cp_rank self.cp_interleave = cp_interleave + if group_cp_sizes is None: + group_cp_sizes = [cp_size] * len(block_sizes) + assert len(group_cp_sizes) == len(block_sizes) + self.group_cp_sizes = torch.tensor( + group_cp_sizes, dtype=torch.int32, device=device + ) self.num_kv_cache_groups = len(self.block_sizes) assert len(max_num_blocks_per_group) == self.num_kv_cache_groups @@ -205,8 +212,11 @@ def compute_slot_mappings( positions, self.block_table_ptrs, self.block_table_strides, + self.num_blocks.gpu, + self.num_blocks.gpu.stride(0), self.block_sizes_tensor, self.kernel_block_sizes_tensor, + self.group_cp_sizes, slot_mappings, slot_mappings.stride(0), self.cp_rank, @@ -268,6 +278,10 @@ def _gather_block_tables_kernel( block_ids = tl.load(src_row_ptr + offset, mask=offset < num_blocks) tl.store(dst_row_ptr + offset, block_ids, mask=offset < num_blocks) + for i in tl.range(num_blocks, max_num_blocks, BLOCK_SIZE): + offset = i + tl.arange(0, BLOCK_SIZE) + tl.store(dst_row_ptr + offset, 0, mask=offset < max_num_blocks) + @triton.jit def _compute_slot_mappings_kernel( @@ -277,8 +291,11 @@ def _compute_slot_mappings_kernel( pos, # [num_tokens] block_table_ptrs, # [num_kv_cache_groups] block_table_strides, # [num_kv_cache_groups] + num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs] + num_blocks_stride, block_sizes, # [num_kv_cache_groups] kernel_block_sizes, # [num_kv_cache_groups] + group_cp_sizes, # [num_kv_cache_groups] slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] slot_mappings_stride, cp_rank, @@ -305,40 +322,52 @@ def _compute_slot_mappings_kernel( block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32) block_table_stride = tl.load(block_table_strides + group_id) + group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride kv_block_size = tl.load(block_sizes + group_id) kernel_block_size = tl.load(kernel_block_sizes + group_id) + group_cp_size = tl.load(group_cp_sizes + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) + num_blocks = tl.load(group_num_blocks_ptr + req_state_idx) start_idx = tl.load(query_start_loc + batch_idx) end_idx = tl.load(query_start_loc + batch_idx + 1) for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE): offset = i + tl.arange(0, TRITON_BLOCK_SIZE) - positions = tl.load(pos + offset, mask=offset < end_idx, other=0) + token_mask = offset < end_idx + positions = tl.load(pos + offset, mask=token_mask, other=0) if CP_SIZE == 1: # Common case: Context parallelism is not used. local_positions = positions - is_local = True + is_local = token_mask else: # Context parallelism is used. virtual_block_size = kv_block_size * CP_SIZE virtual_block_indices = positions // virtual_block_size virtual_block_offsets = positions % virtual_block_size - is_local = virtual_block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank + sharded_is_local = ( + virtual_block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank + ) rounds = virtual_block_offsets // (CP_INTERLEAVE * CP_SIZE) remainder = virtual_block_offsets % CP_INTERLEAVE local_offsets = rounds * CP_INTERLEAVE + remainder - local_positions = virtual_block_indices * kv_block_size + local_offsets + sharded_positions = virtual_block_indices * kv_block_size + local_offsets + # Replicated draft groups store every token locally. Express the + # runtime group choice with vector selects so Triton sees the same + # types on both paths during kernel compilation. + is_replicated = group_cp_size == 1 + is_local = token_mask & (is_replicated | sharded_is_local) + local_positions = tl.where(is_replicated, positions, sharded_positions) block_indices = local_positions // kernel_block_size block_offsets = local_positions % kernel_block_size + valid_block = token_mask & (block_indices < num_blocks) block_numbers = tl.load( block_table_ptr + req_state_idx * block_table_stride + block_indices, - mask=is_local, + mask=is_local & valid_block, other=0, ) slot_ids = block_numbers * kernel_block_size + block_offsets - if CP_SIZE != 1: - slot_ids = tl.where(is_local, slot_ids, PAD_ID) + slot_ids = tl.where(is_local & valid_block, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 32fbd63fe817..539c76f6843e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -541,9 +541,13 @@ def initialize_kv_cache( block_sizes = [] max_num_blocks_per_group = [] + group_cp_sizes = [] for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec block_sizes.append(spec.block_size) + group_cp_sizes.append( + 1 if getattr(spec, "dcp_replicated", False) else self.dcp_size + ) # Let each cache type account for CP. Attention KV is DCP-sharded, # while Mamba/GDN recurrent state is replicated across DCP ranks. max_num_blocks = spec.max_num_blocks_per_req( @@ -601,6 +605,7 @@ def initialize_kv_cache( cp_size=self.dcp_size, cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, + group_cp_sizes=group_cp_sizes, ) self.pcp_manager = pcp.maybe_build_pcp_manager( self.vllm_config, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 01555920c214..7ffb4d408778 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -419,9 +419,11 @@ def propose( seeds, self.block_tables.input_block_tables[gid], self.block_tables.kernel_block_sizes[gid], - self.block_tables.cp_rank, - self.block_tables.cp_size, - self.block_tables.cp_interleave, + # Every DFlash draft cache group is replicated under DCP, so + # draft context/query slots are ordinary local DCP1 slots. + 0, + 1, + 1, self.parallel_drafting_token_id, self.num_query_per_req, self.num_speculative_steps, diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 7ed3e5bcd97a..ba822a53539e 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -608,6 +608,7 @@ def bind_kv_cache( def unbind_kv_cache(forward_context: dict[str, Any]) -> None: """Release cache references retained by attention-like layers.""" + reset_impl_types: set[type] = set() for layer in forward_context.values(): unbind = getattr(layer, "unbind_kv_cache", None) if callable(unbind): @@ -620,6 +621,11 @@ def unbind_kv_cache(forward_context: dict[str, Any]) -> None: impl = getattr(layer, "impl", None) if impl is not None: + impl_type = type(impl) + reset_binding_state = getattr(impl, "reset_kv_cache_binding_state", None) + if callable(reset_binding_state) and impl_type not in reset_impl_types: + reset_binding_state() + reset_impl_types.add(impl_type) if hasattr(impl, "_k_scale_cache"): impl._k_scale_cache = None if hasattr(impl, "_v_scale_cache"):