| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758 |
- --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
- +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
- @@ -32,6 +32,13 @@
- class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]):
- """SM120 FlashInfer sparse-MLA implementation."""
-
- + # glm52-dcp patch: FlashInfer's trtllm_batch_decode_with_kv_cache_mla
- + # accepts return_lse (verified in flashinfer 0.6.14), so this impl CAN
- + # return the decode LSE that DCP's cross-rank merge requires. Upstream
- + # never plumbed it through; see forward_mqa below. Mirrors the parent
- + # FlashInferMLASparseImpl's LSE handling.
- + can_return_lse_for_decode = True
- +
- def __init__(
- self,
- num_heads: int,
- @@ -126,8 +133,12 @@
- ),
- )
-
- + # glm52-dcp patch: size the out buffer from the query's actual head
- + # count, not self.num_heads. Under DCP the MLA decode all-gathers
- + # queries across ranks, so each rank sees the FULL head count (64)
- + # rather than the TP-local slice (16) — self.num_heads is stale here.
- output = q.new_empty(
- - (num_actual_toks, self.num_heads, self.kv_lora_rank),
- + (num_actual_toks, q.shape[1], self.kv_lora_rank),
- dtype=q.dtype,
- )
-
- @@ -153,5 +164,27 @@
- bmm2_scale=1.0,
- sparse_mla_top_k=attn_metadata.topk_tokens,
- kv_scale_format=self.kv_scale_format,
- + # glm52-dcp patch: ask the kernel for the decode LSE when DCP
- + # needs it for the cross-rank merge (base class sets this flag
- + # only when dcp_world_size > 1).
- + return_lse=self.need_to_return_lse_for_decode,
- )
- + if self.need_to_return_lse_for_decode:
- + # With return_lse=True flashinfer returns (out, lse).
- + assert isinstance(out, tuple), (
- + "expected (out, lse) tuple from flashinfer when return_lse=True"
- + )
- + o, lse = out
- + o = o.view(-1, o.shape[-2], o.shape[-1])
- + # Reuse the parent impl's LSE normalization + empty-row masking
- + # conventions (the shared DCP reducer expects (num_tokens, num_heads)
- + # and -inf LSE on rows with no valid top-k indices).
- + from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
- + FlashInferMLASparseImpl as _ParentImpl,
- + )
- + lse = _ParentImpl._normalize_lse(lse, o.shape[0], o.shape[1])
- + empty_rows = (topk_indices_physical == -1).all(dim=-1)
- + o.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
- + lse.masked_fill_(empty_rows.view(-1, 1), float("-inf"))
- + return o, lse
- return out.squeeze(1), None
|
