feat: enable trigger_completion_at_end=False for MNNVL backend

csrc/trtllm_mnnvl_allreduce.cu

@@ -27,16 +27,14 @@ using tvm::ffi::Optional;
     }                                                                               \
   }()
 
-void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_ptr,
-                                   int64_t buffer_ptrs_dev, int64_t buffer_ptr_local,
-                                   TensorView buffer_flags_mnnvl, int64_t nranks, int64_t rank,
-                                   bool rmsnorm_fusion, bool launch_with_pdl, bool use_oneshot,
-                                   Optional<TensorView> output, Optional<TensorView> residual_out,
-                                   Optional<TensorView> residual_in, Optional<TensorView> gamma,
-                                   Optional<double> epsilon, Optional<double> weight_bias,
-                                   Optional<int64_t> quant_type, Optional<TensorView> quant_out,
-                                   Optional<TensorView> sf_out, Optional<TensorView> output_scale,
-                                   Optional<int64_t> layout_code) {
+void trtllm_mnnvl_allreduce_fusion(
+    TensorView input, int64_t multicast_buffer_ptr, int64_t buffer_ptrs_dev,
+    int64_t buffer_ptr_local, TensorView buffer_flags_mnnvl, int64_t nranks, int64_t rank,
+    bool rmsnorm_fusion, bool launch_with_pdl, bool trigger_completion_at_end, bool use_oneshot,
+    Optional<TensorView> output, Optional<TensorView> residual_out,
+    Optional<TensorView> residual_in, Optional<TensorView> gamma, Optional<double> epsilon,
+    Optional<double> weight_bias, Optional<int64_t> quant_type, Optional<TensorView> quant_out,
+    Optional<TensorView> sf_out, Optional<TensorView> output_scale, Optional<int64_t> layout_code) {
   ffi::CUDADeviceGuard device_guard(input.device().device_id);
   auto stream = get_stream(input.device());
 
@@ -152,6 +150,7 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt
     params.bufferFlags = reinterpret_cast<uint32_t*>(buffer_flags_mnnvl.data_ptr());
     params.rmsNormFusion = rmsnorm_fusion;
     params.launchWithPdl = launch_with_pdl;
+    params.triggerCompletionAtEnd = trigger_completion_at_end;
     params.sfLayout = sf_layout;
     params.quantType = quant_type_enum;
 

flashinfer/comm/allreduce.py

@@ -555,11 +555,15 @@ def allreduce_fusion(
     launch_with_pdl : bool
         Use Programmatic Dependent Launch.
     trigger_completion_at_end : bool
-        TRT-LLM only. Controls when PDL completion is signaled. ``True``
-        (default) signals after the kernel finishes (safe, no overlap).
-        ``False`` signals early, allowing the next PDL-aware kernel to
-        overlap with this one. Only safe when the next kernel also calls
-        ``cudaGridDependencySynchronize()``. Ignored by the MNNVL backend.
+        Controls when PDL completion is signaled. ``True`` (default)
+        signals after the kernel finishes (safe, no overlap). ``False``
+        signals early, allowing the next PDL-aware kernel to overlap with
+        this one. Only safe when the next kernel also calls
+        ``cudaGridDependencySynchronize()``. For the MNNVL backend with
+        RMSNorm fusion (two-shot), this controls when the trailing RMSNorm
+        kernel signals completion; the intermediate all-reduce kernel
+        always uses the early trigger so the RMSNorm kernel can be
+        launched eagerly.
     output : Optional[torch.Tensor]
         Pre-allocated AllReduce output buffer, shape
         ``[token_num, hidden_dim]``.
@@ -946,6 +950,7 @@ def _flatten_checked(t, name):
                 launch_with_pdl=launch_with_pdl,
                 output=output,
                 strategy=strategy,
+                trigger_completion_at_end=trigger_completion_at_end,
             )
             return output
 
@@ -975,6 +980,7 @@ def _flatten_checked(t, name):
                 launch_with_pdl=launch_with_pdl,
                 strategy=strategy,
                 weight_bias=weight_bias,
+                trigger_completion_at_end=trigger_completion_at_end,
             )
             return norm_result
 
@@ -1010,6 +1016,7 @@ def _flatten_checked(t, name):
                 launch_with_pdl=launch_with_pdl,
                 strategy=strategy,
                 weight_bias=weight_bias,
+                trigger_completion_at_end=trigger_completion_at_end,
             )
             return quant_result
 

flashinfer/comm/trtllm_mnnvl_ar.py

@@ -282,6 +282,7 @@ def get_trtllm_mnnvl_comm_module():
             "rank",
             "rmsnorm_fusion",
             "launch_with_pdl",
+            "trigger_completion_at_end",
             "use_oneshot",
             "output",
             "residual_out",
