[eldritch] attention: support DCP in SM120 FlashInfer sparse MLA

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

@@ -325,19 +325,32 @@ def build(
         )
 
 
-# Global workspace buffer (lazily initialized)
-_fi_sparse_workspace: torch.Tensor | None = None
+# Global FlashInfer sparse-MLA workspaces, one per CUDA device.
+#
+# vLLM workers are single-device in normal serving, but tests and warmup helpers
+# can construct this backend on different CUDA devices in one process. Keying
+# the scratch buffer by device avoids reusing a tensor allocated on another GPU.
+_fi_sparse_workspace_by_device: dict[torch.device, torch.Tensor] = {}
+
+
+def _normalize_workspace_device(device: torch.device) -> torch.device:
+    device = torch.device(device)
+    if device.type == "cuda" and device.index is None:
+        device = torch.device(f"cuda:{torch.cuda.current_device()}")
+    return device
 
 
 def _get_workspace_buffer(device: torch.device) -> torch.Tensor:
-    global _fi_sparse_workspace
-    if _fi_sparse_workspace is None:
-        _fi_sparse_workspace = torch.zeros(
+    device = _normalize_workspace_device(device)
+    workspace = _fi_sparse_workspace_by_device.get(device)
+    if workspace is None:
+        workspace = torch.zeros(
             FLASHINFER_MLA_SPARSE_WORKSPACE_BUFFER_SIZE,
             dtype=torch.uint8,
             device=device,
         )
-    return _fi_sparse_workspace
+        _fi_sparse_workspace_by_device[device] = workspace
+    return workspace
 
 
 class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]):

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

@@ -16,6 +16,7 @@
     _get_workspace_buffer,
 )
 from vllm.v1.attention.backends.mla.sparse_utils import (
+    triton_convert_dcp_global_index_to_local_index,
     triton_convert_req_index_to_global_index,
 )
 
@@ -32,6 +33,8 @@ def _kv_scale_format_for_model(model_type: str | None) -> str:
 class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]):
     """SM120 FlashInfer sparse-MLA implementation."""
 
