From 2c58118f3b1181f739ffe6c525b855f81bc68764 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Fri, 5 Jun 2026 06:41:02 -0700 Subject: [PATCH 01/10] add support for THD CUDA graph Signed-off-by: HaochenYuan --- .../attention/run_attention_with_cp.py | 13 +++++ .../dot_product_attention/backends.py | 37 ++++++++++++-- .../dot_product_attention/context_parallel.py | 51 +++++++++++++++++-- .../dot_product_attention.py | 11 ++-- .../pytorch/cpp_extensions/fused_attn.py | 26 ++++++++++ 5 files changed, 125 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 52c7dc067c..3cec42e015 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/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: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 8f42983553..baa1127717 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1679,9 +1679,17 @@ def backward(ctx, d_out, *_args): rest = [None] if ctx.use_FAv2_bwd: softmax_lse, rng_state = aux_ctx_tensors - dq = torch.empty_like(q) - dk = torch.empty_like(k) - dv = torch.empty_like(v) + # During CUDA graph capture, allocate with zeros so the memset is baked into + # the captured graph and replay buffers start clean. Outside capture, allocate + # with empty for perf and rely on the explicit tail zero-fill below. + if torch.cuda.is_current_stream_capturing(): + dq = torch.zeros_like(q) + dk = torch.zeros_like(k) + dv = torch.zeros_like(v) + else: + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) d_out, q, k, v, out = [dpa_utils.maybe_contiguous(x) for x in (d_out, q, k, v, out)] # from transformer_engine.pytorch.attention.dot_product_attention import flash_attn_cuda_bwd flash_attn_cuda_bwd( @@ -1708,6 +1716,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 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 36847e40ed..e374e19e4e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -2245,6 +2245,22 @@ def forward( ctx.S_quantizer.scale = S_quantizer.scale.clone() nvtx_range_pop(f"{nvtx_label}") + + # Zero-fill output at positions beyond the actual sequence end (THD CUDA Graph). + # cu_seqlens_q_padded is already local to this CP rank in the THD path, so its + # last entry is the local extent. Skip when the caller did not provide padded + # cu_seqlens -- no static-shape padding can be present in that case. + if ( + qkv_format == "thd" + and cu_seqlens_q_padded is not None + and isinstance(out_ret, torch.Tensor) + and out_ret.shape[0] > 0 + ): + local_actual_t = cu_seqlens_q_padded[-1] + pad_mask = torch.arange(out_ret.shape[0], device=out_ret.device) >= local_actual_t + out_ret.data[pad_mask] = 0 + out.data[pad_mask.view(-1, *([1] * (out.dim() - 1))).expand_as(out)] = 0 + if return_max_logit: return out_ret, max_logit return out_ret @@ -2897,10 +2913,20 @@ 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]. + # Use arange+mask indexing rather than `tensor[scalar_tensor:]` slicing -- the + # latter forces a GPU->CPU sync that is forbidden during CUDA graph capture. + 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 + ): + 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 if ctx.fp8 and ctx.is_input_fp8: dq, dk, dv, _, _ = combine_and_quantize(ctx.qkv_layout, dq, dk, dv, ctx.dQKV_quantizer) @@ -2955,6 +2981,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, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index a244c8f814..72db6ea2c6 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1536,15 +1536,14 @@ def forward( False ), "core_attention_bias must be in one of {bhss, 1hss, b1ss, 11ss, 111s} shapes" - # check if there is padding between sequences when qkv_format='thd' + # Default pad_between_seqs auto-detect. For THD, infer presence of + # inter-sequence padding from whether padded cu_seqlens were supplied -- + # sync-free, and stable across eager and CUDA graph capture (the auto-detect + # must return the same value in both modes for backend selection to match). if pad_between_seqs is None: if qkv_format == "thd": pad_between_seqs = ( - cu_seqlens_q_padded is not None - and not torch.equal(cu_seqlens_q_padded[:-1], cu_seqlens_q[:-1]) - ) or ( - cu_seqlens_kv_padded is not None - and not torch.equal(cu_seqlens_kv_padded[:-1], cu_seqlens_kv[:-1]) + cu_seqlens_q_padded is not None or cu_seqlens_kv_padded is not None ) else: pad_between_seqs = False diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 62dcaadc96..924c3636db 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -340,6 +340,16 @@ def fused_attn_fwd( cuda_graph, ) + # Zero-fill output at positions beyond the actual sequence end (THD CUDA Graph). + # Uses arange+mask indexing instead of `tensor[scalar_tensor:]` slicing -- the + # latter forces a GPU->CPU sync that is forbidden during CUDA graph capture. + if qkv_layout in ("t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"): + out_tensor = output_tensors[0] + actual_t = cu_seqlens_q[-1] if cu_seqlens_q_padded is None else cu_seqlens_q_padded[-1] + if isinstance(out_tensor, torch.Tensor) and out_tensor.shape[0] > 0: + pad_mask = torch.arange(out_tensor.shape[0], device=out_tensor.device) >= actual_t + out_tensor[pad_mask] = 0 + if return_max_logit: qkv_format = qkv_layout.replace("3", "").replace("2", "").split("_")[0] # thd (newer cuDNN runtimes, non-sm120): output_tensors: out [tq, h, d], Stats [tq, h, 1], Max [tq, h, 1] @@ -596,4 +606,20 @@ def fused_attn_bwd( cuda_graph, ) + # Zero-fill dQ/dK/dV at positions beyond the actual sequence end (THD CUDA Graph). + # Use Q's padded boundary for dQ and KV's padded boundary for dK/dV. + if qkv_layout in ("t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"): + q_actual_t = cu_seqlens_q[-1] if cu_seqlens_q_padded is None else cu_seqlens_q_padded[-1] + kv_actual_t = ( + cu_seqlens_kv[-1] if cu_seqlens_kv_padded is None else cu_seqlens_kv_padded[-1] + ) + dq_tensor = output_tensors[0] + if isinstance(dq_tensor, torch.Tensor) and dq_tensor.shape[0] > 0: + q_pad_mask = torch.arange(dq_tensor.shape[0], device=dq_tensor.device) >= q_actual_t + dq_tensor[q_pad_mask] = 0 + for d_tensor in output_tensors[1:3]: # 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 output_tensors From b6aa597788ef2f472dff3510dea8706162ef4ffc Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Fri, 5 Jun 2026 23:31:18 -0700 Subject: [PATCH 02/10] modify comment Signed-off-by: HaochenYuan --- .../pytorch/attention/dot_product_attention/backends.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index baa1127717..0f2c6f5bba 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1679,9 +1679,9 @@ def backward(ctx, d_out, *_args): rest = [None] if ctx.use_FAv2_bwd: softmax_lse, rng_state = aux_ctx_tensors - # During CUDA graph capture, allocate with zeros so the memset is baked into - # the captured graph and replay buffers start clean. Outside capture, allocate - # with empty for perf and rely on the explicit tail zero-fill below. + # Keep capture replay buffers zero-initialized; outside capture, use + # empty_like to avoid the extra memset. The THD tail zero-fill below + # clears tail positions in both modes. if torch.cuda.is_current_stream_capturing(): dq = torch.zeros_like(q) dk = torch.zeros_like(k) From 805746a5a9ca74952cc63130c7c7165141ba14d3 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Tue, 9 Jun 2026 02:22:18 -0700 Subject: [PATCH 03/10] address @timmoon10: drop FAv2-bwd alloc gate, rely on THD tail zero-fill Signed-off-by: HaochenYuan --- .../attention/dot_product_attention/backends.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 0f2c6f5bba..785e438cda 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1679,17 +1679,9 @@ def backward(ctx, d_out, *_args): rest = [None] if ctx.use_FAv2_bwd: softmax_lse, rng_state = aux_ctx_tensors - # Keep capture replay buffers zero-initialized; outside capture, use - # empty_like to avoid the extra memset. The THD tail zero-fill below - # clears tail positions in both modes. - if torch.cuda.is_current_stream_capturing(): - dq = torch.zeros_like(q) - dk = torch.zeros_like(k) - dv = torch.zeros_like(v) - else: - dq = torch.empty_like(q) - dk = torch.empty_like(k) - dv = torch.empty_like(v) + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) d_out, q, k, v, out = [dpa_utils.maybe_contiguous(x) for x in (d_out, q, k, v, out)] # from transformer_engine.pytorch.attention.dot_product_attention import flash_attn_cuda_bwd flash_attn_cuda_bwd( From a2cc7ef88290a394f3e0d7e639c7a8364a109679 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Thu, 11 Jun 2026 10:08:39 -0700 Subject: [PATCH 04/10] Support graph-safe MoE aux loss token count Signed-off-by: HaochenYuan --- tests/pytorch/test_fused_router.py | 64 +++++++++ .../common/fused_router/fused_moe_aux_loss.cu | 123 ++++++++++++++++++ .../include/transformer_engine/fused_router.h | 42 ++++-- .../dot_product_attention/context_parallel.py | 15 --- .../dot_product_attention.py | 15 ++- transformer_engine/pytorch/csrc/extensions.h | 7 + .../pytorch/csrc/extensions/pybind.cpp | 6 +- .../pytorch/csrc/extensions/router.cpp | 34 +++++ transformer_engine/pytorch/router.py | 49 ++++--- 9 files changed, 310 insertions(+), 45 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 974ccee19c..63d5df39aa 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -462,6 +462,70 @@ 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. The captured graph must 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() + 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): + _ = 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.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, + ) + + 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 = aux_loss_pytorch( + probs=probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=new_total, + topk=topk, + num_experts=num_experts, + moe_aux_loss_coeff=coeff, + ) + torch.testing.assert_close(out, ref, 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. diff --git a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu index cc5e5e3bcc..b425e36a1c 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -132,6 +132,115 @@ void fused_moe_aux_loss_forward(const Tensor& probs, const Tensor& tokens_per_ex reinterpret_cast(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. A single-thread + * prelude kernel reads it on device and writes C_coeff into Coeff_buf[0]; + * the main reduction kernel then reads C_coeff from Coeff_buf[0] instead + * of receiving it as a launch parameter, so the value stays dynamic. + * ------------------------------------------------------------------------- */ +__global__ void fused_moe_aux_loss_compute_coeff_kernel(const int64_t* total_num_tokens_ptr, + int num_experts, int topk, float coeff, + float* Coeff_buf) { + if (threadIdx.x == 0 && blockIdx.x == 0) { + const int64_t T = *total_num_tokens_ptr; + const float T_f = static_cast(T); + Coeff_buf[0] = (static_cast(num_experts) * coeff) / static_cast(topk) / T_f / T_f; + } +} + +template +__global__ void fused_moe_aux_loss_forward_kernel_tensor(const DataType* probs, + const IndexType* tokens_per_expert, + int num_rows, int num_cols, + const float* Coeff_buf) { + // Read C_coeff from Coeff_buf[0], written by the prelude kernel that ran + // earlier on the same stream. Backward reads it from the same slot. + const float C_coeff = Coeff_buf[0]; + + // 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(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(shmem_block, static_cast(blockDim.x), + ReduceFuncType::SUM, lane_id); + if (lane_id == 0) { + atomicAdd(const_cast(&Coeff_buf[1]), static_cast(block_sum * C_coeff)); + } + } +} + +template +void fused_moe_aux_loss_forward_kernel_launcher_tensor(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(kThreadsPerWarp) - 1) / + static_cast(kThreadsPerWarp)) * + static_cast(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 (Coeff_buf[1]); the prelude kernel writes + // Coeff_buf[0] = C_coeff derived from device-side total_num_tokens. + NVTE_CHECK_CUDA(cudaMemsetAsync(Coeff_buf + 1, 0, sizeof(float), stream)); + fused_moe_aux_loss_compute_coeff_kernel<<<1, 1, 0, stream>>>(total_num_tokens_dev, num_experts, + topk, coeff, Coeff_buf); + NVTE_CHECK_CUDA(cudaGetLastError()); + + fused_moe_aux_loss_forward_kernel_tensor + <<>>(probs, tokens_per_expert, num_rows, num_cols, + Coeff_buf); + NVTE_CHECK_CUDA(cudaGetLastError()); + + convert_accum_to_output<<<1, 1, 0, stream>>>(Coeff_buf, aux_loss); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +void fused_moe_aux_loss_forward(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(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_tensor( + reinterpret_cast(probs.data.dptr), + reinterpret_cast(tokens_per_expert.data.dptr), + reinterpret_cast(total_num_tokens.data.dptr), num_experts, num_rows, + num_cols, topk, coeff, reinterpret_cast(aux_loss.data.dptr), + reinterpret_cast(Coeff_buf.data.dptr), stream););); +} + template __global__ void fused_moe_aux_loss_backward_kernel(const float* Const_buf, const IndexType* tokens_per_expert, int num_rows, @@ -195,6 +304,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_tensor(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_tensor); + using namespace transformer_engine; + fused_router::fused_moe_aux_loss_forward( + *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, diff --git a/transformer_engine/common/include/transformer_engine/fused_router.h b/transformer_engine/common/include/transformer_engine/fused_router.h index 6cee10bd39..22de780919 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -190,25 +190,45 @@ 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. Adds one extra single-thread + * kernel launch versus the host-int path; prefer this only when the caller + * needs CUDA-graph-safe semantics. + * + * \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_tensor(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. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index e374e19e4e..af3ed49506 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -2246,21 +2246,6 @@ def forward( nvtx_range_pop(f"{nvtx_label}") - # Zero-fill output at positions beyond the actual sequence end (THD CUDA Graph). - # cu_seqlens_q_padded is already local to this CP rank in the THD path, so its - # last entry is the local extent. Skip when the caller did not provide padded - # cu_seqlens -- no static-shape padding can be present in that case. - if ( - qkv_format == "thd" - and cu_seqlens_q_padded is not None - and isinstance(out_ret, torch.Tensor) - and out_ret.shape[0] > 0 - ): - local_actual_t = cu_seqlens_q_padded[-1] - pad_mask = torch.arange(out_ret.shape[0], device=out_ret.device) >= local_actual_t - out_ret.data[pad_mask] = 0 - out.data[pad_mask.view(-1, *([1] * (out.dim() - 1))).expand_as(out)] = 0 - if return_max_logit: return out_ret, max_logit return out_ret diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 3d2687e4a5..aa1481a384 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1562,11 +1562,20 @@ def forward( # inter-sequence padding from whether padded cu_seqlens were supplied -- # sync-free, and stable across eager and CUDA graph capture (the auto-detect # must return the same value in both modes for backend selection to match). + # If padded cu_seqlens are the *same object* as the unpadded ones, no real + # inter-sequence padding exists (only THD tail padding) -- treat as False so + # FlashAttention v2/v4 remain eligible. if pad_between_seqs is None: if qkv_format == "thd": - pad_between_seqs = ( - cu_seqlens_q_padded is not None or cu_seqlens_kv_padded is not None - ) + if ( + cu_seqlens_q_padded is cu_seqlens_q + and cu_seqlens_kv_padded is cu_seqlens_kv + ): + pad_between_seqs = False + else: + pad_between_seqs = ( + cu_seqlens_q_padded is not None or cu_seqlens_kv_padded is not None + ) else: pad_between_seqs = False diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 2b4f899e1d..9a1d349b2c 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -55,6 +55,13 @@ std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, int num_rows, int num_cols, int topk, float coeff); +std::tuple fused_moe_aux_loss_fwd_tensor(at::Tensor probs, + at::Tensor tokens_per_expert, + at::Tensor total_num_tokens, + int num_experts, int num_rows, + int num_cols, int topk, + float coeff); + at::Tensor fused_moe_aux_loss_bwd(at::Tensor Const_buf, at::Tensor tokens_per_expert, int num_rows, int num_cols, at::Tensor grad_aux_loss); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index d6089b1e01..2127ef1c24 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -158,7 +158,11 @@ void init_router_bindings(pybind11::module &m) { m.def("fused_moe_aux_loss_fwd", &fused_moe_aux_loss_fwd, py::arg("probs"), py::arg("tokens_per_expert"), py::arg("total_num_tokens"), py::arg("num_experts"), py::arg("num_rows"), py::arg("num_cols"), py::arg("topk"), py::arg("coeff"), - "Fused aux loss fwd"); + "Fused aux loss fwd (host-int total_num_tokens, host-folded C_coeff)"); + m.def("fused_moe_aux_loss_fwd_tensor", &fused_moe_aux_loss_fwd_tensor, py::arg("probs"), + py::arg("tokens_per_expert"), py::arg("total_num_tokens"), py::arg("num_experts"), + py::arg("num_rows"), py::arg("num_cols"), py::arg("topk"), py::arg("coeff"), + "Fused aux loss fwd (device-tensor total_num_tokens, CUDA-graph-safe)"); m.def("fused_moe_aux_loss_bwd", &fused_moe_aux_loss_bwd, py::arg("Const_buf"), py::arg("tokens_per_expert"), py::arg("num_rows"), py::arg("num_cols"), py::arg("grad_aux_loss"), "Fused aux loss bwd"); diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 5dc5c7fe86..a54653785f 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -238,6 +238,40 @@ std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, return std::make_tuple(aux_loss, Const_buf); } +std::tuple fused_moe_aux_loss_fwd_tensor(at::Tensor probs, + at::Tensor tokens_per_expert, + at::Tensor total_num_tokens, + int num_experts, int num_rows, + int num_cols, int topk, + float coeff) { + TORCH_CHECK(topk > 0, "topk must be greater than 0"); + TORCH_CHECK(num_experts > 0, "num_experts must be greater than 0"); + // Device-tensor path: keep total_num_tokens dynamic across CUDA Graph replays. + // Validate shape and dtype on the host; do not read its value (avoids a sync). + TORCH_CHECK(total_num_tokens.is_cuda(), "total_num_tokens must be a CUDA tensor"); + TORCH_CHECK(total_num_tokens.numel() == 1, + "total_num_tokens must contain exactly one element; got ", total_num_tokens.numel()); + TORCH_CHECK(total_num_tokens.scalar_type() == at::kLong, + "total_num_tokens must be int64; got ", total_num_tokens.scalar_type()); + + // Create the output tensor + at::Tensor aux_loss = at::empty({}, at::dtype(probs.scalar_type()).device(at::kCUDA)); + at::Tensor Const_buf = at::empty({2}, at::dtype(at::kFloat).device(at::kCUDA)); + + auto probs_cu = makeTransformerEngineTensor(probs); + auto tokens_per_expert_cu = makeTransformerEngineTensor(tokens_per_expert); + auto total_num_tokens_cu = makeTransformerEngineTensor(total_num_tokens); + auto aux_loss_cu = makeTransformerEngineTensor(aux_loss); + auto Const_buf_cu = makeTransformerEngineTensor(Const_buf); + + nvte_fused_moe_aux_loss_forward_tensor( + probs_cu.data(), tokens_per_expert_cu.data(), total_num_tokens_cu.data(), num_experts, + num_rows, num_cols, topk, coeff, aux_loss_cu.data(), Const_buf_cu.data(), + at::cuda::getCurrentCUDAStream()); + + return std::make_tuple(aux_loss, Const_buf); +} + at::Tensor fused_moe_aux_loss_bwd(at::Tensor Const_buf, at::Tensor tokens_per_expert, int num_rows, int num_cols, at::Tensor grad_aux_loss) { // Create the output tensor diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 667dee6872..c77c3fcf33 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -266,7 +266,9 @@ def fused_compute_score_for_moe_aux_loss( class FusedAuxLoss(torch.autograd.Function): """ - Fused MoE aux loss. + Fused MoE aux loss. ``total_num_tokens`` may be either a Python int + (host-folded coefficient, original fast path) or a 0-dim int64 CUDA + tensor (device-folded coefficient, CUDA-graph-safe path). """ @staticmethod @@ -274,7 +276,7 @@ def forward( ctx, probs: torch.Tensor, tokens_per_expert: torch.Tensor, - total_num_tokens: int, + total_num_tokens: Union[int, torch.Tensor], num_experts: int, topk: int, coeff: float, @@ -282,16 +284,28 @@ def forward( # pylint: disable=missing-function-docstring num_rows = probs.size(0) num_cols = probs.size(1) - aux_loss, Const_buf = tex.fused_moe_aux_loss_fwd( - probs=probs, - tokens_per_expert=tokens_per_expert, - total_num_tokens=total_num_tokens, - num_experts=num_experts, - num_rows=num_rows, - num_cols=num_cols, - topk=topk, - coeff=coeff, - ) + if isinstance(total_num_tokens, torch.Tensor): + aux_loss, Const_buf = tex.fused_moe_aux_loss_fwd_tensor( + probs=probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=total_num_tokens, + num_experts=num_experts, + num_rows=num_rows, + num_cols=num_cols, + topk=topk, + coeff=coeff, + ) + else: + aux_loss, Const_buf = tex.fused_moe_aux_loss_fwd( + probs=probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=int(total_num_tokens), + num_experts=num_experts, + num_rows=num_rows, + num_cols=num_cols, + topk=topk, + coeff=coeff, + ) ctx.save_for_backward(Const_buf, tokens_per_expert) ctx.num_rows = num_rows ctx.num_cols = num_cols @@ -314,7 +328,7 @@ def backward(ctx, grad_aux_loss): def fused_moe_aux_loss( probs: torch.Tensor, tokens_per_expert: torch.Tensor, - total_num_tokens: int, + total_num_tokens: Union[int, torch.Tensor], num_experts: int, topk: int, coeff: float, @@ -326,8 +340,13 @@ def fused_moe_aux_loss( probs : torch.Tensor in fp32/bf16/fp16 tokens_per_expert : torch.Tensor in int32/int64/fp32/bf16 the number of tokens per expert. - total_num_tokens : int - the total number of tokens used in the aux loss calculation. + total_num_tokens : int or 0-dim int64 CUDA torch.Tensor + the total number of tokens used in the aux loss calculation. Pass a + Python int for the fastest path (coefficient folded on the host). + Pass a 0-dim int64 CUDA tensor when the call is captured into a + CUDA Graph and the value must stay dynamic across replays; this + path adds one extra single-thread kernel launch to compute the + coefficient on device. num_experts : int topk : int coeff : float From 21d8d3a3a82a3adb28ecab68aa03b7240ff59bc6 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Thu, 11 Jun 2026 10:17:41 -0700 Subject: [PATCH 05/10] add graph guard for one zero fill Signed-off-by: HaochenYuan --- .../dot_product_attention/context_parallel.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index af3ed49506..ce1bc395ec 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -2899,19 +2899,24 @@ def backward(ctx, dout, *_args): dq, dk, dv = [x.view(*x.shape[:dim], -1, *x.shape[dim + 2 :]) for x in [dq, dk, dv]] # Zero-fill dQ/dK/dV at positions beyond cu_seqlens_*_padded[-1]. - # Use arange+mask indexing rather than `tensor[scalar_tensor:]` slicing -- the - # latter forces a GPU->CPU sync that is forbidden during CUDA graph capture. 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 ): - 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 + if is_graph_capturing(): + 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: + 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) From cee7b6ad9f1a89af242d4177d06aa01c95a98b56 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Thu, 11 Jun 2026 10:46:32 -0700 Subject: [PATCH 06/10] remove redundant zero-fill Signed-off-by: HaochenYuan --- .../dot_product_attention/context_parallel.py | 3 +++ .../pytorch/cpp_extensions/fused_attn.py | 26 ------------------- 2 files changed, 3 insertions(+), 26 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index ce1bc395ec..a77d9db6a1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -2906,6 +2906,8 @@ def backward(ctx, dout, *_args): 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] @@ -2914,6 +2916,7 @@ def backward(ctx, dout, *_args): 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) diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 924c3636db..62dcaadc96 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -340,16 +340,6 @@ def fused_attn_fwd( cuda_graph, ) - # Zero-fill output at positions beyond the actual sequence end (THD CUDA Graph). - # Uses arange+mask indexing instead of `tensor[scalar_tensor:]` slicing -- the - # latter forces a GPU->CPU sync that is forbidden during CUDA graph capture. - if qkv_layout in ("t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"): - out_tensor = output_tensors[0] - actual_t = cu_seqlens_q[-1] if cu_seqlens_q_padded is None else cu_seqlens_q_padded[-1] - if isinstance(out_tensor, torch.Tensor) and out_tensor.shape[0] > 0: - pad_mask = torch.arange(out_tensor.shape[0], device=out_tensor.device) >= actual_t - out_tensor[pad_mask] = 0 - if return_max_logit: qkv_format = qkv_layout.replace("3", "").replace("2", "").split("_")[0] # thd (newer cuDNN runtimes, non-sm120): output_tensors: out [tq, h, d], Stats [tq, h, 1], Max [tq, h, 1] @@ -606,20 +596,4 @@ def fused_attn_bwd( cuda_graph, ) - # Zero-fill dQ/dK/dV at positions beyond the actual sequence end (THD CUDA Graph). - # Use Q's padded boundary for dQ and KV's padded boundary for dK/dV. - if qkv_layout in ("t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"): - q_actual_t = cu_seqlens_q[-1] if cu_seqlens_q_padded is None else cu_seqlens_q_padded[-1] - kv_actual_t = ( - cu_seqlens_kv[-1] if cu_seqlens_kv_padded is None else cu_seqlens_kv_padded[-1] - ) - dq_tensor = output_tensors[0] - if isinstance(dq_tensor, torch.Tensor) and dq_tensor.shape[0] > 0: - q_pad_mask = torch.arange(dq_tensor.shape[0], device=dq_tensor.device) >= q_actual_t - dq_tensor[q_pad_mask] = 0 - for d_tensor in output_tensors[1:3]: # 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 output_tensors From 823805dc516a6ff861ea4b90f8eb8d133c1d2b17 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 11 Jun 2026 17:48:56 +0000 Subject: [PATCH 07/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/csrc/extensions.h | 9 +++----- .../pytorch/csrc/extensions/router.cpp | 21 ++++++++----------- 2 files changed, 12 insertions(+), 18 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 9a1d349b2c..ddf858fd18 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -55,12 +55,9 @@ std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, int num_rows, int num_cols, int topk, float coeff); -std::tuple fused_moe_aux_loss_fwd_tensor(at::Tensor probs, - at::Tensor tokens_per_expert, - at::Tensor total_num_tokens, - int num_experts, int num_rows, - int num_cols, int topk, - float coeff); +std::tuple fused_moe_aux_loss_fwd_tensor( + at::Tensor probs, at::Tensor tokens_per_expert, at::Tensor total_num_tokens, int num_experts, + int num_rows, int num_cols, int topk, float coeff); at::Tensor fused_moe_aux_loss_bwd(at::Tensor Const_buf, at::Tensor tokens_per_expert, int num_rows, int num_cols, at::Tensor grad_aux_loss); diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index a54653785f..88d983dcbd 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -238,12 +238,9 @@ std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, return std::make_tuple(aux_loss, Const_buf); } -std::tuple fused_moe_aux_loss_fwd_tensor(at::Tensor probs, - at::Tensor tokens_per_expert, - at::Tensor total_num_tokens, - int num_experts, int num_rows, - int num_cols, int topk, - float coeff) { +std::tuple fused_moe_aux_loss_fwd_tensor( + at::Tensor probs, at::Tensor tokens_per_expert, at::Tensor total_num_tokens, int num_experts, + int num_rows, int num_cols, int topk, float coeff) { TORCH_CHECK(topk > 0, "topk must be greater than 0"); TORCH_CHECK(num_experts > 0, "num_experts must be greater than 0"); // Device-tensor path: keep total_num_tokens dynamic across CUDA Graph replays. @@ -251,8 +248,8 @@ std::tuple fused_moe_aux_loss_fwd_tensor(at::Tensor prob TORCH_CHECK(total_num_tokens.is_cuda(), "total_num_tokens must be a CUDA tensor"); TORCH_CHECK(total_num_tokens.numel() == 1, "total_num_tokens must contain exactly one element; got ", total_num_tokens.numel()); - TORCH_CHECK(total_num_tokens.scalar_type() == at::kLong, - "total_num_tokens must be int64; got ", total_num_tokens.scalar_type()); + TORCH_CHECK(total_num_tokens.scalar_type() == at::kLong, "total_num_tokens must be int64; got ", + total_num_tokens.scalar_type()); // Create the output tensor at::Tensor aux_loss = at::empty({}, at::dtype(probs.scalar_type()).device(at::kCUDA)); @@ -264,10 +261,10 @@ std::tuple fused_moe_aux_loss_fwd_tensor(at::Tensor prob auto aux_loss_cu = makeTransformerEngineTensor(aux_loss); auto Const_buf_cu = makeTransformerEngineTensor(Const_buf); - nvte_fused_moe_aux_loss_forward_tensor( - probs_cu.data(), tokens_per_expert_cu.data(), total_num_tokens_cu.data(), num_experts, - num_rows, num_cols, topk, coeff, aux_loss_cu.data(), Const_buf_cu.data(), - at::cuda::getCurrentCUDAStream()); + nvte_fused_moe_aux_loss_forward_tensor(probs_cu.data(), tokens_per_expert_cu.data(), + total_num_tokens_cu.data(), num_experts, num_rows, + num_cols, topk, coeff, aux_loss_cu.data(), + Const_buf_cu.data(), at::cuda::getCurrentCUDAStream()); return std::make_tuple(aux_loss, Const_buf); } From bce6bc50090c6bf873514e73810d085c7228e163 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Tue, 30 Jun 2026 07:13:19 -0700 Subject: [PATCH 08/10] rename & remove prelude kernel Signed-off-by: HaochenYuan --- tests/pytorch/test_fused_router.py | 16 ++-- .../common/fused_router/fused_moe_aux_loss.cu | 85 ++++++++----------- .../include/transformer_engine/fused_router.h | 17 ++-- transformer_engine/pytorch/csrc/extensions.h | 2 +- .../pytorch/csrc/extensions/pybind.cpp | 2 +- .../pytorch/csrc/extensions/router.cpp | 10 +-- transformer_engine/pytorch/router.py | 7 +- 7 files changed, 64 insertions(+), 75 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 63d5df39aa..a3aca0a276 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -464,8 +464,8 @@ def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk, expert_multipl 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. The captured graph must observe the new value - via the device-side coefficient computation.""" + 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 @@ -477,7 +477,7 @@ def test_fused_moe_aux_loss_cuda_graph_capture(): 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() + 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") @@ -487,7 +487,7 @@ def test_fused_moe_aux_loss_cuda_graph_capture(): s.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(s): for _ in range(3): - _ = fused_moe_aux_loss( + warmup_out = fused_moe_aux_loss( probs=probs, tokens_per_expert=tokens_per_expert, total_num_tokens=total_num_tokens_dev, @@ -495,6 +495,8 @@ def test_fused_moe_aux_loss_cuda_graph_capture(): topk=topk, coeff=coeff, ) + torch.autograd.grad(warmup_out, probs) + del warmup_out torch.cuda.current_stream().wait_stream(s) g = torch.cuda.CUDAGraph() @@ -507,6 +509,7 @@ def test_fused_moe_aux_loss_cuda_graph_capture(): 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 @@ -515,15 +518,18 @@ def test_fused_moe_aux_loss_cuda_graph_capture(): 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=probs, + 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: diff --git a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu index b425e36a1c..5a0f761f3f 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -134,30 +134,14 @@ void fused_moe_aux_loss_forward(const Tensor& probs, const Tensor& tokens_per_ex /* ------------------------------------------------------------------------- * CUDA-graph-safe variant: total_num_tokens lives in a 0-dim int64 device - * tensor whose value can change between graph replays. A single-thread - * prelude kernel reads it on device and writes C_coeff into Coeff_buf[0]; - * the main reduction kernel then reads C_coeff from Coeff_buf[0] instead - * of receiving it as a launch parameter, so the value stays dynamic. + * 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. * ------------------------------------------------------------------------- */ -__global__ void fused_moe_aux_loss_compute_coeff_kernel(const int64_t* total_num_tokens_ptr, - int num_experts, int topk, float coeff, - float* Coeff_buf) { - if (threadIdx.x == 0 && blockIdx.x == 0) { - const int64_t T = *total_num_tokens_ptr; - const float T_f = static_cast(T); - Coeff_buf[0] = (static_cast(num_experts) * coeff) / static_cast(topk) / T_f / T_f; - } -} - template -__global__ void fused_moe_aux_loss_forward_kernel_tensor(const DataType* probs, - const IndexType* tokens_per_expert, - int num_rows, int num_cols, - const float* Coeff_buf) { - // Read C_coeff from Coeff_buf[0], written by the prelude kernel that ran - // earlier on the same stream. Backward reads it from the same slot. - const float C_coeff = Coeff_buf[0]; - +__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) { @@ -180,18 +164,22 @@ __global__ void fused_moe_aux_loss_forward_kernel_tensor(const DataType* probs, CompType block_sum = warp_reduce_on_shmem(shmem_block, static_cast(blockDim.x), ReduceFuncType::SUM, lane_id); if (lane_id == 0) { - atomicAdd(const_cast(&Coeff_buf[1]), static_cast(block_sum * C_coeff)); + const float total_num_tokens = static_cast(*total_num_tokens_ptr); + const float C_coeff = (static_cast(num_experts) * coeff) / static_cast(topk) / + total_num_tokens / total_num_tokens; + if (blockIdx.x == 0) { + Coeff_buf[0] = C_coeff; + } + atomicAdd(&Coeff_buf[1], static_cast(block_sum * C_coeff)); } } } template -void fused_moe_aux_loss_forward_kernel_launcher_tensor(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) { +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, @@ -204,15 +192,11 @@ void fused_moe_aux_loss_forward_kernel_launcher_tensor(const DataType* probs, const size_t smem_size = block_size * sizeof(CompType); check_shared_memory_capacity_num_experts(smem_size, num_cols); - // Zero the float accumulator (Coeff_buf[1]); the prelude kernel writes - // Coeff_buf[0] = C_coeff derived from device-side total_num_tokens. + // 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_compute_coeff_kernel<<<1, 1, 0, stream>>>(total_num_tokens_dev, num_experts, - topk, coeff, Coeff_buf); - NVTE_CHECK_CUDA(cudaGetLastError()); - - fused_moe_aux_loss_forward_kernel_tensor - <<>>(probs, tokens_per_expert, num_rows, num_cols, + fused_moe_aux_loss_forward_kernel_graph_safe + <<>>(probs, tokens_per_expert, total_num_tokens_dev, + num_experts, num_rows, num_cols, topk, coeff, Coeff_buf); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -220,10 +204,11 @@ void fused_moe_aux_loss_forward_kernel_launcher_tensor(const DataType* probs, NVTE_CHECK_CUDA(cudaGetLastError()); } -void fused_moe_aux_loss_forward(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) { +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(total_num_tokens.data.dtype)); @@ -233,7 +218,7 @@ void fused_moe_aux_loss_forward(const Tensor& probs, const Tensor& tokens_per_ex probs.data.dtype, DataType, TE_ROUTER_INDEX_TYPE_SWITCH_ALL( tokens_per_expert.data.dtype, IndexType, - fused_moe_aux_loss_forward_kernel_launcher_tensor( + fused_moe_aux_loss_forward_kernel_launcher_graph_safe( reinterpret_cast(probs.data.dptr), reinterpret_cast(tokens_per_expert.data.dptr), reinterpret_cast(total_num_tokens.data.dptr), num_experts, num_rows, @@ -304,15 +289,15 @@ void nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor to *convertNVTETensorCheck(Coeff_buf), stream); } -void nvte_fused_moe_aux_loss_forward_tensor(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_tensor); +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( + 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); diff --git a/transformer_engine/common/include/transformer_engine/fused_router.h b/transformer_engine/common/include/transformer_engine/fused_router.h index 22de780919..5242ddd80b 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -215,19 +215,18 @@ void nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor to /*! \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. Adds one extra single-thread - * kernel launch versus the host-int path; prefer this only when the caller - * needs CUDA-graph-safe semantics. + * 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_tensor(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); +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. * diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index ddf858fd18..a28abcd503 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -55,7 +55,7 @@ std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, int num_rows, int num_cols, int topk, float coeff); -std::tuple fused_moe_aux_loss_fwd_tensor( +std::tuple fused_moe_aux_loss_fwd_graph_safe( at::Tensor probs, at::Tensor tokens_per_expert, at::Tensor total_num_tokens, int num_experts, int num_rows, int num_cols, int topk, float coeff); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 2127ef1c24..1224065218 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -159,7 +159,7 @@ void init_router_bindings(pybind11::module &m) { py::arg("tokens_per_expert"), py::arg("total_num_tokens"), py::arg("num_experts"), py::arg("num_rows"), py::arg("num_cols"), py::arg("topk"), py::arg("coeff"), "Fused aux loss fwd (host-int total_num_tokens, host-folded C_coeff)"); - m.def("fused_moe_aux_loss_fwd_tensor", &fused_moe_aux_loss_fwd_tensor, py::arg("probs"), + m.def("fused_moe_aux_loss_fwd_graph_safe", &fused_moe_aux_loss_fwd_graph_safe, py::arg("probs"), py::arg("tokens_per_expert"), py::arg("total_num_tokens"), py::arg("num_experts"), py::arg("num_rows"), py::arg("num_cols"), py::arg("topk"), py::arg("coeff"), "Fused aux loss fwd (device-tensor total_num_tokens, CUDA-graph-safe)"); diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 88d983dcbd..8801bb8d9c 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -238,7 +238,7 @@ std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, return std::make_tuple(aux_loss, Const_buf); } -std::tuple fused_moe_aux_loss_fwd_tensor( +std::tuple fused_moe_aux_loss_fwd_graph_safe( at::Tensor probs, at::Tensor tokens_per_expert, at::Tensor total_num_tokens, int num_experts, int num_rows, int num_cols, int topk, float coeff) { TORCH_CHECK(topk > 0, "topk must be greater than 0"); @@ -261,10 +261,10 @@ std::tuple fused_moe_aux_loss_fwd_tensor( auto aux_loss_cu = makeTransformerEngineTensor(aux_loss); auto Const_buf_cu = makeTransformerEngineTensor(Const_buf); - nvte_fused_moe_aux_loss_forward_tensor(probs_cu.data(), tokens_per_expert_cu.data(), - total_num_tokens_cu.data(), num_experts, num_rows, - num_cols, topk, coeff, aux_loss_cu.data(), - Const_buf_cu.data(), at::cuda::getCurrentCUDAStream()); + nvte_fused_moe_aux_loss_forward_graph_safe(probs_cu.data(), tokens_per_expert_cu.data(), + total_num_tokens_cu.data(), num_experts, num_rows, + num_cols, topk, coeff, aux_loss_cu.data(), + Const_buf_cu.data(), at::cuda::getCurrentCUDAStream()); return std::make_tuple(aux_loss, Const_buf); } diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index c77c3fcf33..7ac51494c0 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -285,7 +285,7 @@ def forward( num_rows = probs.size(0) num_cols = probs.size(1) if isinstance(total_num_tokens, torch.Tensor): - aux_loss, Const_buf = tex.fused_moe_aux_loss_fwd_tensor( + aux_loss, Const_buf = tex.fused_moe_aux_loss_fwd_graph_safe( probs=probs, tokens_per_expert=tokens_per_expert, total_num_tokens=total_num_tokens, @@ -344,9 +344,8 @@ def fused_moe_aux_loss( the total number of tokens used in the aux loss calculation. Pass a Python int for the fastest path (coefficient folded on the host). Pass a 0-dim int64 CUDA tensor when the call is captured into a - CUDA Graph and the value must stay dynamic across replays; this - path adds one extra single-thread kernel launch to compute the - coefficient on device. + CUDA Graph and the value must stay dynamic across replays; the + coefficient is computed on device by the main reduction kernel. num_experts : int topk : int coeff : float From a69fb1bad56782f6d70d0c867f9ab00dd878244f Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Tue, 30 Jun 2026 13:16:17 -0700 Subject: [PATCH 09/10] Update warp reduction function Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- transformer_engine/common/fused_router/fused_moe_aux_loss.cu | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu index 99b205b73e..b873d966cc 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -161,8 +161,7 @@ __global__ void fused_moe_aux_loss_forward_kernel_graph_safe( 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(shmem_block, static_cast(blockDim.x), - ReduceFuncType::SUM, lane_id); + CompType block_sum = warp_reduce_on_shmem(shmem_block, static_cast(blockDim.x), lane_id); if (lane_id == 0) { const float total_num_tokens = static_cast(*total_num_tokens_ptr); const float C_coeff = (static_cast(num_experts) * coeff) / static_cast(topk) / From 1bc41e5704192e4a34e609d419fb50b2ff58502f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 30 Jun 2026 20:17:36 +0000 Subject: [PATCH 10/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/fused_router/fused_moe_aux_loss.cu | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu index b873d966cc..9fb5adcc1f 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -161,7 +161,8 @@ __global__ void fused_moe_aux_loss_forward_kernel_graph_safe( 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(shmem_block, static_cast(blockDim.x), lane_id); + CompType block_sum = warp_reduce_on_shmem( + shmem_block, static_cast(blockDim.x), lane_id); if (lane_id == 0) { const float total_num_tokens = static_cast(*total_num_tokens_ptr); const float C_coeff = (static_cast(num_experts) * coeff) / static_cast(topk) /