[DSV4.1] Record side-stream work right before its join to keep CUDA-graph replay on one stream

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

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)

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)