# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """DeepSeek V4 FlashInfer sparse MLA backend.""" from typing import TYPE_CHECKING, ClassVar, cast import torch from vllm.config import VllmConfig 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.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 if TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata _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]: return [128] @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 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 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 = DeepseekSparseSWAFlashInferBackend use_fp8_ds_mla_layout: ClassVar[bool] = True @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) if not has_flashinfer_sparse_mla_sm120_config(self.padded_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}). " "Install a FlashInfer build containing " "flashinfer-ai/flashinfer#4380." ) 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() 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." ) # Local indices filled by the index-source layer's indexer. 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( 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." ) 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 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 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], ) ) 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 = 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." ) 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, )