@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 63ff748f65

[Attention][MLA] FlashMLA sparse: DCP on the fp8_ds_mla mixed-batch path + MTP (#46514)

Signed-off-by: Mikhail Kostryukov <mike@triptrack.net> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

drakosha committed Aug 19, 2026 at 07:02 UTC 63ff748f657af0e2a95a712fd2cc994f4299ff8c

3 files changed +295 -34

tests/v1/attention/test_sparse_mla_backends.py

+154

@@ -1492,6 +1492,8 @@ def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: b

1492 topk_indices_buffer=topk_indices,

1493 num_heads=2,

1494 kv_lora_rank=1,

1495 + dcp_world_size=1,

1496 + need_to_return_lse_for_decode=False,

1497 _fp8_flash_mla_kernel=run_kernel,

1498 )

1499 impl._forward_fp8_kv_mixed_batch = MethodType(

@@ -1512,3 +1514,155 @@ def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: b

1514 assert kernel_q_shapes == [(1, num_decode_tokens, 2, 3)]

1515 assert output.shape == (num_decode_tokens, 2, 1)

1516 assert lse is None

1517 +

1518 +

1519 +def _build_sparse_dcp_vllm_config(

1520 + local_heads: int,

1521 + dcp_world_size: int,

1522 + comm_backend: str = "ag_rs",

1523 +):

1524 + """Minimal sparse-MLA VllmConfig for the FlashMLASparse DCP head-envelope

1525 + guard. TP is simulated by mocking ``get_num_attention_heads`` to return the

1526 + per-rank head count, as the decode-correctness test above does.

1527 + """

1528 + kv_lora_rank = 512

1529 + qk_nope_head_dim = 128

1530 + qk_rope_head_dim = 64

1531 + v_head_dim = 128

1532 + head_size = kv_lora_rank + qk_rope_head_dim

1533 + topk_tokens = 128

1534 +

1535 + vllm_config = create_vllm_config(

1536 + model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",

1537 + tensor_parallel_size=1,

1538 + max_model_len=4096,

1539 + block_size=64,

1540 + hf_config_override={

1541 + "index_topk": topk_tokens,

1542 + "attn_module_list_cfg": [{"topk_tokens": topk_tokens}],

1543 + },

1544 + )

1545 + model_config = vllm_config.model_config

1546 + model_config.dtype = torch.bfloat16

1547 + model_config.hf_text_config = SimpleNamespace(

1548 + q_lora_rank=None,

1549 + kv_lora_rank=kv_lora_rank,

1550 + qk_nope_head_dim=qk_nope_head_dim,

1551 + qk_rope_head_dim=qk_rope_head_dim,

1552 + v_head_dim=v_head_dim,

1553 + model_type="deepseek_v2",

1554 + )

1555 + model_config.get_num_attention_heads = MethodType(

1556 + lambda self, parallel_config: local_heads, model_config

1557 + )

1558 + model_config.get_num_kv_heads = MethodType(

1559 + lambda self, parallel_config: 1, model_config

1560 + )

1561 + model_config.get_head_size = MethodType(lambda self: head_size, model_config)

1562 + model_config.get_sliding_window = MethodType(lambda self: None, model_config)

1563 +

1564 + vllm_config.cache_config.cache_dtype = "fp8_ds_mla"

1565 + vllm_config.parallel_config.decode_context_parallel_size = dcp_world_size

1566 + vllm_config.parallel_config.dcp_comm_backend = comm_backend

1567 + # The base builder clones the layer's dense-MHA prefill backend from

1568 + # static_forward_context; the guard tests never run prefill.

1569 + vllm_config.compilation_config.static_forward_context["placeholder"] = (

1570 + SimpleNamespace(prefill_backend=None)

1571 + )

1572 + return vllm_config

1573 +

1574 +

1575 +@pytest.mark.skipif(

1576 + torch.cuda.get_device_capability() < (9, 0),

1577 + reason="FlashMLASparseBackend requires CUDA 9.0 or higher",

1578 +)

1579 +@pytest.mark.parametrize(

1580 + "local_heads,dcp_world_size,should_raise",

1581 + [

1582 + (16, 8, True),

1583 + (24, 4, True),

1584 + (16, 4, False),

1585 + (16, 1, False),

1586 + ],

1587 +)

1588 +def test_fp8_dcp_head_envelope_guard(local_heads, dcp_world_size, should_raise):

1589 + """The fp8 decode envelope (head padding + tile-scheduler metadata) is

1590 + sized from the local head count while the kernel runs on the DCP-gathered

1591 + heads, so the builder must reject configs where the two pad differently.

1592 + """

1593 + device = torch.device(DEVICE_TYPE)

1594 + vllm_config = _build_sparse_dcp_vllm_config(local_heads, dcp_world_size)

1595 + kv_cache_spec = create_standard_kv_cache_spec(vllm_config)

1596 + builder_cls = FlashMLASparseBackend.get_builder_cls()

1597 +

1598 + if should_raise:

1599 + with pytest.raises(NotImplementedError, match="envelope"):

1600 + builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)

