diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 8bb1237e77d5..e05a3d59ee2a 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1001,15 +1001,6 @@ def forward_normal_dual_stream( # quant-once fp8 pair when that is on (also fed to the shared expert), # otherwise the MXFP8 pre-quant issued on routed_quant_stream, whose # layout the shared expert cannot take, or None. - routed_pre_quant_input = pre_quant_input - if should_quant_routed_input_mxfp8: - with torch.cuda.stream(self.routed_quant_stream): - x_q, x_sf = self.experts.quant_method.quantize_routed_input( - hidden_states, routed_hidden_size(self.experts) - ) - ready = self.routed_quant_stream.record_event() - routed_pre_quant_input = Mxfp8RoutedInputPreQuant(x_q, x_sf, ready) - if use_flashinfer_trtllm_bypass: topk_output = BypassedTopKOutput( hidden_states=hidden_states, @@ -1037,6 +1028,17 @@ def forward_normal_dual_stream( expert_location_dispatch_info=dispatch_info, **topk_kwargs, ) + # Recorded after the router so the routed MoE's first kernel, which + # joins this side stream, keeps the main chain on the main stream at + # CUDA-graph replay (the fork point above is unchanged). + routed_pre_quant_input = pre_quant_input + if should_quant_routed_input_mxfp8: + with torch.cuda.stream(self.routed_quant_stream): + x_q, x_sf = self.experts.quant_method.quantize_routed_input( + hidden_states, routed_hidden_size(self.experts) + ) + ready = self.routed_quant_stream.record_event() + routed_pre_quant_input = Mxfp8RoutedInputPreQuant(x_q, x_sf, ready) # The mHC post-split consumes the reduced row without an RMSNorm. use_fused_finalize_all_reduce = ( self._fuse_finalize_all_reduce diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 9e3aeef8a201..1c30ce2f95d4 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -2753,37 +2753,32 @@ def forward( # cross-layer fusion, and the final layer is completed in DeepseekV4Model. return hidden_states, residual, post, comb - def _hc_mix_and_combine( + def _hc_combine( self, x: torch.Tensor, - hc_fn: torch.Tensor, - hc_scale: torch.Tensor, - hc_base: torch.Tensor, apply_pre: Optional[torch.Tensor], norm: RMSNorm, stats_stream: Optional[torch.cuda.Stream] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Mixing coefficients come from x; the sublayer input is x collapsed with - apply_pre (None selects copy 0), then RMS-normalized. - Returns (y, pre, post, comb).""" - from sglang.kernels.ops.layernorm.mhc import ( - hc_combine, - hc_mix_stats, - hc_mix_stats_sinkhorn, - ) - - dtype = x.dtype - x_flat = x.flatten(1) + ) -> torch.Tensor: + """The sublayer input: x collapsed with apply_pre (None selects copy 0), + then RMS-normalized, on the current stream. Forks ``stats_stream`` for + :meth:`_hc_mix_stats`: after the combine for a tiny input so the two do + not compete for SMs, before it otherwise so they overlap.""" + from sglang.kernels.ops.layernorm.mhc import hc_combine - def combine_and_norm(): - if apply_pre is None: - return norm(x[:, 0, :].contiguous()) + tiny = 0 < x.shape[0] <= 8 + if stats_stream is not None and not tiny: + stats_stream.wait_stream(torch.cuda.current_stream()) + if apply_pre is None: + y = norm(x[:, 0, :].contiguous()) + else: from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled + x_flat = x.flatten(1) if ( x.is_cuda and get_platform().is_blackwell - and 0 < x.shape[0] <= 8 + and tiny and self.hc_mult == 4 and x_flat.shape[1] == 20480 and x.dtype == norm.weight.dtype == torch.bfloat16 @@ -2794,11 +2789,36 @@ def combine_and_norm(): ): from sglang.kernels.ops.layernorm.hc_combine_norm import hc_combine_norm - return hc_combine_norm( + y = hc_combine_norm( x_flat, apply_pre, norm.weight, norm.variance_epsilon ) - return norm(hc_combine(x_flat, apply_pre, self.hc_mult, dtype)) + else: + y = norm(hc_combine(x_flat, apply_pre, self.hc_mult, x.dtype)) + if stats_stream is not None and tiny: + stats_stream.wait_stream(torch.cuda.current_stream()) + return y + def _hc_mix_stats( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + stats_stream: Optional[torch.cuda.Stream] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Mixing coefficients (pre, post, comb) from x, on ``stats_stream`` when + one is given (forked by :meth:`_hc_combine`). + + Callers record this right before they join the stream: at CUDA-graph + replay a join continues on the lane of its first-recorded parent, so + recording the side stream ahead of the main-stream kernels it joins + would carry the main chain onto a fresh stream at every layer.""" + from sglang.kernels.ops.layernorm.mhc import ( + hc_mix_stats, + hc_mix_stats_sinkhorn, + ) + + x_flat = x.flatten(1) if ( x.is_cuda and torch.version.cuda is not None @@ -2808,18 +2828,10 @@ def combine_and_norm(): ) and x.dtype == torch.bfloat16 ): - # The split-K partial fixes the reduction order; - # fusing the reduction and sinkhorn preserves batch invariance. + # The split-K partial fixes the reduction order; fusing the + # reduction and sinkhorn preserves batch invariance. main_stream = torch.cuda.current_stream() - # Avoid competing with the statistics projection for a tiny input; - # the coefficients still overlap the attention or FFN that follows. - y = ( - combine_and_norm() - if stats_stream is not None and 0 < x.shape[0] <= 8 - else None - ) if stats_stream is not None: - stats_stream.wait_stream(main_stream) x.record_stream(stats_stream) with ( torch.cuda.stream(stats_stream) @@ -2841,12 +2853,11 @@ def combine_and_norm(): # after the caller joins it, on the main stream. for coefficient in (pre, post, comb): coefficient.record_stream(main_stream) - if y is None: - y = combine_and_norm() - return y, pre, post, comb + return pre, post, comb if x.is_cuda and torch.version.cuda is not None: - # Keep mixing and RMS reductions batch-invariant; cuBLAS/torch reductions can - # change order with num_tokens. Kernel upcasts let x_flat remain a bf16 view. + # Keep mixing and RMS reductions batch-invariant; cuBLAS/torch + # reductions can change order with num_tokens. Kernel upcasts let + # x_flat remain a bf16 view. mixes = hc_mix_stats(x_flat, hc_fn, self.rms_norm_eps).unsqueeze(1) else: x_flat = x_flat.float() @@ -2862,8 +2873,7 @@ def combine_and_norm(): self.hc_sinkhorn_iters, self.hc_eps, ) - y = combine_and_norm() - return y, pre.squeeze(1), post.squeeze(1), comb.squeeze(1) + return pre.squeeze(1), post.squeeze(1), comb.squeeze(1) def _get_hc_stats_stream(self, hidden_states, forward_batch): # Verify batches can also compute coefficients beside the @@ -2894,35 +2904,37 @@ def forward_hc_pre_from_prev( the FFN consumes this attention's. Returns (hidden_states, ffn_pre).""" stats_stream = self._get_hc_stats_stream(hidden_states, forward_batch) residual = hidden_states - x, attn_pre, attn_post, attn_comb = self._hc_mix_and_combine( - hidden_states, - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - apply_pre=prev_pre, - norm=self.input_layernorm, - stats_stream=stats_stream, + x = self._hc_combine( + hidden_states, prev_pre, self.input_layernorm, stats_stream ) with self.self_attn.maybe_use_decode_attn_tp(forward_batch): x = self.self_attn( x=x, positions=positions, forward_batch=forward_batch, x_quant=None ) + attn_pre, attn_post, attn_comb = self._hc_mix_stats( + hidden_states, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + stats_stream, + ) if stats_stream is not None: torch.cuda.current_stream().wait_stream(stats_stream) hidden_states = self.hc_post(x, residual, attn_post, attn_comb) residual = hidden_states - x, ffn_pre, ffn_post, ffn_comb = self._hc_mix_and_combine( + x = self._hc_combine( + hidden_states, attn_pre, self.post_attention_layernorm, stats_stream + ) + x = self._run_moe_ffn_dp_sync( + x, forward_batch, input_ids=input_ids, input_ids_global=input_ids_global + ) + ffn_pre, ffn_post, ffn_comb = self._hc_mix_stats( hidden_states, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base, - apply_pre=attn_pre, - norm=self.post_attention_layernorm, - stats_stream=stats_stream, - ) - x = self._run_moe_ffn_dp_sync( - x, forward_batch, input_ids=input_ids, input_ids_global=input_ids_global + stats_stream, ) if stats_stream is not None: torch.cuda.current_stream().wait_stream(stats_stream) diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index dca2d9308a8d..de83a00a9457 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -696,32 +696,34 @@ def _forward_hc_pre_from_prev( ) -> Tuple[torch.Tensor, torch.Tensor]: stats_stream = self._get_hc_stats_stream(hidden_states, forward_batch) residual = hidden_states - x, attn_pre, attn_post, attn_comb = self._hc_mix_and_combine( + x = self._hc_combine( + hidden_states, prev_pre, self.input_layernorm, stats_stream + ) + with self.self_attn.maybe_use_decode_attn_tp(forward_batch): + x = self.self_attn(positions, x, forward_batch) + attn_pre, attn_post, attn_comb = self._hc_mix_stats( hidden_states, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base, - apply_pre=prev_pre, - norm=self.input_layernorm, - stats_stream=stats_stream, + stats_stream, ) - with self.self_attn.maybe_use_decode_attn_tp(forward_batch): - x = self.self_attn(positions, x, forward_batch) if stats_stream is not None: torch.cuda.current_stream().wait_stream(stats_stream) hidden_states = self.hc_post(x, residual, attn_post, attn_comb) residual = hidden_states - x, ffn_pre, ffn_post, ffn_comb = self._hc_mix_and_combine( + x = self._hc_combine( + hidden_states, attn_pre, self.post_attention_layernorm, stats_stream + ) + x = self._run_ffn(x, forward_batch) + ffn_pre, ffn_post, ffn_comb = self._hc_mix_stats( hidden_states, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base, - apply_pre=attn_pre, - norm=self.post_attention_layernorm, - stats_stream=stats_stream, + stats_stream, ) - x = self._run_ffn(x, forward_batch) if stats_stream is not None: torch.cuda.current_stream().wait_stream(stats_stream) hidden_states = self.hc_post(x, residual, ffn_post, ffn_comb)