diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e6c526451..6cf331007 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index a57215e50..523d05569 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -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, diff --git a/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py b/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py index 6b644e72d..f710996b0 100644 --- a/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py +++ b/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py @@ -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):