1601 + else:

1602 + builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)

1603 + gathered_heads = local_heads * dcp_world_size

1604 + local_pad = 64 if local_heads <= 64 else 128

1605 + gathered_pad = 64 if gathered_heads <= 64 else 128

1606 + assert builder.fp8_decode_padded_heads == local_pad

1607 + assert local_pad == gathered_pad

1608 +

1609 +

1610 +def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch):

1611 + """A decode row whose top-k shard holds no local candidates (all -1) has

1612 + undefined kernel out/lse; it must come back as (0, -inf), the identity of

1613 + the cross-rank LSE merge, or a NaN would survive the merge even at zero

1614 + weight (0 * NaN = NaN)."""

1615 + num_tokens, num_heads, head_dim = 3, 2, 3

1616 + q = torch.empty(num_tokens, num_heads, head_dim, device=DEVICE_TYPE)

1617 + local_indices = torch.tensor(

1618 + [[0, 1, -1, -1], [-1, -1, -1, -1], [2, -1, 3, -1]],

1619 + dtype=torch.int32,

1620 + device=DEVICE_TYPE,

1621 + )

1622 +

1623 + monkeypatch.setattr(

1624 + "vllm.v1.attention.backends.mla.flashmla_sparse."

1625 + "triton_filter_and_convert_dcp_index",

1626 + lambda *args, **kwargs: local_indices,

1627 + )

1628 +

1629 + def run_kernel(**kwargs):

1630 + out = torch.full(

1631 + (1, num_tokens, num_heads, 1), float("nan"), device=DEVICE_TYPE

1632 + )

1633 + lse = torch.full((1, num_heads, num_tokens), float("nan"), device=DEVICE_TYPE)

1634 + for token_id in (0, 2): # rows with local candidates get real values

1635 + out[0, token_id] = float(token_id + 1)

1636 + lse[0, :, token_id] = float(token_id + 1)

1637 + return out, lse

1638 +

1639 + metadata = SimpleNamespace(

1640 + fp8_extra_metadata=FlashMLASparseMetadata.FP8KernelMetadata(

1641 + scheduler_metadata=object(), # type: ignore[arg-type]

1642 + dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),

1643 + cache_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),

1644 + ),

1645 + req_id_per_token=torch.empty(num_tokens, dtype=torch.int32, device=DEVICE_TYPE),

1646 + block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),

1647 + block_size=64,

1648 + cp_kv_cache_interleave_size=1,

1649 + )

1650 + impl = SimpleNamespace(

1651 + dcp_world_size=2,

1652 + dcp_rank=0,

1653 + need_to_return_lse_for_decode=True,

1654 + _fp8_flash_mla_kernel=run_kernel,

1655 + )

1656 +

1657 + out, lse = FlashMLASparseImpl._forward_fp8_kv_mixed_batch(

1658 + impl, q, torch.empty(0, device=DEVICE_TYPE), local_indices, metadata

1659 + )

1660 +

1661 + assert torch.equal(out[1], torch.zeros_like(out[1]))

1662 + assert torch.isneginf(lse[1]).all()

1663 + for token_id in (0, 2):

1664 + assert torch.equal(out[token_id], torch.full_like(out[token_id], token_id + 1))

1665 + assert torch.equal(lse[token_id], torch.full_like(lse[token_id], token_id + 1))

1666 + assert out.is_contiguous()

1667 + assert not out.isnan().any()

1668 + assert not lse.isnan().any()

vllm/v1/attention/backends/mla/flashmla.py

+12

@@ -343,4 +343,16 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

343

344 o = reshape_attn_output_for_spec_decode(o)

345

346 + if self.need_to_return_lse_for_decode:

347 + # FlashMLA returns LSE as [batch, heads, seq_len]; the DCP reducer

348 + # consumes [tokens, heads]. Flattening matters under spec-decode,

349 + # where seq_len > 1. Only DCP consumes lse, so skip the copy

350 + # otherwise.

