from __future__ import annotations import logging from dataclasses import dataclass from typing import ( TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, TypeAlias, ) import torch from sglang.srt.configs.model_config import get_dsa_index_topk, is_deepseek_dsa from sglang.srt.runtime_context import get_parallel logger = logging.getLogger(__name__) from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.dsa.dequant_k_cache import dequantize_k_cache_paged from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import ( DeepseekSparseAttnBackendMTPPrecomputeMixin, PrecomputedMetadata, compute_cu_seqlens, ) from sglang.srt.layers.attention.dsa.dsa_indexer import BaseIndexerMetadata from sglang.srt.layers.attention.dsa.dsa_topk_backend import ( DSATopKBackend, TopkTransformMethod, ) from sglang.srt.layers.attention.dsa.quant_k_cache import quantize_k_cache from sglang.srt.layers.attention.dsa.transform_index import ( transform_index_page_table_decode, transform_index_page_table_prefill, ) from sglang.srt.layers.attention.dsa.utils import ( can_dsa_prefill_cp_round_robin_split, compute_dsa_seqlens, dsa_cp_round_robin_split_data, dsa_cp_round_robin_split_q_seqs, dsa_use_prefill_cp, is_dsa_enable_prefill_cp, is_dsa_prefill_cp_in_seq_split, pad_dsa_cache_seqlens, ) from sglang.srt.layers.attention.utils import ( concat_mla_absorb_q_general, mla_quantize_and_rope_for_fp8, seqlens_expand_triton, ) from sglang.srt.layers.utils.cp_utils import ( cp_all_gather_rerange_output, cp_split_and_rebuild_position, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_buffer from sglang.srt.utils import ( get_bool_env_var, is_cuda, is_gfx95_supported, is_hip, is_sm100_supported, ) # Opt-in (default off): route the fp8 sparse-MLA prefill path through the Triton # per-query flash kernel instead of TileLang. Validated on gfx950 (GLM-5.1 @ # TP4: 16 heads, d_v=512, tail=64). Reads q_nope/q_rope directly (skips the # concat). Enable with SGLANG_DSA_TRITON_PREFILL=1. Decode stays on TileLang. _DSA_TRITON_PREFILL = get_bool_env_var("SGLANG_DSA_TRITON_PREFILL") _IS_GFX95 = is_gfx95_supported() if is_cuda(): import deep_gemm if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.speculative.spec_info import SpecInput def _all_gather_dsa_trtllm_fp8_kv( forward_batch: ForwardBatch, k: torch.Tensor, k_rope: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: kv_lora_rank = k.shape[-1] qk_rope_head_dim = k_rope.shape[-1] kv_dtype = k.dtype kv = torch.cat((k, k_rope), dim=-1).view(torch.uint8) kv = cp_all_gather_rerange_output( kv, get_parallel().attn_cp_size, forward_batch, torch.cuda.current_stream(), ).view(kv_dtype) return kv.split((kv_lora_rank, qk_rope_head_dim), dim=-1) _is_hip = is_hip() if _is_hip: from sglang.srt.layers.attention.dsa.triton_kernel import get_valid_kv_indices from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype try: from aiter import ( # noqa: F401 flash_attn_varlen_func, get_mla_metadata_info_v1, get_mla_metadata_v1, mha_batch_prefill_func, paged_attention_ragged, ) from aiter.mla import mla_decode_fwd, mla_prefill_fwd # noqa: F401 except ImportError: print( "aiter is AMD specific kernel library. Please make sure aiter is installed on your AMD device." ) else: from sglang.jit_kernel.flash_attention import ( flash_attn_varlen_func, flash_attn_with_kvcache, ) def _to_2d_context_lens(seqlens_32: torch.Tensor, batch_size: int) -> torch.Tensor: # Always normalize to (N_total, 1) layout, to avoid deadlock at deep_gemm.fp8_paged_mqa_logits if seqlens_32.dim() == 2: if seqlens_32.size(1) == 1: return seqlens_32 # Fall through and re-flatten if the caller already gave us a (bs, next_n) # view — we want (N_total, 1) regardless. seqlens_32 = seqlens_32.reshape(-1) return seqlens_32.contiguous().view(-1, 1) # Reuse this workspace buffer across all DSA backend instances # Control whether to use fused metadata copy kernel for cuda graph replay (default: enabled) # Set SGLANG_USE_FUSED_METADATA_COPY=0 or false to disable _USE_FUSED_METADATA_COPY = envs.SGLANG_USE_FUSED_METADATA_COPY.get() and not _is_hip _USE_FUSED_METADATA_GENERATION = ( envs.SGLANG_DSA_USE_FUSED_METADATA_GENERATION.get() and not _is_hip ) @dataclass(frozen=True) class DSAFlashMLAMetadata: """Metadata only needed by FlashMLA""" flashmla_metadata: torch.Tensor num_splits: torch.Tensor def slice(self, sli): return DSAFlashMLAMetadata( flashmla_metadata=self.flashmla_metadata, num_splits=self.num_splits[sli], ) def copy_(self, other: DSAFlashMLAMetadata): self.flashmla_metadata.copy_(other.flashmla_metadata) self.num_splits.copy_(other.num_splits) @dataclass(frozen=True) class DSAMetadata: page_size: int # Sequence lengths for the forward batch cache_seqlens_int32: torch.Tensor # Maximum sequence length for query max_seq_len_q: int # Maximum sequence length for key max_seq_len_k: int # Cumulative sequence lengths for query cu_seqlens_q: torch.Tensor # Cumulative sequence lengths for key cu_seqlens_k: torch.Tensor # Page table, the index of KV Cache Tables/Blocks # this table is always with page_size = 1. # None for fused-decode CUDA graphs where the wide [bs, max_ctx_len] table is # never read (attention uses topk_indices, indexer uses real_page_table); the # graph then only materializes the compact real_page_table. See # `dsa_drop_wide_page_table`. page_table_1: Optional[torch.Tensor] # NOTE(dark): This will property be used in: # 1. dense decode/prefill, we use paged flash attention, need real_page_table # 2. sparse decode/prefill, indexer need real_page_table to compute the score real_page_table: torch.Tensor # DSA metadata (dsa prefill are expanded) dsa_cache_seqlens_int32: torch.Tensor # this seqlens is clipped to `topk` dsa_cu_seqlens_q: torch.Tensor # must be arange(0, len(dsa_cu_seqlens_k)) dsa_cu_seqlens_k: torch.Tensor # cumsum of `dsa_cache_seqlens_int32` dsa_extend_seq_lens_list: List[int] dsa_seqlens_expanded: torch.Tensor # expanded, unclipped `seqlens` dsa_max_seqlen_q: Literal[1] = 1 # always 1 for decode, variable for extend flashmla_metadata: Optional[DSAFlashMLAMetadata] = None # DeepGEMM schedule metadata for paged MQA logits (decode/target_verify/draft_extend only). # Precomputed once per forward batch and reused across layers. paged_mqa_schedule_metadata: Optional[torch.Tensor] = None # 2D context_lens used to build the schedule above; the indexer reuses it # as DG's `context_lens` arg so the broadcast doesn't rebuild per layer. paged_mqa_ctx_lens_2d: Optional[torch.Tensor] = None # Precomputed once per forward batch and reused across layers: the # DeepSeek-V4 top-k v2 plan (cluster-threshold metadata) for the folded # decode top-k transform. None unless SGLANG_OPT_USE_TOPK_V2 and decode. topk_v2_plan: Optional[torch.Tensor] = None # The sum of sequence lengths for key, prefill only seq_lens_sum: Optional[int] = None # The flattened 1D page table with shape (seq_lens_sum,), prefill only # this table is always with page_size = 1 page_table_1_flattened: Optional[torch.Tensor] = None # The offset of topk indices in ragged kv, prefill only # shape: (seq_lens_sum,) topk_indices_offset: Optional[torch.Tensor] = None # k_start and k_end in kv cache for each token. indexer_k_start_end: Optional[Tuple[torch.Tensor, torch.Tensor]] = None # seq lens for each batch. indexer_seq_lens_cpu: Optional[torch.Tensor] = None # seq lens for each batch. indexer_seq_lens: Optional[torch.Tensor] = None # batch index for each token. token_to_batch_idx: Optional[torch.Tensor] = None @torch.compile def _compiled_cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor: return torch.cat(tensors, dim=dim) def _cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor: """ Concatenate two tensors along the last dimension. Use this function to concatenate q_nope and q_rope or k_nope and k_rope. """ assert len(tensors) == 2 qk_nope, qk_rope = tensors assert qk_nope.ndim == 3 and qk_rope.ndim == 3 torch._dynamo.mark_dynamic(qk_nope, 0) torch._dynamo.mark_dynamic(qk_rope, 0) return _compiled_cat([qk_nope, qk_rope], dim=dim) @dataclass(frozen=True) class DSAIndexerMetadata(BaseIndexerMetadata): attn_metadata: DSAMetadata topk_transform_method: TopkTransformMethod topk_backend: DSATopKBackend = DSATopKBackend.SGL_KERNEL paged_mqa_schedule_metadata: Optional[torch.Tensor] = None paged_mqa_ctx_lens_2d: Optional[torch.Tensor] = None force_unfused_topk: bool = False def get_seqlens_int32(self) -> torch.Tensor: return self.attn_metadata.cache_seqlens_int32 def get_page_table_64(self) -> torch.Tensor: return self.attn_metadata.real_page_table def get_page_table_1(self) -> torch.Tensor: return self.attn_metadata.page_table_1 def get_seqlens_expanded(self) -> torch.Tensor: return self.attn_metadata.dsa_seqlens_expanded def get_cu_seqlens_k(self) -> torch.Tensor: return self.attn_metadata.cu_seqlens_k def get_indexer_kvcache_range(self) -> Tuple[torch.Tensor, torch.Tensor]: return self.attn_metadata.indexer_k_start_end def get_indexer_seq_len(self) -> torch.Tensor: return self.attn_metadata.indexer_seq_lens def get_indexer_seq_len_cpu(self) -> torch.Tensor: return self.attn_metadata.indexer_seq_lens_cpu def get_dsa_extend_len_cpu(self) -> List[int]: return self.attn_metadata.dsa_extend_seq_lens_list def get_token_to_batch_idx(self) -> torch.Tensor: return self.attn_metadata.token_to_batch_idx def topk_transform( self, logits: torch.Tensor, topk: int, ks: Optional[torch.Tensor] = None, cu_seqlens_q: Optional[torch.Tensor] = None, ke_offset: Optional[torch.Tensor] = None, batch_idx_list: Optional[List[int]] = None, topk_indices_offset_override: Optional[torch.Tensor] = None, ) -> torch.Tensor: if topk_indices_offset_override is not None: cu_topk_indices_offset = topk_indices_offset_override cu_seqlens_q_topk = None elif cu_seqlens_q is not None: cu_seqlens_q = cu_seqlens_q.to(torch.int32) cu_seqlens_q_topk = compute_cu_seqlens(cu_seqlens_q) cu_topk_indices_offset = torch.repeat_interleave( cu_seqlens_q_topk[:-1], cu_seqlens_q, ) else: cu_seqlens_q_topk = self.attn_metadata.cu_seqlens_q cu_topk_indices_offset = self.attn_metadata.topk_indices_offset if ke_offset is not None: seq_lens_topk = ke_offset else: seq_lens_topk = self.get_seqlens_expanded() return self.topk_backend.topk_transform( logits=logits, lengths=seq_lens_topk, topk=topk, topk_transform_method=self.topk_transform_method, attn_metadata=self.attn_metadata, cu_seqlens_q_topk=cu_seqlens_q_topk, topk_indices_offset=cu_topk_indices_offset, row_starts=ks, batch_idx_list=batch_idx_list, force_unfused_topk=self.force_unfused_topk, ) _DSA_IMPL_T: TypeAlias = Literal[ "flashmla_sparse", "flashmla_kv", "fa3", "tilelang", "trtllm" ] class DeepseekSparseAttnBackend( DeepseekSparseAttnBackendMTPPrecomputeMixin, AttentionBackend ): # Decode/verify/draft graph replay rebuilds metadata from static buffers # (page-table width) and never reads seq_lens_cpu / seq_lens_sum; opt out of # the D2H sync. The eager fallback derives lengths from GPU seq_lens. needs_cpu_seq_lens: bool = False def __init__( self, model_runner: ModelRunner, skip_prefill: bool = False, speculative_step_id=0, topk=0, speculative_num_steps=0, ): super().__init__() self.forward_metadata: DSAMetadata self.device = model_runner.device assert isinstance(model_runner.page_size, int) self.real_page_size = model_runner.page_size self.num_splits = ( 1 if model_runner.server_args.enable_deterministic_inference else 0 ) self.use_dsa = is_deepseek_dsa(model_runner.model_config.hf_config) assert self.use_dsa, "DSA backend only supports DeepSeek DSA" self.dsa_kv_cache_store_fp8 = ( model_runner.token_to_kv_pool.dsa_kv_cache_store_fp8 ) self.dsa_index_topk = get_dsa_index_topk(model_runner.model_config.hf_config) self.max_context_len = model_runner.model_config.context_len self.num_q_heads = ( model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size ) self.kv_cache_dim = model_runner.token_to_kv_pool.kv_cache_dim self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim self.kv_lora_rank = model_runner.model_config.kv_lora_rank self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim assert model_runner.req_to_token_pool is not None self.req_to_token_pool = model_runner.req_to_token_pool self.token_to_kv_pool = model_runner.token_to_kv_pool self.hisparse_coordinator = model_runner.hisparse_coordinator self.req_to_token = model_runner.req_to_token_pool.req_to_token self.use_mha: bool = False self.dsa_prefill_impl: _DSA_IMPL_T = ( model_runner.server_args.dsa_prefill_backend ) self.dsa_decode_impl: _DSA_IMPL_T = model_runner.server_args.dsa_decode_backend self.dsa_topk_backend: DSATopKBackend = DSATopKBackend( model_runner.server_args.dsa_topk_backend ) if self.num_q_heads <= 64: self.flashmla_kv_num_q_heads = 64 elif self.num_q_heads <= 128: self.flashmla_kv_num_q_heads = 128 else: # Keep original head count if it exceeds current padded variants. self.flashmla_kv_num_q_heads = self.num_q_heads self.enable_auto_select_prefill_impl = self.dsa_prefill_impl == "flashmla_auto" self._arange_buf = torch.arange(16384, device=self.device, dtype=torch.int32) if _is_hip: max_bs = model_runner.req_to_token_pool.size self.kv_indptr = torch.zeros( (max_bs + 1,), dtype=torch.int32, device=model_runner.device ) self.kv_indices = torch.zeros( max_bs * self.dsa_index_topk, dtype=torch.int32, device=self.device, ) # Aiter mla_decode_fwd supports num_heads multiples of 16 in range [16, 128]. # For models with fewer heads per GPU (e.g. GLM-5 64 heads / TP8 = 8), need to pad the heads to 16. self.need_pad_heads = self.num_q_heads < 16 self.head_repeat_factor = ( 16 // self.num_q_heads if self.num_q_heads < 16 else 1 ) self.num_head_padded = self.num_q_heads * self.head_repeat_factor self.aiter_dsa_max_split_per_batch = 64 self.aiter_dsa_metadata_capacity = 0 self.aiter_dsa_metadata_max_seqlen_q = 0 self.aiter_dsa_metadata_q_dtype = None self.aiter_dsa_metadata_kv_dtype = None self.aiter_dsa_kv_last_page_lens = None self.aiter_dsa_work_metadata = None if ( self.dsa_prefill_impl == "aiter" or self.dsa_decode_impl == "aiter" ) and model_runner.kv_cache_dtype == fp8_dtype: self._ensure_aiter_dsa_decode_metadata_buffer( max_seqlen_q=1, batch_size=max_bs, q_dtype=torch.bfloat16, kv_dtype=fp8_dtype, ) # Speculative decoding self.topk = model_runner.server_args.speculative_eagle_topk or 0 self.speculative_num_steps = speculative_num_steps self.speculative_num_draft_tokens = ( model_runner.server_args.speculative_num_draft_tokens ) self.speculative_step_id = speculative_step_id self.device_capability = torch.cuda.get_device_capability() self.device_sm_major = self.device_capability[0] self.kv_cache_dtype = model_runner.kv_cache_dtype # Allocate global workspace buffer for TRT-LLM kernels (ragged attention on SM100/B200, or trtllm decode) if self.device_sm_major >= 10 or self.dsa_decode_impl == "trtllm": self.workspace_buffer = get_buffer( "dsa_trtllm_workspace", lambda: torch.empty( envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(), dtype=torch.uint8, device=model_runner.device, ), ) else: self.workspace_buffer = None def _make_aiter_dsa_decode_metadata_buffer( self, max_seqlen_q: int, batch_size: int, q_dtype: torch.dtype, kv_dtype: torch.dtype, ): ( (work_metadata_size, work_metadata_type), (work_indptr_size, work_indptr_type), (work_info_set_size, work_info_set_type), (reduce_indptr_size, reduce_indptr_type), (reduce_final_map_size, reduce_final_map_type), (reduce_partial_map_size, reduce_partial_map_type), ) = get_mla_metadata_info_v1( batch_size, max_seqlen_q, self.num_head_padded, q_dtype, kv_dtype, is_sparse=True, fast_mode=False, num_kv_splits=self.aiter_dsa_max_split_per_batch, intra_batch_mode=True, ) return ( torch.empty( work_metadata_size, dtype=work_metadata_type, device=self.device ), torch.empty(work_indptr_size, dtype=work_indptr_type, device=self.device), torch.empty( work_info_set_size, dtype=work_info_set_type, device=self.device ), torch.empty( reduce_indptr_size, dtype=reduce_indptr_type, device=self.device ), torch.empty( reduce_final_map_size, dtype=reduce_final_map_type, device=self.device ), torch.empty( reduce_partial_map_size, dtype=reduce_partial_map_type, device=self.device, ), ) def _ensure_aiter_dsa_decode_metadata_buffer( self, max_seqlen_q: int, batch_size: int, q_dtype: torch.dtype, kv_dtype: torch.dtype, ) -> None: if ( self.aiter_dsa_work_metadata is not None and self.aiter_dsa_metadata_capacity >= batch_size and self.aiter_dsa_metadata_max_seqlen_q == max_seqlen_q and self.aiter_dsa_metadata_q_dtype == q_dtype and self.aiter_dsa_metadata_kv_dtype == kv_dtype ): return ( self.aiter_dsa_work_metadata, self.aiter_dsa_work_indptr, self.aiter_dsa_work_info_set, self.aiter_dsa_reduce_indptr, self.aiter_dsa_reduce_final_map, self.aiter_dsa_reduce_partial_map, ) = self._make_aiter_dsa_decode_metadata_buffer( max_seqlen_q=max_seqlen_q, batch_size=batch_size, q_dtype=q_dtype, kv_dtype=kv_dtype, ) self.aiter_dsa_kv_last_page_lens = torch.ones( (batch_size,), dtype=torch.int32, device=self.device ) self.aiter_dsa_metadata_capacity = batch_size self.aiter_dsa_metadata_max_seqlen_q = max_seqlen_q self.aiter_dsa_metadata_q_dtype = q_dtype self.aiter_dsa_metadata_kv_dtype = kv_dtype def _prepare_aiter_dsa_decode_metadata( self, qo_indptr: torch.Tensor, kv_indptr: torch.Tensor, bs: int, max_seqlen_q: int, q_dtype: torch.dtype, kv_dtype: torch.dtype, ) -> dict: self._ensure_aiter_dsa_decode_metadata_buffer( max_seqlen_q=max_seqlen_q, batch_size=bs, q_dtype=q_dtype, kv_dtype=kv_dtype, ) self.aiter_dsa_kv_last_page_lens[:bs].fill_(1) kv_last_page_lens = self.aiter_dsa_kv_last_page_lens[:bs] get_mla_metadata_v1( qo_indptr, kv_indptr, kv_last_page_lens, self.num_head_padded, 1, False, self.aiter_dsa_work_metadata, self.aiter_dsa_work_info_set, self.aiter_dsa_work_indptr, self.aiter_dsa_reduce_indptr, self.aiter_dsa_reduce_final_map, self.aiter_dsa_reduce_partial_map, page_size=1, kv_granularity=16, max_seqlen_qo=max_seqlen_q, uni_seqlen_qo=max_seqlen_q, fast_mode=False, topk=self.dsa_index_topk, max_split_per_batch=self.aiter_dsa_max_split_per_batch, intra_batch_mode=True, dtype_q=q_dtype, dtype_kv=kv_dtype, ) return { "kv_last_page_lens": kv_last_page_lens, "work_meta_data": self.aiter_dsa_work_metadata, "work_indptr": self.aiter_dsa_work_indptr, "work_info_set": self.aiter_dsa_work_info_set, "reduce_indptr": self.aiter_dsa_reduce_indptr, "reduce_final_map": self.aiter_dsa_reduce_final_map, "reduce_partial_map": self.aiter_dsa_reduce_partial_map, "intra_batch_mode": True, "num_kv_splits": self.aiter_dsa_max_split_per_batch, } def _build_paged_mqa_schedule_2d_ctx_lens( self, forward_mode: ForwardMode, cache_seqlens_int32: torch.Tensor, seqlens_expanded: torch.Tensor, batch_size: int, ) -> torch.Tensor: # target_verify with next_n>=2 uses DG-native q=[B,next_n,H,D] which # needs a [B, next_n] schedule; everything else stays per-token. # TODO: SM90 supports DG-native next_n in {1,2} too — enable once # validated; for now DG-native is SM100+ only. next_n = self.speculative_num_draft_tokens if ( forward_mode.is_target_verify() and next_n and next_n >= 2 and is_sm100_supported() ): return cache_seqlens_int32.view(-1, 1).expand(-1, next_n).contiguous() if forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2(): return _to_2d_context_lens(seqlens_expanded, batch_size) return _to_2d_context_lens(cache_seqlens_int32, batch_size) def _refresh_paged_mqa_schedule_metadata( self, metadata: DSAMetadata, seqlens_32_2d: torch.Tensor, ) -> None: new_schedule = deep_gemm.get_paged_mqa_logits_metadata( seqlens_32_2d, 64, deep_gemm.get_num_sms() ) if metadata.paged_mqa_schedule_metadata is None: object.__setattr__(metadata, "paged_mqa_schedule_metadata", new_schedule) else: metadata.paged_mqa_schedule_metadata.copy_(new_schedule) def _build_topk_v2_plan( self, seqlens_expanded: torch.Tensor ) -> Optional[torch.Tensor]: # Preprocess the folded top-k v2 plan once per forward (shared across # layers), at metadata-build time, from the same seqlens the transform # receives as `lengths` (dsa_seqlens_expanded). This must cover EVERY shape # that dispatches to `_topk_transform_v2_paged` -- decode AND MTP # target-verify / draft-extend, whose expanded row count is exactly what v2 # sees -- otherwise the helper's plan-present assertion fires. None only # when the fold is disabled; such metadata is never dispatched to v2. if not envs.SGLANG_OPT_USE_TOPK_V2.get(): return None from sglang.jit_kernel.dsv4.topk import plan_topk_v2 return plan_topk_v2(seqlens_expanded) def _refresh_topk_v2_plan(self, metadata: DSAMetadata) -> None: # Refresh the plan in-place under CUDA graph replay so the captured # read sees fresh cluster metadata for the replay's decode seq lengths. # `copy_` preserves the buffer's data_ptr captured by the graph. None # means it was not built (fold disabled / non-decode shape), and such a # metadata object is never dispatched to the v2 helper, so there is # nothing to refresh. if metadata.topk_v2_plan is None: return from sglang.jit_kernel.dsv4.topk import plan_topk_v2 metadata.topk_v2_plan.copy_(plan_topk_v2(metadata.dsa_seqlens_expanded)) def _get_fused_topk_page_table(self, topk_indices: torch.Tensor) -> torch.Tensor: if ( self.dsa_topk_backend.is_sgl_kernel() or self.dsa_topk_backend.is_flashinfer() ): return topk_indices raise RuntimeError( f"Unsupported {self.dsa_topk_backend = } for SGLANG_DSA_FUSE_TOPK." ) def get_device_int32_arange(self, length: int) -> torch.Tensor: if length > len(self._arange_buf): next_pow_of_2 = 1 << (length - 1).bit_length() self._arange_buf = torch.arange( next_pow_of_2, device=self.device, dtype=torch.int32 ) return self._arange_buf[:length] def _graph_page_table_width(self, metadata: DSAMetadata) -> int: """Column count to scan req_to_token during graph replay. Reads the wide page_table_1 width when present, else req_to_token's width (the wide table is dropped for fused decode graphs, see `dsa_drop_wide_page_table`).""" if metadata.page_table_1 is not None: return metadata.page_table_1.shape[1] return self.req_to_token.shape[1] def _transform_table_1_to_real(self, page_table: torch.Tensor) -> torch.Tensor: page_size = self.real_page_size if page_size == 1: return page_table max_seqlen_k = page_table.shape[1] strided_indices = torch.arange( 0, max_seqlen_k, page_size, device=page_table.device, dtype=torch.int32 ) return page_table[:, strided_indices] // page_size def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, in_capture: bool = False, ): seq_lens_cpu = ( forward_batch.seq_lens.cpu() if in_capture else forward_batch.seq_lens_cpu ) self._apply_cuda_graph_metadata( bs=forward_batch.batch_size, req_pool_indices=forward_batch.req_pool_indices, seq_lens=forward_batch.seq_lens, seq_lens_cpu=seq_lens_cpu, forward_mode=forward_batch.forward_mode, spec_info=forward_batch.spec_info, out_cache_loc=getattr(forward_batch, "out_cache_loc", None), actual_forward_mode=getattr(forward_batch, "actual_forward_mode", None), ) def init_forward_metadata(self, forward_batch: ForwardBatch): """Init the metadata for a forward pass.""" batch_size = forward_batch.batch_size device = forward_batch.seq_lens.device if forward_batch.forward_mode.is_target_verify(): draft_token_num = self.speculative_num_draft_tokens else: draft_token_num = 0 cache_seqlens_int32 = (forward_batch.seq_lens + draft_token_num).to(torch.int32) cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) if forward_batch.seq_lens_cpu is not None: max_seqlen_k = int( forward_batch.seq_lens_cpu.max().item() + draft_token_num ) else: # needs_cpu_seq_lens=False nulls the host mirror for spec-v2 relay # batches; graph replay uses the static page-table width, so only this # eager (e.g. over-capture-bs) fallback needs a length here. max_seqlen_k = int(forward_batch.seq_lens.max().item()) + draft_token_num # [b, max_seqlen_k] page_table = self.req_to_token_pool.req_to_token[ forward_batch.req_pool_indices, :max_seqlen_k ] page_table_1_flattened = None topk_indices_offset = None # Centralized dispatch: decide all strategies for this batch self.set_dsa_prefill_impl(forward_batch) dsa_impl_for_batch = ( self.dsa_decode_impl if ( forward_batch.forward_mode.is_decode_or_idle() or forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend_v2() ) else self.dsa_prefill_impl ) use_flashmla_kv = (not self.use_mha) and dsa_impl_for_batch == "flashmla_kv" topk_transform_method = self.get_topk_transform_method( forward_batch.forward_mode ) # Batch indices selected when cp enabled: After splitting multiple sequences, # a certain cp rank may not have some of these sequences. # We use bs_idx_cpu to mark which sequences are finally selected by the current cp rank, # a default value of None indicates that all sequences are selected. bs_idx_cpu = None # seq_len_cpu of selected sequences indexer_seq_lens_cpu = forward_batch.seq_lens_cpu indexer_seq_lens = forward_batch.seq_lens if forward_batch.forward_mode.is_decode_or_idle(): extend_seq_lens_cpu = [1] * batch_size max_seqlen_q = 1 cu_seqlens_q = self.get_device_int32_arange(batch_size + 1) seqlens_expanded = cache_seqlens_int32 elif forward_batch.forward_mode.is_target_verify(): max_seqlen_q = 1 cu_seqlens_q = torch.arange( 0, batch_size * self.speculative_num_draft_tokens + 1, 1, dtype=torch.int32, device=device, ) extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * batch_size forward_batch.extend_seq_lens_cpu = extend_seq_lens_cpu seqlens_expanded = seqlens_expand_triton( torch.tensor(extend_seq_lens_cpu, dtype=torch.int32, device=device), cache_seqlens_int32, self.speculative_num_draft_tokens * batch_size, self.speculative_num_draft_tokens, ) page_table = torch.repeat_interleave( page_table, repeats=self.speculative_num_draft_tokens, dim=0 ) elif forward_batch.forward_mode.is_draft_extend_v2(): if forward_batch.extend_prefix_lens_cpu is None: assert forward_batch.extend_prefix_lens is not None forward_batch.extend_prefix_lens_cpu = ( forward_batch.extend_prefix_lens.cpu().tolist() ) if forward_batch.seq_lens_cpu is None: forward_batch.seq_lens_cpu = forward_batch.seq_lens.cpu() forward_batch.seq_lens_sum = int(forward_batch.seq_lens_cpu.sum()) assert ( forward_batch.extend_seq_lens_cpu is not None and forward_batch.extend_seq_lens is not None and forward_batch.extend_prefix_lens_cpu is not None ), "All of them must not be None" extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu assert forward_batch.extend_seq_lens is not None max_seqlen_q = 1 cu_seqlens_q = torch.arange( 0, forward_batch.extend_num_tokens + 1, 1, dtype=torch.int32, device=device, ) seqlens_expanded = seqlens_expand_triton( forward_batch.extend_seq_lens, cache_seqlens_int32, sum(extend_seq_lens_cpu), self.speculative_num_draft_tokens, ) if forward_batch.forward_mode.is_draft_extend_v2(): # DRAFT_EXTEND_V2: V2 worker pre-fills draft KV cache with ALL speculated # tokens upfront. All requests extend by the same fixed # (speculative_num_draft_tokens). Use scalar to avoid GPU sync. page_table = torch.repeat_interleave( page_table, repeats=self.speculative_num_draft_tokens, dim=0 ) else: # DRAFT_EXTEND: the draft worker extends by (num_correct_drafts + 1) # per request after verification. Lengths vary per request based on # how many tokens were accepted. page_table = torch.repeat_interleave( page_table, repeats=forward_batch.extend_seq_lens, dim=0 ) elif forward_batch.forward_mode.is_extend(): assert ( forward_batch.extend_seq_lens_cpu is not None and forward_batch.extend_seq_lens is not None and forward_batch.extend_prefix_lens_cpu is not None ), "All of them must not be None" extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu assert forward_batch.extend_seq_lens is not None extend_seq_lens = forward_batch.extend_seq_lens seqlens_expanded = torch.cat( [ torch.arange( kv_len - qo_len + 1, kv_len + 1, dtype=torch.int32, device=device, ) for qo_len, kv_len in zip( forward_batch.extend_seq_lens_cpu, forward_batch.seq_lens_cpu.tolist(), strict=True, ) ] ) if can_dsa_prefill_cp_round_robin_split(forward_batch): seqlens_expanded = dsa_cp_round_robin_split_data(seqlens_expanded) extend_seq_lens_cpu, extend_seq_lens, bs_idx_cpu, bs_idx = ( dsa_cp_round_robin_split_q_seqs( extend_seq_lens_cpu, extend_seq_lens ) ) indexer_seq_lens_cpu = indexer_seq_lens_cpu[bs_idx_cpu] indexer_seq_lens = indexer_seq_lens[bs_idx] cache_seqlens_int32 = cache_seqlens_int32[bs_idx] cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) max_seqlen_k = ( int(indexer_seq_lens_cpu.max().item() + draft_token_num) if len(indexer_seq_lens_cpu) != 0 else 0 ) page_table = page_table[bs_idx, :max_seqlen_k] if any(forward_batch.extend_prefix_lens_cpu) or bs_idx_cpu is not None: max_seqlen_q = ( max(extend_seq_lens_cpu) if len(extend_seq_lens_cpu) != 0 else 1 ) cu_seqlens_q = compute_cu_seqlens(extend_seq_lens.to(torch.int32)) else: max_seqlen_q = max_seqlen_k cu_seqlens_q = cu_seqlens_k # Check if MHA FP8 dequantization is needed mha_dequantize_needed = ( self.use_mha and self.token_to_kv_pool.dtype == torch.float8_e4m3fn ) forward_batch.using_mha_one_shot_fp8_dequant = mha_dequantize_needed # page_table_1_flattened is only used when prefix sharing is enabled: has_prefix_sharing = any(forward_batch.extend_prefix_lens_cpu) if has_prefix_sharing and ( topk_transform_method == TopkTransformMethod.RAGGED or mha_dequantize_needed ): page_table_1_flattened = torch.cat( [ page_table[i, :kv_len] for i, kv_len in enumerate( indexer_seq_lens_cpu.tolist(), ) ] ) assert page_table_1_flattened.shape[0] == sum( indexer_seq_lens_cpu ), f"{page_table_1_flattened.shape[0] = } must be the same as {sum(indexer_seq_lens_cpu) = }" # Validate indices when logical tokens exceed physical capacity # This is likely to be triggered by PP with high kv reuse & parallelism kv_cache_capacity = ( self.token_to_kv_pool.size + self.token_to_kv_pool.page_size ) if forward_batch.seq_lens_sum > kv_cache_capacity: max_idx = page_table_1_flattened.max().item() assert max_idx < kv_cache_capacity, ( f"Invalid page table index: max={max_idx}, " f"kv_cache_capacity={kv_cache_capacity}" ) if topk_transform_method == TopkTransformMethod.RAGGED: topk_indices_offset = torch.repeat_interleave( cu_seqlens_k[:-1], extend_seq_lens, ) else: assert False, f"Unsupported {forward_batch.forward_mode = }" indexer_k_start_end, token_to_batch_idx = self._cal_indexer_k_start_end( forward_batch, bs_idx_cpu ) # 1D, expanded seqlens (1D means cheap to compute, so always compute it) dsa_cache_seqlens_int32 = compute_dsa_seqlens( original_seq_lens=seqlens_expanded, dsa_index_topk=self.dsa_index_topk, ) dsa_cache_seqlens_int32 = pad_dsa_cache_seqlens( forward_batch, dsa_cache_seqlens_int32 ) dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens_int32) dsa_cu_seqlens_q = self.get_device_int32_arange(len(dsa_cu_seqlens_k)) paged_mqa_schedule_metadata = None paged_mqa_ctx_lens_2d = None if is_cuda() and ( forward_batch.forward_mode.is_decode_or_idle() or forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend_v2() ): paged_mqa_ctx_lens_2d = self._build_paged_mqa_schedule_2d_ctx_lens( forward_batch.forward_mode, cache_seqlens_int32, seqlens_expanded, forward_batch.batch_size, ) # NOTE: block_kv arg must be 64 here — DG computes SPLIT_KV = # block_kv * 4 and both DG's and the indexer's compute kernels # require SPLIT_KV = 256; this is independent of the cache page size. paged_mqa_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( paged_mqa_ctx_lens_2d, 64, deep_gemm.get_num_sms() ) metadata = DSAMetadata( page_size=self.real_page_size, cache_seqlens_int32=cache_seqlens_int32, max_seq_len_q=max_seqlen_q, max_seq_len_k=max_seqlen_k, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, seq_lens_sum=forward_batch.seq_lens_sum, page_table_1=page_table, page_table_1_flattened=page_table_1_flattened, flashmla_metadata=( self._compute_flashmla_metadata( cache_seqlens=dsa_cache_seqlens_int32, seq_len_q=1, ) if use_flashmla_kv else None ), paged_mqa_schedule_metadata=paged_mqa_schedule_metadata, paged_mqa_ctx_lens_2d=paged_mqa_ctx_lens_2d, dsa_cache_seqlens_int32=dsa_cache_seqlens_int32, dsa_cu_seqlens_q=dsa_cu_seqlens_q, dsa_cu_seqlens_k=dsa_cu_seqlens_k, dsa_seqlens_expanded=seqlens_expanded, dsa_extend_seq_lens_list=extend_seq_lens_cpu, real_page_table=self._transform_table_1_to_real(page_table), dsa_max_seqlen_q=1, topk_indices_offset=topk_indices_offset, indexer_k_start_end=indexer_k_start_end, indexer_seq_lens_cpu=indexer_seq_lens_cpu, indexer_seq_lens=indexer_seq_lens, token_to_batch_idx=token_to_batch_idx, topk_v2_plan=self._build_topk_v2_plan(seqlens_expanded), ) self.forward_metadata = metadata def _cal_indexer_k_start_end( self, forward_batch: ForwardBatch, bs_idx: Optional[List[int]] = None, ): if not forward_batch.forward_mode.is_extend_without_speculative(): return None, None if forward_batch.batch_size == 0 or (bs_idx is not None and len(bs_idx) == 0): empty_t = torch.empty(0, dtype=torch.int32, device=self.device) return (empty_t, empty_t), empty_t # Suppose there are two requests, with extend_seq_len = [3, 2] # and seq_lens = [10, 4] # The logits matrix looks like this, with * representing the valid logits # and - representing the invalid logits: # # ********--|---- # *********-|---- # **********|---- # ----------|***- # ----------|**** # # ks = [0, 0, 0, 10, 10] # ke = [8, 9, 10, 13, 14] ks_list = [] ke_list = [] token_to_batch_idx = [] q_offset = 0 k_offset = 0 assert ( forward_batch.seq_lens_cpu is not None and forward_batch.extend_seq_lens_cpu is not None ) for i in range(forward_batch.batch_size): seq_len = forward_batch.seq_lens_cpu[i].item() assert isinstance(seq_len, int) extend_seq_len = forward_batch.extend_seq_lens_cpu[i] ks = torch.full( (extend_seq_len,), k_offset, dtype=torch.int32, device=self.device ) kv_len = seq_len if forward_batch.forward_mode.is_target_verify(): kv_len += self.speculative_num_draft_tokens seq_lens_expanded = torch.arange( kv_len - extend_seq_len + 1, kv_len + 1, dtype=torch.int32, device=self.device, ) ke = ks + seq_lens_expanded ks_list.append(ks) ke_list.append(ke) # bi: The index within the selected batch bs_idx. Entries that were not selected are ignored. bi = bs_idx.index(i) if (bs_idx is not None and i in bs_idx) else i tb = torch.full( (extend_seq_len,), bi, dtype=torch.int32, device=self.device ) token_to_batch_idx.append(tb) if bs_idx is None or i in bs_idx: # skip batch not included in bs_idx q_offset += extend_seq_len k_offset += seq_len ks = torch.cat(ks_list, dim=0) ke = torch.cat(ke_list, dim=0) token_to_batch_idx = torch.cat(token_to_batch_idx, dim=0) if bs_idx is not None: assert can_dsa_prefill_cp_round_robin_split(forward_batch) ks = dsa_cp_round_robin_split_data(ks) ke = dsa_cp_round_robin_split_data(ke) token_to_batch_idx = dsa_cp_round_robin_split_data(token_to_batch_idx) return (ks, ke), token_to_batch_idx def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): """Initialize CUDA graph state for the attention backend. Args: max_bs (int): Maximum batch size to support in CUDA graphs This creates fixed-size tensors that will be reused during CUDA graph replay to avoid memory allocations. """ # Whether we can skip the wide [max_num_tokens, max_ctx_len] page_size=1 # page table in the decode CUDA graph. It is dead weight there only when the # decode top-k routes to the fused v2 kernel: attention reads topk_indices # and the indexer reads the compact real_page_table, so nothing reads the # page_size=1 table. This MUST match the exact condition under which # `DSATopKBackend.topk_transform` dispatches decode PAGED to # `_topk_transform_v2_paged` -- otherwise the legacy transform would read a # dropped (None) table. Hence: fused top-k AND v2 enabled AND index_topk in # the kernel's supported range, on CUDA with page_size>1. Excludes HIP (its # indexer reads page_table_1), hisparse (needs page_size=1 loc translation), # and spec decoding (MTP precompute fast-path + target-verify/draft-extend # still consume the wide table). Computed once from stable config; the graph # is captured once per process. self.dsa_drop_wide_page_table = ( is_cuda() and not _is_hip and self.real_page_size > 1 and self.hisparse_coordinator is None and not self.speculative_num_draft_tokens and envs.SGLANG_DSA_FUSE_TOPK.get() and envs.SGLANG_OPT_USE_TOPK_V2.get() and self.dsa_index_topk is not None and self.dsa_index_topk <= 2048 ) max_ctx_len = self.req_to_token.shape[1] self.decode_cuda_graph_metadata: Dict = { "cache_seqlens": torch.ones( max_num_tokens, dtype=torch.int32, device=self.device ), "cu_seqlens_q": torch.arange( 0, max_bs + 1, dtype=torch.int32, device=self.device ), "cu_seqlens_k": torch.zeros( max_bs + 1, dtype=torch.int32, device=self.device ), # fake page_table for sparse_prefill # Match req_to_token's width exactly. It is over-allocated beyond # context_len because spec decoding lets seq_len transiently overshoot. # When dropping the wide table (fused decode), allocate only the compact # page_size=64 real table; else allocate the wide page_size=1 table and # derive real from it per batch size. "real_page_table": ( torch.zeros( max_num_tokens, (max_ctx_len + self.real_page_size - 1) // self.real_page_size, dtype=torch.int32, device=self.device, ) if self.dsa_drop_wide_page_table else None ), "page_table": ( None if self.dsa_drop_wide_page_table else torch.zeros( max_num_tokens, max_ctx_len, dtype=torch.int32, device=self.device, ) ), "flashmla_metadata": ( self._compute_flashmla_metadata( cache_seqlens=torch.ones( max_num_tokens, dtype=torch.int32, device=self.device ), seq_len_q=1, ) if self.dsa_decode_impl == "flashmla_kv" else None ), } def _build_forward_metadata_cuda_graph( self, bs: int, num_tokens: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, seq_lens_cpu: Optional[torch.Tensor], forward_mode: ForwardMode, spec_info: Optional[SpecInput], out_cache_loc: Optional[torch.Tensor] = None, actual_forward_mode: Optional[ForwardMode] = None, ): """Create and store DSAMetadata for a new batch size during CUDA graph capture.""" self.set_dsa_prefill_impl(forward_batch=None) if forward_mode.is_decode_or_idle(): # Normal Decode # Get sequence information cache_seqlens_int32 = seq_lens.to(torch.int32) cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) # Use max context length for seq_len_k real_rows = bs if self.dsa_drop_wide_page_table: page_table_1 = None max_seqlen_k = self.req_to_token.shape[1] else: page_table_1 = self.decode_cuda_graph_metadata["page_table"][:bs, :] max_seqlen_k = page_table_1.shape[1] max_seqlen_q = 1 # Precompute page table # Precompute cumulative sequence lengths # NOTE(dark): this is always arange, since we are decoding cu_seqlens_q = self.decode_cuda_graph_metadata["cu_seqlens_q"][: bs + 1] dsa_cache_seqlens_int32 = compute_dsa_seqlens( cache_seqlens_int32, dsa_index_topk=self.dsa_index_topk ) seqlens_expanded = cache_seqlens_int32 dsa_extend_seq_lens_list = [1] * bs if self.dsa_decode_impl == "flashmla_kv": flashmla_metadata = self.decode_cuda_graph_metadata[ "flashmla_metadata" ].slice(slice(0, bs + 1)) flashmla_metadata.copy_( self._compute_flashmla_metadata( cache_seqlens=dsa_cache_seqlens_int32, seq_len_q=1, ) ) else: flashmla_metadata = None elif forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2(): cache_seqlens_int32 = (seq_lens + self.speculative_num_draft_tokens).to( torch.int32 ) cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) max_seqlen_q = 1 real_rows = bs * self.speculative_num_draft_tokens if self.dsa_drop_wide_page_table: page_table_1 = None max_seqlen_k = self.req_to_token.shape[1] else: page_table_1 = self.decode_cuda_graph_metadata["page_table"][ :real_rows, : ] max_seqlen_k = page_table_1.shape[1] cu_seqlens_q = torch.arange( 0, bs * self.speculative_num_draft_tokens + 1, 1, dtype=torch.int32, device=self.device, ) extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * bs seqlens_int32_cpu = [ self.speculative_num_draft_tokens + kv_len for kv_len in seq_lens.tolist() ] seqlens_expanded = torch.cat( [ torch.arange( kv_len - qo_len + 1, kv_len + 1, dtype=torch.int32, device=self.device, ) for qo_len, kv_len in zip( extend_seq_lens_cpu, seqlens_int32_cpu, strict=True, ) ] ) dsa_cache_seqlens_int32 = compute_dsa_seqlens( seqlens_expanded, dsa_index_topk=self.dsa_index_topk ) dsa_extend_seq_lens_list = [1] * bs * self.speculative_num_draft_tokens if self.dsa_decode_impl == "flashmla_kv": flashmla_metadata = self.decode_cuda_graph_metadata[ "flashmla_metadata" ].slice(slice(0, bs * self.speculative_num_draft_tokens + 1)) flashmla_metadata.copy_( self._compute_flashmla_metadata( cache_seqlens=dsa_cache_seqlens_int32, seq_len_q=1, ) ) else: flashmla_metadata = None dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens_int32) dsa_cu_seqlens_q = self.get_device_int32_arange(len(dsa_cu_seqlens_k)) if self.dsa_drop_wide_page_table: # Compact page_size=64 static buffer; filled per-replay by the fused # metadata kernel straight from req_to_token (no wide table needed). real_page_table = self.decode_cuda_graph_metadata["real_page_table"][ :real_rows, : ] else: real_page_table = self._transform_table_1_to_real(page_table_1) paged_mqa_schedule_metadata = None paged_mqa_ctx_lens_2d = None if is_cuda() and ( forward_mode.is_decode_or_idle() or forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2() ): paged_mqa_ctx_lens_2d = self._build_paged_mqa_schedule_2d_ctx_lens( forward_mode, cache_seqlens_int32, seqlens_expanded, bs ) paged_mqa_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( paged_mqa_ctx_lens_2d, 64, deep_gemm.get_num_sms() ) metadata = DSAMetadata( page_size=self.real_page_size, cache_seqlens_int32=cache_seqlens_int32, max_seq_len_q=max_seqlen_q, max_seq_len_k=max_seqlen_k, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, page_table_1=page_table_1, flashmla_metadata=flashmla_metadata, paged_mqa_schedule_metadata=paged_mqa_schedule_metadata, paged_mqa_ctx_lens_2d=paged_mqa_ctx_lens_2d, dsa_cache_seqlens_int32=dsa_cache_seqlens_int32, dsa_cu_seqlens_q=dsa_cu_seqlens_q, dsa_cu_seqlens_k=dsa_cu_seqlens_k, dsa_seqlens_expanded=seqlens_expanded, real_page_table=real_page_table, dsa_extend_seq_lens_list=dsa_extend_seq_lens_list, topk_v2_plan=self._build_topk_v2_plan(seqlens_expanded), ) self.decode_cuda_graph_metadata[bs] = metadata self.forward_metadata = metadata def _apply_cuda_graph_metadata( self, bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, seq_lens_cpu: torch.Tensor, forward_mode: ForwardMode, spec_info: Optional[SpecInput], out_cache_loc: Optional[torch.Tensor] = None, actual_forward_mode: Optional[ForwardMode] = None, ): """Shared capture+replay body for the cuda-graph init path. Public entry: :py:meth:`init_forward_metadata_out_graph`. Spec runners also call this directly via _apply_cuda_graph_metadata when they need to pass out_cache_loc / actual_forward_mode explicitly. """ if bs not in self.decode_cuda_graph_metadata: self._build_forward_metadata_cuda_graph( bs, None, req_pool_indices, seq_lens, seq_lens_cpu, forward_mode, spec_info, out_cache_loc, actual_forward_mode, ) return self.set_dsa_prefill_impl(forward_batch=None) seq_lens = seq_lens[:bs] req_pool_indices = req_pool_indices[:bs] # Normal Decode metadata: DSAMetadata = self.decode_cuda_graph_metadata[bs] used_fused_metadata_generation = False target_verify_ctx_lens_written = False if forward_mode.is_decode_or_idle(): # Normal Decode max_len = self._graph_page_table_width(metadata) if _USE_FUSED_METADATA_GENERATION and is_cuda() and not _is_hip: from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_decode_metadata, ) fused_dsa_decode_metadata( seq_lens=seq_lens, req_pool_indices=req_pool_indices, req_to_token=self.req_to_token, cache_seqlens=metadata.cache_seqlens_int32, cu_seqlens_k=metadata.cu_seqlens_k, page_table_1=metadata.page_table_1, dsa_cache_seqlens=metadata.dsa_cache_seqlens_int32, dsa_cu_seqlens_k=metadata.dsa_cu_seqlens_k, real_page_table=metadata.real_page_table, bs=bs, max_len=max_len, dsa_index_topk=self.dsa_index_topk, real_page_size=self.real_page_size, ) cache_seqlens = metadata.cache_seqlens_int32 dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32 seqlens_expanded = cache_seqlens page_indices = None used_fused_metadata_generation = True if not used_fused_metadata_generation: cache_seqlens = seq_lens.to(torch.int32) metadata.cache_seqlens_int32.copy_(cache_seqlens) metadata.cu_seqlens_k[1:].copy_( torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) ) page_indices = self.req_to_token[req_pool_indices, :max_len] metadata.page_table_1[:, :max_len].copy_(page_indices) dsa_cache_seqlens = compute_dsa_seqlens( cache_seqlens, dsa_index_topk=self.dsa_index_topk ) metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens) seqlens_expanded = cache_seqlens elif forward_mode.is_target_verify(): max_seqlen_k = self._graph_page_table_width(metadata) if _USE_FUSED_METADATA_GENERATION and is_cuda() and not _is_hip: from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_target_verify_metadata, ) paged_mqa_ctx_lens_2d = None if ( self.speculative_num_draft_tokens >= 2 and is_sm100_supported() and metadata.paged_mqa_ctx_lens_2d is not None and metadata.paged_mqa_ctx_lens_2d.dim() == 2 and metadata.paged_mqa_ctx_lens_2d.size(0) == bs and metadata.paged_mqa_ctx_lens_2d.size(1) == self.speculative_num_draft_tokens ): paged_mqa_ctx_lens_2d = metadata.paged_mqa_ctx_lens_2d fused_dsa_target_verify_metadata( seq_lens=seq_lens, req_pool_indices=req_pool_indices, req_to_token=self.req_to_token, cache_seqlens=metadata.cache_seqlens_int32, cu_seqlens_k=metadata.cu_seqlens_k, page_table_1=metadata.page_table_1, seqlens_expanded=metadata.dsa_seqlens_expanded, dsa_cache_seqlens=metadata.dsa_cache_seqlens_int32, dsa_cu_seqlens_k=metadata.dsa_cu_seqlens_k, real_page_table=metadata.real_page_table, bs=bs, max_seqlen_k=max_seqlen_k, dsa_index_topk=self.dsa_index_topk, real_page_size=self.real_page_size, next_n=self.speculative_num_draft_tokens, paged_mqa_ctx_lens_2d=paged_mqa_ctx_lens_2d, ) target_verify_ctx_lens_written = paged_mqa_ctx_lens_2d is not None cache_seqlens = metadata.cache_seqlens_int32 seqlens_expanded = metadata.dsa_seqlens_expanded[ : self.speculative_num_draft_tokens * bs ] dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32[ : self.speculative_num_draft_tokens * bs ] page_indices = None used_fused_metadata_generation = True if not used_fused_metadata_generation: cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to( torch.int32 ) metadata.cache_seqlens_int32.copy_(cache_seqlens) metadata.cu_seqlens_k[1:].copy_( torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) ) page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k] page_indices = torch.repeat_interleave( page_indices, repeats=self.speculative_num_draft_tokens, dim=0 ) metadata.page_table_1[:, :max_seqlen_k].copy_(page_indices) # Fill the constant per-req qo lengths on-device; torch.tensor(list, # device=cuda) does a pageable H2D copy that blocks the host. extend_seq_lens = torch.full( (bs,), self.speculative_num_draft_tokens, dtype=torch.int32, device=self.device, ) seqlens_expanded = seqlens_expand_triton( extend_seq_lens, cache_seqlens, self.speculative_num_draft_tokens * bs, self.speculative_num_draft_tokens, ) metadata.dsa_seqlens_expanded.copy_(seqlens_expanded) dsa_cache_seqlens = compute_dsa_seqlens( seqlens_expanded, self.dsa_index_topk ) metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens) elif forward_mode.is_draft_extend_v2(): # V2 draft-extend processes the full padded tree width # (speculative_num_draft_tokens) per req -- a static shape, like # target-verify -- so graph replay stays host-sync-free. seq_lens # already includes the draft KV written by prepare_for_draft_extend; # the per-req accept length is handled downstream by output # selection, not by reshaping the page table here. max_seqlen_k = self._graph_page_table_width(metadata) total_extend_len = self.speculative_num_draft_tokens * bs # See target-verify note: fill on-device to avoid the blocking # pageable H2D from torch.tensor(list, device=cuda). extend_seq_lens = torch.full( (bs,), self.speculative_num_draft_tokens, dtype=torch.int32, device=self.device, ) if _USE_FUSED_METADATA_GENERATION and is_cuda() and not _is_hip: from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_draft_extend_metadata, ) fused_dsa_draft_extend_metadata( seq_lens=seq_lens, extend_seq_lens=extend_seq_lens, req_pool_indices=req_pool_indices, req_to_token=self.req_to_token, cache_seqlens=metadata.cache_seqlens_int32, cu_seqlens_k=metadata.cu_seqlens_k, page_table_1=metadata.page_table_1, seqlens_expanded=metadata.dsa_seqlens_expanded, dsa_cache_seqlens=metadata.dsa_cache_seqlens_int32, dsa_cu_seqlens_k=metadata.dsa_cu_seqlens_k, real_page_table=metadata.real_page_table, bs=bs, total_len=total_extend_len, max_seqlen_k=max_seqlen_k, dsa_index_topk=self.dsa_index_topk, real_page_size=self.real_page_size, max_extend_len=self.speculative_num_draft_tokens, max_total_len=bs * self.speculative_num_draft_tokens, static_extend_len=True, ) cache_seqlens = metadata.cache_seqlens_int32 seqlens_expanded = metadata.dsa_seqlens_expanded[:total_extend_len] dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32[:total_extend_len] page_indices = None used_fused_metadata_generation = True if not used_fused_metadata_generation: cache_seqlens = seq_lens.to(torch.int32) metadata.cache_seqlens_int32.copy_(cache_seqlens) metadata.cu_seqlens_k[1:].copy_( torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) ) page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k] page_indices = torch.repeat_interleave( page_indices, repeats=self.speculative_num_draft_tokens, dim=0 ) metadata.page_table_1[:, :max_seqlen_k].copy_(page_indices) seqlens_expanded = seqlens_expand_triton( extend_seq_lens, cache_seqlens, total_extend_len, self.speculative_num_draft_tokens, ) metadata.dsa_seqlens_expanded[: seqlens_expanded.shape[0]].copy_( seqlens_expanded ) dsa_cache_seqlens = compute_dsa_seqlens( seqlens_expanded, self.dsa_index_topk ) metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens) # Update DeepGEMM paged MQA schedule metadata outside the captured graph. if is_cuda() and ( forward_mode.is_decode_or_idle() or forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2() ): if forward_mode.is_draft_extend_v2(): schedule_seqlens_expanded = metadata.dsa_seqlens_expanded else: schedule_seqlens_expanded = seqlens_expanded if target_verify_ctx_lens_written: seqlens_32_2d = metadata.paged_mqa_ctx_lens_2d else: seqlens_32_2d = self._build_paged_mqa_schedule_2d_ctx_lens( forward_mode, metadata.cache_seqlens_int32, schedule_seqlens_expanded, bs, ) self._refresh_paged_mqa_schedule_metadata(metadata, seqlens_32_2d) self._refresh_topk_v2_plan(metadata) # `copy_` preserves the buffer's data_ptr that the captured graph captured. if not target_verify_ctx_lens_written: if metadata.paged_mqa_ctx_lens_2d is None: object.__setattr__(metadata, "paged_mqa_ctx_lens_2d", seqlens_32_2d) else: metadata.paged_mqa_ctx_lens_2d.copy_(seqlens_32_2d) seqlens_expanded_size = seqlens_expanded.shape[0] assert ( metadata.dsa_cache_seqlens_int32 is not None and metadata.dsa_cu_seqlens_k is not None and self.dsa_index_topk is not None ) if not used_fused_metadata_generation: metadata.dsa_cu_seqlens_k[1 : 1 + seqlens_expanded_size].copy_( torch.cumsum(dsa_cache_seqlens, dim=0, dtype=torch.int32) ) # NOTE(dark): (dsa-) cu_seqlens_q is always arange, no need to copy assert self.real_page_size == metadata.page_size if self.real_page_size > 1: if not used_fused_metadata_generation: real_table = self._transform_table_1_to_real(page_indices) new_rows = real_table.shape[0] new_cols = real_table.shape[1] metadata.real_page_table[:new_rows, :new_cols].copy_(real_table) else: assert metadata.real_page_table is metadata.page_table_1 if self.dsa_decode_impl == "flashmla_kv": flashmla_metadata = metadata.flashmla_metadata.slice( slice(0, seqlens_expanded_size + 1) ) flashmla_metadata.copy_( self._compute_flashmla_metadata( cache_seqlens=dsa_cache_seqlens, seq_len_q=1, ) ) self.forward_metadata = metadata def init_forward_metadata_replay_cuda_graph_from_precomputed( self, bs: int, precomputed: PrecomputedMetadata, forward_mode: ForwardMode, ): """Fast path: copy precomputed metadata to this backend's metadata. This function only performs copy operations, no computation. Args: bs: Batch size precomputed: Precomputed metadata to copy from forward_mode: Forward mode """ self.set_dsa_prefill_impl(forward_batch=None) metadata = self.decode_cuda_graph_metadata[bs] # Track whether fused kernel succeeded fused_kernel_succeeded = False # Use fused CUDA kernel for all copy operations if not _is_hip: try: from sglang.jit_kernel.fused_metadata_copy import ( fused_metadata_copy_cuda, ) # Map forward_mode to integer enum if forward_mode.is_decode_or_idle(): mode_int = 0 # DECODE elif forward_mode.is_target_verify(): mode_int = 1 # TARGET_VERIFY else: raise ValueError(f"Unsupported forward_mode: {forward_mode}") # Prepare FlashMLA tensors if needed flashmla_num_splits_src = None flashmla_num_splits_dst = None flashmla_metadata_src = None flashmla_metadata_dst = None if precomputed.flashmla_metadata is not None: flashmla_num_splits_src = precomputed.flashmla_metadata.num_splits flashmla_num_splits_dst = metadata.flashmla_metadata.num_splits flashmla_metadata_src = ( precomputed.flashmla_metadata.flashmla_metadata ) flashmla_metadata_dst = metadata.flashmla_metadata.flashmla_metadata # Call fused kernel fused_metadata_copy_cuda( # Source tensors precomputed.cache_seqlens, precomputed.cu_seqlens_k, precomputed.page_indices, precomputed.dsa_cache_seqlens, precomputed.seqlens_expanded, precomputed.dsa_cu_seqlens_k, precomputed.real_page_table, flashmla_num_splits_src, flashmla_metadata_src, # Destination tensors metadata.cache_seqlens_int32, metadata.cu_seqlens_k, metadata.page_table_1, metadata.dsa_cache_seqlens_int32, metadata.dsa_seqlens_expanded, metadata.dsa_cu_seqlens_k, ( metadata.real_page_table if precomputed.real_page_table is not None else None ), flashmla_num_splits_dst, flashmla_metadata_dst, # Parameters mode_int, bs, precomputed.max_len, precomputed.max_seqlen_k, precomputed.seqlens_expanded_size, ) # Successfully used fused kernel fused_kernel_succeeded = True except ImportError: print( "Warning: Fused metadata copy kernel not available, falling back to individual copies." ) except Exception as e: print( f"Warning: Fused metadata copy kernel failed with error: {e}, falling back to individual copies." ) # Fallback to individual copy operations if the fused kernel is unavailable # or fails at runtime. if not fused_kernel_succeeded: # Copy basic seqlens metadata.cache_seqlens_int32.copy_(precomputed.cache_seqlens) metadata.cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:]) # Mode-specific copy logic if forward_mode.is_decode_or_idle(): # Decode mode metadata.page_table_1[:, : precomputed.max_len].copy_( precomputed.page_indices ) metadata.dsa_cache_seqlens_int32.copy_(precomputed.dsa_cache_seqlens) # seqlens_expanded is same as cache_seqlens (already copied) elif forward_mode.is_target_verify(): # Target verify mode metadata.page_table_1[:, : precomputed.max_seqlen_k].copy_( precomputed.page_indices ) metadata.dsa_seqlens_expanded.copy_(precomputed.seqlens_expanded) metadata.dsa_cache_seqlens_int32.copy_(precomputed.dsa_cache_seqlens) # Copy DSA cu_seqlens size = precomputed.seqlens_expanded_size metadata.dsa_cu_seqlens_k[1 : 1 + size].copy_( precomputed.dsa_cu_seqlens_k[1 : 1 + size] ) # Copy real page table if precomputed.real_page_table is not None: rows, cols = precomputed.real_page_table.shape metadata.real_page_table[:rows, :cols].copy_( precomputed.real_page_table ) # Copy FlashMLA metadata in fallback path if precomputed.flashmla_metadata is not None: size = precomputed.seqlens_expanded_size flashmla_metadata = metadata.flashmla_metadata.slice(slice(0, size + 1)) flashmla_metadata.copy_(precomputed.flashmla_metadata) # Refresh DeepGEMM paged MQA schedule metadata for the actual seqlens of # this replay (the captured graph holds stale data otherwise, which can # deadlock the kernel when the runtime work decomposition diverges from # the captured one). if is_cuda(): if forward_mode.is_decode_or_idle(): seqlens_32_2d = _to_2d_context_lens(metadata.cache_seqlens_int32, bs) else: seqlens_32_2d = self._build_paged_mqa_schedule_2d_ctx_lens( forward_mode, metadata.cache_seqlens_int32, metadata.dsa_seqlens_expanded, bs, ) self._refresh_paged_mqa_schedule_metadata(metadata, seqlens_32_2d) self._refresh_topk_v2_plan(metadata) if metadata.paged_mqa_ctx_lens_2d is None: object.__setattr__(metadata, "paged_mqa_ctx_lens_2d", seqlens_32_2d) else: metadata.paged_mqa_ctx_lens_2d.copy_(seqlens_32_2d) self.forward_metadata = metadata def forward_extend( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer: RadixAttention, forward_batch: ForwardBatch, save_kv_cache=True, # For multi-head latent attention q_rope: Optional[torch.Tensor] = None, k_rope: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, cos_sin_cache: Optional[torch.Tensor] = None, is_neox: Optional[bool] = False, llama_4_scaling: Optional[torch.Tensor] = None, ) -> torch.Tensor: causal = not layer.is_cross_attention metadata = self.forward_metadata assert causal, "DSA is causal only" dsa_impl = ( self.dsa_decode_impl if ( forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend_v2() ) else self.dsa_prefill_impl ) if dsa_impl == "trtllm" and not self.use_mha: return self._forward_trtllm( q, k, v, layer, forward_batch, metadata.dsa_cache_seqlens_int32, save_kv_cache, q_rope, k_rope, topk_indices, cos_sin_cache, is_neox, llama_4_scaling, is_prefill=True, ) if k is not None: assert v is not None if save_kv_cache: cache_loc = ( forward_batch.out_cache_loc if not layer.is_cross_attention else forward_batch.encoder_out_cache_loc ) self.token_to_kv_pool.set_mla_kv_buffer( # type: ignore layer, cache_loc, k, k_rope, ) # Use MHA kernel if in MHA_ONE_SHOT mode if self.use_mha: assert k is not None and v is not None assert q_rope is None, "MHA_ONE_SHOT path should not pass q_rope" assert ( layer.tp_k_head_num == layer.tp_q_head_num > 1 ), "MHA_ONE_SHOT requires dense multi-head config" return self._forward_standard_mha( q=q, k=k, v=v, layer=layer, forward_batch=forward_batch, metadata=metadata, ) # Do absorbed multi-latent attention (MLA path) assert q_rope is not None kv_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id) if q_rope is not None: q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) q_rope = q_rope.view( -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim ) else: q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) q_nope = q_all[:, :, : layer.v_head_dim] q_rope = q_all[:, :, layer.v_head_dim :] # Align topk_indices with q dimensions # This handles cases where q is padded (TP + partial DP attention) if topk_indices is not None: topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0]) # NOTE(dark): here, we use page size = 1 topk_transform_method = self.get_topk_transform_method( forward_batch.forward_mode ) if envs.SGLANG_DSA_FUSE_TOPK.get(): page_table_1 = self._get_fused_topk_page_table(topk_indices) else: if topk_transform_method == TopkTransformMethod.RAGGED: topk_indices_offset = metadata.topk_indices_offset assert topk_indices_offset is not None mask = topk_indices != -1 topk_indices_offset = ( topk_indices_offset.unsqueeze(1) if topk_indices_offset.ndim == 1 else topk_indices_offset ) topk_indices = torch.where( mask, topk_indices + topk_indices_offset, topk_indices ) elif topk_transform_method == TopkTransformMethod.PAGED: assert metadata.dsa_extend_seq_lens_list is not None page_table_1 = transform_index_page_table_prefill( page_table=metadata.page_table_1, topk_indices=topk_indices, extend_lens_cpu=metadata.dsa_extend_seq_lens_list, page_size=1, ) # todo hisparse: to cover more backends if self.hisparse_coordinator is not None: # flash_mla_sparse_fwd / tilelang require int32 page indices. page_table_1 = self.token_to_kv_pool.translate_loc_to_hisparse_device( page_table_1 ).to(torch.int32) if dsa_impl == "tilelang": if q_rope is not None: # Triton prefill kernel reads q_nope/q_rope directly, skipping # the concat (it splits q into main/tail internally anyway). # Gated to gfx950 + the validated shape (16 heads, d_v=512, # tail=64, topk=2048); everything else uses TileLang. if ( _DSA_TRITON_PREFILL and _IS_GFX95 and kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) and layer.tp_q_head_num == 16 and layer.v_head_dim == 512 and (layer.head_dim - layer.v_head_dim) == 64 and page_table_1.shape[-1] == 2048 and q_nope.shape[0] >= 512 ): from sglang.srt.layers.attention.dsa.triton_sparse_mla import ( triton_sparse_mla_fwd, ) return triton_sparse_mla_fwd( q_nope=q_nope, q_rope=q_rope, kv=kv_cache, indices=page_table_1.unsqueeze(1), sm_scale=layer.scaling, d_v=layer.v_head_dim, ) q_all = concat_mla_absorb_q_general(q_nope, q_rope) return self._forward_tilelang( q_all=q_all, kv_cache=kv_cache, page_table_1=page_table_1, sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) elif dsa_impl == "flashmla_sparse": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) if topk_transform_method == TopkTransformMethod.RAGGED: if any(forward_batch.extend_prefix_lens_cpu): page_table_1_flattened = ( self.forward_metadata.page_table_1_flattened ) assert page_table_1_flattened is not None kv_cache = dequantize_k_cache_paged( kv_cache, page_table_1_flattened ) else: kv_cache = _cat([k, k_rope], dim=-1) page_table_1 = topk_indices return self._forward_flashmla_sparse( q_all=q_all, kv_cache=kv_cache, page_table_1=page_table_1, sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) elif dsa_impl == "flashmla_kv": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) return self._forward_flashmla_kv( q_all=q_all, kv_cache=kv_cache, sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, # TODO optimize args layer=layer, metadata=metadata, page_table_1=page_table_1, ) elif dsa_impl == "fa3": return self._forward_fa3( q_rope=q_rope, kv_cache=kv_cache, v_head_dim=layer.v_head_dim, q_nope=q_nope, page_table=page_table_1, cache_seqlens=metadata.dsa_cache_seqlens_int32, cu_seqlens_q=metadata.dsa_cu_seqlens_q, cu_seqlens_k=metadata.dsa_cu_seqlens_k, max_seqlen_q=metadata.dsa_max_seqlen_q, sm_scale=layer.scaling, logit_cap=layer.logit_cap, page_size=1, ) elif dsa_impl == "aiter": if q_rope is not None: q_all = torch.cat([q_nope, q_rope], dim=-1) return self._forward_aiter_extend( q_all=q_all, kv_cache=kv_cache, page_table_1=page_table_1, layer=layer, ) else: raise ValueError( f"Unsupported {dsa_impl = } for forward_extend. Consider using an other attention backend." ) def forward_decode( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer: RadixAttention, forward_batch: ForwardBatch, save_kv_cache=True, # For multi-head latent attention q_rope: Optional[torch.Tensor] = None, k_rope: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, cos_sin_cache: Optional[torch.Tensor] = None, is_neox: Optional[bool] = False, llama_4_scaling: Optional[torch.Tensor] = None, ) -> torch.Tensor: causal = not layer.is_cross_attention metadata = self.forward_metadata assert causal, "DSA is causal only" if self.dsa_decode_impl == "trtllm": return self._forward_trtllm( q, k, v, layer, forward_batch, metadata.cache_seqlens_int32, save_kv_cache, q_rope, k_rope, topk_indices, cos_sin_cache, is_neox, llama_4_scaling, ) if k is not None: assert v is not None if save_kv_cache: cache_loc = ( forward_batch.out_cache_loc if not layer.is_cross_attention else forward_batch.encoder_out_cache_loc ) self.token_to_kv_pool.set_mla_kv_buffer( # type: ignore layer, cache_loc, k, k_rope, ) # Do absorbed multi-latent attention kv_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id) if q_rope is not None: q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) q_rope = q_rope.view( -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim ) # Caller passed split q_nope / q_rope; we'll need to concat below if # the chosen impl wants q_all. q_all = None else: # Caller passed already-concatenated q (q_all = q). Reuse it directly # via a zero-copy view; the impl-specific blocks below will skip the # otherwise redundant concat_mla_absorb_q_general call. q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) q_nope = q_all[:, :, : layer.v_head_dim] q_rope = q_all[:, :, layer.v_head_dim :] # Align topk_indices with q dimensions if topk_indices is not None: topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0]) if self.hisparse_coordinator is not None: page_table_1 = self.hisparse_coordinator.swap_in_selected_pages( forward_batch.req_pool_indices, forward_batch.seq_lens, topk_indices, layer.layer_id, ) elif envs.SGLANG_DSA_FUSE_TOPK.get(): page_table_1 = self._get_fused_topk_page_table(topk_indices) else: page_table_1 = transform_index_page_table_decode( page_table=metadata.page_table_1, topk_indices=topk_indices, page_size=1, ) if self.dsa_decode_impl == "flashmla_sparse": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) return self._forward_flashmla_sparse( q_all=q_all, kv_cache=kv_cache, page_table_1=page_table_1, sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) elif self.dsa_decode_impl == "flashmla_kv": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) return self._forward_flashmla_kv( q_all=q_all, kv_cache=kv_cache, sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, # TODO optimize args layer=layer, metadata=metadata, page_table_1=page_table_1, ) elif self.dsa_decode_impl == "tilelang": # Cat-skip (HIP-only): when caller passes q_rope=None on HIP, q_all # has already been set to a zero-copy view of q in the else branch # above and we can reuse it directly. The `not _is_hip` clause keeps # CUDA / MUSA paths byte-identical to pre-patch by always re-cat. if q_all is None or not _is_hip: q_all = concat_mla_absorb_q_general(q_nope, q_rope) return self._forward_tilelang( q_all=q_all, kv_cache=kv_cache, page_table_1=page_table_1, sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) elif self.dsa_decode_impl == "fa3": return self._forward_fa3( q_rope=q_rope, kv_cache=kv_cache, v_head_dim=layer.v_head_dim, q_nope=q_nope, page_table=page_table_1, cache_seqlens=metadata.dsa_cache_seqlens_int32, cu_seqlens_q=metadata.dsa_cu_seqlens_q, cu_seqlens_k=metadata.dsa_cu_seqlens_k, max_seqlen_q=metadata.dsa_max_seqlen_q, sm_scale=layer.scaling, logit_cap=layer.logit_cap, page_size=1, ) elif self.dsa_decode_impl == "aiter": if q_all is None or not _is_hip: q_all = torch.cat([q_nope, q_rope], dim=-1) return self._forward_aiter( q_all=q_all, kv_cache=kv_cache, page_table_1=page_table_1, layer=layer, metadata=metadata, bs=forward_batch.batch_size, ) else: assert False, f"Unsupported {self.dsa_decode_impl = }" def _forward_fa3( self, q_rope: torch.Tensor, kv_cache: torch.Tensor, v_head_dim: int, q_nope: torch.Tensor, page_table: torch.Tensor, cache_seqlens: torch.Tensor, cu_seqlens_q: torch.Tensor, cu_seqlens_k: torch.Tensor, max_seqlen_q: int, sm_scale: float, logit_cap: float, page_size: int, ) -> torch.Tensor: k_rope_cache = kv_cache[:, :, v_head_dim:] c_kv_cache = kv_cache[:, :, :v_head_dim] qk_rope_dim = k_rope_cache.shape[-1] k_rope_cache = k_rope_cache.view(-1, page_size, 1, qk_rope_dim) c_kv_cache = c_kv_cache.view(-1, page_size, 1, v_head_dim) o = flash_attn_with_kvcache( q=q_rope, k_cache=k_rope_cache, v_cache=c_kv_cache, qv=q_nope, page_table=page_table, cache_seqlens=cache_seqlens, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k_new=cu_seqlens_k, max_seqlen_q=max_seqlen_q, softmax_scale=sm_scale, causal=True, softcap=logit_cap, return_softmax_lse=False, num_splits=self.num_splits, ) return o # type: ignore def _forward_flashmla_sparse( self, q_all: torch.Tensor, kv_cache: torch.Tensor, v_head_dim: int, page_table_1: torch.Tensor, sm_scale: float, ) -> torch.Tensor: from sgl_kernel.flash_mla import flash_mla_sparse_fwd # FlashMLA sparse kernel requires num_heads to be a multiple of 64 (Hopper) or 128 (Blackwell) # When using TP, num_heads might be smaller (e.g., 256//8=32) num_tokens, num_heads, head_dim = q_all.shape # Determine required padding based on GPU architecture (use cached value) required_padding = 128 if self.device_sm_major >= 10 else 64 need_padding = num_heads % required_padding != 0 if need_padding: assert required_padding % num_heads == 0, ( f"num_heads {num_heads} cannot be padded to {required_padding}. " f"TP size may be too large for this model." ) # Pad q to required size q_padded = q_all.new_zeros((num_tokens, required_padding, head_dim)) q_padded[:, :num_heads, :] = q_all q_input = q_padded else: q_input = q_all # indices shape must be (s_q, h_kv=1, topk), keep h_kv=1 unchanged indices_input = page_table_1.unsqueeze(1) o, _, _ = flash_mla_sparse_fwd( q=q_input, kv=kv_cache, indices=indices_input, sm_scale=sm_scale, d_v=v_head_dim, ) # Trim output back to original num_heads if we padded if need_padding: o = o[:, :num_heads, :] return o def _forward_flashmla_kv( self, q_all: torch.Tensor, kv_cache: torch.Tensor, v_head_dim: int, sm_scale: float, layer, metadata: DSAMetadata, page_table_1, ) -> torch.Tensor: from sgl_kernel.flash_mla import flash_mla_with_kvcache cache_seqlens = metadata.dsa_cache_seqlens_int32 assert metadata.flashmla_metadata is not None # TODO the 2nd dim is seq_len_q, need to be >1 when MTP q_all = q_all.view(-1, 1, layer.tp_q_head_num, layer.head_dim) num_q_heads = q_all.shape[2] target_q_heads = self.flashmla_kv_num_q_heads if target_q_heads != num_q_heads: # Pad q heads to match FlashMLA decode supported head-count variants. q_input = q_all.new_zeros( q_all.shape[0], q_all.shape[1], target_q_heads, q_all.shape[3] ) q_input[:, :, :num_q_heads, :] = q_all else: q_input = q_all kv_cache = kv_cache.view(-1, self.real_page_size, 1, self.kv_cache_dim) assert self.real_page_size == 64, "only page size 64 is supported" if not self.dsa_kv_cache_store_fp8: # inefficiently quantize the whole cache kv_cache = quantize_k_cache(kv_cache) indices = page_table_1.unsqueeze(1) assert ( indices.shape[-1] == self.dsa_index_topk ) # requirement of FlashMLA decode kernel o, _ = flash_mla_with_kvcache( q=q_input, k_cache=kv_cache, cache_seqlens=cache_seqlens, head_dim_v=v_head_dim, tile_scheduler_metadata=metadata.flashmla_metadata.flashmla_metadata, num_splits=metadata.flashmla_metadata.num_splits, softmax_scale=sm_scale, indices=indices, # doc says it is not used, but if pass in None then error block_table=torch.empty( (q_all.shape[0], 0), dtype=torch.int32, device=q_all.device ), is_fp8_kvcache=True, ) if target_q_heads != num_q_heads: o = o[:, :, :num_q_heads, :] return o def _forward_standard_mha( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer: RadixAttention, forward_batch: ForwardBatch, metadata: DSAMetadata, ) -> torch.Tensor: """Standard MHA using FlashAttention varlen for MHA_ONE_SHOT mode.""" q = q.view(-1, layer.tp_q_head_num, layer.head_dim) k = k.view(-1, layer.tp_k_head_num, layer.head_dim) v = v.view(-1, layer.tp_v_head_num, layer.v_head_dim) # MHA_ONE_SHOT: k/v include all tokens (prefix + current) cu_seqlens_q = metadata.cu_seqlens_q cu_seqlens_k = metadata.cu_seqlens_k max_seqlen_k = metadata.max_seq_len_k causal = True # Verify batch sizes match (length of cu_seqlens should be batch_size + 1) assert len(cu_seqlens_q) == len(cu_seqlens_k), ( f"batch_size mismatch: cu_seqlens_q has {len(cu_seqlens_q)-1} requests, " f"cu_seqlens_k has {len(cu_seqlens_k)-1} requests" ) # Use TRTLLm ragged attention for SM100 (Blackwell/B200) to avoid FA4 accuracy issues if self.device_sm_major >= 10: import flashinfer seq_lens = metadata.cache_seqlens_int32 return flashinfer.prefill.trtllm_ragged_attention_deepseek( query=q, key=k, value=v, workspace_buffer=self.workspace_buffer, seq_lens=seq_lens, max_q_len=metadata.max_seq_len_q, max_kv_len=max_seqlen_k, bmm1_scale=layer.scaling, bmm2_scale=1.0, o_sf_scale=1.0, batch_size=forward_batch.batch_size, window_left=-1, cum_seq_lens_q=cu_seqlens_q, cum_seq_lens_kv=cu_seqlens_k, enable_pdl=False, is_causal=causal, return_lse=False, skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(), ) # Use FA3 for SM90 (Hopper/H200) return flash_attn_varlen_func( q=q, k=k, v=v, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=metadata.max_seq_len_q, max_seqlen_k=max_seqlen_k, softmax_scale=layer.scaling, causal=causal, ) def _forward_tilelang( self, q_all: torch.Tensor, kv_cache: torch.Tensor, v_head_dim: int, page_table_1: torch.Tensor, sm_scale: float, ) -> torch.Tensor: from sglang.srt.layers.attention.dsa.tilelang_kernel import tilelang_sparse_fwd return tilelang_sparse_fwd( q=q_all, kv=kv_cache, indices=page_table_1.unsqueeze(1), sm_scale=sm_scale, d_v=v_head_dim, ) def _forward_aiter( self, q_all: torch.Tensor, kv_cache: torch.Tensor, page_table_1: torch.Tensor, layer: RadixAttention, metadata: DSAMetadata, bs: int, ) -> torch.Tensor: q = q_all.reshape(-1, layer.tp_q_head_num * layer.head_dim) if layer.head_dim != layer.v_head_dim: o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) else: o = torch.empty_like(q) if self.need_pad_heads: q_kernel = q.view( -1, layer.tp_q_head_num, layer.head_dim ).repeat_interleave(self.head_repeat_factor, dim=1) o_kernel = q.new_empty( ( q.shape[0], layer.tp_q_head_num * self.head_repeat_factor, layer.v_head_dim, ) ) else: q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim) o_kernel = o.view(-1, layer.tp_q_head_num, layer.v_head_dim) q_scale = None kv_scale = None aiter_persistent_kwargs = {} if kv_cache.dtype == fp8_dtype: kv_scale = torch.ones((), dtype=torch.float32, device=q_kernel.device) kv_indptr = self.kv_indptr non_minus1_mask = page_table_1 != -1 non_minus1_counts = non_minus1_mask.sum(dim=1) kv_indptr[1 : bs + 1] = torch.cumsum(non_minus1_counts, dim=0) kv_indices = self.kv_indices get_valid_kv_indices(page_table_1, kv_indptr, kv_indices, bs) kv_last_page_lens = metadata.cu_seqlens_q if kv_cache.dtype == fp8_dtype: aiter_persistent_kwargs = self._prepare_aiter_dsa_decode_metadata( metadata.cu_seqlens_q, kv_indptr, bs, metadata.max_seq_len_q, q_kernel.dtype, kv_cache.dtype, ) kv_last_page_lens = aiter_persistent_kwargs.pop("kv_last_page_lens") mla_decode_fwd( q_kernel, kv_cache.view(-1, 1, 1, layer.head_dim), o_kernel, metadata.cu_seqlens_q, kv_indptr, kv_indices, kv_last_page_lens, metadata.max_seq_len_q, sm_scale=layer.scaling, logit_cap=layer.logit_cap, q_scale=q_scale, kv_scale=kv_scale, **aiter_persistent_kwargs, ) if self.need_pad_heads: o = o_kernel[:, :: self.head_repeat_factor, :] return o def _forward_aiter_extend( self, q_all: torch.Tensor, kv_cache: torch.Tensor, page_table_1: torch.Tensor, layer: RadixAttention, ) -> torch.Tensor: num_tokens = q_all.shape[0] q = q_all.reshape(-1, layer.tp_q_head_num * layer.head_dim) if layer.head_dim != layer.v_head_dim: o = q.new_empty((num_tokens, layer.tp_q_head_num * layer.v_head_dim)) else: o = torch.empty_like(q) if self.need_pad_heads: q_kernel = q.view( -1, layer.tp_q_head_num, layer.head_dim ).repeat_interleave(self.head_repeat_factor, dim=1) o_kernel = q.new_empty( ( num_tokens, layer.tp_q_head_num * self.head_repeat_factor, layer.v_head_dim, ) ) else: q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim) o_kernel = o.view(-1, layer.tp_q_head_num, layer.v_head_dim) q_scale = None kv_scale = None aiter_persistent_kwargs = {} if kv_cache.dtype == fp8_dtype: kv_scale = torch.ones((), dtype=torch.float32, device=q_kernel.device) non_minus1_mask = page_table_1 != -1 non_minus1_counts = non_minus1_mask.sum(dim=1) kv_indptr = torch.zeros(num_tokens + 1, dtype=torch.int32, device=self.device) kv_indptr[1:] = torch.cumsum(non_minus1_counts, dim=0) # Allocate kv_indices with upper-bound size (num_tokens * topk) topk = page_table_1.shape[1] kv_indices = torch.zeros( num_tokens * topk, dtype=torch.int32, device=self.device ) # Use get_valid_kv_indices kernel to extract valid indices get_valid_kv_indices(page_table_1, kv_indptr, kv_indices, num_tokens) # Build cu_seqlens_q for extend: each token is treated as seq_len_q=1 cu_seqlens_q = torch.arange( 0, num_tokens + 1, dtype=torch.int32, device=self.device ) kv_last_page_lens = cu_seqlens_q if kv_cache.dtype == fp8_dtype: aiter_persistent_kwargs = self._prepare_aiter_dsa_decode_metadata( cu_seqlens_q, kv_indptr, num_tokens, 1, q_kernel.dtype, kv_cache.dtype, ) kv_last_page_lens = aiter_persistent_kwargs.pop("kv_last_page_lens") # TODO support more forward_mode mla_decode_fwd( q_kernel, kv_cache.view(-1, 1, 1, layer.head_dim), o_kernel, cu_seqlens_q, kv_indptr, kv_indices, kv_last_page_lens, 1, # max_seq_len_q = 1 for per-token attention sm_scale=layer.scaling, logit_cap=layer.logit_cap, q_scale=q_scale, kv_scale=kv_scale, **aiter_persistent_kwargs, ) if self.need_pad_heads: o = o_kernel[:, :: self.head_repeat_factor, :] return o def _forward_trtllm( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer: RadixAttention, forward_batch: ForwardBatch, seq_lens: torch.Tensor, save_kv_cache=True, # For multi-head latent attention q_rope: Optional[torch.Tensor] = None, k_rope: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, cos_sin_cache: Optional[torch.Tensor] = None, is_neox: Optional[bool] = False, llama_4_scaling: Optional[torch.Tensor] = None, is_prefill: bool = False, ) -> torch.Tensor: """Forward using TRT-LLM sparse MLA kernel.""" import flashinfer.decode metadata = self.forward_metadata merge_query = q_rope is not None if self.kv_cache_dtype == torch.float8_e4m3fn: # For FP8 path, we quantize the query and rope parts and merge them into a single tensor # Note: rope application in deepseek_v2.py:forward_absorb_prepare is skipped for FP8 decode path of this trtllm_mla backend assert q_rope is not None, "For FP8 path q_rope should not be None." assert k_rope is not None, "For FP8 path k_rope should not be None." assert ( cos_sin_cache is not None ), "For FP8 path cos_sin_cache should not be None." rope_positions = forward_batch.positions if dsa_use_prefill_cp(forward_batch): rope_positions = cp_split_and_rebuild_position( forward_batch, rope_positions ) q, k, k_rope = mla_quantize_and_rope_for_fp8( q, q_rope, k.squeeze(1), k_rope.squeeze(1), rope_positions, cos_sin_cache, is_neox, self.kv_lora_rank, self.qk_rope_head_dim, ) if save_kv_cache and dsa_use_prefill_cp(forward_batch): k, k_rope = _all_gather_dsa_trtllm_fp8_kv(forward_batch, k, k_rope) merge_query = False # Save KV cache if requested if save_kv_cache: assert ( k is not None and k_rope is not None ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." cache_loc = ( forward_batch.out_cache_loc if not layer.is_cross_attention else forward_batch.encoder_out_cache_loc ) self.token_to_kv_pool.set_mla_kv_buffer(layer, cache_loc, k, k_rope) k_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id) kv_cache = k_cache.view(-1, self.real_page_size, self.kv_cache_dim).unsqueeze(1) if merge_query: q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) q_rope_reshaped = q_rope.view( -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim ) q_all = concat_mla_absorb_q_general(q_nope, q_rope_reshaped) else: q_all = q.view(-1, layer.tp_q_head_num, layer.head_dim) # Align topk_indices with q dimensions if topk_indices is not None: topk_indices = self._pad_topk_indices(topk_indices, q.shape[0]) if envs.SGLANG_DSA_FUSE_TOPK.get(): page_table_1 = self._get_fused_topk_page_table(topk_indices) elif is_prefill: page_table_1 = transform_index_page_table_prefill( page_table=metadata.page_table_1, topk_indices=topk_indices, extend_lens_cpu=metadata.dsa_extend_seq_lens_list, page_size=1, ) else: page_table_1 = transform_index_page_table_decode( page_table=metadata.page_table_1, topk_indices=topk_indices, page_size=1, ) q_scale = 1.0 k_scale = ( layer.k_scale_float if getattr(layer, "k_scale_float", None) is not None else 1.0 ) bmm1_scale = q_scale * k_scale * layer.scaling batch_size = page_table_1.shape[0] _, num_heads, head_dim = q_all.shape q = q_all.view(batch_size, 1, num_heads, head_dim) kv = kv_cache.view(-1, 1, self.real_page_size, self.kv_cache_dim) block_tables = page_table_1.unsqueeze(1) seq_lens = metadata.cache_seqlens_int32 if seq_lens is None else seq_lens if ( dsa_use_prefill_cp(forward_batch) and is_dsa_prefill_cp_in_seq_split() and forward_batch.attn_cp_metadata is not None ): cp_meta = forward_batch.attn_cp_metadata seq_chunks = list(torch.split(seq_lens, cp_meta.split_list, dim=0)) seq_lens = torch.cat([seq_chunks[i] for i in cp_meta.zigzag_index], dim=0) out = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( query=q, kv_cache=kv, workspace_buffer=self.workspace_buffer, qk_nope_head_dim=self.qk_nope_head_dim, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, block_tables=block_tables, seq_lens=seq_lens, max_seq_len=metadata.max_seq_len_k, sparse_mla_top_k=self.dsa_index_topk, bmm1_scale=bmm1_scale, backend="trtllm-gen", skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), ) return out def _pad_topk_indices( self, topk_indices: torch.Tensor, num_tokens: int ) -> torch.Tensor: current_tokens = topk_indices.shape[0] if current_tokens == num_tokens: return topk_indices assert current_tokens <= num_tokens, ( f"topk_indices rows ({current_tokens}) > num_tokens ({num_tokens}); " "this indicates a mismatch between indexer output and q layout." ) pad_size = num_tokens - current_tokens padding = torch.full( (pad_size, topk_indices.shape[1]), -1, dtype=topk_indices.dtype, device=topk_indices.device, ) return torch.cat([topk_indices, padding], dim=0) def get_cuda_graph_seq_len_fill_value(self): """Get the fill value for sequence length in CUDA graph.""" return 1 def set_dsa_prefill_impl(self, forward_batch: Optional[ForwardBatch] = None): """ Decide all attention prefill dispatch strategies for this batch. """ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( is_in_breakable_cuda_graph, ) from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, ) from sglang.srt.utils import get_device_sm, is_blackwell # Decide MHA vs MLA if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph(): # Can't branch on seq_lens_cpu in graph replay, force MHA off to # guarantee correctness. self.use_mha = False elif ( forward_batch and forward_batch.forward_mode.is_extend_without_speculative() ): # Check if sequence meets criteria for MHA_ONE_SHOT assert forward_batch.seq_lens_cpu is not None max_kv_len = forward_batch.seq_lens_cpu.max().item() sum_seq_lens = sum(forward_batch.seq_lens_cpu) device_sm = get_device_sm() # Requirements: H200/B200, short sequences, supported dtype, fits in chunk self.use_mha = ( ( device_sm == 90 or (device_sm >= 100 and device_sm < 110) ) # SM90/SM100 only and max_kv_len <= envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get() # Short enough for MHA and self.token_to_kv_pool.dtype in [torch.bfloat16, torch.float8_e4m3fn] and sum_seq_lens <= forward_batch.get_max_chunk_capacity() # Fits in chunk and (not is_dsa_enable_prefill_cp()) # CP not enabled and (self.hisparse_coordinator is None) ) else: self.use_mha = False # Decode/verify always use MLA # Set MLA implementation only if not using MHA if not self.use_mha and self.enable_auto_select_prefill_impl: if self.dsa_kv_cache_store_fp8: if ( is_blackwell() and forward_batch is not None and forward_batch.forward_mode == ForwardMode.EXTEND ): total_kv_tokens = forward_batch.seq_lens_sum total_q_tokens = forward_batch.extend_num_tokens # Heuristic based on benchmarking flashmla_kv vs flashmla_sparse + dequantize_k_cache_paged if total_kv_tokens < total_q_tokens * 512: self.dsa_prefill_impl = "flashmla_sparse" return self.dsa_prefill_impl = "flashmla_kv" else: # bf16 kv cache self.dsa_prefill_impl = "flashmla_sparse" def get_topk_transform_method( self, forward_mode: Optional[ForwardMode] = None ) -> TopkTransformMethod: """ SGLANG_DSA_FUSE_TOPK controls whether to fuse the topk transform into the topk kernel. This method is used to select the topk transform method which can be fused or unfused. """ if ( # disable for MTP self.dsa_kv_cache_store_fp8 and self.dsa_prefill_impl == "flashmla_sparse" and forward_mode == ForwardMode.EXTEND ): topk_transform_method = TopkTransformMethod.RAGGED else: topk_transform_method = TopkTransformMethod.PAGED return topk_transform_method def get_indexer_metadata( self, layer_id: int, forward_batch: ForwardBatch ) -> DSAIndexerMetadata: force_unfused = ( self.hisparse_coordinator is not None and forward_batch.forward_mode.is_decode_or_idle() ) return DSAIndexerMetadata( attn_metadata=self.forward_metadata, topk_transform_method=self.get_topk_transform_method( forward_batch.forward_mode ), topk_backend=self.dsa_topk_backend, paged_mqa_schedule_metadata=self.forward_metadata.paged_mqa_schedule_metadata, paged_mqa_ctx_lens_2d=self.forward_metadata.paged_mqa_ctx_lens_2d, force_unfused_topk=force_unfused, ) def _compute_flashmla_metadata(self, cache_seqlens: torch.Tensor, seq_len_q: int): from sgl_kernel.flash_mla import get_mla_metadata num_heads_q = self.flashmla_kv_num_q_heads flashmla_metadata, num_splits = get_mla_metadata( cache_seqlens=cache_seqlens, # TODO doc says `num_q_tokens_per_q_seq * num_heads_q // num_heads_k` # but the name looks like need seq_len_q? num_q_tokens_per_head_k=seq_len_q * num_heads_q // 1, num_heads_k=1, num_heads_q=num_heads_q, is_fp8_kvcache=True, topk=self.dsa_index_topk, ) return DSAFlashMLAMetadata( flashmla_metadata=flashmla_metadata, num_splits=num_splits, ) class DeepseekSparseAttnMultiStepBackend: # Per-step draft decode replays from precomputed GPU metadata; opt out so # decide_needs_cpu_seq_lens' OR over the backends stays False. needs_cpu_seq_lens: bool = False def __init__( self, model_runner: ModelRunner, topk: int, speculative_num_steps: int ): self.topk = topk self.speculative_num_steps = speculative_num_steps self.attn_backends = [] for i in range(self.speculative_num_steps - 1): self.attn_backends.append( DeepseekSparseAttnBackend( model_runner, speculative_step_id=i, topk=self.topk, speculative_num_steps=self.speculative_num_steps, ) ) def init_forward_metadata(self, forward_batch: ForwardBatch): for i in range(self.speculative_num_steps - 1): self.attn_backends[i].init_forward_metadata(forward_batch) def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): for i in range(self.speculative_num_steps - 1): self.attn_backends[i].init_cuda_graph_state(max_bs, max_num_tokens) def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, in_capture: bool = False, ): from sglang.srt.model_executor.forward_batch_info import build_inner_fb_view if in_capture: inner_fb = build_inner_fb_view( forward_batch, bs=forward_batch.batch_size, forward_mode=ForwardMode.DECODE, ) for i in range(self.speculative_num_steps - 1): self.attn_backends[i].init_forward_metadata_out_graph( inner_fb, in_capture=True ) return bs = forward_batch.batch_size # Precompute metadata once (shared across all backends) precomputed = self.attn_backends[0]._precompute_replay_metadata( bs=bs, req_pool_indices=forward_batch.req_pool_indices, seq_lens=forward_batch.seq_lens, seq_lens_cpu=forward_batch.seq_lens_cpu, forward_mode=ForwardMode.DECODE, ) # Use multi-backend fused copy when we have 3 or more backends # This is 3x faster than calling the single-backend copy 3 times if self.speculative_num_steps > 3: try: from sglang.jit_kernel.fused_metadata_copy import ( fused_metadata_copy_multi_cuda, ) metadata0 = self.attn_backends[0].decode_cuda_graph_metadata[bs] metadata1 = self.attn_backends[1].decode_cuda_graph_metadata[bs] metadata2 = self.attn_backends[2].decode_cuda_graph_metadata[bs] # Set dsa_prefill_impl for first 3 backends (required by the method) for i in range(3): self.attn_backends[i].set_dsa_prefill_impl(forward_batch=None) # Prepare FlashMLA tensors if needed flashmla_num_splits_src = None flashmla_metadata_src = None flashmla_num_splits_dst0 = None flashmla_num_splits_dst1 = None flashmla_num_splits_dst2 = None flashmla_metadata_dst0 = None flashmla_metadata_dst1 = None flashmla_metadata_dst2 = None if precomputed.flashmla_metadata is not None: flashmla_num_splits_src = precomputed.flashmla_metadata.num_splits flashmla_metadata_src = ( precomputed.flashmla_metadata.flashmla_metadata ) flashmla_num_splits_dst0 = metadata0.flashmla_metadata.num_splits flashmla_num_splits_dst1 = metadata1.flashmla_metadata.num_splits flashmla_num_splits_dst2 = metadata2.flashmla_metadata.num_splits flashmla_metadata_dst0 = ( metadata0.flashmla_metadata.flashmla_metadata ) flashmla_metadata_dst1 = ( metadata1.flashmla_metadata.flashmla_metadata ) flashmla_metadata_dst2 = ( metadata2.flashmla_metadata.flashmla_metadata ) # Call the multi-backend fused kernel for first 3 backends fused_metadata_copy_multi_cuda( # Source tensors precomputed.cache_seqlens, precomputed.cu_seqlens_k, precomputed.page_indices, precomputed.dsa_cache_seqlens, precomputed.dsa_cu_seqlens_k, precomputed.real_page_table, flashmla_num_splits_src, flashmla_metadata_src, # Destination tensors for backend 0 metadata0.cache_seqlens_int32, metadata0.cu_seqlens_k, metadata0.page_table_1, metadata0.dsa_cache_seqlens_int32, metadata0.dsa_cu_seqlens_k, ( metadata0.real_page_table if precomputed.real_page_table is not None else None ), flashmla_num_splits_dst0, flashmla_metadata_dst0, # Destination tensors for backend 1 metadata1.cache_seqlens_int32, metadata1.cu_seqlens_k, metadata1.page_table_1, metadata1.dsa_cache_seqlens_int32, metadata1.dsa_cu_seqlens_k, ( metadata1.real_page_table if precomputed.real_page_table is not None else None ), flashmla_num_splits_dst1, flashmla_metadata_dst1, # Destination tensors for backend 2 metadata2.cache_seqlens_int32, metadata2.cu_seqlens_k, metadata2.page_table_1, metadata2.dsa_cache_seqlens_int32, metadata2.dsa_cu_seqlens_k, ( metadata2.real_page_table if precomputed.real_page_table is not None else None ), flashmla_num_splits_dst2, flashmla_metadata_dst2, # Parameters bs, precomputed.max_len, precomputed.seqlens_expanded_size, ) # Copy remaining backends one by one (if > 3 backends) for i in range(3, self.speculative_num_steps - 1): self.attn_backends[ i ].init_forward_metadata_replay_cuda_graph_from_precomputed( bs=bs, precomputed=precomputed, forward_mode=ForwardMode.DECODE, ) except (ImportError, Exception) as e: # Fallback to loop if multi-backend kernel not available or fails if isinstance(e, ImportError): print( "Warning: Multi-backend fused metadata copy kernel not available, falling back to loop." ) else: print( f"Warning: Multi-backend fused metadata copy kernel failed with error: {e}, falling back to loop." ) for i in range(self.speculative_num_steps - 1): self.attn_backends[ i ].init_forward_metadata_replay_cuda_graph_from_precomputed( bs=bs, precomputed=precomputed, forward_mode=ForwardMode.DECODE, ) else: # Less than 3 backends: copy to each backend individually for i in range(self.speculative_num_steps - 1): self.attn_backends[ i ].init_forward_metadata_replay_cuda_graph_from_precomputed( bs=bs, precomputed=precomputed, forward_mode=ForwardMode.DECODE, ) def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None: for i in range(self.speculative_num_steps - 1): self.attn_backends[i].init_forward_metadata_in_graph(forward_batch) # Backward-compat aliases (deprecated: use DSA class names) DeepseekSparseAttnBackend = DeepseekSparseAttnBackend DeepseekSparseAttnMultiStepBackend = DeepseekSparseAttnMultiStepBackend DSAMetadata = DSAMetadata DSAFlashMLAMetadata = DSAFlashMLAMetadata DSAIndexerMetadata = DSAIndexerMetadata