# SPDX-License-Identifier: Apache-2.0 # 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 ( compute_fp8_einsum_recipe, deep_gemm_fp8_o_proj, ) from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention from vllm.models.deepseek_v4_1.common.ops import ( build_flashinfer_mixed_sparse_indices, compute_global_topk_indices_and_lens, ) from vllm.models.deepseek_v4_1.sparse_mla import ( DeepseekV4FlashMLAMetadata, DeepseekV4SparseMLABackend, DeepseekV4SparseMLAMetadataBuilder, DeepseekV41SparseSWAMetadataBuilder, ) from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4 from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf from vllm.v1.attention.backends.mla.compressor_utils import ( 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] = {} def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor: workspace = _flashinfer_dsv4_workspace_by_device.get(device) if workspace is None: workspace = torch.zeros( _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device=device, ) _flashinfer_dsv4_workspace_by_device[device] = workspace return workspace def _packed_block_span(pool: torch.Tensor) -> int: """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): == block_size for unpacked KV, larger when packed (#44577). Raises if not token-aligned.""" block_stride = pool.stride(0) token_stride = pool.stride(-2) if block_stride % token_stride != 0: raise NotImplementedError( "FLASHINFER_MLA_SPARSE_DSV4 packed KV requires the per-block stride " f"({block_stride}) to be a multiple of the per-token stride " f"({token_stride}); this layout is not supported yet." ) return block_stride // token_stride # Sparse MLA h_q counts accepted natively (flashinfer>=0.6.14, #3545). _SPARSE_MLA_SUPPORTED_Q_HEADS = (8, 16, 32, 64, 128) def _pad_to_supported_q_heads(num_heads: int) -> int: for supported in _SPARSE_MLA_SUPPORTED_Q_HEADS: if num_heads <= supported: return supported raise ValueError( f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads " "(sparse MLA kernel requires h_q in {8, 16, 32, 64, 128})." ) def _required_sm120_sparse_topk(vllm_config: VllmConfig, window_size: int) -> int: """Return the SM120 DSV4 SWA specialization needed by this model.""" if not vllm_config.attention_config.use_non_causal: return window_size speculative_config = vllm_config.speculative_config if speculative_config is None: return window_size return get_dspark_swa_index_width( window_size, speculative_config.num_speculative_tokens, ) class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend): """FlashInfer backend using the DSv4 sparse metadata/cache layout. Inherits the base and backend reuses its``DeepseekV4SparseMLAMetadataBuilder`` """ supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ "auto", "bfloat16", "fp8", "fp8_e4m3", "fp8_ds_mla", ] @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: # Tech2Wild/Kai 2026-09-10 (SM12x): 64-state compressed pages, so a # ratio-1 kv-source layer needs 64-token pages (128 stays first so the # cache-config default remains 128 for ratio-2 layers). if current_platform.is_cuda() and current_platform.is_device_capability_family(120): return [128, 64] return [128] @classmethod def get_compressed_block_size(cls, compress_ratio: int) -> int | None: """Tech2Wild/Kai 2026-09-10: tokens per compressed-KV page for a kv-source layer on SM12x (FlashInfer SM120 sparse-MLA prefill/decode kernels are instantiated for 64-state pages); None = use cache_config.block_size.""" if compress_ratio > 0 and current_platform.is_cuda() and current_platform.is_device_capability_family(120): return 64 * compress_ratio return None @staticmethod def get_name() -> str: return "FLASHINFER_MLA_SPARSE_DSV41" @classmethod def get_supported_head_sizes(cls) -> list[int]: return [512] @classmethod def supports_sink(cls) -> bool: return True @classmethod def is_sparse(cls) -> bool: return True @classmethod def supports_compute_capability(cls, capability: DeviceCapability) -> bool: return capability.major in [10, 12] @classmethod def supports_combination( cls, head_size: int, dtype: torch.dtype, kv_cache_dtype: CacheDType | None, block_size: int | None, use_mla: bool, has_sink: bool, use_sparse: bool, use_mm_prefix: bool, device_capability: DeviceCapability, ) -> str | None: if device_capability.major == 10: if kv_cache_dtype == "fp8_ds_mla": return ( "FLASHINFER_MLA_SPARSE_DSV4 SM10x uses the plain " "per-tensor FP8 KV layout, not fp8_ds_mla" ) if kv_cache_dtype not in (None, "auto", "bfloat16", "fp8", "fp8_e4m3"): return "kv_cache_dtype not supported" return None if device_capability.major == 12: if kv_cache_dtype not in ("fp8", "fp8_e4m3", "fp8_ds_mla"): return "kv_cache_dtype not supported" from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 if not has_flashinfer_sparse_mla_sm120(): return ( "FLASHINFER_MLA_SPARSE_DSV4 SM120 requires FlashInfer's " "sparse MLA decode API" ) return None return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x" @staticmethod def get_builder_cls() -> type["DeepseekV4FlashInferSparseMLAMetadataBuilder"]: return DeepseekV4FlashInferSparseMLAMetadataBuilder class DeepseekV4FlashInferSparseMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): """Varlen-capable metadata builder for the FlashInfer sparse MLA backend.""" _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).""" _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend): @staticmethod def get_builder_cls() -> type[DeepseekSparseSWAFlashInferMetadataBuilder]: return DeepseekSparseSWAFlashInferMetadataBuilder class DeepseekSparseSWAFlashInferSM120Backend(DeepseekSparseSWAFlashInferBackend): """Tech2Wild/Kai 2026-09-10: SM12x. FlashInfer's SM120 sparse-MLA decode and prefill kernels are instantiated only for 64-token pages (_DECODE_DSV4_PAGE_BLOCK_SIZE = 64); the base backend's "any multiple of 32" made vLLM build the SWA cache at 32-token pages, which the SM120 path rejects ("no decode kernel for this shape ... page_block_size=32").""" @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [64] @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: return 64 @classmethod def get_swa_block_size(cls) -> int: return 64 class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4.""" backend_cls = DeepseekV4FlashInferMLASparseBackend swa_backend_cls = DeepseekSparseSWAFlashInferBackend use_fp8_ds_mla_layout: ClassVar[bool] = False @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: return _pad_to_supported_q_heads(num_heads) def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: return deep_gemm_fp8_o_proj( o, positions, self.rotary_emb.cos_sin_cache, self.wo_a, self.wo_b, n_groups=self.n_local_groups, heads_per_group=self.n_local_heads // self.n_local_groups, nope_dim=self.nope_head_dim, rope_dim=self.rope_head_dim, o_lora_rank=self.o_lora_rank, einsum_recipe=self._einsum_recipe, tma_aligned_scales=self._tma_aligned_scales, ) def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe( self._o_proj_block_size ) # Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``. if self.kv_cache_torch_dtype != torch.float8_e4m3fn: return fp8_q_scale = 1.0 fp8_kv_scale = 1.0 self.register_buffer( "_flashinfer_fp8_q_scale", torch.tensor([fp8_q_scale], dtype=torch.float32), persistent=False, ) self.register_buffer( "_flashinfer_fp8_q_scale_inv", torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32), persistent=False, ) self.register_buffer( "_flashinfer_fp8_kv_scale", torch.tensor([fp8_kv_scale], dtype=torch.float32), persistent=False, ) # TRTLLM-gen takes scalar scale args on a distinct C++ path vs # one-element tensors, so these are Python floats. self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale self._flashinfer_fp8_bmm2_scale = fp8_kv_scale def forward_mqa( self, q: torch.Tensor, kv: torch.Tensor, positions: torch.Tensor, output: torch.Tensor, ) -> None: # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output buffer # is allocated at the padded head count while q arrives at the local # head count; _forward pads q to match before the launcher. assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], ( f"output buffer shape {output.shape} incompatible with q shape {q.shape}" ) assert output.shape[1] >= q.shape[1], ( f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}" ) # Per-tensor FP8 q produces a bf16 attention output. expected_output_dtype = ( torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype ) assert output.dtype == expected_output_dtype, ( f"output dtype {output.dtype} must match expected {expected_output_dtype} " f"for q dtype {q.dtype}" ) forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata if attn_metadata is None: # Warmup dummy run: FlashInfer reads the cache directly and lazily # allocates its workspace, so nothing to reserve here. output.zero_() return assert isinstance(attn_metadata, dict) # Compressed-cache metadata lives on the kv-source layer's prefix; # consumers share that cache and its block table. flashmla_metadata = cast( DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.compressed_cache_prefix) if self.compressed_cache_prefix is not None else None, ) swa_metadata = cast( "DeepseekSparseSWAMetadata | None", attn_metadata.get(self.swa_cache_layer.prefix), ) assert swa_metadata is not None swa_only = self.compress_ratio == 0 # SWA-only layers have no compressed KV cache; consumers read the kv # source's cache. self_kv_cache = None if swa_only else self._compressed_kv_cache() swa_kv_cache = self.swa_cache_layer.kv_cache self._forward( q=q, kv_cache=self_kv_cache, swa_k_cache=swa_kv_cache, swa_metadata=swa_metadata, attn_metadata=flashmla_metadata, swa_only=swa_only, output=output, ) def _build_sparse_index_metadata( self, kv_cache: torch.Tensor | None, swa_k_cache: torch.Tensor, swa_metadata: "DeepseekSparseSWAMetadata", attn_metadata: DeepseekV4FlashMLAMetadata | None, swa_only: bool, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Build the combined sparse-index tensors for the mixed batch. Returns ``(compressed_kv_cache, seq_lens, sparse_indices, sparse_topk_lens)``. """ num_decodes = swa_metadata.num_decodes num_prefills = swa_metadata.num_prefills num_decode_tokens = swa_metadata.num_decode_tokens num_prefill_tokens = swa_metadata.num_prefill_tokens num_reqs = num_decodes + num_prefills num_tokens = num_decode_tokens + num_prefill_tokens assert swa_metadata.seq_lens is not None assert swa_metadata.query_start_loc is not None assert swa_metadata.token_to_req_indices is not None assert swa_metadata.decode_swa_indices is not None assert swa_metadata.block_table is not None decode_swa_indices = swa_metadata.decode_swa_indices.reshape( num_decode_tokens, swa_metadata.decode_swa_width ) decode_compressed_topk_lens = None decode_compressed_indices_are_local = False decode_is_valid_token = None if swa_only: assert self.topk_indices_buffer is not None compressed_kv_cache = swa_k_cache decode_compressed_indices = None prefill_topk_indices = self.topk_indices_buffer[ num_decode_tokens:num_tokens, :0 ] compressed_block_table = None compressed_block_size = swa_metadata.block_size top_k = 0 else: assert kv_cache is not None assert attn_metadata is not None assert self.topk_indices_buffer is not None assert swa_metadata.is_valid_token is not None compressed_kv_cache = kv_cache compressed_block_table = attn_metadata.block_table[:num_reqs] compressed_block_size = attn_metadata.block_size // self.compress_ratio # Local indices filled by the index-source layer's indexer. if num_prefill_tokens > 0: prefill_topk_indices = self.topk_indices_buffer[ num_decode_tokens:num_tokens ] top_k = prefill_topk_indices.shape[-1] else: prefill_topk_indices = self.topk_indices_buffer[:0, :0] top_k = 0 decode_compressed_indices_are_local = True decode_is_valid_token = swa_metadata.is_valid_token[:num_decode_tokens] if num_decode_tokens > 0: decode_compressed_indices = self.topk_indices_buffer[:num_decode_tokens] else: # Keep the logical width aligned with the mixed-batch case so # pure-prefill steps reuse the same Triton specialization. decode_compressed_indices = prefill_topk_indices[:0] query_start_loc = swa_metadata.query_start_loc[: num_reqs + 1] seq_lens = swa_metadata.seq_lens[:num_reqs] assert seq_lens.dtype == torch.int32 # SWA-only layers all build the same mixed sparse indices, so the first # one caches them for the step; indexer layers depend on their own topk # indices and stay uncached. cached_sparse = ( swa_metadata.flashinfer_sparse_index_cache.get("swa_only") if swa_only else None ) if cached_sparse is None: swa_block_span = _packed_block_span(swa_k_cache) compressed_block_span = _packed_block_span(compressed_kv_cache) sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices( decode_swa_indices, decode_compressed_indices, decode_compressed_topk_lens, prefill_topk_indices[:num_prefill_tokens], query_start_loc, seq_lens, swa_metadata.token_to_req_indices[:num_tokens], swa_metadata.block_table[:num_reqs], swa_metadata.block_size, compressed_block_table, compressed_block_size, self.window_size, self.compress_ratio, top_k, decode_compressed_indices_are_local=decode_compressed_indices_are_local, decode_is_valid_token=decode_is_valid_token, swa_block_span=swa_block_span, compressed_block_span=compressed_block_span, prefill_left_visible=swa_metadata.prefill_left_visible, prefill_right_visible=swa_metadata.prefill_right_visible, # getattr for tests that bypass __init__ via object.__new__. max_image_tokens=getattr(self, "max_image_tokens", 0), ) if swa_only: swa_metadata.flashinfer_sparse_index_cache["swa_only"] = ( sparse_indices, sparse_topk_lens, ) else: sparse_indices, sparse_topk_lens = cached_sparse return compressed_kv_cache, seq_lens, sparse_indices, sparse_topk_lens def _forward( self, q: torch.Tensor, kv_cache: torch.Tensor | None, swa_k_cache: torch.Tensor, swa_metadata: "DeepseekSparseSWAMetadata", attn_metadata: DeepseekV4FlashMLAMetadata | None, swa_only: bool, output: torch.Tensor, ) -> None: assert self.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn) num_decodes = swa_metadata.num_decodes num_prefills = swa_metadata.num_prefills num_decode_tokens = swa_metadata.num_decode_tokens num_prefill_tokens = swa_metadata.num_prefill_tokens num_reqs = num_decodes + num_prefills num_tokens = num_decode_tokens + num_prefill_tokens if num_tokens == 0: return ( compressed_kv_cache, seq_lens, sparse_indices, sparse_topk_lens, ) = self._build_sparse_index_metadata( kv_cache=kv_cache, swa_k_cache=swa_k_cache, swa_metadata=swa_metadata, attn_metadata=attn_metadata, swa_only=swa_only, ) # CUDA graph execution can pad q/output past the scheduled token count; # restrict to the real tokens (the launcher validates sparse indices). query = q[:num_tokens] output = output[:num_tokens] bmm1_scale: float | torch.Tensor = self.scale bmm2_scale: float | torch.Tensor = 1.0 if self.kv_cache_torch_dtype == torch.float8_e4m3fn: assert query.dtype == torch.float8_e4m3fn bmm1_scale = self._flashinfer_fp8_bmm1_scale bmm2_scale = self._flashinfer_fp8_bmm2_scale else: assert query.dtype == torch.bfloat16 query = query.contiguous() # The TRTLLM-gen sparse-MLA kernel requires h_q in {64, 128}; zero-pad # the query heads to the allocated output head count. Padded heads attend # to the shared KV and are sliced off downstream (output is padded too). padded_heads = output.shape[1] if query.shape[1] < padded_heads: padded_query = query.new_zeros( (query.shape[0], padded_heads, query.shape[2]) ) padded_query[:, : query.shape[1], :] = query query = padded_query workspace = _get_flashinfer_dsv4_workspace(q.device) query_start_loc = swa_metadata.query_start_loc query_start_loc_cpu = swa_metadata.query_start_loc_cpu assert query_start_loc is not None and query_start_loc_cpu is not None # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for # uniform-q batches, and this avoids flattening mixed batches into one call. if num_decode_tokens > 0: decode_cu = query_start_loc[: num_decodes + 1] flashinfer_trtllm_batch_decode_sparse_mla_dsv4( query=query[:num_decode_tokens], swa_kv_cache=swa_k_cache, workspace_buffer=workspace, sparse_indices=sparse_indices[:num_decode_tokens], compressed_kv_cache=compressed_kv_cache, sparse_topk_lens=sparse_topk_lens[:num_decode_tokens], seq_lens=seq_lens[:num_decodes], out=output[:num_decode_tokens], bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, sinks=self.attn_sink, cum_seq_lens_q=decode_cu, max_q_len=swa_metadata.max_decode_query_len, ) if num_prefill_tokens > 0: # The prefill query view re-anchors at offset 0, so rebase the # cumulative query offsets to start at 0. prefill_cu = ( query_start_loc[num_decodes : num_reqs + 1] - query_start_loc[num_decodes] ) prefill_cu_cpu = query_start_loc_cpu[num_decodes : num_reqs + 1] prefill_lens_cpu = prefill_cu_cpu[1:] - prefill_cu_cpu[:-1] flashinfer_trtllm_batch_decode_sparse_mla_dsv4( query=query[num_decode_tokens:num_tokens], swa_kv_cache=swa_k_cache, workspace_buffer=workspace, sparse_indices=sparse_indices[num_decode_tokens:num_tokens], compressed_kv_cache=compressed_kv_cache, sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens], seq_lens=seq_lens[num_decodes:num_reqs], out=output[num_decode_tokens:num_tokens], bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, sinks=self.attn_sink, cum_seq_lens_q=prefill_cu, max_q_len=int(prefill_lens_cpu.max().item()), ) class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention): """DeepSeek V4 sparse MLA attention through FlashInfer's SM120 kernels.""" 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: return _get_flashinfer_dsv4_workspace(device) @staticmethod def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor: if kv_cache.dtype == torch.float8_e4m3fn: kv_cache = kv_cache.view(torch.uint8) if kv_cache.dim() == 4: return kv_cache return kv_cache.unsqueeze(-2) @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: return _pad_to_supported_q_heads(num_heads) def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: return deep_gemm_fp8_o_proj( o, positions, self.rotary_emb.cos_sin_cache, self.wo_a, self.wo_b, n_groups=self.n_local_groups, heads_per_group=self.n_local_heads // self.n_local_groups, nope_dim=self.nope_head_dim, rope_dim=self.rope_head_dim, o_lora_rank=self.o_lora_rank, einsum_recipe=self._einsum_recipe, tma_aligned_scales=self._tma_aligned_scales, ) def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None: super().__init__(vllm_config, *args, **kwargs) from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120_config required_topk = _required_sm120_sparse_topk(vllm_config, self.window_size) # 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={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 ) # Per-tensor FP8 cache path scales. if self.kv_cache_torch_dtype != torch.float8_e4m3fn: return fp8_q_scale = 1.0 fp8_kv_scale = 1.0 self.register_buffer( "_flashinfer_fp8_q_scale", torch.tensor([fp8_q_scale], dtype=torch.float32), persistent=False, ) self.register_buffer( "_flashinfer_fp8_q_scale_inv", torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32), persistent=False, ) self.register_buffer( "_flashinfer_fp8_kv_scale", torch.tensor([fp8_kv_scale], dtype=torch.float32), persistent=False, ) # FlashInfer expects scalar scale arguments for this path. self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale self._flashinfer_fp8_bmm2_scale = fp8_kv_scale def _reserve_empty_forward_workspace(self) -> None: self._get_workspace( torch.device("cuda", torch.accelerator.current_device_index()) ) def _forward_sparse_impl( self, q: torch.Tensor, output: torch.Tensor, flashmla_metadata: DeepseekV4FlashMLAMetadata | None, swa_metadata: "DeepseekSparseSWAMetadata", self_kv_cache: torch.Tensor | None, swa_kv_cache: torch.Tensor, swa_only: bool, ) -> None: num_decode_tokens = swa_metadata.num_decode_tokens if swa_metadata.num_prefills > 0: self._forward_prefill( q=q[num_decode_tokens:], compressed_k_cache=self_kv_cache, swa_k_cache=swa_kv_cache, output=output[num_decode_tokens:], attn_metadata=flashmla_metadata, swa_metadata=swa_metadata, ) if swa_metadata.num_decodes > 0: self._forward_decode( q=q[:num_decode_tokens], kv_cache=self_kv_cache, swa_metadata=swa_metadata, attn_metadata=flashmla_metadata, swa_only=swa_only, output=output[:num_decode_tokens], ) def forward_mqa( self, q: torch.Tensor, kv: torch.Tensor, positions: torch.Tensor, output: torch.Tensor, ) -> None: # Output may be padded to backend-supported head counts. assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], ( f"output buffer shape {output.shape} incompatible with q shape {q.shape}" ) assert output.shape[1] >= q.shape[1], ( f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}" ) # Per-tensor FP8 q produces a bf16 attention output. expected_output_dtype = ( torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype ) assert output.dtype == expected_output_dtype, ( f"output dtype {output.dtype} must match expected {expected_output_dtype} " f"for q dtype {q.dtype}" ) forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata if attn_metadata is None: self._reserve_empty_forward_workspace() output.zero_() return assert isinstance(attn_metadata, dict) # Compressed-cache metadata lives on the kv-source layer's prefix; # consumers share that cache and its block table. flashmla_metadata = cast( DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.compressed_cache_prefix) if self.compressed_cache_prefix is not None else None, ) swa_metadata = cast( "DeepseekSparseSWAMetadata | None", attn_metadata.get(self.swa_cache_layer.prefix), ) assert swa_metadata is not None swa_only = self.compress_ratio == 0 # SWA-only layers have no compressed KV cache; consumers read the kv # source's cache. self_kv_cache = None if swa_only else self._compressed_kv_cache() swa_kv_cache = self.swa_cache_layer.kv_cache self._forward_sparse_impl( q=q, output=output, flashmla_metadata=flashmla_metadata, swa_metadata=swa_metadata, self_kv_cache=self_kv_cache, swa_kv_cache=swa_kv_cache, swa_only=swa_only, ) def _prepare_query(self, q: torch.Tensor, output: torch.Tensor) -> torch.Tensor: if self.kv_cache_torch_dtype == torch.float8_e4m3fn: assert q.dtype == torch.float8_e4m3fn q = q.to(torch.bfloat16) else: assert q.dtype == torch.bfloat16 padded_heads = output.shape[1] if q.shape[1] < padded_heads: padded_query = q.new_zeros((q.shape[0], padded_heads, q.shape[2])) padded_query[:, : q.shape[1], :] = q 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, kv_cache: torch.Tensor | None, swa_metadata: "DeepseekSparseSWAMetadata", attn_metadata: DeepseekV4FlashMLAMetadata | None, swa_only: bool, output: torch.Tensor, ) -> None: num_decodes = swa_metadata.num_decodes num_decode_tokens = swa_metadata.num_decode_tokens extra_sparse_indices = None extra_sparse_lengths = None if not swa_only: if attn_metadata is None: raise RuntimeError( "Sparse MLA metadata is required for compressed layers." ) if swa_metadata.is_valid_token is None: raise RuntimeError( "SWA validity metadata is required for compressed layers." ) if self.topk_indices_buffer is None: raise RuntimeError( "Compressed-layer decode requires top-k indices from the 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 = self._compressed_topk_to_slots( self.topk_indices_buffer[:num_decode_tokens], swa_metadata.token_to_req_indices, attn_metadata.block_table[:num_decodes], block_size, is_valid, ) extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1) swa_indices = swa_metadata.decode_swa_indices swa_lens = swa_metadata.decode_swa_lens assert swa_indices is not None assert swa_lens is not None q = self._prepare_query(q, output) swa_cache = self._as_sparse_cache(self.swa_cache_layer.kv_cache) extra_cache = self._as_sparse_cache(kv_cache) if kv_cache is not None else None if extra_cache is not None and extra_sparse_indices is None: 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, workspace_buffer=self._get_workspace(q.device), sparse_indices=swa_indices, compressed_kv_cache=extra_cache, out=output, bmm1_scale=self.scale, sinks=self.attn_sink, kv_layout="NHD", swa_topk_lens=swa_lens, extra_sparse_indices=extra_sparse_indices, extra_sparse_topk_lens=extra_sparse_lengths, ) def _forward_prefill( self, q: torch.Tensor, compressed_k_cache: torch.Tensor | None, swa_k_cache: torch.Tensor, output: torch.Tensor, attn_metadata: DeepseekV4FlashMLAMetadata | None, swa_metadata: "DeepseekSparseSWAMetadata", ) -> None: swa_only = self.compress_ratio == 0 num_prefills = swa_metadata.num_prefills num_decodes = swa_metadata.num_decodes num_decode_tokens = swa_metadata.num_decode_tokens num_prefill_tokens = swa_metadata.num_prefill_tokens query_start_loc_cpu = swa_metadata.query_start_loc_cpu assert query_start_loc_cpu is not None prefill_token_base = query_start_loc_cpu[num_decodes] 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( "Compressed-layer prefill requires top-k indices from the indexer." ) if attn_metadata is None: raise RuntimeError("Compressed-layer prefill metadata is missing.") if swa_metadata.token_to_req_indices is None: raise RuntimeError( "Compressed-layer prefill request mapping is missing." ) if swa_metadata.is_valid_token is None: raise RuntimeError( "Compressed-layer prefill validity metadata is missing." ) # Local indices filled by the index-source layer's indexer. local_topk_indices = self.topk_indices_buffer[ num_decode_tokens : num_decode_tokens + num_prefill_tokens ] prefill_token_slice = slice( num_decode_tokens, num_decode_tokens + num_prefill_tokens ) block_size = attn_metadata.block_size // self.compress_ratio 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 q = self._prepare_query(q, output) swa_kv_paged = self._as_sparse_cache(swa_k_cache) if swa_only: extra_kv_paged = None else: if compressed_k_cache is None: raise RuntimeError( "Compressed sparse MLA layers require their compressed KV 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 ) // self.PREFILL_CHUNK_SIZE for chunk_idx in range(num_chunks): chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills) query_start = ( query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base ) query_end = ( query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base ) extra_sparse_indices_chunk = ( extra_sparse_indices[query_start:query_end] if extra_sparse_indices is not None else None ) extra_sparse_lengths_chunk = ( extra_sparse_lengths[query_start:query_end] if extra_sparse_lengths is not None else None ) q_chunk = q[query_start:query_end] swa_indices_chunk = swa_metadata.prefill_swa_indices[query_start:query_end] swa_lens_chunk = swa_metadata.prefill_swa_lens[query_start:query_end] if extra_kv_paged is not None and extra_sparse_indices_chunk is None: raise RuntimeError( "Compressed sparse MLA prefill requires compressed sparse indices." ) 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, )