351 + num_decodes, q_num_heads, seq_len = lse.shape

352 + lse = (

353 + lse.permute(0, 2, 1)

354 + .reshape(num_decodes * seq_len, q_num_heads)

355 + .contiguous()

356 + )

357 +

358 return o, lse

vllm/v1/attention/backends/mla/flashmla_sparse.py

+129 -34

@@ -28,6 +28,7 @@ from vllm.v1.attention.backend import (

28 )

29 from vllm.v1.attention.backends.mla.sparse_utils import (

30 triton_convert_req_index_to_global_index,

31 + triton_filter_and_convert_dcp_index,

32 )

33 from vllm.v1.attention.backends.utils import (

34 reshape_attn_output_for_spec_decode,

@@ -254,9 +255,12 @@ class FlashMLASparseMetadataBuilder(

255 )

256 else:

257 threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256)

258 + # Varlen decodes are safe under DCP: causality comes from the

259 + # indexer's top-k indices, not from the kernel metadata.

260 self._init_reorder_batch_threshold(

261 threshold,

262 supports_spec_as_decode=True,

263 + supports_dcp_with_varlen=(parallel_config.cp_kv_cache_interleave_size == 1),

264 )

265

266 sm_count = num_compute_units(device.index)

@@ -313,6 +317,42 @@ class FlashMLASparseMetadataBuilder(

317 device=device,

318 )

319

320 + self.fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL

321 +

322 + if parallel_config.decode_context_parallel_size > 1:

323 + if parallel_config.dcp_comm_backend != "ag_rs":

324 + raise NotImplementedError(

325 + "DCP for FlashMLA sparse is only validated with the "

326 + "default 'ag_rs' DCP comm backend; got "

327 + f"'{parallel_config.dcp_comm_backend}'"

328 + )

329 + if not self.fp8_use_mixed_batch:

330 + raise NotImplementedError(

331 + "DCP for FlashMLA sparse is only supported on the "

332 + "mixed-batch fp8 path (num_heads < "

333 + f"{MIN_HEADS_FOR_BF16_PREFILL}); the separate "

334 + "prefill/decode path returns the LSE for decode tokens "

335 + "only, while the DCP merge needs it for every token"

336 + )

337 + # Head padding (and the tile-scheduler metadata sized from it) is

338 + # computed from the local head count, but the kernel runs on the

339 + # DCP-gathered heads.

340 + gathered_num_heads = (

341 + self.num_heads * parallel_config.decode_context_parallel_size

342 + )

343 + gathered_padded_heads = FlashMLASparseImpl._compute_fp8_decode_padded_heads(

344 + gathered_num_heads

345 + )

346 + if self.fp8_decode_padded_heads != gathered_padded_heads:

347 + raise NotImplementedError(

348 + "DCP for FlashMLA sparse requires the local and "

349 + "DCP-gathered head counts to pad to the same fp8 decode "

350 + f"kernel envelope; got {self.num_heads} local heads "

351 + f"(pad to {self.fp8_decode_padded_heads}) vs "

352 + f"{gathered_num_heads} gathered heads (pad to "

353 + f"{gathered_padded_heads})"

354 + )

355 +

356 def _build_fp8_mixed_decode_prefill(

357 self,

358 common_attn_metadata: CommonAttentionMetadata,

@@ -496,10 +536,9 @@ class FlashMLASparseMetadataBuilder(

536 ) -> FlashMLASparseMetadata:

537 metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)

538

499 - fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL

500 - metadata.fp8_use_mixed_batch = fp8_use_mixed_batch

539 + metadata.fp8_use_mixed_batch = self.fp8_use_mixed_batch

540 if self.use_fp8_kv_cache:

502 - if fp8_use_mixed_batch:

541 + if self.fp8_use_mixed_batch:

542 metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill(

543 common_attn_metadata

544 )

@@ -512,6 +551,8 @@ class FlashMLASparseMetadataBuilder(

551

552

553 class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

554 + can_return_lse_for_decode: bool = True

555 +

556 @staticmethod

557 def _compute_fp8_decode_padded_heads(num_heads: int) -> int:

558 # FP8 decode kernel only supports h_q = 64 or 128

@@ -566,6 +607,14 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

607 "fp8_ds_mla kv-cache dtype"

608 )

609

610 + if self.need_to_return_lse_for_decode and not is_quantized_kv_cache(

611 + kv_cache_dtype

612 + ):

613 + raise NotImplementedError(

614 + "DCP for FlashMLA sparse requires an fp8_ds_mla kv-cache; "

615 + "the bf16 sparse path is not supported under DCP."

616 + )

