[PyT] [Common] add support for enabling cuda graph under thd format in megatron.

tests/pytorch/attention/run_attention_with_cp.py

@@ -411,6 +411,13 @@ def run_dpa_with_cp(
             cu_seqlens_kv=cu_seqlens_kv,
             cu_seqlens_q_padded=cu_seqlens_q_padded,
             cu_seqlens_kv_padded=cu_seqlens_kv_padded,
+            # Test runner sets cu_seqlens_q == cu_seqlens_q_padded for the
+            # FlashAttention path, i.e. no inter-sequence padding. Declare this
+            # explicitly so the sync-free auto-detect (which conservatively
+            # picks True when padded cu_seqlens are present) does not disable FA.
+            pad_between_seqs=(
+                (kernel_backend != "FlashAttention") if qkv_format == "thd" else None
+            ),
             fp8_output=fp8_mha,
         )
         if config.return_max_logit:
@@ -528,6 +535,12 @@ def run_dpa_with_cp(
             cu_seqlens_kv=cu_seqlens_kv,
             cu_seqlens_q_padded=cu_seqlens_q_padded,
             cu_seqlens_kv_padded=cu_seqlens_kv_padded,
+            # See note above (non-CP branch): same explicit declaration so
+            # FlashAttention isn't disabled by the conservative sync-free
+            # auto-detect when this test path constructs no inter-seq padding.
+            pad_between_seqs=(
+                (kernel_backend != "FlashAttention") if qkv_format == "thd" else None
+            ),
             fp8_output=fp8_mha,
         )
         if config.return_max_logit:

tests/pytorch/test_fused_router.py

@@ -523,6 +523,76 @@ def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk, expert_multipl
     torch.testing.assert_close(probs.grad, probs_clone.grad, atol=atol, rtol=rtol)
 
 
