From e0026f7c92c91f7c039ab7b823caf65207c8cbb2 Mon Sep 17 00:00:00 2001 From: Johnsonms Date: Thu, 18 Dec 2025 13:48:27 -0800 Subject: [PATCH] [Performance] optimize NSA backend metadata computation for multi-step speculative decoding (#14781) --- .../nsa/nsa_backend_mtp_precompute.py | 324 ++++++++++++++++++ .../sglang/srt/layers/attention/nsa/utils.py | 5 + .../srt/layers/attention/nsa_backend.py | 127 ++++++- 3 files changed, 440 insertions(+), 16 deletions(-) create mode 100644 python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py diff --git a/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py b/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py new file mode 100644 index 000000000..b9450ce09 --- /dev/null +++ b/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py @@ -0,0 +1,324 @@ +"""Multi-step precompute utilities for Native Sparse Attention backend. + +This module provides optimization utilities for multi-step speculative decoding +by precomputing shared metadata once and copying it to multiple backend instances. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.layers.attention.nsa.utils import compute_nsa_seqlens + +if TYPE_CHECKING: + from sglang.srt.model_executor.forward_batch_info import ForwardMode + from sglang.srt.speculative.spec_info import SpecInput + + +@dataclass +class PrecomputedMetadata: + """Precomputed metadata shared across multiple backend instances. + + Used for multi-step speculative decoding where multiple backends + need identical metadata. Precomputing once and copying N times + is much faster than computing N times. + + """ + + # Basic seqlens + cache_seqlens: torch.Tensor # int32, [bs] + cu_seqlens_k: torch.Tensor # int32, [bs+1] + + # Page table + page_indices: torch.Tensor # int32, [bs, max_len] or [expanded_bs, max_len] + real_page_table: Optional[torch.Tensor] # int32, transformed version + + # NSA seqlens + seqlens_expanded: torch.Tensor # int32, [expanded_size] + nsa_cache_seqlens: torch.Tensor # int32, [expanded_size] + nsa_cu_seqlens_k: torch.Tensor # int32, [expanded_size+1] + seqlens_expanded_size: int + + # Dimensions + max_len: int # for decode/draft_extend + max_seqlen_k: int # for target_verify + + # FlashMLA (optional) + flashmla_metadata: Optional[torch.Tensor] = None + + +def compute_cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor: + """Compute cumulative sequence lengths with padding.""" + assert seqlens.dtype == torch.int32 + return torch.nn.functional.pad( + torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0) + ) + + +class NativeSparseAttnBackendMTPPrecomputeMixin: + """Mixin class providing metadata precomputation for multi-step speculative decoding. + + This mixin provides the _precompute_replay_metadata method and its helpers, + which are used to optimize CUDA graph replay in multi-step scenarios. + """ + + def _precompute_replay_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"], + ) -> PrecomputedMetadata: + """Precompute all shared metadata for multi-step backends. + + This function extracts and computes all operations that are + identical across different backend instances in multi-step + speculative decoding. + + Args: + bs: Batch size + req_pool_indices: Request pool indices [bs] + seq_lens: Sequence lengths [bs] + seq_lens_cpu: Sequence lengths on CPU [bs] + forward_mode: Forward mode (decode/target_verify/draft_extend) + spec_info: Speculative decoding info (for draft_extend mode) + + Returns: + PrecomputedMetadata containing all shared intermediate results + """ + # Slice inputs to batch size + seq_lens = seq_lens[:bs] + seq_lens_cpu = seq_lens_cpu[:bs] + req_pool_indices = req_pool_indices[:bs] + + # Dispatch to mode-specific precomputation + if forward_mode.is_decode_or_idle(): + return self._precompute_decode_mode( + bs, req_pool_indices, seq_lens, seq_lens_cpu + ) + elif forward_mode.is_target_verify(): + return self._precompute_target_verify_mode( + bs, req_pool_indices, seq_lens, seq_lens_cpu + ) + elif forward_mode.is_draft_extend(): + return self._precompute_draft_extend_mode( + bs, req_pool_indices, seq_lens, seq_lens_cpu, spec_info + ) + else: + raise ValueError(f"Unsupported forward mode: {forward_mode}") + + def _precompute_decode_mode( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + ) -> PrecomputedMetadata: + """Precompute metadata for normal decode mode.""" + max_len = int(seq_lens_cpu.max().item()) + + # Convert to int32 and compute cumsum + cache_seqlens = seq_lens.to(torch.int32) + cu_seqlens_k = compute_cu_seqlens(cache_seqlens) + + # Get page indices from cache + page_indices = self.req_to_token[req_pool_indices, :max_len] + + # Compute NSA seqlens + nsa_cache_seqlens = compute_nsa_seqlens( + cache_seqlens, nsa_index_topk=self.nsa_index_topk + ) + seqlens_expanded = cache_seqlens + seqlens_expanded_size = seqlens_expanded.shape[0] + + # Compute NSA cumsum + nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens) + + # Transform page table if needed + if self.real_page_size > 1: + real_page_table = self._transform_table_1_to_real(page_indices) + else: + real_page_table = None # Will use page_indices directly + + # Compute FlashMLA metadata if needed + flashmla_metadata = None + if self.nsa_decode_impl == "flashmla_kv": + flashmla_metadata = self._compute_flashmla_metadata( + cache_seqlens=nsa_cache_seqlens, + seq_len_q=1, + ) + + return PrecomputedMetadata( + cache_seqlens=cache_seqlens, + cu_seqlens_k=cu_seqlens_k, + page_indices=page_indices, + real_page_table=real_page_table, + seqlens_expanded=seqlens_expanded, + nsa_cache_seqlens=nsa_cache_seqlens, + nsa_cu_seqlens_k=nsa_cu_seqlens_k, + seqlens_expanded_size=seqlens_expanded_size, + max_len=max_len, + max_seqlen_k=max_len, + flashmla_metadata=flashmla_metadata, + ) + + def _precompute_target_verify_mode( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + ) -> PrecomputedMetadata: + """Precompute metadata for target verify mode.""" + max_seqlen_k = int( + seq_lens_cpu.max().item() + self.speculative_num_draft_tokens + ) + + # Cache seqlens with draft tokens + cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to(torch.int32) + cu_seqlens_k = compute_cu_seqlens(cache_seqlens) + + # Page indices (repeated for each draft token) + 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 + ) + + # Generate expanded seqlens + 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_cpu.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, + ) + ] + ) + + # Compute NSA seqlens + nsa_cache_seqlens = compute_nsa_seqlens(seqlens_expanded, self.nsa_index_topk) + seqlens_expanded_size = seqlens_expanded.shape[0] + + # NSA cumsum + nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens) + + # Transform page table + if self.real_page_size > 1: + real_page_table = self._transform_table_1_to_real(page_indices) + else: + real_page_table = None + + # FlashMLA metadata + flashmla_metadata = None + if self.nsa_decode_impl == "flashmla_kv": + flashmla_metadata = self._compute_flashmla_metadata( + cache_seqlens=nsa_cache_seqlens, + seq_len_q=1, + ) + + return PrecomputedMetadata( + cache_seqlens=cache_seqlens, + cu_seqlens_k=cu_seqlens_k, + page_indices=page_indices, + real_page_table=real_page_table, + seqlens_expanded=seqlens_expanded, + nsa_cache_seqlens=nsa_cache_seqlens, + nsa_cu_seqlens_k=nsa_cu_seqlens_k, + seqlens_expanded_size=seqlens_expanded_size, + max_len=-1, # Not used in this mode + max_seqlen_k=max_seqlen_k, + flashmla_metadata=flashmla_metadata, + ) + + def _precompute_draft_extend_mode( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + spec_info: "SpecInput", + ) -> PrecomputedMetadata: + """Precompute metadata for draft extend mode.""" + max_seqlen_k = int(seq_lens_cpu.max().item()) + + # Cache seqlens + cache_seqlens = seq_lens.to(torch.int32) + cu_seqlens_k = compute_cu_seqlens(cache_seqlens) + + # Extend seqlens from spec_info + extend_seq_lens = spec_info.accept_length[:bs] + extend_seq_lens_cpu = extend_seq_lens.tolist() + + # Page indices (repeated per accept length) + page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k] + page_indices = torch.repeat_interleave( + page_indices, repeats=extend_seq_lens, dim=0 + ) + + # Generate expanded seqlens + 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, + seq_lens_cpu.tolist(), + strict=True, + ) + ] + ) + + # Compute NSA seqlens + nsa_cache_seqlens = compute_nsa_seqlens(seqlens_expanded, self.nsa_index_topk) + seqlens_expanded_size = seqlens_expanded.shape[0] + + # NSA cumsum + nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens) + + # Transform page table + if self.real_page_size > 1: + real_page_table = self._transform_table_1_to_real(page_indices) + else: + real_page_table = None + + # FlashMLA metadata + flashmla_metadata = None + if self.nsa_decode_impl == "flashmla_kv": + flashmla_metadata = self._compute_flashmla_metadata( + cache_seqlens=nsa_cache_seqlens, + seq_len_q=1, + ) + + return PrecomputedMetadata( + cache_seqlens=cache_seqlens, + cu_seqlens_k=cu_seqlens_k, + page_indices=page_indices, + real_page_table=real_page_table, + seqlens_expanded=seqlens_expanded, + nsa_cache_seqlens=nsa_cache_seqlens, + nsa_cu_seqlens_k=nsa_cu_seqlens_k, + seqlens_expanded_size=seqlens_expanded_size, + max_len=max_seqlen_k, + max_seqlen_k=max_seqlen_k, + flashmla_metadata=flashmla_metadata, + ) diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 24371cd4e..9cbaa6c0a 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -19,6 +19,11 @@ NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8 = get_bool_env_var( NSA_QUANT_K_CACHE_FAST = get_bool_env_var("SGLANG_NSA_QUANT_K_CACHE_FAST", "true") NSA_DEQUANT_K_CACHE_FAST = get_bool_env_var("SGLANG_NSA_DEQUANT_K_CACHE_FAST", "true") +# Environment variable to control mtp precomputing of metadata for multi-step speculative decoding +NSA_ENABLE_MTP_PRECOMPUTE_METADATA = get_bool_env_var( + "SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA", "true" +) + def print_nsa_bool_env_vars(): msg = "" diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 7ea4322c2..18b1b9daf 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -10,6 +10,11 @@ from sglang.srt.configs.model_config import get_nsa_index_topk, is_deepseek_nsa from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.nsa.dequant_k_cache import dequantize_k_cache_paged +from sglang.srt.layers.attention.nsa.nsa_backend_mtp_precompute import ( + NativeSparseAttnBackendMTPPrecomputeMixin, + PrecomputedMetadata, + compute_cu_seqlens, +) from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata from sglang.srt.layers.attention.nsa.quant_k_cache import quantize_k_cache from sglang.srt.layers.attention.nsa.transform_index import ( @@ -17,6 +22,7 @@ from sglang.srt.layers.attention.nsa.transform_index import ( transform_index_page_table_prefill, ) from sglang.srt.layers.attention.nsa.utils import ( + NSA_ENABLE_MTP_PRECOMPUTE_METADATA, NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8, NSA_FUSE_TOPK, compute_nsa_seqlens, @@ -224,17 +230,12 @@ class NSAIndexerMetadata(BaseIndexerMetadata): assert False, f"Unsupported {self.topk_transform_method = }" -def compute_cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor: - assert seqlens.dtype == torch.int32 - return torch.nn.functional.pad( - torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0) - ) - - _NSA_IMPL_T: TypeAlias = Literal["flashmla_sparse", "flashmla_kv", "fa3", "tilelang"] -class NativeSparseAttnBackend(AttentionBackend): +class NativeSparseAttnBackend( + NativeSparseAttnBackendMTPPrecomputeMixin, AttentionBackend +): def __init__( self, model_runner: ModelRunner, @@ -886,6 +887,77 @@ class NativeSparseAttnBackend(AttentionBackend): 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_nsa_prefill_impl(forward_batch=None) + + metadata = self.decode_cuda_graph_metadata[bs] + + # 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.nsa_cache_seqlens_int32.copy_(precomputed.nsa_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.nsa_seqlens_expanded.copy_(precomputed.seqlens_expanded) + metadata.nsa_cache_seqlens_int32.copy_(precomputed.nsa_cache_seqlens) + + elif forward_mode.is_draft_extend(): + # Draft extend mode + rows = precomputed.page_indices.shape[0] + cols = precomputed.max_seqlen_k + metadata.page_table_1[:rows, :cols].copy_(precomputed.page_indices) + + size = precomputed.seqlens_expanded_size + metadata.nsa_seqlens_expanded[:size].copy_(precomputed.seqlens_expanded) + metadata.nsa_cache_seqlens_int32[:size].copy_(precomputed.nsa_cache_seqlens) + + # Copy NSA cu_seqlens + size = precomputed.seqlens_expanded_size + metadata.nsa_cu_seqlens_k[1 : 1 + size].copy_( + precomputed.nsa_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) + else: + # real_page_table is same as page_table_1 (already copied) + pass + + # Copy FlashMLA metadata + if precomputed.flashmla_metadata is not None: + flashmla_metadata = metadata.flashmla_metadata.slice(slice(0, size + 1)) + flashmla_metadata.copy_(precomputed.flashmla_metadata) + + self.forward_metadata = metadata + def forward_extend( self, q: torch.Tensor, @@ -1587,14 +1659,37 @@ class NativeSparseAttnMultiStepBackend: def init_forward_metadata_replay_cuda_graph( self, forward_batch: ForwardBatch, bs: int ): - for i in range(self.speculative_num_steps): - self.attn_backends[i].init_forward_metadata_replay_cuda_graph( - bs, - forward_batch.req_pool_indices, - forward_batch.seq_lens, - seq_lens_sum=-1, - encoder_lens=None, + if NSA_ENABLE_MTP_PRECOMPUTE_METADATA: + # 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, spec_info=forward_batch.spec_info, - seq_lens_cpu=forward_batch.seq_lens_cpu, ) + + # Fast copy to each backend (1-2x faster than computing N times) + for i in range(self.speculative_num_steps): + self.attn_backends[ + i + ].init_forward_metadata_replay_cuda_graph_from_precomputed( + bs=bs, + precomputed=precomputed, + forward_mode=ForwardMode.DECODE, + ) + else: + # Fallback: compute metadata separately for each backend + for i in range(self.speculative_num_steps): + self.attn_backends[i].init_forward_metadata_replay_cuda_graph( + bs=bs, + req_pool_indices=forward_batch.req_pool_indices, + seq_lens=forward_batch.seq_lens, + seq_lens_sum=forward_batch.seq_lens_sum, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + seq_lens_cpu=forward_batch.seq_lens_cpu, + out_cache_loc=None, + )