[Feature] Support DCP with FP8 KV cache in MLA decode path (#44044)
Signed-off-by: shivampr <shivampr.dev@gmail.com> Signed-off-by: Shivam <shivamprasad91@gmail.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
shivampr committed
Jun 24, 2026 at 12:28 UTC
49f2104c53c910033aae05524af120e7c39c5c48
4 files changed
+359
-15
tests/kernels/attention/test_cache.py
+72
@@ -1015,6 +1015,78 @@ def test_gather_and_maybe_dequant_cache_mla(
1015
torch.testing.assert_close(dst, expected)
1016
1017
1018
+@pytest.mark.parametrize("kv_lora_rank", [512])
1019
+@pytest.mark.parametrize("qk_rope_head_dim", [64])
1020
+@pytest.mark.parametrize("block_size", [16])
1021
+@pytest.mark.parametrize("num_blocks", [128])
1022
+@pytest.mark.parametrize("dtype", [torch.float32])
1023
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
1024
+@pytest.mark.parametrize("device", CUDA_DEVICES)
1025
+@torch.inference_mode()
1026
+def test_gather_and_maybe_dequant_cache_mla_with_seq_starts(
1027
+ kv_lora_rank,
1028
+ qk_rope_head_dim,
1029
+ block_size,
1030
+ num_blocks,
1031
+ dtype,
1032
+ kv_cache_dtype,
1033
+ device,
1034
+):
1035
+ entry_size = kv_lora_rank + qk_rope_head_dim
1036
+ scale = torch.tensor(0.1, dtype=torch.float32, device=device)
1037
+ src_cache = _create_mla_cache(
1038
+ num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device
1039
+ )
1040
+ _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)
1041
+
1042
+ seq_starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device)
1043
+ seq_lens = torch.tensor([20, 10, 16], dtype=torch.int32, device=device)
1044
+ batch_size = seq_lens.shape[0]
1045
+ total_tokens = seq_lens.sum().item()
1046
+ cu_seq_lens = torch.empty((batch_size + 1), dtype=torch.int32, device=device)
1047
+ cu_seq_lens[0] = 0
1048
+ cu_seq_lens[1:] = seq_lens.cumsum(dim=0)
1049
+ token_to_seq = torch.repeat_interleave(
1050
+ torch.arange(batch_size, dtype=torch.int32, device=device), seq_lens
1051
+ )
1052
+
1053
+ block_table = torch.empty(
1054
+ (batch_size, num_blocks), dtype=torch.int32, device=device
1055
+ )
1056
+ for b in range(batch_size):
1057
+ block_table[b, :] = torch.randperm(num_blocks, device=device)
1058
+
1059
+ if kv_cache_dtype == "fp8":
1060
+ dequant_src_cache = torch.empty_like(src_cache, dtype=dtype)
1061
+ ops.convert_fp8(dequant_src_cache, src_cache, scale.item())
1062
+ else:
1063
+ dequant_src_cache = src_cache
1064
+
1065
+ expected_rows = []
1066
+ for b in range(batch_size):
1067
+ start = seq_starts[b].item()
1068
+ length = seq_lens[b].item()
1069
+ for offset in range(start, start + length):
1070
+ block_id = block_table[b, offset // block_size]
1071
+ slot = offset % block_size
1072
+ expected_rows.append(dequant_src_cache[block_id, slot])
1073
+ expected = torch.stack(expected_rows)
1074
+
1075
+ dst = torch.zeros((total_tokens, entry_size), dtype=dtype, device=device)
1076
+ ops.gather_and_maybe_dequant_cache(
1077
+ src_cache,
1078
+ dst,
1079
+ block_table,
1080
+ cu_seq_lens,
1081
+ token_to_seq,
1082
+ total_tokens,
1083
+ kv_cache_dtype,
1084
+ scale,
1085
+ seq_starts,
1086
+ )
1087
+ torch.testing.assert_close(dst, expected)
1088
+
1089
+
1090
@pytest.mark.parametrize("kv_lora_rank", [512])
1091
@pytest.mark.parametrize("qk_rope_head_dim", [64])
1092
@pytest.mark.parametrize("block_size", [16])
tests/v1/attention/test_mla_backends.py
+215
@@ -19,6 +19,7 @@ from tests.v1.attention.utils import (
19
)
20
from vllm import _custom_ops as ops
21
from vllm.config.vllm import set_current_vllm_config
22
+from vllm.model_executor.layers.attention import mla_attention as mla_attention_module
23
from vllm.model_executor.layers.attention.mla_attention import (
24
MLAAttention,
25
QueryLenSupport,
@@ -30,6 +31,7 @@ from vllm.utils.math_utils import cdiv
31
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
32
from vllm.v1.attention.backend import CommonAttentionMetadata
33
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
34
+from vllm.v1.attention.backends.mla import flashmla as flashmla_module
35
from vllm.v1.attention.backends.mla.prefill import (
36
MLAPrefillBackendEnum,
37
get_mla_prefill_backend,
@@ -552,6 +554,10 @@ class MockMLAAttentionLayer(MLAAttention):
554
)
555
else:
556
mqa_q = (mqa_ql_nope, mqa_q_pe)
557
+ if self.impl.dcp_world_size > 1:
558
+ if isinstance(mqa_q, tuple):
559
+ mqa_q = torch.cat(mqa_q, dim=-1)
560
+ mqa_q = mla_attention_module.get_dcp_group().all_gather(mqa_q, dim=1)
561
562
attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
563
@@ -569,6 +575,215 @@ class MockMLAAttentionLayer(MLAAttention):
575
return output
576
577
578
+def test_mock_mla_dcp_fp8_decode_gathers_quantized_query(
579
+ monkeypatch, default_vllm_config
580
+):
581
+ if not torch.cuda.is_available():
582
+ pytest.skip("CUDA is required for FP8 decode query quantization path.")
583
+
584
+ device = torch.device(f"{DEVICE_TYPE}:0")
585
+ num_tokens = 2
586
+ num_heads = 2
587
+ qk_nope_head_dim = 4
588
+ qk_rope_head_dim = 2
589
+ v_head_dim = 3
590
+ kv_lora_rank = 5
591
+
592
+ class _DummyKVProj:
593
+ def __init__(self):
594
+ # Shape expected by MockMLAAttentionLayer.__init__
595
+ self.weight = torch.randn(
596
+ num_heads * (qk_nope_head_dim + v_head_dim),
597
+ kv_lora_rank,
598
+ device=device,
599
+ dtype=torch.float32,
600
+ )
601
+
602
+ class _FakeImpl:
603
+ def __init__(self):
604
+ self.kv_cache_dtype = "fp8"
605
+ self.supports_quant_query_input = True
606
+ self.dcp_world_size = 2
607
+ self.forward_q = None
608
+
609
+ def forward_mha(self, *args, **kwargs):
610
+ return None
611
+
612
+ def forward_mqa(self, q, kv_cache, attn_metadata, layer):
613
+ self.forward_q = q
614
+ assert isinstance(q, torch.Tensor)
615
+ bsz, _, _ = q.shape
616
+ return (
617
+ torch.zeros(
618
+ bsz,
619
+ num_heads,
620
+ kv_lora_rank,
621
+ device=q.device,
622
+ dtype=torch.float32,
623
+ ),
624
+ None,
625
+ )
626
+
627
+ class _FakeDCPGroup:
628
+ def __init__(self):
629
+ self.calls = 0
630
+ self.input_dtype = None
631
+ self.input_shape = None
632
+
633
+ def all_gather(self, x, dim=1):
634
+ self.calls += 1
635
+ self.input_dtype = x.dtype
636
+ self.input_shape = tuple(x.shape)
637
+ return torch.cat([x, x], dim=dim)
638
+
639
+ fake_group = _FakeDCPGroup()
640
+ monkeypatch.setattr(mla_attention_module, "get_dcp_group", lambda: fake_group)
641
+
642
+ impl = _FakeImpl()
643
+ with set_current_vllm_config(default_vllm_config):
644
+ layer = MockMLAAttentionLayer(
645
+ impl=impl,
646
+ num_heads=num_heads,
647
+ qk_nope_head_dim=qk_nope_head_dim,
648
+ qk_rope_head_dim=qk_rope_head_dim,
649
+ v_head_dim=v_head_dim,
650
+ kv_lora_rank=kv_lora_rank,
651
+ device=device,
652
+ kv_b_proj=_DummyKVProj(),
653
+ q_scale=1.0,
654
+ k_scale=1.0,
655
+ )
656
+
657
+ q = torch.randn(
658
+ num_tokens,
659
+ num_heads,
660
+ qk_nope_head_dim + qk_rope_head_dim,
661
+ device=device,
662
+ dtype=torch.float32,
663
+ )
664
+ kv_c = torch.randn(num_tokens, kv_lora_rank, device=device, dtype=torch.float32)
665
+ k_pe = torch.randn(
666
+ num_tokens, 1, qk_rope_head_dim, device=device, dtype=torch.float32
667
+ )
668
+ kv_cache = torch.empty(0, device=device, dtype=torch.float32)
669
+ output = torch.empty(
670
+ num_tokens, num_heads * v_head_dim, device=device, dtype=torch.float32
671
+ )
672
+
673
+ class _AttnMeta:
674
+ num_decode_tokens = num_tokens
675
+ num_decodes = 1
676
+ num_prefills = 0
677
+ slot_mapping = torch.empty(0, dtype=torch.long, device=device)
678
+
679
+ layer.forward_impl(q, kv_c, k_pe, kv_cache, _AttnMeta(), output)
680
+
681
+ assert fake_group.calls == 1
682
+ assert fake_group.input_dtype == current_platform.fp8_dtype()
683
+ assert fake_group.input_shape == (
684
+ num_tokens,
685
+ num_heads,
686
+ kv_lora_rank + qk_rope_head_dim,
687
+ )
688
+ assert isinstance(impl.forward_q, torch.Tensor)
689
+ assert tuple(impl.forward_q.shape) == (
690
+ num_tokens,
691
+ num_heads * impl.dcp_world_size,
692
+ kv_lora_rank + qk_rope_head_dim,
693
+ )
694
+
695
+
696
+@pytest.mark.parametrize("is_fp8_kvcache", [False, True], ids=["bf16", "fp8"])
697
+def test_flashmla_dcp_decode_metadata_uses_gathered_query_heads(
698
+ monkeypatch, is_fp8_kvcache
699
+):
700
+ class _FakeSchedulerMetadata:
701
+ tile_scheduler_metadata = None
702
+ num_splits = None
703
+
704
+ base_call: tuple[torch.Tensor, int, int, bool] | None = None
705
+ fp8_call: tuple[torch.Tensor, int, int] | None = None
706
+
707
+ def fake_get_mla_metadata(
708
+ seq_lens_device,
709
+ num_q_tokens_per_head_k,
710
+ num_heads_k,
711
+ is_fp8_kvcache=False,
712
+ ):
713
+ nonlocal base_call
714
+ base_call = (
715
+ seq_lens_device,
716
+ num_q_tokens_per_head_k,
717
+ num_heads_k,
718
+ is_fp8_kvcache,
719
+ )
720
+ return _FakeSchedulerMetadata(), None
721
+
722
+ def fake_get_mla_metadata_dense_fp8(
723
+ seq_lens_device, num_q_tokens_per_head_k, num_heads_k
724
+ ):
725
+ nonlocal fp8_call
726
+ fp8_call = (
727
+ seq_lens_device,
728
+ num_q_tokens_per_head_k,
729
+ num_heads_k,
730
+ )
731
+ return (
732
+ torch.empty((0, 8), dtype=torch.int32),
733
+ torch.empty((0,), dtype=torch.int32),
734
+ )
735
+
736
+ monkeypatch.setattr(flashmla_module, "get_mla_metadata", fake_get_mla_metadata)
737
+ monkeypatch.setattr(
738
+ flashmla_module,
739
+ "get_mla_metadata_dense_fp8",
740
+ fake_get_mla_metadata_dense_fp8,
741
+ )
742
+
743
+ builder = object.__new__(flashmla_module.FlashMLAMetadataBuilder)
744
+ builder.num_q_heads = 4
745
+ builder.dcp_world_size = 2
746
+ builder.is_fp8_kvcache = is_fp8_kvcache
747
+ builder.compilation_config = type(
748
+ "_CompilationConfig",
749
+ (),
750
+ {
751
+ "cudagraph_mode": type(
752
+ "_CudaGraphMode",
753
+ (),
754
+ {"has_full_cudagraphs": lambda self: False},
755
+ )()
756
+ },
757
+ )()
758
+
759
+ seq_lens = torch.tensor([16, 24], dtype=torch.int32)
760
+ query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32)
761
+
762
+ metadata = builder._build_decode(
763
+ block_table_tensor=torch.empty((2, 1), dtype=torch.int32),
764
+ seq_lens_device=seq_lens,
765
+ max_seq_len=24,
766
+ query_start_loc_cpu=query_start_loc,
767
+ query_start_loc_device=query_start_loc,
768
+ num_decode_tokens=2,
769
+ dcp_tot_seq_lens_device=None,
770
+ )
771
+
772
+ assert base_call is not None
773
+ assert base_call[0] is seq_lens
774
+ assert base_call[1:] == (8, 1, is_fp8_kvcache)
775
+ if is_fp8_kvcache:
776
+ assert metadata.scheduler_metadata.tile_scheduler_metadata is not None
777
+ assert metadata.scheduler_metadata.num_splits is not None
778
+ assert fp8_call is not None
779
+ assert fp8_call[0] is seq_lens
780
+ assert fp8_call[1:] == (8, 1)
781
+ else:
782
+ assert metadata.scheduler_metadata.tile_scheduler_metadata is None
783
+ assert metadata.scheduler_metadata.num_splits is None
784
+ assert fp8_call is None
785
+
786
+
787
def run_attention_backend(
788
backend: AttentionBackendEnum,
789
kv_cache_spec: MLAAttentionSpec,
vllm/model_executor/layers/attention/mla_attention.py
+68
-14
@@ -800,9 +800,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):
800
else:
801
mqa_q = (mqa_ql_nope, mqa_q_pe)
802
if self.impl.dcp_world_size > 1:
803
- assert not fp8_attention, "DCP not support fp8 kvcache now."
804
- # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P)
805
- mqa_q = torch.cat(mqa_q, dim=-1)
803
+ if isinstance(mqa_q, tuple):
804
+ # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P)
805
+ mqa_q = torch.cat(mqa_q, dim=-1)
806
# mqa_q do allgather in head dim.
807
mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)
808
@@ -1263,6 +1263,7 @@ class MLACommonPrefillMetadata:
1263
padded_local_chunk_seq_lens: list[list[int]] | None = None
1264
local_context_lens_allranks: list[list[int]] | None = None
1265
padded_local_cu_seq_lens: torch.Tensor | None = None
1266
+ padded_local_token_to_seq: torch.Tensor | None = None
1267
cu_seq_lens_lst: list[list[int]] | None = None
1268
chunk_size: int | None = None
1269
prefill_tokens_with_context: int | None = None
@@ -1787,6 +1788,21 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1788
out=padded_local_cu_chunk_seq_lens_cpu[:, 1:],
1789
dtype=torch.int32,
1790
)
1791
+ max_padded_local_tokens_over_chunk = (
1792
+ padded_local_cu_chunk_seq_lens_cpu[:, -1].max().item()
1793
+ )
1794
+ padded_local_token_to_seq_tensor_cpu = torch.zeros(
1795
+ [num_chunks, max_padded_local_tokens_over_chunk],
1796
+ dtype=torch.int32,
1797
+ )
1798
+ for i in range(num_chunks):
1799
+ chunk_token_to_seq_tensor = torch.repeat_interleave(
1800
+ range_idx, padded_local_chunk_seq_lens[i]
1801
+ )
1802
+ chunk_len = chunk_token_to_seq_tensor.shape[0]
1803
+ padded_local_token_to_seq_tensor_cpu[i, :chunk_len] = (
1804
+ chunk_token_to_seq_tensor
1805
+ )
1806
1807
prefill_tokens_with_context = None
1808
if num_prefills_with_context_cpu > 0:
@@ -1811,6 +1827,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1827
padded_local_cu_seq_lens=padded_local_cu_chunk_seq_lens_cpu.to(
1828
device, non_blocking=True
1829
),
1830
+ padded_local_token_to_seq=padded_local_token_to_seq_tensor_cpu.to(
1831
+ device, non_blocking=True
1832
+ ),
1833
cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),
1834
chunk_size=padded_local_max_context_chunk_across_ranks,
1835
prefill_tokens_with_context=prefill_tokens_with_context,
@@ -2187,7 +2206,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2206
k_scale: torch.Tensor,
2207
dcp_world_size: int,
2208
):
2190
- assert k_scale is None, "DCP not support scaled kvcache now."
2209
assert attn_metadata.prefill is not None
2210
prefill_metadata = attn_metadata.prefill
2211
assert prefill_metadata.prefill_backend is not None
@@ -2195,9 +2213,11 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2213
assert prefill_metadata.chunked_context.padded_local_chunk_seq_lens is not None
2214
assert prefill_metadata.chunked_context.local_context_lens_allranks is not None
2215
assert prefill_metadata.chunked_context.padded_local_cu_seq_lens is not None
2216
+ assert prefill_metadata.chunked_context.padded_local_token_to_seq is not None
2217
assert prefill_metadata.chunked_context.cu_seq_lens_lst is not None
2218
assert prefill_metadata.chunked_context.chunk_size is not None
2219
2220
+ use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
2221
output = None
2222
merge_output = None
2223
iters = len(prefill_metadata.chunked_context.seq_tot)
@@ -2205,16 +2225,37 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2225
2226
for i in range(iters):
2227
toks = prefill_metadata.chunked_context.seq_tot[i]
2208
- ops.cp_gather_cache(
2209
- src_cache=kv_c_and_k_pe_cache,
2210
- dst=workspace,
2211
- block_table=prefill_metadata.block_table,
2212
- cu_seq_lens=prefill_metadata.chunked_context.padded_local_cu_seq_lens[
2213
- i
2214
- ],
2215
- batch_size=attn_metadata.num_prefills,
2216
- seq_starts=prefill_metadata.chunked_context.starts[i],
2228
+ if toks == 0:
2229
+ continue
2230
+ padded_local_cu_seq_lens = (
2231
+ prefill_metadata.chunked_context.padded_local_cu_seq_lens[i]
2232
)
2233
+ if is_quantized_kv_cache(self.kv_cache_dtype) and (
2234
+ self.kv_cache_dtype != "fp8_ds_mla"
2235
+ ):
2236
+ assert k_scale is not None
2237
+ ops.gather_and_maybe_dequant_cache(
2238
+ src_cache=kv_c_and_k_pe_cache,
2239
+ dst=workspace,
2240
+ block_table=prefill_metadata.block_table,
2241
+ cu_seq_lens=padded_local_cu_seq_lens,
2242
+ token_to_seq=prefill_metadata.chunked_context.padded_local_token_to_seq[
2243
+ i
2244
+ ],
2245
+ num_tokens=toks,
2246
+ kv_cache_dtype=self.kv_cache_dtype,
2247
+ scale=k_scale,
2248
+ seq_starts=prefill_metadata.chunked_context.starts[i],
2249
+ )
2250
+ else:
2251
+ ops.cp_gather_cache(
2252
+ src_cache=kv_c_and_k_pe_cache,
2253
+ dst=workspace,
2254
+ block_table=prefill_metadata.block_table,
2255
+ cu_seq_lens=padded_local_cu_seq_lens,
2256
+ batch_size=attn_metadata.num_prefills,
2257
+ seq_starts=prefill_metadata.chunked_context.starts[i],
2258
+ )
2259
# workspace
2260
# |------- N tokens --------|--------- N*dcp_size tokens ----------|
2261
# |<- use for local_gather ->|<--------- use for allgather -------->|
@@ -2252,9 +2293,22 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2293
toks=toks,
2294
)
2295
2296
+ kv_b_proj_w_dtype = (
2297
+ self.kv_b_proj.weight.dtype
2298
+ if hasattr(self.kv_b_proj, "weight")
2299
+ else self.kv_b_proj.params_dtype
2300
+ )
2301
+ if (
2302
+ use_fp8_prefill or kv_b_proj_w_dtype != current_platform.fp8_dtype()
2303
+ ) and kv_b_proj_w_dtype != torch.uint8:
2304
+ kv_c_normed = kv_c_normed.to(kv_b_proj_w_dtype)
2305
+
2306
kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
2307
-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
2308
)
2309
+ if use_fp8_prefill:
2310
+ kv_nope = kv_nope.to(prefill_metadata.q_data_type)
2311
+ k_pe = k_pe.to(prefill_metadata.q_data_type)
2312
k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
2313
k = self._concat_k_nope_k_pe(k_nope, k_pe)
2314
@@ -2346,7 +2400,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2400
q,
2401
kv_c_and_k_pe_cache,
2402
attn_metadata,
2349
- k_scale=None,
2403
+ k_scale=k_scale,
2404
dcp_world_size=self.dcp_world_size,
2405
)
2406
)
vllm/v1/attention/backends/mla/flashmla.py
+4
-1
@@ -171,7 +171,10 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
171
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
172
# we use the max but all should be the same due to uniform length requirement
173
max_query_len = query_lens_cpu.max().item()
174
- num_q_tokens_per_head_k = max_query_len * self.num_q_heads // 1
174
+ num_q_heads = self.num_q_heads
175
+ if self.dcp_world_size > 1:
176
+ num_q_heads *= self.dcp_world_size
177
+ num_q_tokens_per_head_k = max_query_len * num_q_heads // 1
178
scheduler_metadata, _ = get_mla_metadata(
179
seq_lens_device,
180
num_q_tokens_per_head_k,