--- baked-ours/models/deepseek_v4_1/nvidia/flashinfer_sparse.py 2026-09-10 20:13:20.000000000 +0200 +++ baked/models/deepseek_v4_1/nvidia/flashinfer_sparse.py 2026-09-11 10:13:47.000000000 +0200 @@ -2,11 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """DeepSeek V4 FlashInfer sparse MLA backend.""" +import os from typing import TYPE_CHECKING, ClassVar, cast import torch from vllm.config import VllmConfig +from vllm.distributed import get_dcp_group +from vllm.logger import init_logger from vllm.config.cache import CacheDType from vllm.forward_context import get_forward_context from vllm.models.deepseek_v4.nvidia.ops.o_proj import ( @@ -32,10 +35,94 @@ get_dspark_swa_index_width, ) from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWABackend +from vllm.v1.attention.backends.mla.sparse_utils import ( + triton_filter_and_convert_dcp_index, +) +from vllm.v1.attention.ops.dcp import ( + CPTritonContext, + cp_lse_ag_out_rs, + dcp_a2a_lse_reduce, +) if TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata +logger = init_logger(__name__) + +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( # noqa: E402 + _PREFILL_TIMING as _DSV41_PREFILL_TIMING, + _PrefillTimer as _DSV41PrefillTimer, + prefill_acc_add as _dsv41_prefill_acc_add, +) + +# Stage 9: DCP prefill attention via compressed-KV gather (no q gather / combine). +_DCP_PREFILL_GATHER = os.environ.get("DSV41_DCP_PREFILL_GATHER", "0") == "1" + +# FlashInfer's SM120 sparse-MLA kernels write this (not -inf) into the LSE of a +# row with no candidate and no sink; the DCP combine needs -inf there. +_SM120_EMPTY_LSE = -1e29 + + +def _sm120_sparse_attention_with_lse( + query: torch.Tensor, + swa_kv_cache: torch.Tensor, + workspace: torch.Tensor, + swa_indices: torch.Tensor, + swa_lens: torch.Tensor | None, + compressed_kv_cache: torch.Tensor | None, + extra_indices: torch.Tensor | None, + extra_lens: torch.Tensor | None, + out: torch.Tensor, + sm_scale: float, + sinks: torch.Tensor | None, + lse: torch.Tensor, +) -> None: + """FlashInfer's SM120 packed sparse-MLA attention, returning the LSE. + + Mirrors the SM120 branch of ``trtllm_batch_decode_sparse_mla_dsv4`` + (flashinfer 0.7.0rc1, ``flashinfer/mla/_core.py``), whose public entry + hard-codes ``return_lse=False``. ``query``/``out`` are ``[T, H, 512]`` + bf16, ``lse`` is ``[T, H]`` fp32 and receives the LSE in base 2 of the + ``sm_scale``-scaled logits with the sink folded in (``log2(2^lse + + 2^(sink*log2e))``); rows with no candidate and no sink hold ``-1e30``. + """ + from flashinfer.mla._core import ( + _check_sm120_dsv4_kv_cache_layout, + _SparseMLASegment, + _trtllm_batch_decode_sparse_mla_sm120, + ) + + num_tokens = query.shape[0] + swa_kv_cache = _check_sm120_dsv4_kv_cache_layout(swa_kv_cache, "NHD", "swa_kv_cache") + segments = [ + _SparseMLASegment(indices=swa_indices.reshape(num_tokens, -1), lengths=swa_lens) + ] + if extra_indices is not None: + assert compressed_kv_cache is not None + compressed_kv_cache = _check_sm120_dsv4_kv_cache_layout( + compressed_kv_cache, "NHD", "compressed_kv_cache" + ) + segments.append( + _SparseMLASegment( + indices=extra_indices.reshape(num_tokens, -1), + lengths=extra_lens, + kv_cache=compressed_kv_cache, + ) + ) + _trtllm_batch_decode_sparse_mla_sm120( + query=query.unsqueeze(1), + kv_cache=swa_kv_cache, + workspace_buffer=workspace, + sparse_mla_segments=segments, + out=out.unsqueeze(1), + sm_scale=float(sm_scale), + sinks=sinks, + lse=lse, + return_lse=True, + kv_scale_format="auto", + kv_cache_format="fp8", + ) + _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024 _flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {} @@ -192,6 +279,23 @@ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # DCP port staging (DCP2-PLAN.md): the sharded-KV attention path + # (per-rank sparse attention + LSE combine) lands in stage 3. Until + # DSV41_DCP_STAGE >= 3, a DCP>1 boot must fail here -- after the KV + # cache config is computed and logged, before any forward -- rather + # than run the dcp=1 kernels over a sharded cache and serve garbage. + import os + + dcp = self.vllm_config.parallel_config.decode_context_parallel_size + if dcp > 1 and int(os.environ.get("DSV41_DCP_STAGE", "3")) < 3: + raise NotImplementedError( + f"DeepSeek-V4.1 sparse MLA under DCP{dcp}: the attention combine " + "(port stage 3) is disabled by DSV41_DCP_STAGE; unset it or run " + "with --decode-context-parallel-size 1." + ) + class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder): """SWA metadata for the FlashInfer sparse decode path (varlen decode).""" @@ -580,6 +684,11 @@ backend_cls = DeepseekV4FlashInferMLASparseBackend swa_backend_cls = DeepseekSparseSWAFlashInferSM120Backend use_fp8_ds_mla_layout: ClassVar[bool] = True + # DCP: per-rank partial attention with the LSE + cross-rank combine. + can_return_lse_for_decode: ClassVar[bool] = True + # Under DCP the per-token q gather / combine has no cross-token dependency: + # bound the transient [T, 2*heads, 512] buffers inside a prefill chunk. + DCP_PREFILL_TOKEN_CHUNK: ClassVar[int] = 2048 @staticmethod def _get_workspace(device: torch.device) -> torch.Tensor: @@ -618,14 +727,56 @@ from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120_config required_topk = _required_sm120_sparse_topk(vllm_config, self.window_size) - if not has_flashinfer_sparse_mla_sm120_config(self.padded_heads, required_topk): + # DCP: the kernel sees the gathered heads [rank0 heads | rank1 heads | + # ...] in TP order; the local slice must be unpadded so the head + # reduce-scatter hands each rank exactly its own heads back. + self.n_dcp_heads = self.n_local_heads * self.dcp_world_size + if self.dcp_world_size > 1: + if self.padded_heads != self.n_local_heads: + raise NotImplementedError( + "DeepSeek-V4.1 DCP needs an unpadded local head count in " + f"{_SPARSE_MLA_SUPPORTED_Q_HEADS}; got {self.n_local_heads} " + f"(padded to {self.padded_heads})." + ) + if _pad_to_supported_q_heads(self.n_dcp_heads) != self.n_dcp_heads: + raise NotImplementedError( + f"DeepSeek-V4.1 DCP gathered head count {self.n_dcp_heads} is " + f"not a supported sparse-MLA h_q {_SPARSE_MLA_SUPPORTED_Q_HEADS}." + ) + kernel_heads = self.n_dcp_heads if self.dcp_world_size > 1 else self.padded_heads + if not has_flashinfer_sparse_mla_sm120_config(kernel_heads, required_topk): raise RuntimeError( "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires a FlashInfer " "DSV4 sparse MLA decode specialization for " - f"(num_q_heads={self.padded_heads}, top_k={required_topk}). " + f"(num_q_heads={kernel_heads}, top_k={required_topk}). " "Install a FlashInfer build containing " "flashinfer-ai/flashinfer#4380." ) + # DCP state: the group (consecutive TP ranks), the gathered sink + # (built lazily on the first forward, after weight loading) and the + # neutral SWA index rows used on ranks other than 0. + self.dcp_group = get_dcp_group() if self.dcp_world_size > 1 else None + self._dcp_sink: torch.Tensor | None = None + self._dcp_empty_swa: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + # One Triton launch context per layer for the LSE-combine kernel (a + # fresh context per call re-specializes the launch on every step). + self._cp_ctx = CPTritonContext() if self.dcp_world_size > 1 else None + # Combine flavour: "ag_rs" (LSE all-gather + head reduce-scatter, two + # collectives) or "a2a" (one all-to-all of packed output+LSE). + self._dcp_combine_a2a = os.environ.get("DSV41_DCP_COMBINE", "ag_rs") == "a2a" + # Optional per-phase timing (CUDA events, synchronises; diagnostics + # only): DSV41_DCP_TIMING=1 logs gather / kernel / combine ms per call, + # averaged over the last 100 calls, on DCP rank 0. + self._dcp_timing = os.environ.get("DSV41_DCP_TIMING", "0") == "1" + self._dcp_t_acc = [0.0, 0.0, 0.0, 0] + if self.dcp_world_size > 1 and self.compress_ratio > 0: + logger.info_once( + "DeepSeek-V4.1 SM120 attention under DCP%d: %d local heads -> %d " + "gathered, SWA window + sink on DCP rank 0, base-2 LSE combine", + self.dcp_world_size, + self.n_local_heads, + self.n_dcp_heads, + ) self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe( self._o_proj_block_size ) @@ -762,6 +913,262 @@ q = padded_query return q.contiguous() + # ---- DCP helpers ------------------------------------------------------- + + def _compressed_topk_to_slots( + self, + topk_indices: torch.Tensor, + token_to_req_indices: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + is_valid_token: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Indexer top-k (compressed ids) -> physical slots + per-row count. + + DCP1: the request-local ids map through the block table as today. + DCP>1: the ids are GLOBAL compressed ids (indexer merge / short-context + bypass); de-interleave them in state coordinates, drop the states other + ranks own and compact this rank's slots to a prefix (-1 tail, count). + ``block_size`` is the number of states per page on this rank. + """ + num_tokens = topk_indices.shape[0] + if self.dcp_world_size == 1: + return compute_global_topk_indices_and_lens( + topk_indices, + token_to_req_indices, + block_table, + block_size, + is_valid_token, + ) + slots, lens = triton_filter_and_convert_dcp_index( + token_to_req_indices[:num_tokens].contiguous(), + block_table, + topk_indices, + dcp_size=self.dcp_world_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + BLOCK_N=128, + return_valid_counts=True, + ) + # CUDA-graph padding rows attend nothing (parity with the DCP1 kernel). + lens.masked_fill_(~is_valid_token[:num_tokens], 0) + return slots, lens + + # Stage 9: one prefill step's gathered compressed KV per source cache and + # remapped top-k per index buffer, shared by every layer of that step. The + # metadata object is held so the cache resets exactly once per step. + _dcp_gather_cache: ClassVar[dict] = {"meta": None, "kv": {}, "idx": {}, "geom": None} + + def _dcp_prefill_gather( + self, + compressed_k_cache: torch.Tensor, + topk_indices: torch.Tensor, + token_to_req: torch.Tensor, + is_valid_token: torch.Tensor, + attn_metadata: DeepseekV4FlashMLAMetadata, + num_decodes: int, + num_prefills: int, + block_size: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """DCP prefill without q gather / LSE combine: gather the DCP ranks' + compressed pages of every prefill request into one rank-major paged + buffer and remap the GLOBAL top-k ids into it (see stage 9 notes). + Returns (gathered cache [W*P_total, block_size, 1, B], slots, lens).""" + cache = type(self)._dcp_gather_cache + if cache["meta"] is not attn_metadata: + cache["meta"] = attn_metadata + cache["kv"] = {} + cache["idx"] = {} + cache["geom"] = None + W = self.dcp_world_size + K = self.cp_kv_cache_interleave_size + bs = block_size + dev = topk_indices.device + if cache["geom"] is None: + assert attn_metadata.seq_lens is not None + seq_lens = attn_metadata.seq_lens[num_decodes : num_decodes + num_prefills].tolist() + # states that exist after this step (state = pos // ratio) -> virtual + # blocks of bs*W states = one page per rank; at least one page so + # the collective never carries zero bytes. + pages = [max(1, -(-(int(sl) // self.compress_ratio) // (bs * W))) for sl in seq_lens] + base = [0] + for p in pages: + base.append(base[-1] + W * p) + cache["geom"] = ( + pages, + torch.tensor(pages, device=dev, dtype=torch.int64), + torch.tensor(base[:-1], device=dev, dtype=torch.int64), + ) + pages, pages_t, base_t = cache["geom"] + key = compressed_k_cache.data_ptr() + kv = cache["kv"].get(key) + if kv is None: + with _DSV41PrefillTimer("attn_gather", int(topk_indices.shape[0])): + pool = self._as_sparse_cache(compressed_k_cache) + parts = [] + for i, p in enumerate(pages): + page_ids = attn_metadata.block_table[num_decodes + i, :p].long() + local = pool.index_select(0, page_ids).contiguous() + parts.append(self.dcp_group.all_gather(local, dim=0)) + kv = parts[0] if len(parts) == 1 else torch.cat(parts, dim=0) + cache["kv"][key] = kv + ikey = (topk_indices.data_ptr(), tuple(topk_indices.shape)) + idx = cache["idx"].get(ikey) + if idx is None: + n_req = max(len(pages) - 1, 0) + req = (token_to_req.long() - num_decodes).clamp_(0, n_req) + ids = topk_indices.long() + valid = ids >= 0 + st = ids.clamp_min(0) + span = bs * W + vblock = st // span + off = st % span + owner = (off // K) % W + loff = (off // (K * W)) * K + off % K + slot = (base_t[req][:, None] + owner * pages_t[req][:, None] + vblock) * bs + loff + slots = torch.where(valid, slot, torch.full_like(slot, -1)).to(torch.int32) + lens = valid.sum(dim=-1, dtype=torch.int32) + lens.masked_fill_(~is_valid_token[: lens.shape[0]], 0) + idx = (slots, lens) + cache["idx"][ikey] = idx + return kv, idx[0], idx[1] + + def _dcp_attn_sink(self) -> torch.Tensor: + """The sink for the gathered heads, in TP head order (collective).""" + if self._dcp_sink is None: + assert self.dcp_group is not None + local = self.attn_sink.data[: self.n_local_heads].contiguous() + self._dcp_sink = self.dcp_group.all_gather(local, dim=0) + return self._dcp_sink + + def _dcp_neutral_swa( + self, num_tokens: int, width: int, device: torch.device + ) -> tuple[torch.Tensor, torch.Tensor]: + """All -1 SWA indices and zero lengths (ranks other than 0).""" + buf = self._dcp_empty_swa.get(width) + if buf is None or buf[0].shape[0] < num_tokens: + rows = max(num_tokens, self.max_num_batched_tokens) + buf = ( + torch.full((rows, 1, width), -1, dtype=torch.int32, device=device), + torch.zeros(rows, dtype=torch.int32, device=device), + ) + self._dcp_empty_swa[width] = buf + return buf[0][:num_tokens], buf[1][:num_tokens] + + def _run_sm120_dcp( + self, + q: torch.Tensor, + swa_cache: torch.Tensor, + swa_indices: torch.Tensor, + swa_lens: torch.Tensor, + extra_cache: torch.Tensor | None, + extra_indices: torch.Tensor | None, + extra_lens: torch.Tensor | None, + output: torch.Tensor, + ) -> None: + """One DCP attention step for ``q``/``output`` of ``[T, n_local_heads, 512]``. + + Gather the local heads to ``[T, n_local_heads * W, 512]`` (TP order), + run this rank's sparse attention over its compressed-state shard with + the LSE, then combine across the DCP group and reduce-scatter the heads + back. The replicated SWA window and the sink are attended on DCP rank 0 + only, so with ``L_r`` the per-rank base-2 LSEs (sink folded into + ``L_0``) the combine ``L = log2(sum_r 2^L_r)`` and ``o = sum_r + 2^(L_r - L) o_r`` reproduce the DCP1 softmax exactly. + """ + assert self.dcp_group is not None + num_tokens = q.shape[0] + # Timing events cannot be synchronised inside a CUDA-graph capture (the + # eager break runs on the capture stream during warmup/capture). + capturing = torch.cuda.is_current_stream_capturing() + prefill_timing = _DSV41_PREFILL_TIMING and num_tokens >= 1024 and not capturing + timing = (self._dcp_timing or prefill_timing) and not capturing + if self._dcp_timing: + logger.info_once( + "DCP path active on layer %d (capturing=%s, T=%d)", + self.layer_id, + capturing, + num_tokens, + ) + if timing: + ev = [torch.cuda.Event(enable_timing=True) for _ in range(4)] + ev[0].record() + # Collective: every rank builds the gathered sink (lazily, once); only + # DCP rank 0 applies it, together with the replicated SWA window. + sink = self._dcp_attn_sink() + q_gathered = self.dcp_group.all_gather(q, dim=1) + num_heads = q_gathered.shape[1] + if timing: + ev[1].record() + out_gathered = torch.empty( + (num_tokens, num_heads, self.head_dim), dtype=torch.bfloat16, device=q.device + ) + lse = torch.empty((num_tokens, num_heads), dtype=torch.float32, device=q.device) + if self.dcp_rank != 0: + swa_indices, swa_lens = self._dcp_neutral_swa( + num_tokens, swa_indices.shape[-1], q.device + ) + _sm120_sparse_attention_with_lse( + q_gathered, + swa_cache, + self._get_workspace(q.device), + swa_indices, + swa_lens, + extra_cache, + extra_indices, + extra_lens, + out_gathered, + self.scale, + sink if self.dcp_rank == 0 else None, + lse, + ) + if timing: + ev[2].record() + # Rows with no local candidate (and no sink) carry the kernel's -1e30 + # LSE and a zero output; the combine kernel's max-subtracted exp2 makes + # that the merge identity without any masking (test_stage3_combine). + if self._dcp_combine_a2a: + combined = dcp_a2a_lse_reduce( + out_gathered, lse, self.dcp_group, is_lse_base_on_e=False + ) + else: + combined = cp_lse_ag_out_rs( + out_gathered, + lse, + self.dcp_group, + ctx=self._cp_ctx, + is_lse_base_on_e=False, + ) + output.copy_(combined) + if timing: + ev[3].record() + ev[3].synchronize() + if prefill_timing: + _dsv41_prefill_acc_add("attn_gather", ev[0].elapsed_time(ev[1])) + _dsv41_prefill_acc_add("attn_kernel", ev[1].elapsed_time(ev[2])) + _dsv41_prefill_acc_add("attn_combine", ev[2].elapsed_time(ev[3])) + acc = self._dcp_t_acc + acc[0] += ev[0].elapsed_time(ev[1]) + acc[1] += ev[1].elapsed_time(ev[2]) + acc[2] += ev[2].elapsed_time(ev[3]) + acc[3] += 1 + if self._dcp_timing and acc[3] % 100 == 0 and self.dcp_rank == 0: + logger.info( + "DCP timing layer %d (T=%d): gather %.3f ms, kernel %.3f ms, " + "combine %.3f ms per call (avg of %d)", + self.layer_id, + num_tokens, + acc[0] / acc[3], + acc[1] / acc[3], + acc[2] / acc[3], + acc[3], + ) + acc[0] = acc[1] = acc[2] = 0.0 + acc[3] = 0 + def _forward_decode( self, q: torch.Tensor, @@ -789,10 +1196,11 @@ raise RuntimeError( "Compressed-layer decode requires top-k indices from the indexer." ) - # Local indices filled by the index-source layer's indexer. + # Indices filled by the index-source layer's indexer (request-local + # at DCP1, global compressed ids under DCP). is_valid = swa_metadata.is_valid_token[:num_decode_tokens] block_size = attn_metadata.block_size // self.compress_ratio - global_indices, extra_sparse_lengths = compute_global_topk_indices_and_lens( + global_indices, extra_sparse_lengths = self._compressed_topk_to_slots( self.topk_indices_buffer[:num_decode_tokens], swa_metadata.token_to_req_indices, attn_metadata.block_table[:num_decodes], @@ -812,6 +1220,20 @@ raise RuntimeError( "Compressed sparse MLA decode requires compressed sparse indices." ) + if self.dcp_world_size > 1 and not swa_only: + # SWA-only layers keep the DCP1 path: their cache is replicated and + # every rank runs its own heads over the whole window. + self._run_sm120_dcp( + q, + swa_cache, + swa_indices, + swa_lens, + extra_cache, + extra_sparse_indices, + extra_sparse_lengths, + output, + ) + return flashinfer_trtllm_batch_decode_sparse_mla_dsv4( query=q, swa_kv_cache=swa_cache, @@ -849,6 +1271,7 @@ extra_sparse_indices: torch.Tensor | None = None extra_sparse_lengths: torch.Tensor | None = None + gathered_kv: torch.Tensor | None = None if not swa_only: if self.topk_indices_buffer is None: raise RuntimeError( @@ -872,15 +1295,30 @@ num_decode_tokens, num_decode_tokens + num_prefill_tokens ) block_size = attn_metadata.block_size // self.compress_ratio - extra_sparse_indices, extra_sparse_lengths = ( - compute_global_topk_indices_and_lens( - local_topk_indices, - swa_metadata.token_to_req_indices[prefill_token_slice], - attn_metadata.block_table, - block_size, - swa_metadata.is_valid_token[prefill_token_slice], + if self.dcp_world_size > 1 and _DCP_PREFILL_GATHER: + assert compressed_k_cache is not None + gathered_kv, extra_sparse_indices, extra_sparse_lengths = ( + self._dcp_prefill_gather( + compressed_k_cache, + local_topk_indices, + swa_metadata.token_to_req_indices[prefill_token_slice], + swa_metadata.is_valid_token[prefill_token_slice], + attn_metadata, + num_decodes, + num_prefills, + block_size, + ) + ) + else: + extra_sparse_indices, extra_sparse_lengths = ( + self._compressed_topk_to_slots( + local_topk_indices, + swa_metadata.token_to_req_indices[prefill_token_slice], + attn_metadata.block_table, + block_size, + swa_metadata.is_valid_token[prefill_token_slice], + ) ) - ) assert swa_metadata.prefill_swa_indices is not None assert swa_metadata.prefill_swa_lens is not None @@ -894,7 +1332,11 @@ raise RuntimeError( "Compressed sparse MLA layers require their compressed KV cache." ) - extra_kv_paged = self._as_sparse_cache(compressed_k_cache) + extra_kv_paged = ( + gathered_kv + if gathered_kv is not None + else self._as_sparse_cache(compressed_k_cache) + ) num_chunks = ( num_prefills + self.PREFILL_CHUNK_SIZE - 1 @@ -927,17 +1369,40 @@ raise RuntimeError( "Compressed sparse MLA prefill requires compressed sparse indices." ) - flashinfer_trtllm_batch_decode_sparse_mla_dsv4( - query=q_chunk, - swa_kv_cache=swa_kv_paged, - workspace_buffer=self._get_workspace(q.device), - sparse_indices=swa_indices_chunk, - compressed_kv_cache=extra_kv_paged, - out=output[query_start:query_end], - bmm1_scale=self.scale, - sinks=self.attn_sink, - kv_layout="NHD", - swa_topk_lens=swa_lens_chunk, - extra_sparse_indices=extra_sparse_indices_chunk, - extra_sparse_topk_lens=extra_sparse_lengths_chunk, - ) + if self.dcp_world_size > 1 and not swa_only and gathered_kv is None: + # Chunk bounds come from the CPU query_start_loc (identical on + # every rank), so the collectives below line up across ranks. + qs, qe = int(query_start), int(query_end) + for sub_start in range(qs, qe, self.DCP_PREFILL_TOKEN_CHUNK): + sub_end = min(sub_start + self.DCP_PREFILL_TOKEN_CHUNK, qe) + assert extra_sparse_indices is not None + assert extra_sparse_lengths is not None + self._run_sm120_dcp( + q[sub_start:sub_end], + swa_kv_paged, + swa_metadata.prefill_swa_indices[sub_start:sub_end], + swa_metadata.prefill_swa_lens[sub_start:sub_end], + extra_kv_paged, + extra_sparse_indices[sub_start:sub_end], + extra_sparse_lengths[sub_start:sub_end], + output[sub_start:sub_end], + ) + continue + with _DSV41PrefillTimer( + "attn_kernel", + int(query_end - query_start) if gathered_kv is not None else 0, + ): + flashinfer_trtllm_batch_decode_sparse_mla_dsv4( + query=q_chunk, + swa_kv_cache=swa_kv_paged, + workspace_buffer=self._get_workspace(q.device), + sparse_indices=swa_indices_chunk, + compressed_kv_cache=extra_kv_paged, + out=output[query_start:query_end], + bmm1_scale=self.scale, + sinks=self.attn_sink, + kv_layout="NHD", + swa_topk_lens=swa_lens_chunk, + extra_sparse_indices=extra_sparse_indices_chunk, + extra_sparse_topk_lens=extra_sparse_lengths_chunk, + )