minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support piecewise CUDA graph with NSA (#23351)

Browse Source
This commit is contained in:
nvjullin
2026-05-22 14:39:50 -07:00
committed by GitHub
parent 2df9e8b4b3
commit cadfa2d025
12 changed files with 317 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/hadamard.py
+9
View File
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Callable
import torch
from sglang.jit_kernel.utils import KERNEL_PATH, cache_once, load_jit, make_cpp_args
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
from tvm_ffi.module import Module
@@ -56,6 +57,14 @@ def _hadamard_transform_impl(
return out.reshape(shapes_og)
def _hadamard_transform_fake_impl(
x: torch.Tensor,
scale: float = 1.0,
) -> torch.Tensor:
return torch.empty_like(x)
@register_custom_op(fake_impl=_hadamard_transform_fake_impl)
def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
module = _jit_hadamard_module(x.dtype)
return _hadamard_transform_impl(x, scale, 8, module.hadamard_transform)
python/sglang/srt/compilation/piecewise_context_manager.py
+7
View File
@@ -71,6 +71,7 @@ class ForwardContext:
self.quant_config = None
self.moe_layers = None
self.moe_fusions = None
self.dsa_indexers = None
def set_forward_batch(self, forward_batch: ForwardBatch):
self.forward_batch = forward_batch
@@ -87,6 +88,9 @@ class ForwardContext:
def set_moe_fusions(self, fusions: List[Any]):
self.moe_fusions = fusions
def set_dsa_indexers(self, indexers: List[Any]):
self.dsa_indexers = indexers
_forward_context: Optional[ForwardContext] = None
@@ -104,6 +108,7 @@ def set_forward_context(
quant_config: Any,
moe_layers: List[Any],
moe_fusions: List[Any],
dsa_indexers: Optional[List[Any]] = None,
):
global _forward_context
_forward_context = ForwardContext()
@@ -112,6 +117,8 @@ def set_forward_context(
_forward_context.set_quant_config(quant_config)
_forward_context.set_moe_layers(moe_layers)
_forward_context.set_moe_fusions(moe_fusions)
if dsa_indexers is not None:
_forward_context.set_dsa_indexers(dsa_indexers)
try:
yield
finally:
python/sglang/srt/configs/model_config.py
-3
View File
@@ -329,7 +329,6 @@ class ModelConfig:
self.use_ngram_embedding = getattr(self.hf_config, "use_ngram_embedding", False)
self.is_piecewise_cuda_graph_disabled_model = (
is_piecewise_cuda_graph_disabled_model(self.hf_config.architectures)
or is_deepseek_dsa(self.hf_text_config)
)
self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)
@@ -1556,11 +1555,9 @@ multimodal_model_archs = [
]
piecewise_cuda_graph_disabled_model_archs = [
"DeepseekV32ForCausalLM",
"DeepseekV4ForCausalLM",
"DeepseekV4ForCausalLMNextN",
"Qwen3NextForCausalLM",
"GlmMoeDsaForCausalLM",
"BailingMoeV2_5ForCausalLM",
"LLaDAModelLM",
]
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+167 -35
View File
@@ -12,6 +12,10 @@ from sglang.jit_kernel.fused_store_index_cache import (
can_use_dsa_fused_store,
fused_store_index_k_cache,
)
from sglang.srt.compilation.piecewise_context_manager import (
get_forward_context,
is_in_piecewise_cuda_graph,
)
from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.utils import (
aiter_can_use_preshuffle_paged_mqa,
@@ -95,6 +99,80 @@ if TYPE_CHECKING:
DUAL_STREAM_TOKEN_THRESHOLD = 1024 if _is_cuda else 0
if _is_cuda:
from sglang.srt.compilation.compilation_config import register_split_op
from sglang.srt.utils.custom_op import register_custom_op
@register_custom_op(mutates_args=["topk_result"])
@register_split_op()
def k_cache_and_topk_result(
layer_id: int,
key: torch.Tensor,
q_fp8: torch.Tensor,
weights: torch.Tensor,
topk_result: torch.Tensor,
) -> None:
assert (
_is_cuda
), "Internal error: piecewise CUDA graph is only supported on CUDA"
from sglang.srt.layers.attention.dsa.triton_kernel import act_quant
forward_batch = get_forward_context().forward_batch
indexer = get_forward_context().dsa_indexers[layer_id]
metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch)
# slice off padding from piecewise CUDA graph
extend_num_tokens = forward_batch.extend_num_tokens
indexer._store_index_k_cache(
forward_batch=forward_batch,
layer_id=layer_id,
key=key[:extend_num_tokens],
act_quant=act_quant,
out_cache_loc=forward_batch.out_cache_loc[:extend_num_tokens],
)
indexer._get_topk_ragged(
False,
forward_batch,
layer_id,
q_fp8[:extend_num_tokens],
weights,
metadata,
topk_result,
)
def _logits_head_gate_pcg_fake_impl(
x: torch.Tensor,
weight: torch.Tensor,
n_heads_inv_sqrt: float,
softmax_scale: float,
q_scale: torch.Tensor,
) -> torch.Tensor:
return torch.empty(
(x.shape[0], weight.shape[0], q_scale.shape[-1]),
dtype=torch.float32,
device=x.device,
)
@register_custom_op(fake_impl=_logits_head_gate_pcg_fake_impl)
def logits_head_gate_pcg(
x: torch.Tensor,
weight: torch.Tensor,
n_heads_inv_sqrt: float,
softmax_scale: float,
q_scale: torch.Tensor,
) -> torch.Tensor:
from sglang.srt.layers.deep_gemm_wrapper import entrypoint as deep_gemm_wrapper
out = torch.empty(
(x.shape[0], weight.shape[0]), dtype=torch.float32, device=x.device
)
deep_gemm_wrapper.gemm_nt_bf16bf16f32(x, weight, out)
weights = out * n_heads_inv_sqrt
weights = weights.unsqueeze(-1) * q_scale * softmax_scale
return weights
class BaseIndexerMetadata(ABC):
@abstractmethod
def get_seqlens_int32(self) -> torch.Tensor:
@@ -441,7 +519,8 @@ class Indexer(MultiPlatformOp):
def _update_rope_guarded(dst: torch.Tensor, src: torch.Tensor) -> None:
# On AMD with in-place RoPE kernels, self-aliasing can occur;
# skip write-back when src/dst tensors point to a single memory.
if src.data_ptr() == dst.data_ptr():
# data_ptr() is not comparable inside torch.compile, so skip the guard there.
if not torch.compiler.is_compiling() and src.data_ptr() == dst.data_ptr():
return
dst.copy_(src)
@@ -627,6 +706,7 @@ class Indexer(MultiPlatformOp):
q_fp8: torch.Tensor,
weights: torch.Tensor,
metadata: BaseIndexerMetadata,
topk_result: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if TYPE_CHECKING:
assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool)
@@ -669,9 +749,10 @@ class Indexer(MultiPlatformOp):
device_index = device.index
assert device_index is not None, "q_fp8 must be on an indexed CUDA device"
topk_result = torch.full(
(token_nums, self.index_topk), -1, device=device, dtype=torch.int32
)
if topk_result is None:
topk_result = torch.full(
(token_nums, self.index_topk), -1, device=device, dtype=torch.int32
)
if batch_size == 0:
return topk_result
@@ -855,6 +936,9 @@ class Indexer(MultiPlatformOp):
actual_seq_q: int,
cp_index: List[Tuple[int, int, int]] = None,
) -> torch.Tensor:
assert (
not is_in_piecewise_cuda_graph()
), "DSA context parallel (_get_topk_ragged_with_cp) not supported under piecewise CUDA graph"
if TYPE_CHECKING:
assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool)
@@ -1001,6 +1085,9 @@ class Indexer(MultiPlatformOp):
topk: int,
layer_id: int,
) -> Optional[torch.Tensor]:
assert (
not is_in_piecewise_cuda_graph()
), "DSA forward_indexer (non-CUDA loop path) not supported under piecewise CUDA graph"
if not _is_npu:
from sglang.srt.layers.attention.dsa.tilelang_kernel import fp8_index
@@ -1083,21 +1170,26 @@ class Indexer(MultiPlatformOp):
key: torch.Tensor,
*,
act_quant=None, # fallback only
out_cache_loc: Optional[torch.Tensor] = None,
) -> None:
"""
Store DSA indexer K cache for current step.
Preferred: fused_store_index_k_cache(key, cache, out_cache_loc, page_size)
Fallback : act_quant(key) + token_to_kv_pool.set_index_k_scale_buffer(...)
out_cache_loc will default to forward_batch.out_cache_loc if not provided.
"""
# Fast path: JIT fused store (CUDA, page_size=64, non-fnuz)
if out_cache_loc is None:
out_cache_loc = forward_batch.out_cache_loc
if (
_is_cuda
and (not _is_fp8_fnuz)
and can_use_dsa_fused_store(
key.dtype,
forward_batch.out_cache_loc.dtype,
out_cache_loc.dtype,
get_token_to_kv_pool().page_size,
)
):
@@ -1108,7 +1200,7 @@ class Indexer(MultiPlatformOp):
fused_store_index_k_cache(
key,
buf,
forward_batch.out_cache_loc,
out_cache_loc,
get_token_to_kv_pool().page_size,
)
return
@@ -1141,13 +1233,12 @@ class Indexer(MultiPlatformOp):
assert act_quant is not None
k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt)
out_loc = forward_batch.out_cache_loc
if not out_loc.is_contiguous():
out_loc = out_loc.contiguous()
if not out_cache_loc.is_contiguous():
out_cache_loc = out_cache_loc.contiguous()
get_token_to_kv_pool().set_index_k_scale_buffer(
layer_id=layer_id,
loc=out_loc,
loc=out_cache_loc,
index_k=k_fp8,
index_k_scale=k_scale,
)
@@ -1186,7 +1277,15 @@ class Indexer(MultiPlatformOp):
# a tuple like (x_fp8, x_scale[, y]). Use `x_meta` for shape/device queries.
x_meta = x[0] if isinstance(x, tuple) else x
metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch)
# In piecewise CUDA graph mode, metadata is fetched inside custom ops via get_forward_context() to
# prevent Dynamo from guarding on forward_metadata identity (which changes each
# replay when init_forward_metadata creates a new ForwardMetadata object).
if not is_in_piecewise_cuda_graph():
metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch)
if metadata is None:
return None
else:
metadata = None
enable_dual_stream = (
self.alt_stream is not None
@@ -1195,14 +1294,13 @@ class Indexer(MultiPlatformOp):
and q_lora.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD
)
# skip DSA if attention backend choose to skip this batch
if metadata is None:
return None
# Determine if should skip topk based on sequence length
# We can only skip the logits computation if cuda graph is not involved
skip_logits_computation = False
if forward_batch.forward_mode.is_extend_without_speculative():
if (
not is_in_piecewise_cuda_graph()
and forward_batch.forward_mode.is_extend_without_speculative()
):
if forward_batch.seq_lens_cpu is not None:
max_kv_len = forward_batch.seq_lens_cpu.max().item()
skip_logits_computation = max_kv_len <= self.index_topk
@@ -1258,7 +1356,7 @@ class Indexer(MultiPlatformOp):
act_quant=act_quant,
)
current_stream.wait_stream(self.alt_stream)
else:
elif not is_in_piecewise_cuda_graph():
q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
self._store_index_k_cache(
forward_batch=forward_batch,
@@ -1266,6 +1364,10 @@ class Indexer(MultiPlatformOp):
key=key,
act_quant=act_quant,
)
else:
# piecewise CUDA graph need to split graph on store_k_cache and mqa_logits,
# so delay store_k_cache after weights proj.
q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
# aiter (ROCm gfx95): the 3-tuple (fp8, scale, bf16) from
# fused_rms_fp8_group_quant is passed directly to _get_logits_head_gate,
@@ -1308,25 +1410,37 @@ class Indexer(MultiPlatformOp):
else:
x_for_gate = x
weights = self._get_logits_head_gate(x_for_gate, q_scale)
if is_in_piecewise_cuda_graph():
weights = logits_head_gate_pcg(
x_for_gate,
self.weights_proj.weight,
self.n_heads**-0.5,
self.softmax_scale,
q_scale,
)
else:
weights = self._get_logits_head_gate(x_for_gate, q_scale)
if _is_cuda or _is_hip:
assert forward_batch.seq_lens_cpu is not None
if len(forward_batch.seq_lens_cpu) == 0:
# this seems b/c max-pad, no worries?
# if x.shape[0] != 0:
# print(
# "HACK: seq_lens empty but x not empty, hackily return all-invalid topk_result"
# )
return maybe_capture_indexer_topk(
layer_id,
torch.full(
(x_meta.shape[0], self.index_topk),
-1,
dtype=torch.int,
device=x_meta.device,
),
)
# In piecewise CUDA graph, any access to seq_lens_cpu creates a Dynamo shape guard.
# Piecewise CUDA graph never has empty batches.
if not is_in_piecewise_cuda_graph():
assert forward_batch.seq_lens_cpu is not None
if len(forward_batch.seq_lens_cpu) == 0:
# this seems b/c max-pad, no worries?
# if x.shape[0] != 0:
# print(
# "HACK: seq_lens empty but x not empty, hackily return all-invalid topk_result"
# )
return maybe_capture_indexer_topk(
layer_id,
torch.full(
(x_meta.shape[0], self.index_topk),
-1,
dtype=torch.int,
device=x_meta.device,
),
)
if (
forward_batch.forward_mode.is_decode_or_idle()
@@ -1379,6 +1493,24 @@ class Indexer(MultiPlatformOp):
layer_id,
torch.cat([topk_result_prev, topk_result_next], dim=0),
)
elif is_in_piecewise_cuda_graph():
assert (
not enable_dual_stream
), "Internal error: piecewise CUDA graph should not be enabled with dual stream"
topk_result = torch.full(
(q_fp8.shape[0], self.index_topk),
-1,
device=q_fp8.device,
dtype=torch.int32,
)
k_cache_and_topk_result(
layer_id=layer_id,
key=key,
q_fp8=q_fp8,
weights=weights,
topk_result=topk_result,
)
else:
topk_result = self._get_topk_ragged(
enable_dual_stream,
python/sglang/srt/layers/attention/dsa_backend.py
+14 -3
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import logging
from dataclasses import dataclass
from enum import IntEnum, auto
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, TypeAlias
@@ -7,6 +8,8 @@ from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, TypeAlia
import torch
from sglang.srt.configs.model_config import get_dsa_index_topk, is_deepseek_dsa
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
@@ -2169,8 +2172,8 @@ class DeepseekSparseAttnBackend(
backend="trtllm-gen",
skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
)
# Output: [batch, q_len=1, heads, v_dim] -> [batch, heads, v_dim]
return out.squeeze(1)
return out
def _pad_topk_indices(
self, topk_indices: torch.Tensor, num_tokens: int
@@ -2201,10 +2204,18 @@ class DeepseekSparseAttnBackend(
"""
Decide all attention prefill dispatch strategies for this batch.
"""
from sglang.srt.compilation.piecewise_context_manager import (
is_in_piecewise_cuda_graph,
)
from sglang.srt.utils import get_device_sm, is_blackwell
# Decide MHA vs MLA
if forward_batch and forward_batch.forward_mode.is_extend_without_speculative():
if is_in_piecewise_cuda_graph():
# Can't branch on seq_lens_cpu in PCG, 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()
python/sglang/srt/layers/layernorm.py
+20 -1
View File
@@ -53,7 +53,26 @@ _flashinfer_layernorm_available = False
if _is_cuda or _is_xpu or _is_musa:
if _is_flashinfer_available:
try:
from flashinfer.norm import layernorm
import flashinfer.norm
from sglang.srt.utils.custom_op import register_custom_op
def _layernorm_fake_impl(
input: torch.Tensor,
gamma: torch.Tensor,
beta: torch.Tensor,
eps: float = 1e-6,
) -> torch.Tensor:
return torch.empty_like(input)
@register_custom_op(fake_impl=_layernorm_fake_impl)
def layernorm(
input: torch.Tensor,
gamma: torch.Tensor,
beta: torch.Tensor,
eps: float = 1e-6,
) -> torch.Tensor:
return flashinfer.norm.layernorm(input, gamma, beta, eps)
_flashinfer_layernorm_available = True
except (ImportError, AttributeError):
python/sglang/srt/layers/radix_attention.py
+14
View File
@@ -160,6 +160,12 @@ def unified_attention_with_output(
q_rope: Optional[torch.Tensor] = None,
k_rope: Optional[torch.Tensor] = None,
sinks: Optional[torch.Tensor] = None,
# MLA / TRT-LLM / NSA paths pass these through RadixAttention.forward(**kwargs);
# they must appear in the schema when --enforce-piecewise-cuda-graph is on.
cos_sin_cache: Optional[torch.Tensor] = None,
is_neox: Optional[bool] = None,
llama_4_scaling: Optional[torch.Tensor] = None,
topk_indices: Optional[torch.Tensor] = None,
) -> None:
context = get_forward_context()
forward_batch = context.forward_batch
@@ -178,6 +184,14 @@ def unified_attention_with_output(
kwargs["k_rope"] = k_rope[:real_num_tokens]
if sinks is not None:
kwargs["sinks"] = sinks
if cos_sin_cache is not None:
kwargs["cos_sin_cache"] = cos_sin_cache
if is_neox is not None:
kwargs["is_neox"] = is_neox
if llama_4_scaling is not None:
kwargs["llama_4_scaling"] = llama_4_scaling
if topk_indices is not None:
kwargs["topk_indices"] = topk_indices[:real_num_tokens]
original_out_cache_loc = forward_batch.out_cache_loc
# Keep the original ForwardBatch object and only narrow cache locations for
python/sglang/srt/model_executor/model_runner.py
+6
View File
@@ -2850,6 +2850,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.attention_layers = []
self.moe_layers = []
self.moe_fusions = []
self.dsa_indexers = []
for layer in layer_model.layers:
attn_layer = None
if hasattr(layer, "self_attn"):
@@ -2902,6 +2903,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
moe_fusion = layer.mixer
self.moe_layers.append(moe_block)
self.moe_fusions.append(moe_fusion)
# NSA indexers (None for layers without NSA)
dsa_indexer = None
if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "indexer"):
dsa_indexer = layer.self_attn.indexer
self.dsa_indexers.append(dsa_indexer)
if len(self.attention_layers) < self.model_config.num_hidden_layers:
# TODO(yuwei): support Non-Standard GQA
python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py
+4
View File
@@ -298,6 +298,7 @@ class PiecewiseCudaGraphRunner:
self.attention_layers = self.model_runner.attention_layers
self.moe_layers = self.model_runner.moe_layers
self.moe_fusions = self.model_runner.moe_fusions
self.dsa_indexers = getattr(self.model_runner, "dsa_indexers", None)
if get_global_graph_memory_pool() is None:
set_global_graph_memory_pool(self.device_module.graph_pool_handle())
@@ -432,6 +433,7 @@ class PiecewiseCudaGraphRunner:
self.quant_config,
self.moe_layers,
self.moe_fusions,
dsa_indexers=self.dsa_indexers,
):
_ = self.model_runner.model.forward(
forward_batch.input_ids,
@@ -622,6 +624,7 @@ class PiecewiseCudaGraphRunner:
self.quant_config,
self.moe_layers,
self.moe_fusions,
dsa_indexers=self.dsa_indexers,
):
self.model_runner.model.forward(
forward_batch.input_ids,
@@ -795,6 +798,7 @@ class PiecewiseCudaGraphRunner:
self.quant_config,
self.moe_layers,
self.moe_fusions,
dsa_indexers=self.dsa_indexers,
):
# Due to the dispatch kernel for MLA model, we init the metadata with original forward_batch
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
python/sglang/srt/models/deepseek_v2.py
+1
View File
@@ -1981,6 +1981,7 @@ class DeepseekV2DecoderLayer(nn.Module):
if (
isinstance(self.mlp, DeepseekV2MoE)
and not self.mlp.experts.moe_runner_config.inplace
and not torch.compiler.is_compiling()
):
from sglang.srt.layers.moe.moe_runner.base import moe_output_buffer_ctx
python/sglang/srt/server_args.py
+4 -16
View File
@@ -1347,6 +1347,10 @@ class ServerArgs:
# 18. CUDA Graph debug mode
if self.debug_cuda_graph:
self.disable_piecewise_cuda_graph = True
# 19. DSA prefill context parallelism (attn_cp_size is set later in
# _handle_model_specific_adjustments, so check the flag directly here)
if self.enable_dsa_prefill_context_parallel:
self.disable_piecewise_cuda_graph = True
def _handle_multi_item_scoring(self):
"""Setup and validate multi-item scoring constraints.
@@ -3183,22 +3187,6 @@ class ServerArgs:
self.ep_size == 1
), "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1"
# TODO(yuwei): Fix piecewise cuda graph support for bypassed topk MoE backends.
# Exception: GptOssForCausalLM wraps the entire MoE block in its own
# custom op (moe_impl), so bypassed topk is handled inside the op body.
if (
not self.enforce_piecewise_cuda_graph
and self.moe_runner_backend in ("flashinfer_trtllm", "flashinfer_mxfp4")
and self.get_model_config().hf_config.architectures[0]
!= "GptOssForCausalLM"
):
self.disable_piecewise_cuda_graph = True
logger.info(
f"Piecewise cuda graph is disabled for MoE runner backend "
f"'{self.moe_runner_backend}' (bypassed topk is incompatible "
f"with torch.compile)."
)
def _handle_a2a_moe(self):
if self.enable_deepep_waterfill and self.moe_a2a_backend != "deepep":
logger.warning(