+def test_fused_moe_aux_loss_cuda_graph_capture():
+    """CUDA-graph-safe path: total_num_tokens is a device tensor whose value
+    changes between replays. Forward and backward must both observe the new
+    value via the device-side coefficient computation."""
+    dtype = torch.float32
+    num_tokens = 4096
+    num_experts = 128
+    topk = 4
+    num_cols = num_experts
+    coeff = 0.01
+
+    offset = torch.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype, device="cuda") * 1e-4
+    probs = (
+        torch.arange(-num_cols // 2, num_cols // 2, device="cuda", dtype=dtype) * 1e-2
+    ).unsqueeze(0).repeat(num_tokens, 1) + offset.unsqueeze(1)
+    probs = probs.contiguous().requires_grad_(True)
+    tokens_per_expert = torch.randint(1, 1000, (num_cols,), device="cuda", dtype=torch.int32)
+
+    total_num_tokens_dev = torch.tensor(num_tokens, dtype=torch.int64, device="cuda")
+
+    # Warmup on a side stream to satisfy CUDA Graph capture requirements.
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        for _ in range(3):
+            warmup_out = fused_moe_aux_loss(
+                probs=probs,
+                tokens_per_expert=tokens_per_expert,
+                total_num_tokens=total_num_tokens_dev,
+                num_experts=num_experts,
+                topk=topk,
+                coeff=coeff,
+            )
+            torch.autograd.grad(warmup_out, probs)
+        del warmup_out
+    torch.cuda.current_stream().wait_stream(s)
+
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        out = fused_moe_aux_loss(
+            probs=probs,
+            tokens_per_expert=tokens_per_expert,
+            total_num_tokens=total_num_tokens_dev,
+            num_experts=num_experts,
+            topk=topk,
+            coeff=coeff,
+        )
+        (grad_probs,) = torch.autograd.grad(out, probs)
+
+    atol, rtol = _get_tolerances(dtype, num_cols)
+    # Replay with several distinct token counts; the captured graph must pick
+    # up each new value through total_num_tokens_dev.
+    for new_total in (num_tokens, num_tokens // 2, num_tokens * 2 - 17):
+        total_num_tokens_dev.fill_(new_total)
+        g.replay()
+        torch.cuda.synchronize()
+        ref_probs = probs.detach().clone().requires_grad_(True)
+        ref = aux_loss_pytorch(
+            probs=ref_probs,
+            tokens_per_expert=tokens_per_expert,
+            total_num_tokens=new_total,
+            topk=topk,
+            num_experts=num_experts,
+            moe_aux_loss_coeff=coeff,
+        )
+        (ref_grad_probs,) = torch.autograd.grad(ref, ref_probs)
+        torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)
+        torch.testing.assert_close(grad_probs, ref_grad_probs, atol=atol, rtol=rtol)
+
+
 def _bytemap_to_bitmap_u8(bytemap: torch.Tensor) -> torch.Tensor:
     """Reference packer: bool[T, E] -> uint8[T, ceil(E/8)] LSB-first.
 

transformer_engine/common/fused_router/fused_moe_aux_loss.cu

@@ -132,6 +132,100 @@ void fused_moe_aux_loss_forward(const Tensor& probs, const Tensor& tokens_per_ex
               reinterpret_cast<float*>(Coeff_buf.data.dptr), stream);););
 }
 
+/* -------------------------------------------------------------------------
+ *  CUDA-graph-safe variant: total_num_tokens lives in a 0-dim int64 device
+ *  tensor whose value can change between graph replays. Each CTA's reduction
+ *  lane computes C_coeff directly from the device value. The first CTA also
+ *  writes C_coeff into Coeff_buf[0] for backward.
+ * ------------------------------------------------------------------------- */
+template <typename DataType, typename IndexType>
+__global__ void fused_moe_aux_loss_forward_kernel_graph_safe(
+    const DataType* probs, const IndexType* tokens_per_expert, const int64_t* total_num_tokens_ptr,
+    int num_experts, int num_rows, int num_cols, int topk, float coeff, float* Coeff_buf) {
+  // Reduction body matches the scalar-input kernel above.
+  CompType thread_sum = CompType(0);
+  for (int col = threadIdx.x; col < num_cols; col += blockDim.x) {
+    CompType col_sum = CompType(0);
+    for (int row = blockIdx.x; row < num_rows; row += gridDim.x) {
+      col_sum += CompType(probs[row * num_cols + col]);
+    }
+    col_sum *= CompType(tokens_per_expert[col]);
+    thread_sum += col_sum;
+  }
+
+  extern __shared__ float shmem[];
+  CompType* shmem_block = reinterpret_cast<CompType*>(shmem);
+  shmem_block[threadIdx.x] = thread_sum;
+  __syncthreads();
+
+  const int warp_id = threadIdx.x / kThreadsPerWarp;
+  const int lane_id = threadIdx.x % kThreadsPerWarp;
+  if (warp_id == 0) {
+    CompType block_sum = warp_reduce_on_shmem<CompType, ReduceFuncType::SUM>(
+        shmem_block, static_cast<int>(blockDim.x), lane_id);
+    if (lane_id == 0) {
+      const float total_num_tokens = static_cast<float>(*total_num_tokens_ptr);
+      const float C_coeff = (static_cast<float>(num_experts) * coeff) / static_cast<float>(topk) /
+                            total_num_tokens / total_num_tokens;
+      if (blockIdx.x == 0) {
+        Coeff_buf[0] = C_coeff;
+      }
+      atomicAdd(&Coeff_buf[1], static_cast<float>(block_sum * C_coeff));
+    }
+  }
+}
+
+template <typename DataType, typename IndexType>
+void fused_moe_aux_loss_forward_kernel_launcher_graph_safe(
+    const DataType* probs, const IndexType* tokens_per_expert, const int64_t* total_num_tokens_dev,
+    int num_experts, int num_rows, int num_cols, int topk, float coeff, DataType* aux_loss,
+    float* Coeff_buf, cudaStream_t stream) {
+  NVTE_CHECK(num_cols > 0, "num_cols must be positive, got ", num_cols);
+  NVTE_CHECK(num_experts > 0, "num_experts must be positive, got ", num_experts);
+  NVTE_CHECK(num_cols % num_experts == 0, "Number of input columns (", num_cols,
+             ") must be a multiple of number of experts (", num_experts, ").");
+
+  const int block_size = ((std::min(1024, num_cols) + static_cast<int>(kThreadsPerWarp) - 1) /
+                          static_cast<int>(kThreadsPerWarp)) *
+                         static_cast<int>(kThreadsPerWarp);
+  const int grid_size = cuda::sm_count() * 2;
+  const size_t smem_size = block_size * sizeof(CompType);
+  check_shared_memory_capacity_num_experts(smem_size, num_cols);
+
+  // Zero the float accumulator. The main kernel writes Coeff_buf[0] for backward.
+  NVTE_CHECK_CUDA(cudaMemsetAsync(Coeff_buf + 1, 0, sizeof(float), stream));
+  fused_moe_aux_loss_forward_kernel_graph_safe<DataType, IndexType>
+      <<<grid_size, block_size, smem_size, stream>>>(probs, tokens_per_expert, total_num_tokens_dev,
+                                                     num_experts, num_rows, num_cols, topk, coeff,
+                                                     Coeff_buf);
+  NVTE_CHECK_CUDA(cudaGetLastError());
+
+  convert_accum_to_output<DataType><<<1, 1, 0, stream>>>(Coeff_buf, aux_loss);
+  NVTE_CHECK_CUDA(cudaGetLastError());
+}
+
+void fused_moe_aux_loss_forward_graph_safe(const Tensor& probs, const Tensor& tokens_per_expert,
+                                           const Tensor& total_num_tokens, int num_experts,
+                                           int num_rows, int num_cols, int topk, float coeff,
+                                           Tensor& aux_loss, Tensor& Coeff_buf,
+                                           cudaStream_t stream) {
+  NVTE_CHECK(total_num_tokens.data.dtype == DType::kInt64,
+             "total_num_tokens must be a 0-dim int64 tensor; got dtype ",
+             static_cast<int>(total_num_tokens.data.dtype));
+  NVTE_CHECK(total_num_tokens.numel() == 1,
+             "total_num_tokens must contain exactly one element; got ", total_num_tokens.numel());
+  TE_ROUTER_PROBS_TYPE_SWITCH_ALL(
+      probs.data.dtype, DataType,
+      TE_ROUTER_INDEX_TYPE_SWITCH_ALL(
+          tokens_per_expert.data.dtype, IndexType,
+          fused_moe_aux_loss_forward_kernel_launcher_graph_safe<DataType, IndexType>(
+              reinterpret_cast<DataType*>(probs.data.dptr),
+              reinterpret_cast<IndexType*>(tokens_per_expert.data.dptr),
+              reinterpret_cast<const int64_t*>(total_num_tokens.data.dptr), num_experts, num_rows,
+              num_cols, topk, coeff, reinterpret_cast<DataType*>(aux_loss.data.dptr),
+              reinterpret_cast<float*>(Coeff_buf.data.dptr), stream);););
+}
+
 template <typename DataType, typename IndexType>
 __global__ void fused_moe_aux_loss_backward_kernel(const float* Const_buf,
                                                    const IndexType* tokens_per_expert, int num_rows,
@@ -195,6 +289,20 @@ void nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor to
       *convertNVTETensorCheck(Coeff_buf), stream);
 }
 
+void nvte_fused_moe_aux_loss_forward_graph_safe(const NVTETensor probs,
+                                                const NVTETensor tokens_per_expert,
+                                                const NVTETensor total_num_tokens, int num_experts,
+                                                int num_rows, int num_cols, int topk, float coeff,
+                                                NVTETensor aux_loss, NVTETensor Coeff_buf,
+                                                cudaStream_t stream) {
+  NVTE_API_CALL(nvte_fused_moe_aux_loss_forward_graph_safe);
+  using namespace transformer_engine;
+  fused_router::fused_moe_aux_loss_forward_graph_safe(
+      *convertNVTETensorCheck(probs), *convertNVTETensorCheck(tokens_per_expert),
+      *convertNVTETensorCheck(total_num_tokens), num_experts, num_rows, num_cols, topk, coeff,
+      *convertNVTETensorCheck(aux_loss), *convertNVTETensorCheck(Coeff_buf), stream);
+}
+
 void nvte_fused_moe_aux_loss_backward(const NVTETensor Const_buf,
                                       const NVTETensor tokens_per_expert, int num_rows,
                                       int num_cols, NVTETensor grad_aux_loss, NVTETensor grad_probs,

transformer_engine/common/include/transformer_engine/fused_router.h

@@ -210,25 +210,44 @@ void nvte_fused_score_for_moe_aux_loss_backward(const NVTETensor intermediate_ou
                                                 int num_experts, int topk, int score_function,
                                                 NVTETensor grad_logits, cudaStream_t stream);
 
-/*! \brief Forward pass for auxiliary loss.
+/*! \brief Forward pass for auxiliary loss. Host-int total_num_tokens path:
+ *  the coefficient is folded on the host and passed as a kernel argument.
+ *  Prefer this path when total_num_tokens is statically known and the call
+ *  is not captured into a CUDA Graph.
  *
- *  \param[in]     probs           Probabilities from the forward pass.
+ *  \param[in]     probs              Probabilities from the forward pass.
  *  \param[in]     tokens_per_expert  Number of tokens per expert.
- *  \param[in]     total_num_tokens   Number of total tokens. Will be used in seq/global aux loss.
- *  \param[in]     num_experts     Number of experts.
- *  \param[in]     num_rows        Number of rows of probs.
- *  \param[in]     num_cols        Number of columns of probs.
- *  \param[in]     topk            Topk value.
- *  \param[in]     coeff           Coefficient.
- *  \param[out]    aux_loss        Output GPU scalar for auxiliary loss.
- *  \param[out]    Const_buf       Output GPU scalar for temporary constant buffer for backward pass.
- *  \param[in]     stream          CUDA stream used for the operation.
+ *  \param[in]     total_num_tokens   Number of total tokens. Used in seq/global aux loss.
+ *  \param[in]     num_experts        Number of experts.
+ *  \param[in]     num_rows           Number of rows of probs.
+ *  \param[in]     num_cols           Number of columns of probs.
+ *  \param[in]     topk               Topk value.
+ *  \param[in]     coeff              Coefficient.
+ *  \param[out]    aux_loss           Output GPU scalar for auxiliary loss.
+ *  \param[out]    Const_buf          Output GPU scalar for temporary constant buffer for backward
+ *                                    pass.
+ *  \param[in]     stream             CUDA stream used for the operation.
  */
 void nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor tokens_per_expert,
                                      int total_num_tokens, int num_experts, int num_rows,
                                      int num_cols, int topk, float coeff, NVTETensor aux_loss,
                                      NVTETensor Const_buf, cudaStream_t stream);
 
+/*! \brief Forward pass for auxiliary loss. Device-tensor total_num_tokens path:
+ *  the coefficient is computed on device from a 0-dim int64 GPU tensor so its
+ *  value stays dynamic across CUDA Graph replays. Prefer this path when the
+ *  caller needs CUDA-graph-safe semantics with a dynamic token count.
+ *
+ *  \param[in]     total_num_tokens   0-dim int64 GPU tensor with the total token count.
+ *  Other parameters as in :c:func:`nvte_fused_moe_aux_loss_forward`.
+ */
+void nvte_fused_moe_aux_loss_forward_graph_safe(const NVTETensor probs,
+                                                const NVTETensor tokens_per_expert,
+                                                const NVTETensor total_num_tokens, int num_experts,
+                                                int num_rows, int num_cols, int topk, float coeff,
+                                                NVTETensor aux_loss, NVTETensor Const_buf,
+                                                cudaStream_t stream);
+
 /*! \brief Backward pass for auxiliary loss.
  *
  *  \param[in]     Const_buf       Constant buffer from the forward pass.

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -1708,6 +1708,29 @@ def backward(ctx, d_out, *_args):
             dq = dq[..., : d_out.shape[-1]]
             dk = dk[..., : d_out.shape[-1]]
             dv = dv[..., : d_out.shape[-1]]
+            # Zero-fill positions beyond cu_seqlens_*_padded[-1] in dQ/dK/dV for THD.
+            # Use Q's padded boundary for dQ and KV's padded boundary for dK/dV.
+            # Sync-free `arange + mask` so capture and eager paths run the same code.
+            _qkv_format = ctx.qkv_layout.split("_")[0].replace("3", "").replace("2", "")
+            if _qkv_format == "thd":
+                if (
+                    cu_seqlens_q_padded is not None
+                    and isinstance(dq, torch.Tensor)
+                    and dq.shape[0] > 0
+                ):
+                    q_pad_mask = (
+                        torch.arange(dq.shape[0], device=dq.device) >= cu_seqlens_q_padded[-1]
+                    )
+                    dq[q_pad_mask] = 0
+                if cu_seqlens_kv_padded is not None:
+                    kv_actual_t = cu_seqlens_kv_padded[-1]
+                    for d_tensor in (dk, dv):
+                        if isinstance(d_tensor, torch.Tensor) and d_tensor.shape[0] > 0:
+                            kv_pad_mask = (
+                                torch.arange(d_tensor.shape[0], device=d_tensor.device)
+                                >= kv_actual_t
+                            )
+                            d_tensor[kv_pad_mask] = 0
         else:
             with get_nvtx_range_context("FusedAttnFunc.backward"):
                 # get nominal data type of dq, dk, dv

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py

@@ -2194,6 +2194,7 @@ def forward(
                 ctx.S_quantizer.scale = S_quantizer.scale.clone()
 
         nvtx_range_pop(f"{nvtx_label}")
+
         if return_max_logit:
             return out_ret, max_logit
         return out_ret
@@ -2846,10 +2847,28 @@ def backward(ctx, dout, *_args):
             dim = ctx.qkv_format.index("s")
             dq, dk, dv = [x.view(*x.shape[:dim], -1, *x.shape[dim + 2 :]) for x in [dq, dk, dv]]
 
-        if ctx.qkv_format == "thd" and not ctx.use_fused_attention:
-            dq[cu_seqlens_q_padded[-1] :].fill_(0)
-            dk[cu_seqlens_kv_padded[-1] :].fill_(0)
-            dv[cu_seqlens_kv_padded[-1] :].fill_(0)
+        # Zero-fill dQ/dK/dV at positions beyond cu_seqlens_*_padded[-1].
+        if (
+            ctx.qkv_format == "thd"
+            and not ctx.use_fused_attention
+            and cu_seqlens_q_padded is not None
+            and cu_seqlens_kv_padded is not None
+        ):
+            if is_graph_capturing():
+                # arange+mask under capture: `tensor[scalar_tensor:]` slicing would
+                # force a GPU->CPU sync that is forbidden during CUDA graph capture.
+                q_pad_mask = torch.arange(dq.shape[0], device=dq.device) >= cu_seqlens_q_padded[-1]
+                kv_pad_mask = (
+                    torch.arange(dk.shape[0], device=dk.device) >= cu_seqlens_kv_padded[-1]
+                )
+                dq[q_pad_mask] = 0
+                dk[kv_pad_mask] = 0
+                dv[kv_pad_mask] = 0
+            else:
+                # Pre-existing TE eager-mode behaviour.
+                dq[cu_seqlens_q_padded[-1] :].fill_(0)
+                dk[cu_seqlens_kv_padded[-1] :].fill_(0)
+                dv[cu_seqlens_kv_padded[-1] :].fill_(0)
 
         if ctx.fp8 and ctx.is_input_fp8:
             dq, dk, dv, _, _ = combine_and_quantize(ctx.qkv_layout, dq, dk, dv, ctx.dQKV_quantizer)
@@ -2904,6 +2923,23 @@ def backward(ctx, dout, *_args):
 
         nvtx_range_pop(f"{nvtx_label}")
 
+        # Zero-fill dQ/dK/dV at positions beyond the actual sequence end (THD CUDA Graph).
+        # cu_seqlens_*_padded are already local to this CP rank in the THD path.
+        # Use Q's padded boundary for dQ and KV's padded boundary for dK/dV.
+        # Skip the corresponding zero-fill when its padded cu_seqlens is absent.
+        if ctx.qkv_format == "thd":
+            if cu_seqlens_q_padded is not None and isinstance(dq, torch.Tensor) and dq.shape[0] > 0:
+                q_pad_mask = torch.arange(dq.shape[0], device=dq.device) >= cu_seqlens_q_padded[-1]
+                dq[q_pad_mask] = 0
+            if cu_seqlens_kv_padded is not None:
+                kv_actual_t = cu_seqlens_kv_padded[-1]
+                for d_tensor in [dk, dv]:
+                    if isinstance(d_tensor, torch.Tensor) and d_tensor.shape[0] > 0:
+                        kv_pad_mask = (
+                            torch.arange(d_tensor.shape[0], device=d_tensor.device) >= kv_actual_t
+                        )
+                        d_tensor[kv_pad_mask] = 0
+
         return (
             None,
             dq,