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

[Fix] Fix DSA indexer fusion for NeoX RoPE (#30111)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-07-04 03:20:55 -07:00
committed by GitHub
parent 92b800c531
commit e552f6ed75
3 changed files with 19 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+1 -1
View File
@@ -668,7 +668,7 @@ class Envs:
SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True)
SGLANG_DSA_USE_FUSED_METADATA_GENERATION = EnvBool(True)
SGLANG_DSA_TOPK_BROADCAST = EnvBool(False)
SGLANG_DISABLE_DSA_INDEXER_FUSION = EnvBool(True)
SGLANG_DISABLE_DSA_INDEXER_FUSION = EnvBool(False)
# sgl-kernel
SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+17 -13
View File
@@ -53,7 +53,6 @@ logger = logging.getLogger(__name__)
global _use_multi_stream
_is_cuda = is_cuda()
_use_dsa_indexer_fusion = _is_cuda and not envs.SGLANG_DISABLE_DSA_INDEXER_FUSION.get()
_is_hip = is_hip()
_is_npu = is_npu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
@@ -364,6 +363,11 @@ class Indexer(MultiPlatformOp):
self.index_topk = index_topk
self.q_lora_rank = q_lora_rank
self.layer_id = layer_id
self.use_dsa_indexer_fusion = (
_is_cuda
and not envs.SGLANG_DISABLE_DSA_INDEXER_FUSION.get()
and not is_neox_style
)
self.alt_stream = alt_stream
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
@@ -388,7 +392,7 @@ class Indexer(MultiPlatformOp):
prefix=add_prefix("wq_b", prefix),
)
if _use_dsa_indexer_fusion:
if self.use_dsa_indexer_fusion:
self.wk_weights_proj = ReplicatedLinear(
self.hidden_size,
self.head_dim + self.n_heads,
@@ -498,7 +502,7 @@ class Indexer(MultiPlatformOp):
def _maybe_rotate(self, x: torch.Tensor) -> torch.Tensor:
# Fusion drops the (logit-preserving) Hadamard rotation; without it the
# index-K cache here matches the fused path that decode reads back.
return x if _use_dsa_indexer_fusion else rotate_activation(x)
return x if self.use_dsa_indexer_fusion else rotate_activation(x)
def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
if (
@@ -534,7 +538,7 @@ class Indexer(MultiPlatformOp):
)
with torch.cuda.stream(self.alt_stream):
# TODO we should also put DeepGEMM half SM here?
if _use_dsa_indexer_fusion:
if self.use_dsa_indexer_fusion:
key, weights_raw = self._fused_k_weights(x)
else:
key, _ = self.wk(x)
@@ -553,7 +557,7 @@ class Indexer(MultiPlatformOp):
q_rope, _ = torch.split(
query, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1
)
if _use_dsa_indexer_fusion:
if self.use_dsa_indexer_fusion:
key, weights_raw = self._fused_k_weights(x)
else:
key, _ = self.wk(x)
@@ -1204,7 +1208,7 @@ class Indexer(MultiPlatformOp):
# Write the same K representation the decode path reads back: fused
# (no-Hadamard) when fusion is on, else the legacy Hadamard path.
if _use_dsa_indexer_fusion:
if self.use_dsa_indexer_fusion:
key_raw, _ = self._fused_k_weights(x)
if num_tokens is not None:
assert num_tokens <= key_raw.shape[0]
@@ -1655,12 +1659,12 @@ class Indexer(MultiPlatformOp):
# wrapper owns base+delta and no LoRA kernel runs under torch.compile.
# Fusion folds weights_proj into wk_weights_proj, so weights_proj is
# absent then; short-circuit before touching it.
weights_proj_lora = not _use_dsa_indexer_fusion and getattr(
weights_proj_lora = not self.use_dsa_indexer_fusion and getattr(
self.weights_proj, "set_lora", False
)
if (
_use_dsa_indexer_fusion
self.use_dsa_indexer_fusion
and not in_piecewise_or_breakable_cuda_graph
and forward_batch.attn_cp_metadata is None
):
@@ -1708,7 +1712,7 @@ class Indexer(MultiPlatformOp):
elif enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle():
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
if not _use_dsa_indexer_fusion:
if not self.use_dsa_indexer_fusion:
if weights_proj_lora:
weights = self.weights_proj(x)[0].float() * self.n_heads**-0.5
else:
@@ -1725,7 +1729,7 @@ class Indexer(MultiPlatformOp):
act_quant=act_quant,
)
current_stream.wait_stream(self.alt_stream)
if _use_dsa_indexer_fusion:
if self.use_dsa_indexer_fusion:
weights = self._scale_head_gates(weights_raw, q_scale)
else:
weights = self._apply_q_scale_and_softmax_scale(weights, q_scale)
@@ -1804,7 +1808,7 @@ class Indexer(MultiPlatformOp):
x_for_gate = x
if in_piecewise_or_breakable_cuda_graph:
if _use_dsa_indexer_fusion:
if self.use_dsa_indexer_fusion:
weights = scale_head_gate_graph(
weights_raw,
self.n_heads**-0.5,
@@ -1821,7 +1825,7 @@ class Indexer(MultiPlatformOp):
self.softmax_scale,
q_scale,
)
elif _use_dsa_indexer_fusion:
elif self.use_dsa_indexer_fusion:
weights = self._scale_head_gates(weights_raw, q_scale)
elif weights_proj_lora:
weights = self.weights_proj(x_for_gate)[0].float() * self.n_heads**-0.5
@@ -2312,7 +2316,7 @@ def pcg_dsa_indexer_prefill_split(
# Fused path stores K (no-Hadamard) and computes q_fp8 + head gate in the
# fused kernels, sliced to the unpadded count. Single stream: the split op is
# captured, so the dual-stream overlap is disabled.
if _use_dsa_indexer_fusion:
if indexer.use_dsa_indexer_fusion:
q_fp8, weights = indexer._fused_q_prepare_and_store(
x,
q_lora,
test/registered/8-gpu-models/test_deepseek_v32_indexcache.py
+1 -1
View File
@@ -63,7 +63,7 @@ class TestDeepseekV32IndexTopkPattern(CustomTestCase):
write_github_step_summary(
f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.93)
self.assertGreater(metrics["accuracy"], 0.935)
class TestDeepseekV32IndexFreq(CustomTestCase):