+    can_return_lse_for_decode: bool = True
+
     def __init__(
         self,
         num_heads: int,
@@ -100,7 +103,38 @@ def __init__(
         assert self.topk_indices_buffer is not None
 
         self.supports_quant_query_input = False
-        self._workspace_buffer: torch.Tensor | None = None
+        self.supports_dcp_quant_query_input = False
+
+        parallel_config = vllm_config.parallel_config
+        self.pcp_world_size = 1
+        self.pcp_rank = 0
+        self.dcp_world_size = int(parallel_config.decode_context_parallel_size)
+        self.dcp_rank = 0
+        if self.dcp_world_size > 1:
+            from vllm.distributed.parallel_state import get_dcp_group
+
+            dcp_group = get_dcp_group()
+            self.dcp_rank = dcp_group.rank_in_group
+            if dcp_group.world_size != self.dcp_world_size:
+                raise RuntimeError(
+                    "FLASHINFER_MLA_SPARSE_SM120 DCP group size "
+                    f"{dcp_group.world_size} does not match configured "
+                    f"decode_context_parallel_size={self.dcp_world_size}"
+                )
+        self.total_cp_world_size = self.pcp_world_size * self.dcp_world_size
+        self.total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank
+        self.cp_kv_cache_interleave_size = (
+            parallel_config.cp_kv_cache_interleave_size
+        )
+        self.need_to_return_lse_for_decode = self.dcp_world_size > 1
+
+        # Allocate before memory profiling so KV sizing accounts for FlashInfer's
+        # TRTLLM sparse MLA scratch instead of discovering it at first decode.
+        if torch.cuda.is_available():
+            device = torch.device(f"cuda:{torch.cuda.current_device()}")
+            self._workspace_buffer: torch.Tensor | None = _get_workspace_buffer(device)
+        else:
+            self._workspace_buffer = None
 
     def forward_mqa(
         self,
@@ -113,25 +147,50 @@ def forward_mqa(
             q = torch.cat(q, dim=-1)
 
         num_actual_toks = q.shape[0]
+        num_actual_heads = q.shape[1]
 
         assert self.topk_indices_buffer is not None
         topk_indices = self.topk_indices_buffer[:num_actual_toks]
 
-        topk_indices_physical = cast(
-            torch.Tensor,
-            triton_convert_req_index_to_global_index(
-                attn_metadata.req_id_per_token[:num_actual_toks],
-                attn_metadata.block_table,
-                topk_indices,
-                BLOCK_SIZE=attn_metadata.block_size,
-                NUM_TOPK_TOKENS=topk_indices.shape[1],
-            ),
-        )
+        if self.dcp_world_size > 1:
+            seq_lens = torch.empty(
+                num_actual_toks, dtype=torch.int32, device=q.device
+            )
+            topk_indices_physical, seq_lens = (
+                triton_convert_dcp_global_index_to_local_index(
+                    attn_metadata.req_id_per_token[:num_actual_toks],
+                    attn_metadata.block_table,
+                    topk_indices,
+                    dcp_world_size=self.dcp_world_size,
+                    dcp_rank=self.dcp_rank,
+                    cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+                    BLOCK_SIZE=attn_metadata.block_size,
+                    NUM_TOPK_TOKENS=topk_indices.shape[1],
+                    valid_counts=seq_lens,
+                )
+            )
+        else:
+            topk_indices_physical = cast(
+                torch.Tensor,
+                triton_convert_req_index_to_global_index(
+                    attn_metadata.req_id_per_token[:num_actual_toks],
+                    attn_metadata.block_table,
+                    topk_indices,
+                    BLOCK_SIZE=attn_metadata.block_size,
+                    NUM_TOPK_TOKENS=topk_indices.shape[1],
+                ),
+            )
+            seq_lens = None
 
         output = q.new_empty(
-            (num_actual_toks, self.num_heads, self.kv_lora_rank),
+            (num_actual_toks, num_actual_heads, self.kv_lora_rank),
             dtype=q.dtype,
         )
+        lse = (
+            q.new_empty((num_actual_toks, num_actual_heads), dtype=torch.float32)
+            if self.need_to_return_lse_for_decode
+            else None
+        )
 
         if self._workspace_buffer is None:
             self._workspace_buffer = _get_workspace_buffer(q.device)
@@ -140,20 +199,30 @@ def forward_mqa(
             flashinfer_trtllm_batch_decode_with_kv_cache_mla,
         )
 
-        out = flashinfer_trtllm_batch_decode_with_kv_cache_mla(
+        ret = flashinfer_trtllm_batch_decode_with_kv_cache_mla(
             query=q.unsqueeze(1),
             kv_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(1),
             workspace_buffer=self._workspace_buffer,
             qk_nope_head_dim=self.qk_nope_head_dim,
             kv_lora_rank=self.kv_lora_rank,
             qk_rope_head_dim=self.qk_rope_head_dim,
             block_tables=topk_indices_physical.unsqueeze(1),
-            seq_lens=None,
+            seq_lens=seq_lens,
             max_seq_len=attn_metadata.topk_tokens,
             out=output.unsqueeze(1),
             bmm1_scale=self.scale,
             bmm2_scale=1.0,
             sparse_mla_top_k=attn_metadata.topk_tokens,
             kv_scale_format=self.kv_scale_format,
+            lse=None if lse is None else lse.unsqueeze(1),
+            return_lse=self.need_to_return_lse_for_decode,
         )
-        return out.squeeze(1), None
+        if not self.need_to_return_lse_for_decode:
+            return ret.squeeze(1), None
+        if isinstance(ret, tuple):
+            out, lse = ret
+        else:
+            out = ret
+            assert lse is not None
+        lse = lse.reshape(num_actual_toks, -1)[:, :num_actual_heads].contiguous()
+        return out.squeeze(1), lse

vllm/v1/worker/gpu/warmup.py

@@ -45,15 +45,34 @@ def run_mixed_prefill_decode_warmup(
 
     decode_req_id = f"{req_id_prefix}_decode_"
     prefill_req_id = f"{req_id_prefix}_prefill_"
-    decode_prompt_len = 2
     decode_scheduled_tokens = 1
-    prefill_len = num_tokens - decode_scheduled_tokens
-    decode_token_ids = list(range(decode_prompt_len))
-    prefill_token_ids = list(range(prefill_len))
 
     kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
     num_kv_cache_groups = len(kv_cache_groups)
     group_block_sizes = [g.kv_cache_spec.block_size for g in kv_cache_groups]
+
+    # Under DCP, the sparse decode kernels expect every DCP rank to have a
+    # non-empty local KV range. A two-token synthetic decode prompt can occupy
+    # only the first cache block, leaving later DCP ranks with zero local slots
+    # during mixed prefill+decode autotune. Span one block per DCP rank so the
+    # warmup exercises a valid DCP decode shape.
+    dcp_size = max(1, getattr(model_runner, "dcp_size", 1))
+    cp_interleave = max(1, getattr(model_runner, "cp_interleave", 1))
+    min_dcp_decode_prompt_len = max(group_block_sizes) * dcp_size * cp_interleave
+    decode_prompt_len = max(2, min_dcp_decode_prompt_len)
+    if decode_prompt_len > model_runner.scheduler_config.max_num_batched_tokens:
+        logger.warning(
+            "Skipping V2 mixed prefill+decode warmup because DCP decode prompt "
+            "length %d exceeds max_num_batched_tokens=%d.",
+            decode_prompt_len,
+            model_runner.scheduler_config.max_num_batched_tokens,
+        )
+        return False
+
+    prefill_len = num_tokens - decode_scheduled_tokens
+    decode_token_ids = list(range(decode_prompt_len))
+    prefill_token_ids = list(range(prefill_len))
+
     decode_prefill_block_counts = [
         cdiv(decode_prompt_len, block_size) for block_size in group_block_sizes
     ]