From 938612ea4c223c7e1acbd29986a957b1f648d0e0 Mon Sep 17 00:00:00 2001 From: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> Date: Fri, 25 Sep 2026 03:39:49 +0530 Subject: [PATCH 1/4] Guard fused MoE auxiliary loss integer offsets Signed-off-by: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> --- tests/pytorch/test_fused_router.py | 45 +++++++++++++++++++ .../common/fused_router/fused_moe_aux_loss.cu | 9 ++++ 2 files changed, 54 insertions(+) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index bd89ecebc9e..54491568c4c 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -1119,6 +1119,51 @@ 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) +@pytest.mark.parametrize("path", ["forward", "graph_safe_forward", "backward"]) +def test_fused_moe_aux_loss_rejects_int_offset_overflow(path): + num_rows, num_cols = 8_388_609, 256 # Above INT_MAX elements. + bytes_needed = num_rows * num_cols * 2 + if not torch.cuda.is_available(): + pytest.skip("Requires a CUDA GPU") + torch.cuda.empty_cache() + free_bytes, _ = torch.cuda.mem_get_info() + if free_bytes < bytes_needed + (1 << 30): + pytest.skip("Needs 4.0 GiB for BF16 input/output plus 1 GiB headroom") + + tokens_per_expert = torch.ones(num_cols, device="cuda", dtype=torch.int32) + try: + if path == "backward": + tex.fused_moe_aux_loss_bwd( + Const_buf=torch.ones(2, device="cuda", dtype=torch.float32), + tokens_per_expert=tokens_per_expert, + num_rows=num_rows, + num_cols=num_cols, + grad_aux_loss=torch.ones((), device="cuda", dtype=torch.bfloat16), + ) + else: + probs = torch.empty((num_rows, num_cols), device="cuda", dtype=torch.bfloat16) + arguments = dict( + probs=probs, + tokens_per_expert=tokens_per_expert, + num_experts=num_cols, + num_rows=num_rows, + num_cols=num_cols, + topk=1, + coeff=0.01, + ) + if path == "forward": + tex.fused_moe_aux_loss_fwd(total_num_tokens=num_rows, **arguments) + else: + total = torch.tensor(num_rows, device="cuda", dtype=torch.int64) + tex.fused_moe_aux_loss_fwd_graph_safe(total_num_tokens=total, **arguments) + except torch.cuda.OutOfMemoryError: + pytest.skip("Could not allocate the BF16 input/output tensor") + except RuntimeError as exc: + assert "num_rows * num_cols exceeds INT_MAX" in str(exc) + else: + pytest.fail("Expected the int-offset guard to reject the shape") + + 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 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 9fb5adcc1f4..50b1bcb0a62 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -89,6 +89,9 @@ void fused_moe_aux_loss_forward_kernel_launcher(const DataType* probs, 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(static_cast(num_rows) * num_cols <= INT_MAX, + "num_rows * num_cols exceeds INT_MAX (kernel uses int offsets), got ", + static_cast(num_rows) * num_cols); // Sequence aux loss batches independent sequences along the expert dimension. NVTE_CHECK(num_cols % num_experts == 0, "Number of input columns (", num_cols, ") must be a multiple of number of experts (", num_experts, ")."); @@ -182,6 +185,9 @@ void fused_moe_aux_loss_forward_kernel_launcher_graph_safe( 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(static_cast(num_rows) * num_cols <= INT_MAX, + "num_rows * num_cols exceeds INT_MAX (kernel uses int offsets), got ", + static_cast(num_rows) * num_cols); NVTE_CHECK(num_cols % num_experts == 0, "Number of input columns (", num_cols, ") must be a multiple of number of experts (", num_experts, ")."); @@ -252,6 +258,9 @@ void fused_moe_aux_loss_backward_kernel_launcher(const float* Const_buf, const IndexType* tokens_per_expert, int num_rows, int num_cols, DataType* grad_aux_loss, DataType* grad_probs, cudaStream_t stream) { + NVTE_CHECK(static_cast(num_rows) * num_cols <= INT_MAX, + "num_rows * num_cols exceeds INT_MAX (kernel uses int offsets), got ", + static_cast(num_rows) * num_cols); // Meta data for the kernel int block_size = 256; int grid_size = (num_rows + block_size - 1) / block_size; From bd198e338decdd6931507c5d100d3f8e758191e8 Mon Sep 17 00:00:00 2001 From: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> Date: Fri, 2 Oct 2026 22:14:50 +0530 Subject: [PATCH 2/4] Update tests/pytorch/test_fused_router.py Co-authored-by: Alp Dener Signed-off-by: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> --- tests/pytorch/test_fused_router.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 54491568c4c..08670c40dda 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -1123,8 +1123,6 @@ def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk, expert_multipl def test_fused_moe_aux_loss_rejects_int_offset_overflow(path): num_rows, num_cols = 8_388_609, 256 # Above INT_MAX elements. bytes_needed = num_rows * num_cols * 2 - if not torch.cuda.is_available(): - pytest.skip("Requires a CUDA GPU") torch.cuda.empty_cache() free_bytes, _ = torch.cuda.mem_get_info() if free_bytes < bytes_needed + (1 << 30): From 56943b5acc570f0d0d098d9be3dc53bf4cc80f1c Mon Sep 17 00:00:00 2001 From: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> Date: Wed, 7 Oct 2026 06:31:16 +0530 Subject: [PATCH 3/4] Use 64-bit indexing for fused MoE auxiliary loss Signed-off-by: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> --- tests/pytorch/test_fused_router.py | 37 ++++++++++++------- .../common/fused_router/fused_moe_aux_loss.cu | 28 +++++--------- 2 files changed, 34 insertions(+), 31 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 08670c40dda..fafe44917a5 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -1120,7 +1120,7 @@ def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk, expert_multipl @pytest.mark.parametrize("path", ["forward", "graph_safe_forward", "backward"]) -def test_fused_moe_aux_loss_rejects_int_offset_overflow(path): +def test_fused_moe_aux_loss_int64_offsets(path): num_rows, num_cols = 8_388_609, 256 # Above INT_MAX elements. bytes_needed = num_rows * num_cols * 2 torch.cuda.empty_cache() @@ -1131,15 +1131,22 @@ def test_fused_moe_aux_loss_rejects_int_offset_overflow(path): tokens_per_expert = torch.ones(num_cols, device="cuda", dtype=torch.int32) try: if path == "backward": - tex.fused_moe_aux_loss_bwd( - Const_buf=torch.ones(2, device="cuda", dtype=torch.float32), + tokens_per_expert = torch.arange(1, num_cols + 1, device="cuda", dtype=torch.int32) + grad_probs = tex.fused_moe_aux_loss_bwd( + Const_buf=torch.tensor([0.25, 0.0], device="cuda", dtype=torch.float32), tokens_per_expert=tokens_per_expert, num_rows=num_rows, num_cols=num_cols, - grad_aux_loss=torch.ones((), device="cuda", dtype=torch.bfloat16), + grad_aux_loss=torch.tensor(2.0, device="cuda", dtype=torch.bfloat16), ) + expected = (0.5 * tokens_per_expert).to(torch.bfloat16) + for start in range(0, num_rows, 65536): + chunk = grad_probs[start : start + 65536] + torch.testing.assert_close(chunk, expected.expand_as(chunk), atol=0, rtol=0) else: - probs = torch.empty((num_rows, num_cols), device="cuda", dtype=torch.bfloat16) + probs = torch.zeros((num_rows, num_cols), device="cuda", dtype=torch.bfloat16) + # Nonzero values on both sides of the int32 offset boundary. + probs[-2:] = 1 arguments = dict( probs=probs, tokens_per_expert=tokens_per_expert, @@ -1147,19 +1154,23 @@ def test_fused_moe_aux_loss_rejects_int_offset_overflow(path): num_rows=num_rows, num_cols=num_cols, topk=1, - coeff=0.01, + coeff=1.0, ) if path == "forward": - tex.fused_moe_aux_loss_fwd(total_num_tokens=num_rows, **arguments) + aux_loss, const_buf = tex.fused_moe_aux_loss_fwd( + total_num_tokens=num_cols, **arguments + ) else: - total = torch.tensor(num_rows, device="cuda", dtype=torch.int64) - tex.fused_moe_aux_loss_fwd_graph_safe(total_num_tokens=total, **arguments) + total = torch.tensor(num_cols, device="cuda", dtype=torch.int64) + aux_loss, const_buf = tex.fused_moe_aux_loss_fwd_graph_safe( + total_num_tokens=total, **arguments + ) + torch.testing.assert_close(aux_loss, probs.new_tensor(2.0), atol=0, rtol=0) + torch.testing.assert_close( + const_buf, const_buf.new_tensor([1.0 / num_cols, 2.0]), atol=0, rtol=0 + ) except torch.cuda.OutOfMemoryError: pytest.skip("Could not allocate the BF16 input/output tensor") - except RuntimeError as exc: - assert "num_rows * num_cols exceeds INT_MAX" in str(exc) - else: - pytest.fail("Expected the int-offset guard to reject the shape") def test_fused_moe_aux_loss_cuda_graph_capture(): 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 50b1bcb0a62..9d6eceb3cc7 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -37,11 +37,11 @@ __global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs, // Grid-stride over rows so that every row is processed exactly once. // Each thread processes a subset of columns. - for (int col = threadIdx.x; col < num_cols; col += blockDim.x) { + for (int64_t col = threadIdx.x; col < num_cols; col += blockDim.x) { CompType col_sum = CompType(0); // Accumulate probs over the rows assigned to this CTA (grid-stride). - for (int row = blockIdx.x; row < num_rows; row += gridDim.x) { + for (int64_t row = blockIdx.x; row < num_rows; row += gridDim.x) { col_sum += CompType(probs[row * num_cols + col]); } @@ -89,9 +89,6 @@ void fused_moe_aux_loss_forward_kernel_launcher(const DataType* probs, 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(static_cast(num_rows) * num_cols <= INT_MAX, - "num_rows * num_cols exceeds INT_MAX (kernel uses int offsets), got ", - static_cast(num_rows) * num_cols); // Sequence aux loss batches independent sequences along the expert dimension. NVTE_CHECK(num_cols % num_experts == 0, "Number of input columns (", num_cols, ") must be a multiple of number of experts (", num_experts, ")."); @@ -147,9 +144,9 @@ __global__ void fused_moe_aux_loss_forward_kernel_graph_safe( 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) { + for (int64_t 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) { + for (int64_t row = blockIdx.x; row < num_rows; row += gridDim.x) { col_sum += CompType(probs[row * num_cols + col]); } col_sum *= CompType(tokens_per_expert[col]); @@ -185,9 +182,6 @@ void fused_moe_aux_loss_forward_kernel_launcher_graph_safe( 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(static_cast(num_rows) * num_cols <= INT_MAX, - "num_rows * num_cols exceeds INT_MAX (kernel uses int offsets), got ", - static_cast(num_rows) * num_cols); NVTE_CHECK(num_cols % num_experts == 0, "Number of input columns (", num_cols, ") must be a multiple of number of experts (", num_experts, ")."); @@ -237,17 +231,18 @@ __global__ void fused_moe_aux_loss_backward_kernel(const float* Const_buf, const IndexType* tokens_per_expert, int num_rows, int num_cols, DataType* grad_aux_loss, DataType* grad_probs) { - int global_warp_num = gridDim.x * blockDim.x / kThreadsPerWarp; - int global_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / kThreadsPerWarp; + int64_t global_warp_num = static_cast(gridDim.x) * blockDim.x / kThreadsPerWarp; + int64_t global_warp_id = + (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) / kThreadsPerWarp; int lane_id = threadIdx.x % kThreadsPerWarp; // Loop: for all positions in each row - for (int i = lane_id; i < num_cols; i += kThreadsPerWarp) { + for (int64_t i = lane_id; i < num_cols; i += kThreadsPerWarp) { float C_coeff = Const_buf[0]; CompType tokens_per_expert_i = static_cast(tokens_per_expert[i]); CompType grad_aux_loss_value = static_cast(grad_aux_loss[0]); // Loop: for all rows - for (int j = global_warp_id; j < num_rows; j += global_warp_num) { + for (int64_t j = global_warp_id; j < num_rows; j += global_warp_num) { grad_probs[j * num_cols + i] = C_coeff * tokens_per_expert_i * grad_aux_loss_value; } } @@ -258,12 +253,9 @@ void fused_moe_aux_loss_backward_kernel_launcher(const float* Const_buf, const IndexType* tokens_per_expert, int num_rows, int num_cols, DataType* grad_aux_loss, DataType* grad_probs, cudaStream_t stream) { - NVTE_CHECK(static_cast(num_rows) * num_cols <= INT_MAX, - "num_rows * num_cols exceeds INT_MAX (kernel uses int offsets), got ", - static_cast(num_rows) * num_cols); // Meta data for the kernel int block_size = 256; - int grid_size = (num_rows + block_size - 1) / block_size; + int grid_size = (static_cast(num_rows) + block_size - 1) / block_size; fused_moe_aux_loss_backward_kernel<<>>( Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss, grad_probs); NVTE_CHECK_CUDA(cudaGetLastError()); From 198df0eb30c7058b917b89cd6e44113eb229b930 Mon Sep 17 00:00:00 2001 From: arcusbuilds Date: Thu, 8 Oct 2026 23:22:01 +0530 Subject: [PATCH 4/4] Unroll forward row loop in fused MoE aux loss Signed-off-by: arcusbuilds --- transformer_engine/common/fused_router/fused_moe_aux_loss.cu | 2 ++ 1 file changed, 2 insertions(+) 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 9d6eceb3cc7..ddebfce8ad6 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -41,6 +41,7 @@ __global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs, CompType col_sum = CompType(0); // Accumulate probs over the rows assigned to this CTA (grid-stride). +#pragma unroll 4 for (int64_t row = blockIdx.x; row < num_rows; row += gridDim.x) { col_sum += CompType(probs[row * num_cols + col]); } @@ -146,6 +147,7 @@ __global__ void fused_moe_aux_loss_forward_kernel_graph_safe( CompType thread_sum = CompType(0); for (int64_t col = threadIdx.x; col < num_cols; col += blockDim.x) { CompType col_sum = CompType(0); +#pragma unroll 4 for (int64_t row = blockIdx.x; row < num_rows; row += gridDim.x) { col_sum += CompType(probs[row * num_cols + col]); }