diff --git a/tests/v1/attention/test_flashinfer_mla_dcp.py b/tests/v1/attention/test_flashinfer_mla_dcp.py index 4fe545407f..039d30c9aa 100644 --- a/tests/v1/attention/test_flashinfer_mla_dcp.py +++ b/tests/v1/attention/test_flashinfer_mla_dcp.py @@ -84,7 +84,8 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch): @requires_flashinfer_mla -def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch): +@pytest.mark.parametrize("causal", [True, False], ids=["causal", "noncausal"]) +def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal): import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla impl = MagicMock() @@ -106,7 +107,7 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch): seq_lens = torch.tensor([4, 5], dtype=torch.int32) global_causal_seq_lens = torch.tensor([31, 35], dtype=torch.int32) attn_metadata = MagicMock() - attn_metadata.causal = True + attn_metadata.causal = causal attn_metadata.num_decode_tokens = num_tokens attn_metadata.num_decodes = num_reqs attn_metadata.max_seq_len = 9 @@ -116,10 +117,23 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch): query = torch.ones(num_tokens, 24, 576, dtype=torch.bfloat16) kv_cache = torch.ones(2, 128, 576, dtype=torch.bfloat16) + kernel_batch = num_reqs if causal else num_tokens + kernel_query_len = query_len if causal else 1 kernel = MagicMock( return_value=( - torch.ones(num_reqs, query_len, 24, 512, dtype=torch.bfloat16), - torch.ones(num_reqs, query_len, 24, dtype=torch.float32), + torch.ones( + kernel_batch, + kernel_query_len, + 24, + 512, + dtype=torch.bfloat16, + ), + torch.ones( + kernel_batch, + kernel_query_len, + 24, + dtype=torch.float32, + ), ) ) monkeypatch.setattr(flashinfer_mla, "_get_workspace_buffer", MagicMock()) @@ -130,14 +144,25 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch): ) call = kernel.call_args.kwargs - assert call["query"].shape == (num_reqs, query_len, 24, 576) - assert call["block_tables"] is block_table - assert call["seq_lens"] is seq_lens + assert call["query"].shape == (kernel_batch, kernel_query_len, 24, 576) + expected_block_table = ( + block_table if causal else block_table.repeat_interleave(query_len, dim=0) + ) + expected_seq_lens = seq_lens if causal else seq_lens.repeat_interleave(query_len) + expected_global_seq_lens = ( + global_causal_seq_lens + if causal + else global_causal_seq_lens.repeat_interleave(query_len) + ) + torch.testing.assert_close(call["block_tables"], expected_block_table) + torch.testing.assert_close(call["seq_lens"], expected_seq_lens) assert call["backend"] == "cute-dsl" assert call["enable_dcp"] is True assert call["cp_world"] == 8 assert call["cp_rank"] == 3 - assert call["causal_seqlens_kv_global"] is global_causal_seq_lens + torch.testing.assert_close( + call["causal_seqlens_kv_global"], expected_global_seq_lens + ) assert "multi_ctas_kv_counter_buffer" not in call assert flashinfer_mla.FlashInferMLAImpl.lse_base_on_e assert output.shape == (num_tokens, 24, 512) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 455060413f..c096593b4f 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -1021,7 +1021,7 @@ def test_tokenspeed_mla_noncausal_capability(): assert tokenspeed_mla_module.TokenspeedMLABackend.supports_non_causal() -def test_flashinfer_mla_dspark_dcp_supports_target_only(monkeypatch): +def test_flashinfer_mla_dspark_dcp_supports_target_and_draft(monkeypatch): flashinfer_mla_module = pytest.importorskip( "vllm.v1.attention.backends.mla.flashinfer_mla" ) @@ -1053,7 +1053,7 @@ def test_flashinfer_mla_dspark_dcp_supports_target_only(monkeypatch): assert reason is None assert backend.supports_non_causal() assert builder.supports_non_causal_multi_token_decode - assert not backend.supports_non_causal_dcp() + assert backend.supports_non_causal_dcp() @pytest.mark.parametrize( diff --git a/tests/v1/attention/test_mla_noncausal.py b/tests/v1/attention/test_mla_noncausal.py index 7aca8400f9..040a2ad4be 100644 --- a/tests/v1/attention/test_mla_noncausal.py +++ b/tests/v1/attention/test_mla_noncausal.py @@ -109,7 +109,7 @@ def test_noncausal_decode_metadata_keeps_live_request_buffers(): ) -def test_mla_cache_marker_is_promoted_to_group_capability(): +def test_mla_cache_marker_is_a_group_property(): kwargs = { "block_size": 64, "num_kv_heads": 1, @@ -120,7 +120,8 @@ def test_mla_cache_marker_is_promoted_to_group_capability(): unmarked = MLAAttentionSpec(**kwargs) assert MLAAttentionSpec.merge([marked, marked]).non_causal_multi_token_decode - assert MLAAttentionSpec.merge([marked, unmarked]).non_causal_multi_token_decode assert not MLAAttentionSpec.merge( [unmarked, unmarked] ).non_causal_multi_token_decode + with pytest.raises(AssertionError, match="non_causal_multi_token_decode"): + MLAAttentionSpec.merge([marked, unmarked]) diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py index 4ef99badaa..f66d466c53 100644 --- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py +++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py @@ -37,10 +37,14 @@ def _make_fake_speculator( fake_input_buffers = SimpleNamespace( query_start_loc=torch.zeros(max_num_reqs + 1, dtype=torch.int32), seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32), + dcp_local_seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32), ) fake_block_tables = SimpleNamespace( input_block_tables=[torch.zeros(max_num_reqs, 4, dtype=torch.int32)], slot_mappings=torch.zeros(1, max_num_tokens, dtype=torch.int64), + cp_size=1, + cp_rank=0, + cp_interleave=1, ) return SimpleNamespace( arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"), @@ -126,3 +130,32 @@ def test_build_draft_attn_metadata_clamps_to_max_model_len(): bound = captured["seq_lens_cpu_upper_bound"] # 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected. assert torch.equal(bound, torch.tensor([1024, 503], dtype=torch.int32)) + + +def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens(): + fake = _make_fake_speculator() + fake.block_tables.cp_size = 2 + fake.block_tables.cp_rank = 1 + fake.block_tables.cp_interleave = 4 + fake.input_buffers.seq_lens[:3] = torch.tensor([5, 9, 16]) + + def fake_prepare(out, seq_lens, num_reqs, dcp_size, dcp_rank, cp_interleave): + assert seq_lens is fake.input_buffers.seq_lens + assert (num_reqs, dcp_size, dcp_rank, cp_interleave) == (3, 2, 1, 4) + out[:num_reqs].copy_(torch.tensor([1, 4, 8], dtype=torch.int32)) + out[num_reqs:].zero_() + + with patch.object(base_speculator, "prepare_dcp_local_seq_lens", fake_prepare): + captured = _run_build( + fake, + num_reqs=3, + num_reqs_padded=4, + num_tokens_padded=4, + base=torch.tensor([5, 9, 16]), + step=0, + ) + + local = captured["dcp_local_seq_lens"] + assert isinstance(local, torch.Tensor) + assert local.data_ptr() == fake.input_buffers.dcp_local_seq_lens.data_ptr() + assert local.tolist() == [1, 4, 8, 0] diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index fbe5f6844f..7932565197 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py @@ -120,6 +120,7 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa. supports_non_causal_multi_token_decode: ClassVar[bool] = True + supports_non_causal_multi_token_dcp: ClassVar[bool] = True def __init__( self, @@ -288,8 +289,8 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): seq_lens = attn_metadata.decode.seq_lens if not attn_metadata.causal: - # FlashInfer decode has no causal flag. For TP-only DSpark, flatten - # each non-causal query block into independent single-token rows. + # FlashInfer decode has no causal flag. Flatten each non-causal + # query block into independent single-token rows. query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes q = q.unsqueeze(1) if query_len > 1: @@ -326,6 +327,10 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): assert self.cp_kv_cache_interleave_size == 1 causal_seqlens_kv_global = attn_metadata.decode.dcp_tot_seq_lens assert causal_seqlens_kv_global is not None + if not attn_metadata.causal and query_len > 1: + causal_seqlens_kv_global = causal_seqlens_kv_global.repeat_interleave( + query_len + ) extra_kwargs.update( enable_dcp=True, cp_world=self.dcp_world_size, diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index d5f7cff7c1..79c0055edc 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -570,6 +570,11 @@ class MLAAttentionSpec(FullAttentionSpec): "All attention layers in the same KV cache group must use the same " "quantization method, tokens per state, and model version." ) + non_causal_mtd_set = {spec.non_causal_multi_token_decode for spec in specs} + assert len(non_causal_mtd_set) == 1, ( + "All attention layers in the same KV cache group must agree on " + "non_causal_multi_token_decode." + ) merged_spec = cls( block_size=specs[0].block_size, num_kv_heads=specs[0].num_kv_heads, @@ -582,9 +587,7 @@ class MLAAttentionSpec(FullAttentionSpec): cache_dtype_str=cache_dtype_str_set.pop(), tokens_per_state=tokens_per_state_set.pop(), model_version=model_version_set.pop(), - non_causal_multi_token_decode=any( - spec.non_causal_multi_token_decode for spec in specs - ), + non_causal_multi_token_decode=non_causal_mtd_set.pop(), ) for spec in specs: for f in fields(AttentionSpec): diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 2ff7b22d06..fbd2432c93 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -18,7 +18,7 @@ from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.cp_utils import cp_local_slot, prepare_dcp_local_seq_lens +from vllm.v1.worker.gpu.cp_utils import cp_local_slot from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState @@ -291,16 +291,6 @@ class DFlashSpeculator(DraftModelSpeculator): if not self.draft_attn_layer_names: return None assert num_query_per_req is None # Omitted for DFlash, read from self instead - if dcp_local_seq_lens is None and self.block_tables.cp_size > 1: - prepare_dcp_local_seq_lens( - self.input_buffers.dcp_local_seq_lens, - self.input_buffers.seq_lens, - num_reqs, - self.block_tables.cp_size, - self.block_tables.cp_rank, - self.block_tables.cp_interleave, - ) - dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens return super()._build_draft_attn_metadata( num_reqs, num_reqs_padded, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 382132e61e..57af43bc0b 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -21,6 +21,7 @@ from vllm.v1.worker.gpu.attn_utils import ( init_attn_backend, ) from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.dp_utils import DPSyncState from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState @@ -292,6 +293,18 @@ class DraftModelSpeculator(BaseSpeculator): out=draft_seq_lens_cpu_upper_bound[:num_reqs], ) draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len) + if dcp_local_seq_lens is None and self.block_tables.cp_size > 1: + # Draft steps advance and rewind their own global sequence lengths, + # so the target model's DCP-local lengths may already be stale. + prepare_dcp_local_seq_lens( + self.input_buffers.dcp_local_seq_lens, + self.input_buffers.seq_lens, + num_reqs, + self.block_tables.cp_size, + self.block_tables.cp_rank, + self.block_tables.cp_interleave, + ) + dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens attn_metadata = build_attn_metadata( attn_groups=self.attn_groups, num_reqs=num_reqs_padded,