@@ -305,6 +306,7 @@ def trtllm_mnnvl_allreduce_fusion(
         rank: int,
         rmsnorm_fusion: bool,
         launch_with_pdl: bool,
+        trigger_completion_at_end: bool,
         use_oneshot: bool,
         output: Optional[torch.Tensor],
         residual_out: Optional[torch.Tensor],
@@ -330,6 +332,12 @@ def trtllm_mnnvl_allreduce_fusion(
             rank: Current process rank
             rmsnorm_fusion: Whether to perform RMSNorm fusion
             launch_with_pdl: Whether to launch with PDL
+            trigger_completion_at_end: When True, the kernel signals PDL launch
+                completion at the end of the kernel (safe, no overlap with the
+                next PDL-aware kernel). When False, completion is signaled
+                earlier inside the kernel to allow the next PDL-aware kernel
+                to start before this one fully finishes. Only has an effect
+                when launch_with_pdl is True on SM90+.
             use_oneshot: Whether to use one-shot (true) or two-shot (false)
             output: Output tensor
             residual_out: Residual output tensor (if rmsnorm)
@@ -346,6 +354,7 @@ def trtllm_mnnvl_allreduce_fusion(
             rank,
             rmsnorm_fusion,
             launch_with_pdl,
+            trigger_completion_at_end,
             use_oneshot,
             output,
             residual_out,
@@ -371,6 +380,7 @@ def trtllm_mnnvl_allreduce(
     launch_with_pdl: bool,
     output: Optional[torch.Tensor] = None,
     strategy: MNNVLAllreduceFusionStrategy = MNNVLAllreduceFusionStrategy.AUTO,
+    trigger_completion_at_end: bool = True,
 ) -> torch.Tensor:
     """Perform an MNNVL all-reduce sum across tensor-parallel ranks.
 
@@ -415,6 +425,11 @@ def trtllm_mnnvl_allreduce(
         launch_with_pdl: Whether to launch with PDL
         output: Output tensor to store the result, empty tensor will be created if not provided.
         strategy: MNNVLAllreduceFusionStrategy. Internal heuristics will be used if not provided.
+        trigger_completion_at_end: Controls when PDL completion is signaled.
+            True (default): signal completion after the kernel finishes (safe,
+            no overlap). False: signal completion early, allowing the next
+            PDL-aware kernel to overlap with this one. Only has an effect
+            when ``launch_with_pdl`` is True on SM90+.
     Returns:
         output: Reduced tensor [num_tokens, hidden_dim]
     """
@@ -456,6 +471,7 @@ def trtllm_mnnvl_allreduce(
         workspace.rank,
         False,  # No RMSNorm Fusion
         launch_with_pdl,
+        trigger_completion_at_end,
         strategy == MNNVLAllreduceFusionStrategy.ONESHOT,
         output,
         None,
@@ -477,6 +493,7 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm(
     residual_out: Optional[torch.Tensor] = None,
     launch_with_pdl: bool = False,
     strategy: MNNVLAllreduceFusionStrategy = MNNVLAllreduceFusionStrategy.AUTO,
+    trigger_completion_at_end: bool = True,
     weight_bias: float = 0.0,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Performs MNNVL Allreduce + Residual + RMSNorm.
@@ -495,6 +512,13 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm(
         residual_out: Residual output tensor [num_tokens, hidden_dim], empty tensor will be created if not provided.
         launch_with_pdl: Whether to launch with PDL
         strategy: MNNVLAllreduceFusionStrategy. Internal heuristics will be used if not provided.
+        trigger_completion_at_end: Controls when PDL completion is signaled by
+            the *last* kernel in the chain (the rmsNorm-Lamport kernel for the
+            two-shot strategy, or the one-shot fusion kernel itself for the
+            one-shot strategy). True (default): signal completion after the
+            last kernel fully finishes. False: signal completion early to
+            allow the next PDL-aware kernel to overlap. Only has an effect
+            when ``launch_with_pdl`` is True on SM90+.
         weight_bias: Bias added to gamma before scaling. 0.0 (default) for standard
             RMSNorm (gamma * x * rsqrt(...)); 1.0 for Gemma / Qwen3.5 RMSNorm
             ((1 + gamma) * x * rsqrt(...)).
@@ -555,6 +579,7 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm(
         workspace.rank,
         True,  # RMSNorm Fusion
         launch_with_pdl,
+        trigger_completion_at_end,
         strategy == MNNVLAllreduceFusionStrategy.ONESHOT,
         output,
         residual_out,
@@ -582,6 +607,7 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant(
     launch_with_pdl: bool = False,
     strategy: MNNVLAllreduceFusionStrategy = MNNVLAllreduceFusionStrategy.AUTO,
     weight_bias: float = 0.0,
+    trigger_completion_at_end: bool = True,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]:
     """Perform MNNVL AllReduce + Residual + RMSNorm + FP8/NVFP4 quantization.
 
@@ -769,6 +795,7 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant(
         workspace.rank,
         True,
         launch_with_pdl,
+        trigger_completion_at_end,
         strategy == MNNVLAllreduceFusionStrategy.ONESHOT,
         output,
         residual_out,
@@ -913,6 +940,7 @@ def trtllm_mnnvl_all_reduce(
         rank,
         False,  # No RMSNorm Fusion
         launch_with_pdl,
+        False,  # trigger_completion_at_end: keep legacy early-trigger behavior
         False,  # Use two-shot
         out,
         None,
@@ -1010,6 +1038,7 @@ def trtllm_mnnvl_fused_allreduce_rmsnorm(
         rank,
         True,  # RMSNorm Fusion
         launch_with_pdl,
+        False,  # trigger_completion_at_end: keep legacy early-trigger behavior
         False,
         normed_output,
         prenorm_output,