diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 3d2f99b51b..82b9df262f 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/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index ab12216df8..68d3ed9565 100644 --- a/tests/pytorch/test_fused_router.py +++ b/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. 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 a80e90c0fd..9fb5adcc1f 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,100 @@ 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. 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 +__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(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), 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) / + 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_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(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. 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 + <<>>(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<<<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(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( + 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 +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, diff --git a/transformer_engine/common/include/transformer_engine/fused_router.h b/transformer_engine/common/include/transformer_engine/fused_router.h index 08f347c616..03fec2eb2c 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/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. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 8f42983553..785e438cda 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/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 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 61a46a8652..a62ca73187 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/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, 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 03008bb2d7..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 @@ -1558,16 +1558,24 @@ 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 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 - 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]) - ) + 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 69c73fe2fa..6edfbdc00e 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -57,6 +57,10 @@ 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_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); + 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 14272c9ac4..9c9ec36138 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_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)"); 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 70762e3729..d59d6bc415 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -321,6 +321,37 @@ 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_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"); + 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_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); +} + 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 519d06ce23..51350ab2fd 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -280,7 +280,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 @@ -288,7 +290,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, @@ -296,16 +298,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_graph_safe( + 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 @@ -328,7 +342,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, @@ -340,8 +354,12 @@ 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; the + coefficient is computed on device by the main reduction kernel. num_experts : int topk : int coeff : float