chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Performance] optimize NSA backend metadata computation for multi-step speculative decoding (#14781)

Browse Source
This commit is contained in:
Johnsonms
2025-12-18 13:48:27 -08:00
committed by GitHub
parent 9749d3e346
commit e0026f7c92
3 changed files with 440 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

324
python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py Normal file
View File

@@ -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,
)

5
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -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 = ""

127
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -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,
)