617 +

618 if kv_cache_dtype == "fp8_ds_mla":

619 # Reserve workspace during initialization

620 assert vllm_config is not None and vllm_config.model_config is not None

@@ -590,7 +639,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

639 kv_c_and_k_pe_cache: torch.Tensor,

640 topk_indices: torch.Tensor,

641 attn_metadata: FlashMLASparseMetadata,

593 - ) -> torch.Tensor:

642 + ) -> tuple[torch.Tensor, torch.Tensor]:

643 # Convert per-request indices to global slots (decode) or workspace

644 # offsets (prefill). req_id_per_token covers the whole batch; slice it

645 # to the MQA tokens (q may exclude prefill tokens routed to dense MHA).

@@ -712,7 +761,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

761 chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice]

762 chunk_topk_length = topk_length[chunk.tokens_slice]

763

715 - attn_out[chunk.tokens_slice] = self._bf16_flash_mla_kernel(

764 + attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel(

765 chunk_q,

766 chunk_workspace,

767 chunk_topk_indices_workspace,

@@ -727,22 +776,42 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

776 kv_c_and_k_pe_cache: torch.Tensor,

777 topk_indices: torch.Tensor,

778 attn_metadata: FlashMLASparseMetadata,

730 - ) -> torch.Tensor:

779 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

780 """Mixed batch FP8 forward path that treats all tokens as one batch.

781

782 This is equivalent to main branch's approach and avoids the BF16

783 prefill kernel which has head padding overhead when num_heads is small.

784 Used when use_mixed_batch is True.

785 +

786 + The lse is only returned when DCP needs it, otherwise None.

787 """

737 - # Convert per-request indices to global slots (decode) or workspace

738 - # offsets (prefill).

739 - topk_indices = triton_convert_req_index_to_global_index(

740 - attn_metadata.req_id_per_token[: topk_indices.shape[0]],

741 - attn_metadata.block_table,

742 - topk_indices,

743 - BLOCK_SIZE=attn_metadata.block_size,

744 - NUM_TOPK_TOKENS=topk_indices.shape[1],

745 - )

788 + if self.dcp_world_size > 1:

789 + # The indexer emits global token ids; keep this rank's shard and

790 + # convert to local slots. compact_valid_to_front=False keeps the

791 + # scattered -1s, which the fp8 kernel masks natively and the

792 + # empty-row neutralization below relies on. req_id is sliced to

793 + # topk_indices rows (the converter grids from req_id).

794 + topk_indices = triton_filter_and_convert_dcp_index(

795 + attn_metadata.req_id_per_token[: topk_indices.shape[0]],

796 + attn_metadata.block_table,

797 + topk_indices,

798 + dcp_size=self.dcp_world_size,

799 + dcp_rank=self.dcp_rank,

800 + cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size,

801 + BLOCK_SIZE=attn_metadata.block_size,

802 + NUM_TOPK_TOKENS=topk_indices.shape[1],

803 + compact_valid_to_front=False,

804 + )

805 + else:

806 + # Convert per-request indices to global slots (decode) or workspace

807 + # offsets (prefill).

808 + topk_indices = triton_convert_req_index_to_global_index(

809 + attn_metadata.req_id_per_token[: topk_indices.shape[0]],

810 + attn_metadata.block_table,

811 + topk_indices,

812 + BLOCK_SIZE=attn_metadata.block_size,

813 + NUM_TOPK_TOKENS=topk_indices.shape[1],

814 + )

815

816 assert attn_metadata.fp8_extra_metadata is not None

817 assert isinstance(

@@ -750,15 +819,29 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

819 )

820 fp8_metadata = attn_metadata.fp8_extra_metadata

821

753 - _attn_out, _ = self._fp8_flash_mla_kernel(

822 + _attn_out, _lse = self._fp8_flash_mla_kernel(

823 q=q.unsqueeze(0), # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D)

824 kv_c_and_k_pe_cache=kv_c_and_k_pe_cache,

825 topk_indices=topk_indices.unsqueeze(0), # (T, topk) -> (1, T, topk)

826 kernel_metadata=fp8_metadata,

827 )

759 -

828 # Output is (1, T, H, D_v), squeeze back to (T, H, D_v)

761 - return _attn_out.squeeze(0)

829 + out = _attn_out.squeeze(0)

830 +

831 + if not self.need_to_return_lse_for_decode:

832 + return out, None

833 +

834 + # Kernel LSE is (1, H, T); the DCP merge consumes (T, H).

835 + lse = _lse.squeeze(0).transpose(0, 1)

836 + # Rows where this rank owns none of the selected tokens (all indices

837 + # -1) have undefined out/lse; (0, -inf) is the identity element of the

838 + # cross-rank LSE merge, so it drops this rank from those rows.

839 + empty_rows = (topk_indices == -1).all(dim=-1)

840 + out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)

841 + lse.masked_fill_(empty_rows.view(-1, 1), float("-inf"))

842 + # The head-padding slice above can leave `out` non-contiguous, and the

843 + # merge feeds it to reduce_scatter.

844 + return out.contiguous(), lse

845

846 def _fp8_flash_mla_kernel(

847 self,

@@ -793,9 +876,10 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

876 softmax_scale=self.softmax_scale,

877 )

878

796 - # Slice output back to actual head count if we padded

879 + # Slice output and lse back to actual head count if we padded

880 if actual_num_heads < padded_num_heads:

881 out = out[:, :, :actual_num_heads, :]

882 + lse = lse[:, :actual_num_heads, :]

883

884 return out, lse

885

@@ -805,35 +889,42 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

889 kv_c_and_k_pe_cache: torch.Tensor,

890 topk_indices: torch.Tensor,

891 topk_length: torch.Tensor | None = None,

808 - ) -> torch.Tensor:

892 + ) -> tuple[torch.Tensor, torch.Tensor]:

893 num_tokens = q.shape[0]

894 kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(

895 -1, 1, kv_c_and_k_pe_cache.shape[-1]

896 )

897

898 # NOTE(Chen): kernel requires num_local_head to be a multiple of

815 - # 64 on hopper and 128 on blackwell

816 - if self.num_heads % self.prefill_padding != 0:

817 - assert self.prefill_padding % self.num_heads == 0

899 + # 64 on hopper and 128 on blackwell. Pad from q's head count, not

900 + # self.num_heads: under DCP the heads are all-gathered before this.

901 + actual_num_heads = q.shape[1]

902 + padded_num_heads = (

903 + (actual_num_heads + self.prefill_padding - 1)

904 + // self.prefill_padding

905 + * self.prefill_padding

906 + )

907 + if actual_num_heads < padded_num_heads:

908 logger.warning_once(

819 - f"Padding num_heads from {self.num_heads} to "

820 - f"{self.prefill_padding} for BF16 sparse prefill kernel"

909 + f"Padding num_heads from {actual_num_heads} to "

910 + f"{padded_num_heads} for BF16 sparse prefill kernel"

911 )

822 - q_padded = q.new_empty((q.shape[0], self.prefill_padding, q.shape[2]))

823 - q_padded[:, : self.num_heads, :] = q

912 + q_padded = q.new_empty((q.shape[0], padded_num_heads, q.shape[2]))

913 + q_padded[:, :actual_num_heads, :] = q

914 q = q_padded

915

916 topk_indices = topk_indices.view(num_tokens, 1, -1)

827 - output = flash_mla_sparse_fwd(

917 + output, _, lse = flash_mla_sparse_fwd(

918 q,

919 kv_c_and_k_pe_cache,

920 topk_indices,

921 self.softmax_scale,

922 topk_length=topk_length,

833 - )[0]

923 + )

924

835 - output = output[:, : self.num_heads, :]

836 - return output

925 + output = output[:, :actual_num_heads, :]

926 + lse = lse[:, :actual_num_heads]

927 + return output, lse

928

929 def forward_mqa(

930 self,

@@ -859,12 +950,16 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

950

951 use_fp8_cache = self.kv_cache_dtype == "fp8_ds_mla"

952

953 + lse: torch.Tensor | None = None

954 +

955 if not use_fp8_cache:

863 - attn_out = self._forward_bf16_kv(

956 + attn_out, bf16_lse = self._forward_bf16_kv(

957 q, kv_c_and_k_pe_cache, topk_indices, attn_metadata

958 )

959 + if self.need_to_return_lse_for_decode:

960 + lse = bf16_lse

961 elif attn_metadata.fp8_use_mixed_batch:

867 - attn_out = self._forward_fp8_kv_mixed_batch(

962 + attn_out, lse = self._forward_fp8_kv_mixed_batch(

963 q, kv_c_and_k_pe_cache, topk_indices, attn_metadata

964 )

965 else:

@@ -872,4 +967,4 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

967 q, kv_c_and_k_pe_cache, topk_indices, attn_metadata

968 )

969

875 - return attn_out, None

970 + return attn_out, lse