From edf1e35987375cbbdf95b6d0eabe4dc1bd5f4d66 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu Date: Wed, 9 Sep 2026 23:33:05 -0700 Subject: [PATCH 1/2] initial draft Signed-off-by: Zhongbo Zhu Signed-off-by: zhongboz --- tests/pytorch/test_fused_router.py | 282 ++++++++++++++++-- .../common/fused_router/fused_moe_aux_loss.cu | 180 ++++++++--- .../include/transformer_engine/fused_router.h | 54 +++- transformer_engine/pytorch/csrc/extensions.h | 14 +- .../pytorch/csrc/extensions/pybind.cpp | 3 + .../pytorch/csrc/extensions/router.cpp | 56 +++- transformer_engine/pytorch/router.py | 46 ++- 7 files changed, 560 insertions(+), 75 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index c6c82fa3d78..bb35c7dcaf9 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -15,6 +15,7 @@ fused_moe_aux_loss, mark_qb_bin_bounds_validated, ) +import transformer_engine.pytorch.router as te_router import transformer_engine_torch as tex seed = 42 @@ -1183,23 +1184,246 @@ 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(): +def _moe_aux_loss_determinism_inputs( + dtype, num_tokens, num_experts, expert_multiplier, token_count_dtype +): + """Asymmetric positive scores spanning several orders of magnitude.""" + generator = torch.Generator(device="cuda").manual_seed(1234) + scores = torch.rand( + (num_tokens, expert_multiplier, num_experts), device="cuda", generator=generator + ).pow(4) + scores *= torch.logspace(-4, 0, num_experts, device="cuda") + scores /= scores.sum(dim=-1, keepdim=True) + probs = scores.reshape(num_tokens, -1).to(dtype).requires_grad_(True) + tokens_per_expert = torch.randint( + 1, num_tokens + 1, (num_experts * expert_multiplier,), device="cuda", generator=generator + ).to(token_count_dtype) + if tokens_per_expert.is_floating_point(): + # Global aux loss may use averaged, non-integral expert counts. + tokens_per_expert /= 3 + return probs, tokens_per_expert + + +def _assert_bitwise_equal(actual, expected): + # Compare storage bits, including for the scalar output and low-precision dtypes. + assert actual.dtype == expected.dtype + assert actual.shape == expected.shape + assert torch.equal(actual.reshape(-1).view(torch.uint8), expected.reshape(-1).view(torch.uint8)) + + +@pytest.mark.parametrize("tensor_total", [False, True], ids=["host_total", "tensor_total"]) +@pytest.mark.parametrize("deterministic", [None, False, True], ids=["default", "atomic", "fixed"]) +def test_fused_moe_aux_loss_deterministic_dispatch(monkeypatch, tensor_total, deterministic): + """Check Python dispatch and autograd argument count without launching CUDA.""" + probs = torch.ones(3, 2, requires_grad=True) + tokens_per_expert = torch.tensor([1, 2], dtype=torch.int32) + total_num_tokens = torch.tensor(3, dtype=torch.int64) if tensor_total else 3 + constants = torch.tensor([3.0, 4.0]) + workspace = torch.empty(256, dtype=torch.float32) + calls = [] + workspace_requests = [] + + def fake_workspace(device): + workspace_requests.append(device) + return workspace + + def fake_forward(path, **kwargs): + calls.append((path, kwargs["deterministic"])) + assert kwargs["probs"] is probs + assert kwargs["num_rows"] == 3 and kwargs["num_cols"] == 2 + assert kwargs["workspace"] is (workspace if deterministic is True else None) + return torch.tensor(2.0), constants + + def fake_scalar_forward(**kwargs): + return fake_forward("host", **kwargs) + + def fake_tensor_forward(**kwargs): + return fake_forward("tensor", **kwargs) + + def fake_backward(**kwargs): + assert kwargs["Const_buf"] is constants + assert kwargs["tokens_per_expert"] is tokens_per_expert + return torch.full_like(probs, 7.0) * kwargs["grad_aux_loss"] + + monkeypatch.setattr(tex, "fused_moe_aux_loss_fwd", fake_scalar_forward) + monkeypatch.setattr(tex, "fused_moe_aux_loss_fwd_graph_safe", fake_tensor_forward) + monkeypatch.setattr(tex, "fused_moe_aux_loss_bwd", fake_backward) + monkeypatch.setattr(te_router, "_get_moe_aux_loss_workspace", fake_workspace) + kwargs = {} if deterministic is None else {"deterministic": deterministic} + out = fused_moe_aux_loss(probs, tokens_per_expert, total_num_tokens, 2, 1, 0.01, **kwargs) + out.backward(torch.tensor(0.5)) + + assert calls == [("tensor" if tensor_total else "host", deterministic is True)] + assert workspace_requests == ([probs.device] if deterministic is True else []) + torch.testing.assert_close(probs.grad, torch.full_like(probs, 3.5)) + + +def test_fused_moe_aux_loss_workspace_reuse(): + """Eager calls share scratch only when ordered on the same device/stream.""" + device = torch.device("cuda", torch.cuda.current_device()) + workspaces = [] + streams = [torch.cuda.Stream(), torch.cuda.Stream()] + for stream in streams: + with torch.cuda.stream(stream): + workspace = te_router._get_moe_aux_loss_workspace(device) + again = te_router._get_moe_aux_loss_workspace(device) + assert workspace.data_ptr() == again.data_ptr() + assert workspace.device == device and workspace.dtype == torch.float32 + workspaces.append(workspace) + assert workspaces[0].data_ptr() != workspaces[1].data_ptr() + + +def test_fused_moe_aux_loss_cuda_graph_workspace_isolation(monkeypatch): + """Graphs captured on one stream need private scratch for concurrent replay.""" + num_tokens, num_experts, topk, coeff = 1025, 128, 4, 0.01 + probs_a, counts_a = _moe_aux_loss_determinism_inputs( + torch.float32, num_tokens, num_experts, 1, torch.int32 + ) + inputs = [ + (probs_a, counts_a, torch.tensor(num_tokens, dtype=torch.int64, device="cuda")), + ( + probs_a.detach().flip((1,)).requires_grad_(True), + counts_a.roll(3) + 1, + torch.tensor(num_tokens * 2 + 1, dtype=torch.int64, device="cuda"), + ), + ] + + def run(probs, counts, total): + out = fused_moe_aux_loss(probs, counts, total, num_experts, topk, coeff, deterministic=True) + (grad,) = torch.autograd.grad(out, probs) + return out, grad + + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + expected = [] + with torch.cuda.stream(capture_stream): + eager_workspace = te_router._get_moe_aux_loss_workspace(probs_a.device) + for args in inputs: + for _ in range(3): + warmup_out, warmup_grad = run(*args) + expected.append((warmup_out.detach().clone(), warmup_grad.clone())) + del warmup_out, warmup_grad + torch.cuda.current_stream().wait_stream(capture_stream) + + # Record addresses only, so the test does not extend scratch tensor lifetimes. + captured_workspace_ptrs = [] + get_workspace = te_router._get_moe_aux_loss_workspace + + def trace_workspace(device): + workspace = get_workspace(device) + if torch.cuda.is_current_stream_capturing(): + captured_workspace_ptrs.append(workspace.data_ptr()) + return workspace + + monkeypatch.setattr(te_router, "_get_moe_aux_loss_workspace", trace_workspace) + graphs, captured_outputs = [], [] + for args in inputs: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + outputs = run(*args) + graphs.append(graph) + captured_outputs.append(outputs) + assert len(captured_workspace_ptrs) == 2 + assert captured_workspace_ptrs[0] != captured_workspace_ptrs[1] + assert eager_workspace.data_ptr() not in captured_workspace_ptrs + with torch.cuda.stream(capture_stream): + assert get_workspace(probs_a.device).data_ptr() == eager_workspace.data_ptr() + + replay_streams = [torch.cuda.Stream(), torch.cuda.Stream()] + for stream in replay_streams: + stream.wait_stream(torch.cuda.current_stream()) + for _ in range(8): + for stream, graph in zip(replay_streams, graphs): + with torch.cuda.stream(stream): + graph.replay() + for stream in replay_streams: + torch.cuda.current_stream().wait_stream(stream) + + for (out, grad), (expected_out, expected_grad), (probs, counts, total) in zip( + captured_outputs, expected, inputs + ): + _assert_bitwise_equal(out, expected_out) + _assert_bitwise_equal(grad, expected_grad) + ref_probs = probs.detach().clone().requires_grad_(True) + ref = aux_loss_pytorch(ref_probs, counts, int(total.item()), topk, num_experts, coeff) + (ref_grad,) = torch.autograd.grad(ref, ref_probs) + torch.testing.assert_close(out, ref, atol=1e-6, rtol=2e-5) + torch.testing.assert_close(grad, ref_grad, atol=1e-6, rtol=2e-5) + + +@pytest.mark.parametrize("device_total", [False, True], ids=["host_total", "device_total"]) +@pytest.mark.parametrize( + "dtype,num_tokens,num_experts,expert_multiplier,token_count_dtype", + [ + pytest.param(torch.float32, 4097, 128, 1, torch.int32, id="multi_cta"), + pytest.param(torch.float32, 257, 128, 3, torch.float32, id="sequence_global"), + pytest.param(torch.float32, 3, 7, 1, torch.int64, id="small_partial_count"), + pytest.param(torch.float32, 129, 1025, 1, torch.int32, id="wide"), + pytest.param(torch.float16, 4097, 128, 1, torch.int32, id="fp16"), + pytest.param(torch.bfloat16, 257, 128, 1, torch.bfloat16, id="bf16"), + ], +) +def test_fused_moe_aux_loss_deterministic( + dtype, num_tokens, num_experts, expert_multiplier, token_count_dtype, device_total +): + """Each launch must produce identical loss/gradient bits on the same GPU.""" + probs, tokens_per_expert = _moe_aux_loss_determinism_inputs( + dtype, num_tokens, num_experts, expert_multiplier, token_count_dtype + ) + total_num_tokens = ( + torch.tensor(num_tokens, dtype=torch.int64, device="cuda") if device_total else num_tokens + ) + topk, coeff = 4, 0.01 + grad_output = torch.tensor(0.37, dtype=dtype, device="cuda") + + # Accumulate the reference in FP32, then compare in the public output dtype. + ref_probs = probs.detach().float().requires_grad_(True) + ref = aux_loss_pytorch( + ref_probs, tokens_per_expert.float(), num_tokens, topk, num_experts, coeff + ) + (ref_grad,) = torch.autograd.grad(ref, ref_probs, grad_outputs=grad_output.float()) + finfo = torch.finfo(dtype) + rtol = max(2e-5, 2 * finfo.eps) + atol = finfo.tiny * finfo.eps + expected_out = expected_grad = None + for _ in range(32): + out = fused_moe_aux_loss( + probs, tokens_per_expert, total_num_tokens, num_experts, topk, coeff, deterministic=True + ) + (grad,) = torch.autograd.grad(out, probs, grad_outputs=grad_output) + if expected_out is None: + torch.testing.assert_close(out, ref.to(dtype), atol=atol, rtol=rtol) + torch.testing.assert_close(grad, ref_grad.to(dtype), atol=atol, rtol=rtol) + expected_out, expected_grad = out.detach().clone(), grad.clone() + else: + _assert_bitwise_equal(out, expected_out) + _assert_bitwise_equal(grad, expected_grad) + + +@pytest.mark.parametrize("deterministic", [False, True], ids=["atomic", "fixed"]) +@pytest.mark.parametrize( + "num_tokens,num_experts,expert_multiplier,token_count_dtype", + [ + pytest.param(4096, 128, 1, torch.int32, id="multi_cta"), + pytest.param(3, 7, 2, torch.float32, id="small_sequence_global"), + ], +) +def test_fused_moe_aux_loss_cuda_graph_capture( + num_tokens, num_experts, expert_multiplier, token_count_dtype, deterministic +): """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 + num_cols = num_experts * expert_multiplier 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) + probs, tokens_per_expert = _moe_aux_loss_determinism_inputs( + dtype, num_tokens, num_experts, expert_multiplier, token_count_dtype + ) + original_probs = probs.detach().clone() + original_counts = tokens_per_expert.clone() total_num_tokens_dev = torch.tensor(num_tokens, dtype=torch.int64, device="cuda") @@ -1215,6 +1439,7 @@ def test_fused_moe_aux_loss_cuda_graph_capture(): num_experts=num_experts, topk=topk, coeff=coeff, + deterministic=deterministic, ) torch.autograd.grad(warmup_out, probs) del warmup_out @@ -1229,16 +1454,25 @@ def test_fused_moe_aux_loss_cuda_graph_capture(): num_experts=num_experts, topk=topk, coeff=coeff, + deterministic=deterministic, ) (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): + # Replay A -> B -> A in the same graph. Restore both inputs and the device + # token count so stale partial sums or coefficients cannot pass unnoticed. + original_out = original_grad = None + for changed_inputs, new_total in ( + (False, num_tokens), + (True, num_tokens * 2 + 1), + (False, num_tokens), + ): + with torch.no_grad(): + probs.copy_(original_probs.flip((1,)) if changed_inputs else original_probs) + tokens_per_expert.copy_( + original_counts.roll(3) + 1 if changed_inputs else original_counts + ) 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, @@ -1249,8 +1483,22 @@ def test_fused_moe_aux_loss_cuda_graph_capture(): 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) + expected_out = expected_grad = None + for _ in range(16): + g.replay() + if expected_out is None or not deterministic: + torch.testing.assert_close(out, ref, atol=atol, rtol=rtol) + torch.testing.assert_close(grad_probs, ref_grad_probs, atol=atol, rtol=rtol) + expected_out, expected_grad = out.detach().clone(), grad_probs.clone() + else: + _assert_bitwise_equal(out, expected_out) + _assert_bitwise_equal(grad_probs, expected_grad) + if deterministic and not changed_inputs: + if original_out is None: + original_out, original_grad = expected_out, expected_grad + else: + _assert_bitwise_equal(expected_out, original_out) + _assert_bitwise_equal(expected_grad, original_grad) 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 9fb5adcc1f4..0d93bc88c76 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -8,6 +8,8 @@ #include #include +#include + #include "../common.h" #include "../util/logging.h" #include "../utils.cuh" @@ -17,11 +19,34 @@ namespace transformer_engine { namespace fused_router { -template +// Fixed launch geometry for each input shape bounds scratch to at most 1 KiB. +// Each CTA contributes one FP32 partial sum. +constexpr int kAuxLossMaxReductionBlocks = 256; + +int aux_loss_reduction_blocks(int num_rows) { + NVTE_CHECK(num_rows >= 0, "num_rows must be nonnegative, got ", num_rows); + return std::max(1, std::min(num_rows, kAuxLossMaxReductionBlocks)); +} + +void check_aux_loss_workspace(const Tensor& workspace, int num_rows) { + const size_t num_partials = aux_loss_reduction_blocks(num_rows); + NVTE_CHECK(workspace.data.dptr != nullptr, "Aux-loss workspace must be preallocated."); + NVTE_CHECK(workspace.data.dtype == DType::kFloat32, "Aux-loss workspace must have FP32 dtype."); + NVTE_CHECK(workspace.numel() >= num_partials, "Aux-loss workspace needs at least ", num_partials, + " FP32 elements, got ", workspace.numel()); +} + +void check_aux_loss_const_buffer(const Tensor& const_buf) { + NVTE_CHECK(const_buf.data.dtype == DType::kFloat32 && const_buf.numel() >= 2, + "Aux-loss constant buffer must contain at least two FP32 elements."); +} + +template __global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs, const IndexType* tokens_per_expert, int total_num_tokens, int num_rows, int num_cols, - int topk, float coeff, float* Coeff_buf) { + int topk, float coeff, float* Coeff_buf, + float* partial_sums) { // ----------------------------------------------------------------------- // 1) Write the CPU-computed coefficient into a device buffer to re-use in BWD // ----------------------------------------------------------------------- @@ -66,7 +91,13 @@ __global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs, CompType block_sum = warp_reduce_on_shmem( shmem_block, static_cast(blockDim.x), lane_id); if (lane_id == 0) { - atomicAdd(&Coeff_buf[1], static_cast(block_sum * coeff)); + if constexpr (kDeterministic) { + // Every CTA writes its own slot, including empty row partitions. + partial_sums[blockIdx.x] = static_cast(block_sum * coeff); + } else { + // Retained only for callers of the original C API. + atomicAdd(&Coeff_buf[1], static_cast(block_sum * coeff)); + } } } } @@ -77,16 +108,34 @@ __global__ void convert_accum_to_output(const float* Coeff_buf, DataType* aux_lo aux_loss[0] = static_cast(Coeff_buf[1]); } +// One full warp reads fixed, strided slices and combines them with a fixed tree. +// This replaces the legacy cast kernel, so deterministic forward still uses two +// launches. The final FP32 sum remains in Coeff_buf[1] for compatibility. +template +__global__ void reduce_aux_loss_partials(const float* partial_sums, int num_partials, + float* Coeff_buf, DataType* aux_loss) { + const int lane_id = threadIdx.x; + float sum = 0.0f; + for (int i = lane_id; i < num_partials; i += kThreadsPerWarp) { + sum += partial_sums[i]; + } + for (int offset = kThreadsPerWarp / 2; offset > 0; offset /= 2) { + sum += __shfl_down_sync(0xffffffff, sum, offset); + } + if (lane_id == 0) { + Coeff_buf[1] = sum; + aux_loss[0] = static_cast(sum); + } +} + /* ------------------------------------------------------------------------- * Kernel launcher -- simplified (no cluster launch). * ------------------------------------------------------------------------- */ -template -void fused_moe_aux_loss_forward_kernel_launcher(const DataType* probs, - const IndexType* tokens_per_expert, - int total_num_tokens, int num_experts, int num_rows, - int num_cols, int topk, float coeff, - DataType* aux_loss, float* Coeff_buf, - cudaStream_t stream) { +template +void fused_moe_aux_loss_forward_kernel_launcher( + const DataType* probs, const IndexType* tokens_per_expert, int total_num_tokens, + int num_experts, int num_rows, int num_cols, int topk, float coeff, DataType* aux_loss, + float* Coeff_buf, cudaStream_t stream, float* partial_sums = nullptr) { 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); // Sequence aux loss batches independent sequences along the expert dimension. @@ -97,39 +146,47 @@ void fused_moe_aux_loss_forward_kernel_launcher(const DataType* probs, 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 int grid_size = kDeterministic ? aux_loss_reduction_blocks(num_rows) : cuda::sm_count() * 2; // One CompType per thread in shared memory. const size_t smem_size = block_size * sizeof(CompType); check_shared_memory_capacity_num_experts(smem_size, num_cols); - // Compute final coefficient and zero the float accumulator (Coeff_buf[1]) before launch. + // Keep coefficient placement unchanged: scale each partial before summing. const float C_coeff = (num_experts * coeff) / topk / total_num_tokens / total_num_tokens; - NVTE_CHECK_CUDA(cudaMemsetAsync(Coeff_buf + 1, 0, sizeof(float), stream)); - fused_moe_aux_loss_forward_kernel + if constexpr (!kDeterministic) { + NVTE_CHECK_CUDA(cudaMemsetAsync(Coeff_buf + 1, 0, sizeof(float), stream)); + } + fused_moe_aux_loss_forward_kernel <<>>(probs, tokens_per_expert, total_num_tokens, - num_rows, num_cols, topk, C_coeff, Coeff_buf); + num_rows, num_cols, topk, C_coeff, Coeff_buf, + partial_sums); NVTE_CHECK_CUDA(cudaGetLastError()); - // Convert the float accumulator to the output DataType. - convert_accum_to_output<<<1, 1, 0, stream>>>(Coeff_buf, aux_loss); + if constexpr (kDeterministic) { + reduce_aux_loss_partials + <<<1, kThreadsPerWarp, 0, stream>>>(partial_sums, grid_size, Coeff_buf, aux_loss); + } else { + convert_accum_to_output<<<1, 1, 0, stream>>>(Coeff_buf, aux_loss); + } NVTE_CHECK_CUDA(cudaGetLastError()); } +template void fused_moe_aux_loss_forward(const Tensor& probs, const Tensor& tokens_per_expert, int 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) { + cudaStream_t stream, float* partial_sums = nullptr) { 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( + fused_moe_aux_loss_forward_kernel_launcher( reinterpret_cast(probs.data.dptr), reinterpret_cast(tokens_per_expert.data.dptr), total_num_tokens, num_experts, num_rows, num_cols, topk, coeff, reinterpret_cast(aux_loss.data.dptr), - reinterpret_cast(Coeff_buf.data.dptr), stream););); + reinterpret_cast(Coeff_buf.data.dptr), stream, partial_sums););); } /* ------------------------------------------------------------------------- @@ -138,10 +195,11 @@ void fused_moe_aux_loss_forward(const Tensor& probs, const Tensor& tokens_per_ex * lane computes C_coeff directly from the device value. The first CTA also * writes C_coeff into Coeff_buf[0] for backward. * ------------------------------------------------------------------------- */ -template +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) { + int num_experts, int num_rows, int num_cols, int topk, float coeff, float* Coeff_buf, + float* partial_sums) { // Reduction body matches the scalar-input kernel above. CompType thread_sum = CompType(0); for (int col = threadIdx.x; col < num_cols; col += blockDim.x) { @@ -170,16 +228,20 @@ __global__ void fused_moe_aux_loss_forward_kernel_graph_safe( if (blockIdx.x == 0) { Coeff_buf[0] = C_coeff; } - atomicAdd(&Coeff_buf[1], static_cast(block_sum * C_coeff)); + if constexpr (kDeterministic) { + partial_sums[blockIdx.x] = static_cast(block_sum * C_coeff); + } else { + atomicAdd(&Coeff_buf[1], static_cast(block_sum * C_coeff)); + } } } } -template +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) { + float* Coeff_buf, cudaStream_t stream, float* partial_sums = nullptr) { 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, @@ -188,27 +250,35 @@ void fused_moe_aux_loss_forward_kernel_launcher_graph_safe( 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 int grid_size = kDeterministic ? aux_loss_reduction_blocks(num_rows) : 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 + // The main kernel refreshes Coeff_buf[0] for backward on every graph replay. + if constexpr (!kDeterministic) { + 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); + Coeff_buf, partial_sums); NVTE_CHECK_CUDA(cudaGetLastError()); - convert_accum_to_output<<<1, 1, 0, stream>>>(Coeff_buf, aux_loss); + if constexpr (kDeterministic) { + reduce_aux_loss_partials + <<<1, kThreadsPerWarp, 0, stream>>>(partial_sums, grid_size, Coeff_buf, aux_loss); + } else { + convert_accum_to_output<<<1, 1, 0, stream>>>(Coeff_buf, aux_loss); + } NVTE_CHECK_CUDA(cudaGetLastError()); } +template 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) { + Tensor& aux_loss, Tensor& Coeff_buf, cudaStream_t stream, + float* partial_sums = nullptr) { 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)); @@ -218,12 +288,13 @@ void fused_moe_aux_loss_forward_graph_safe(const Tensor& probs, const Tensor& to 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( + 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););); + reinterpret_cast(Coeff_buf.data.dptr), stream, partial_sums););); } template @@ -303,6 +374,45 @@ void nvte_fused_moe_aux_loss_forward_graph_safe(const NVTETensor probs, *convertNVTETensorCheck(aux_loss), *convertNVTETensorCheck(Coeff_buf), stream); } +size_t nvte_get_moe_aux_loss_workspace_size() { + NVTE_API_CALL(nvte_get_moe_aux_loss_workspace_size); + return transformer_engine::fused_router::kAuxLossMaxReductionBlocks * sizeof(float); +} + +void nvte_fused_moe_aux_loss_forward_v2(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, NVTETensor workspace, + cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_moe_aux_loss_forward_v2); + using namespace transformer_engine; + auto& scratch = *convertNVTETensorCheck(workspace); + fused_router::check_aux_loss_workspace(scratch, num_rows); + auto& constants = *convertNVTETensorCheck(Const_buf); + fused_router::check_aux_loss_const_buffer(constants); + fused_router::fused_moe_aux_loss_forward( + *convertNVTETensorCheck(probs), *convertNVTETensorCheck(tokens_per_expert), total_num_tokens, + num_experts, num_rows, num_cols, topk, coeff, *convertNVTETensorCheck(aux_loss), constants, + stream, reinterpret_cast(scratch.data.dptr)); +} + +void nvte_fused_moe_aux_loss_forward_graph_safe_v2( + 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, NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_moe_aux_loss_forward_graph_safe_v2); + using namespace transformer_engine; + auto& scratch = *convertNVTETensorCheck(workspace); + fused_router::check_aux_loss_workspace(scratch, num_rows); + auto& constants = *convertNVTETensorCheck(Const_buf); + fused_router::check_aux_loss_const_buffer(constants); + 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), constants, stream, + reinterpret_cast(scratch.data.dptr)); +} + 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 323f04905e2..bda9516c49c 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -293,8 +293,8 @@ void nvte_fused_score_for_moe_aux_loss_backward(const NVTETensor intermediate_ou /*! \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. + * This legacy entrypoint uses an atomic reduction and is not deterministic. + * Prefer :c:func:`nvte_fused_moe_aux_loss_forward_v2` for new callers. * * \param[in] probs Probabilities from the forward pass. * \param[in] tokens_per_expert Number of tokens per expert. @@ -305,8 +305,7 @@ void nvte_fused_score_for_moe_aux_loss_backward(const NVTETensor intermediate_ou * \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[out] Const_buf Two FP32 elements: backward coefficient and forward sum. * \param[in] stream CUDA stream used for the operation. */ void nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor tokens_per_expert, @@ -316,8 +315,8 @@ 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. Prefer this path when the - * caller needs CUDA-graph-safe semantics with a dynamic token count. + * value stays dynamic across CUDA Graph replays. This legacy entrypoint uses + * an atomic reduction. Prefer :c:func:`nvte_fused_moe_aux_loss_forward_graph_safe_v2`. * * \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`. @@ -329,6 +328,49 @@ void nvte_fused_moe_aux_loss_forward_graph_safe(const NVTETensor probs, NVTETensor aux_loss, NVTETensor Const_buf, cudaStream_t stream); +/*! \brief Return the maximum workspace size in bytes for deterministic auxiliary loss. + * + * The capacity covers all input shapes. This host-only query does not allocate + * memory or launch kernels. Divide by sizeof(float) to allocate FP32 elements. + * + * \return Maximum required workspace size in bytes. + */ +size_t nvte_get_moe_aux_loss_workspace_size(); + +/*! \brief Deterministic forward pass for auxiliary loss, with a host token count. + * + * Each CTA writes an FP32 partial sum into workspace, then a second kernel + * reduces these values in a fixed order. Repeated execution with identical + * inputs on the same hardware and software is bitwise reproducible; equality + * to other reduction implementations is not guaranteed. + * + * Provide an FP32 workspace with capacity in bytes given by + * :c:func:`nvte_get_moe_aux_loss_workspace_size`. Its contents are overwritten + * and need no initialization. Workspace must remain valid through the queued + * forward kernels, but does not need to be saved for backward. + * + * \param[in,out] workspace Preallocated temporary FP32 buffer. + * Other parameters as in :c:func:`nvte_fused_moe_aux_loss_forward`. + */ +void nvte_fused_moe_aux_loss_forward_v2(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, NVTETensor workspace, + cudaStream_t stream); + +/*! \brief Deterministic forward with a device-tensor token count. + * + * total_num_tokens is a 0-dim int64 GPU tensor and may change between CUDA + * Graph replays. Its value is never read on the host or used to size workspace. + * Both the loss and saved backward coefficient are refreshed on each replay. + * Workspace requirements and lifetime follow :c:func:`nvte_fused_moe_aux_loss_forward_v2`. + * Other parameters as in :c:func:`nvte_fused_moe_aux_loss_forward_graph_safe`. + */ +void nvte_fused_moe_aux_loss_forward_graph_safe_v2( + 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, NVTETensor workspace, 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/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 5a4378ea328..b6ccef07604 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -59,15 +59,17 @@ std::tuple fused_score_for_moe_aux_loss_fwd( void fused_score_for_moe_aux_loss_bwd(at::Tensor intermediate_output, at::Tensor grad_scores, at::Tensor grad_logits, int topk, std::string score_function); -std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, - at::Tensor tokens_per_expert, - int total_num_tokens, int num_experts, - int num_rows, int num_cols, int topk, - float coeff); +size_t get_moe_aux_loss_workspace_size(); + +std::tuple fused_moe_aux_loss_fwd( + at::Tensor probs, at::Tensor tokens_per_expert, int total_num_tokens, int num_experts, + int num_rows, int num_cols, int topk, float coeff, bool deterministic = false, + std::optional workspace = std::nullopt); 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); + int num_rows, int num_cols, int topk, float coeff, bool deterministic = false, + std::optional workspace = std::nullopt); 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 c7f70182f36..0f49ce80935 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -164,13 +164,16 @@ void init_router_bindings(pybind11::module &m) { m.def("fused_score_for_moe_aux_loss_bwd", &fused_score_for_moe_aux_loss_bwd, py::arg("intermediate_output"), py::arg("grad_scores"), py::arg("grad_logits"), py::arg("topk"), py::arg("score_function"), "Fused aux loss with score function bwd"); + m.def("get_moe_aux_loss_workspace_size", &get_moe_aux_loss_workspace_size); 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"), + py::arg("deterministic") = false, py::arg("workspace") = py::none(), "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"), + py::arg("deterministic") = false, py::arg("workspace") = py::none(), "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"), diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 77ee80c0995..fec71b422de 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -417,11 +417,29 @@ void fused_score_for_moe_aux_loss_bwd(at::Tensor intermediate_output, at::Tensor at::cuda::getCurrentCUDAStream()); } +size_t get_moe_aux_loss_workspace_size() { + // Python allocates FP32 elements; the common API returns bytes. + return nvte_get_moe_aux_loss_workspace_size() / sizeof(float); +} + +static TensorWrapper make_aux_loss_workspace(const std::optional &workspace, + const at::Tensor &probs) { + TORCH_CHECK(workspace.has_value(), "deterministic aux loss requires a workspace"); + TORCH_CHECK(workspace->is_cuda() && workspace->device() == probs.device(), + "aux-loss workspace must be on the same CUDA device as probs"); + TORCH_CHECK(workspace->scalar_type() == at::kFloat && workspace->is_contiguous(), + "aux-loss workspace must be contiguous FP32 storage"); + TORCH_CHECK(workspace->numel() > 0, "aux-loss workspace must not be empty"); + return makeTransformerEngineTensor(*workspace); +} + std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, at::Tensor tokens_per_expert, int total_num_tokens, int num_experts, int num_rows, int num_cols, int topk, - float coeff) { + float coeff, bool deterministic, + std::optional workspace) { + at::cuda::CUDAGuard device_guard(probs.device()); TORCH_CHECK(topk > 0, "topk must be greater than 0"); TORCH_CHECK(total_num_tokens > 0, "total_num_tokens must be greater than 0"); TORCH_CHECK(num_experts > 0, "num_experts must be greater than 0"); @@ -435,16 +453,27 @@ std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, auto aux_loss_cu = makeTransformerEngineTensor(aux_loss); auto Const_buf_cu = makeTransformerEngineTensor(Const_buf); - nvte_fused_moe_aux_loss_forward(probs_cu.data(), tokens_per_expert_cu.data(), total_num_tokens, - num_experts, num_rows, num_cols, topk, coeff, aux_loss_cu.data(), - Const_buf_cu.data(), at::cuda::getCurrentCUDAStream()); + const auto stream = at::cuda::getCurrentCUDAStream(); + if (!deterministic) { + nvte_fused_moe_aux_loss_forward(probs_cu.data(), tokens_per_expert_cu.data(), total_num_tokens, + num_experts, num_rows, num_cols, topk, coeff, + aux_loss_cu.data(), Const_buf_cu.data(), stream); + return std::make_tuple(aux_loss, Const_buf); + } + + auto workspace_cu = make_aux_loss_workspace(workspace, probs); + nvte_fused_moe_aux_loss_forward_v2( + probs_cu.data(), tokens_per_expert_cu.data(), total_num_tokens, num_experts, num_rows, + num_cols, topk, coeff, aux_loss_cu.data(), Const_buf_cu.data(), workspace_cu.data(), stream); 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) { + int num_rows, int num_cols, int topk, float coeff, bool deterministic, + std::optional workspace) { + at::cuda::CUDAGuard device_guard(probs.device()); 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. @@ -465,10 +494,19 @@ std::tuple fused_moe_aux_loss_fwd_graph_safe( 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()); + const auto stream = at::cuda::getCurrentCUDAStream(); + if (!deterministic) { + 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(), stream); + return std::make_tuple(aux_loss, Const_buf); + } + + auto workspace_cu = make_aux_loss_workspace(workspace, probs); + nvte_fused_moe_aux_loss_forward_graph_safe_v2(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(), workspace_cu.data(), stream); return std::make_tuple(aux_loss, Const_buf); } diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 898310b0d61..59d35681221 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -14,6 +14,7 @@ global memory. For example, the gradient is required to have the same dtype as the input. """ +import functools import math from typing import Optional, Union @@ -466,6 +467,32 @@ def fused_compute_score_for_moe_aux_loss( return FusedComputeScoresForMoEAuxLoss.apply(logits, topk, score_function, routing_map_format) +@functools.lru_cache(maxsize=None) +def _get_moe_aux_loss_workspace_size() -> int: + """Maximum FP32 scratch capacity, queried once from the CUDA implementation.""" + return tex.get_moe_aux_loss_workspace_size() + + +@functools.lru_cache(maxsize=None) +def _get_cached_moe_aux_loss_workspace(device_index: int, stream_handle: int) -> torch.Tensor: + """Persistent scratch for eager calls, with independent storage for each stream.""" + # stream_handle is part of the cache key: concurrent streams must not share scratch. + return torch.empty(_get_moe_aux_loss_workspace_size(), dtype=torch.float32, device=device_index) + + +def _get_moe_aux_loss_workspace(device: torch.device) -> torch.Tensor: + """Reuse eager scratch; give captured operations private graph-pool storage.""" + with torch.cuda.device(device): + if torch.cuda.is_current_stream_capturing(): + # Allocate only during capture, never on replay. Do not cache graph-pool + # storage: a graph may replay concurrently with eager work or another graph. + return torch.empty( + _get_moe_aux_loss_workspace_size(), dtype=torch.float32, device=device + ) + stream = torch.cuda.current_stream(device) + return _get_cached_moe_aux_loss_workspace(stream.device.index, stream.cuda_stream) + + class FusedAuxLoss(torch.autograd.Function): """ Fused MoE aux loss. ``total_num_tokens`` may be either a Python int @@ -482,10 +509,12 @@ def forward( num_experts: int, topk: int, coeff: float, + deterministic: bool, ): # pylint: disable=missing-function-docstring num_rows = probs.size(0) num_cols = probs.size(1) + workspace = _get_moe_aux_loss_workspace(probs.device) if deterministic else None if isinstance(total_num_tokens, torch.Tensor): aux_loss, Const_buf = tex.fused_moe_aux_loss_fwd_graph_safe( probs=probs, @@ -496,6 +525,8 @@ def forward( num_cols=num_cols, topk=topk, coeff=coeff, + deterministic=deterministic, + workspace=workspace, ) else: aux_loss, Const_buf = tex.fused_moe_aux_loss_fwd( @@ -507,6 +538,8 @@ def forward( num_cols=num_cols, topk=topk, coeff=coeff, + deterministic=deterministic, + workspace=workspace, ) ctx.save_for_backward(Const_buf, tokens_per_expert) ctx.num_rows = num_rows @@ -524,7 +557,7 @@ def backward(ctx, grad_aux_loss): num_cols=ctx.num_cols, grad_aux_loss=grad_aux_loss, ) - return grad_probs, None, None, None, None, None + return grad_probs, None, None, None, None, None, None def fused_moe_aux_loss( @@ -534,6 +567,7 @@ def fused_moe_aux_loss( num_experts: int, topk: int, coeff: float, + deterministic: bool = False, ) -> torch.Tensor: """ Fused MoE aux loss. @@ -552,10 +586,18 @@ def fused_moe_aux_loss( topk : int coeff : float the coefficient of the aux loss. + deterministic : bool, default = False + Use a fixed-order two-stage FP32 reduction instead of atomic accumulation. + Repeated calls with identical inputs on the same hardware and software + are bitwise reproducible. Results may differ from other reduction orders. + Eager calls reuse scratch per device and stream; CUDA Graph capture uses + private scratch that is reused on replay. Returns ------- aux_loss : torch.Tensor. A scalar tensor in the same dtype as the "probs". """ - return FusedAuxLoss.apply(probs, tokens_per_expert, total_num_tokens, num_experts, topk, coeff) + return FusedAuxLoss.apply( + probs, tokens_per_expert, total_num_tokens, num_experts, topk, coeff, deterministic + ) From 07887f8efee91b1bb5f6b14833dd3eb13876b805 Mon Sep 17 00:00:00 2001 From: zhongboz Date: Thu, 8 Oct 2026 02:46:56 -0700 Subject: [PATCH 2/2] Align deterministic MoE aux-loss workspace queries with the common API Query workspace shape and dtype through the deterministic forward APIs, retain stream and CUDA Graph isolation, and cover metadata queries and invalid scratch buffers. Fix the stream cache-key lint warning and upstream import ordering. Validated on GB200: 201 TE tests passed (12 invalid top-k cases skipped), 6 C API tests passed, and 40 MCore checks passed with 1 expected atomic negative-control XPASS. Full PyTorch QA lint, pre-commit, and license checks passed. Signed-off-by: zhongboz --- tests/cpp/operator/CMakeLists.txt | 1 + .../cpp/operator/test_fused_moe_aux_loss.cpp | 72 +++++++++++++++++++ tests/pytorch/test_fused_router.py | 45 +++++++++++- .../common/fused_router/fused_moe_aux_loss.cu | 20 +++--- .../include/transformer_engine/fused_router.h | 22 +++--- .../pytorch/attention/fused_mla_q_uproj.py | 3 +- transformer_engine/pytorch/csrc/extensions.h | 2 +- .../pytorch/csrc/extensions/pybind.cpp | 2 +- .../pytorch/csrc/extensions/router.cpp | 11 ++- transformer_engine/pytorch/router.py | 21 +++--- 10 files changed, 160 insertions(+), 39 deletions(-) create mode 100644 tests/cpp/operator/test_fused_moe_aux_loss.cpp diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index c5acd043363..4b0cf0dab66 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -33,6 +33,7 @@ add_executable(test_operator test_scaled_activation.cu test_normalization.cu test_normalization_mxfp8.cu + test_fused_moe_aux_loss.cpp test_memset.cu test_splits_to_offsets.cu test_multi_cast_transpose.cu diff --git a/tests/cpp/operator/test_fused_moe_aux_loss.cpp b/tests/cpp/operator/test_fused_moe_aux_loss.cpp new file mode 100644 index 00000000000..64ece1db304 --- /dev/null +++ b/tests/cpp/operator/test_fused_moe_aux_loss.cpp @@ -0,0 +1,72 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include + +#include +#include +#include + +namespace { + +using transformer_engine::DType; +using transformer_engine::TensorWrapper; + +// Null input/output handles are allowed only for the metadata query. These +// calls also verify invalid workspaces fail before accessing other tensors. +void call_aux_loss(bool graph_safe, NVTETensor workspace, int num_rows) { + if (graph_safe) { + nvte_fused_moe_aux_loss_forward_graph_safe_v2(nullptr, nullptr, nullptr, 0, num_rows, 0, 0, + 0.0f, nullptr, nullptr, workspace, nullptr); + } else { + nvte_fused_moe_aux_loss_forward_v2(nullptr, nullptr, 0, 0, num_rows, 0, 0, 0.0f, nullptr, + nullptr, workspace, nullptr); + } +} + +class FusedMoeAuxLossWorkspaceTest : public ::testing::TestWithParam {}; + +void expect_workspace_error(bool graph_safe, NVTETensor workspace, int num_rows, + const char* message) { + try { + call_aux_loss(graph_safe, workspace, num_rows); + FAIL() << "Invalid workspace was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_NE(std::string(error.what()).find(message), std::string::npos) << error.what(); + } +} + +TEST_P(FusedMoeAuxLossWorkspaceTest, QueryWithoutAccessingInputs) { + size_t capacity = 0; + for (int num_rows : {0, 3, 257, 4097}) { + // A query replaces stale metadata as well as populating empty metadata. + TensorWrapper workspace(nullptr, std::vector{2, 3}, DType::kFloat16); + ASSERT_NO_THROW(call_aux_loss(GetParam(), workspace.data(), num_rows)); + ASSERT_EQ(workspace.shape().ndim, 1); + EXPECT_EQ(workspace.dtype(), DType::kFloat32); + EXPECT_EQ(nvte_tensor_data(workspace.data()), nullptr); + EXPECT_GT(workspace.shape().data[0], 0); + if (capacity == 0) capacity = workspace.shape().data[0]; + EXPECT_EQ(workspace.shape().data[0], capacity); + } +} + +TEST_P(FusedMoeAuxLossWorkspaceTest, RejectInsufficientCapacity) { + float storage = 0.0f; + TensorWrapper workspace(&storage, std::vector{1}, DType::kFloat32); + expect_workspace_error(GetParam(), workspace.data(), 257, "workspace needs at least"); +} + +TEST_P(FusedMoeAuxLossWorkspaceTest, RejectWrongDtype) { + float storage = 0.0f; + TensorWrapper workspace(&storage, std::vector{1}, DType::kFloat16); + expect_workspace_error(GetParam(), workspace.data(), 1, "workspace must have FP32 dtype"); +} + +INSTANTIATE_TEST_SUITE_P(HostAndDeviceTokenCount, FusedMoeAuxLossWorkspaceTest, ::testing::Bool()); + +} // namespace diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index bb35c7dcaf9..9c18bd4a5db 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -1261,6 +1261,7 @@ def fake_backward(**kwargs): def test_fused_moe_aux_loss_workspace_reuse(): """Eager calls share scratch only when ordered on the same device/stream.""" device = torch.device("cuda", torch.cuda.current_device()) + shape, dtype = te_router._get_moe_aux_loss_workspace_config() workspaces = [] streams = [torch.cuda.Stream(), torch.cuda.Stream()] for stream in streams: @@ -1268,11 +1269,53 @@ def test_fused_moe_aux_loss_workspace_reuse(): workspace = te_router._get_moe_aux_loss_workspace(device) again = te_router._get_moe_aux_loss_workspace(device) assert workspace.data_ptr() == again.data_ptr() - assert workspace.device == device and workspace.dtype == torch.float32 + assert workspace.device == device and workspace.dtype == dtype + assert tuple(workspace.shape) == shape workspaces.append(workspace) assert workspaces[0].data_ptr() != workspaces[1].data_ptr() +@pytest.mark.parametrize("device_total", [False, True], ids=["host_total", "device_total"]) +@pytest.mark.parametrize( + "workspace_kind,match", + [ + ("missing", "requires a workspace"), + ("empty", "must not be empty"), + ("undersized", "needs at least"), + ("wrong_dtype", "contiguous FP32"), + ("noncontiguous", "contiguous FP32"), + ("cpu", "same CUDA device"), + ], +) +def test_fused_moe_aux_loss_invalid_workspace(device_total, workspace_kind, match): + """Reject unsafe scratch before either forward variant launches kernels.""" + num_tokens, num_experts = 257, 8 + probs = torch.ones(num_tokens, num_experts, device="cuda") + counts = torch.ones(num_experts, dtype=torch.int32, device="cuda") + shape, dtype = te_router._get_moe_aux_loss_workspace_config() + workspace = torch.empty(shape, dtype=dtype, device="cuda") + if workspace_kind == "missing": + workspace = None + elif workspace_kind == "empty": + workspace = workspace[:0] + elif workspace_kind == "undersized": + workspace = workspace[:1] + elif workspace_kind == "wrong_dtype": + workspace = workspace.half() + elif workspace_kind == "noncontiguous": + workspace = workspace.repeat_interleave(2)[::2] + elif workspace_kind == "cpu": + workspace = torch.empty(shape, dtype=dtype) + total = ( + torch.tensor(num_tokens, dtype=torch.int64, device="cuda") if device_total else num_tokens + ) + forward = tex.fused_moe_aux_loss_fwd_graph_safe if device_total else tex.fused_moe_aux_loss_fwd + with pytest.raises(RuntimeError, match=match): + forward( + probs, counts, total, num_experts, num_tokens, num_experts, 1, 0.01, True, workspace + ) + + def test_fused_moe_aux_loss_cuda_graph_workspace_isolation(monkeypatch): """Graphs captured on one stream need private scratch for concurrent replay.""" num_tokens, num_experts, topk, coeff = 1025, 128, 4, 0.01 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 0d93bc88c76..814f9cdccdb 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -28,9 +28,14 @@ int aux_loss_reduction_blocks(int num_rows) { return std::max(1, std::min(num_rows, kAuxLossMaxReductionBlocks)); } -void check_aux_loss_workspace(const Tensor& workspace, int num_rows) { +void populate_aux_loss_workspace_config(Tensor& workspace, int num_rows) { const size_t num_partials = aux_loss_reduction_blocks(num_rows); - NVTE_CHECK(workspace.data.dptr != nullptr, "Aux-loss workspace must be preallocated."); + if (workspace.data.dptr == nullptr) { + // The bounded capacity can be reused across all input shapes on a stream. + workspace.data.shape = {kAuxLossMaxReductionBlocks}; + workspace.data.dtype = DType::kFloat32; + return; + } NVTE_CHECK(workspace.data.dtype == DType::kFloat32, "Aux-loss workspace must have FP32 dtype."); NVTE_CHECK(workspace.numel() >= num_partials, "Aux-loss workspace needs at least ", num_partials, " FP32 elements, got ", workspace.numel()); @@ -374,11 +379,6 @@ void nvte_fused_moe_aux_loss_forward_graph_safe(const NVTETensor probs, *convertNVTETensorCheck(aux_loss), *convertNVTETensorCheck(Coeff_buf), stream); } -size_t nvte_get_moe_aux_loss_workspace_size() { - NVTE_API_CALL(nvte_get_moe_aux_loss_workspace_size); - return transformer_engine::fused_router::kAuxLossMaxReductionBlocks * sizeof(float); -} - void nvte_fused_moe_aux_loss_forward_v2(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, @@ -387,7 +387,8 @@ void nvte_fused_moe_aux_loss_forward_v2(const NVTETensor probs, const NVTETensor NVTE_API_CALL(nvte_fused_moe_aux_loss_forward_v2); using namespace transformer_engine; auto& scratch = *convertNVTETensorCheck(workspace); - fused_router::check_aux_loss_workspace(scratch, num_rows); + fused_router::populate_aux_loss_workspace_config(scratch, num_rows); + if (scratch.data.dptr == nullptr) return; auto& constants = *convertNVTETensorCheck(Const_buf); fused_router::check_aux_loss_const_buffer(constants); fused_router::fused_moe_aux_loss_forward( @@ -403,7 +404,8 @@ void nvte_fused_moe_aux_loss_forward_graph_safe_v2( NVTE_API_CALL(nvte_fused_moe_aux_loss_forward_graph_safe_v2); using namespace transformer_engine; auto& scratch = *convertNVTETensorCheck(workspace); - fused_router::check_aux_loss_workspace(scratch, num_rows); + fused_router::populate_aux_loss_workspace_config(scratch, num_rows); + if (scratch.data.dptr == nullptr) return; auto& constants = *convertNVTETensorCheck(Const_buf); fused_router::check_aux_loss_const_buffer(constants); fused_router::fused_moe_aux_loss_forward_graph_safe( diff --git a/transformer_engine/common/include/transformer_engine/fused_router.h b/transformer_engine/common/include/transformer_engine/fused_router.h index bda9516c49c..921e6bd38a9 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -328,15 +328,6 @@ void nvte_fused_moe_aux_loss_forward_graph_safe(const NVTETensor probs, NVTETensor aux_loss, NVTETensor Const_buf, cudaStream_t stream); -/*! \brief Return the maximum workspace size in bytes for deterministic auxiliary loss. - * - * The capacity covers all input shapes. This host-only query does not allocate - * memory or launch kernels. Divide by sizeof(float) to allocate FP32 elements. - * - * \return Maximum required workspace size in bytes. - */ -size_t nvte_get_moe_aux_loss_workspace_size(); - /*! \brief Deterministic forward pass for auxiliary loss, with a host token count. * * Each CTA writes an FP32 partial sum into workspace, then a second kernel @@ -344,12 +335,15 @@ size_t nvte_get_moe_aux_loss_workspace_size(); * inputs on the same hardware and software is bitwise reproducible; equality * to other reduction implementations is not guaranteed. * - * Provide an FP32 workspace with capacity in bytes given by - * :c:func:`nvte_get_moe_aux_loss_workspace_size`. Its contents are overwritten - * and need no initialization. Workspace must remain valid through the queued - * forward kernels, but does not need to be saved for backward. + * To query workspace requirements, pass a valid workspace tensor whose data + * pointer is nullptr. The call populates its shape and dtype and returns without + * launching kernels or accessing other tensors, which may be nullptr during + * this query. The reported capacity covers all input shapes. Allocate storage + * with the reported metadata, then call again to execute the operation. + * Workspace contents are overwritten and need no initialization. Storage must + * remain valid through the queued forward kernels, but is not saved for backward. * - * \param[in,out] workspace Preallocated temporary FP32 buffer. + * \param[in,out] workspace Temporary buffer, or an empty tensor to query its metadata. * Other parameters as in :c:func:`nvte_fused_moe_aux_loss_forward`. */ void nvte_fused_moe_aux_loss_forward_v2(const NVTETensor probs, const NVTETensor tokens_per_expert, diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index b15e53e6ee3..43685f12a03 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -15,9 +15,10 @@ from importlib.metadata import PackageNotFoundError, version as get_pkg_version import torch -import transformer_engine_torch as tex from packaging.version import Version as PkgVersion +import transformer_engine_torch as tex + try: import triton import triton.language as tl diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index b6ccef07604..89c9c263ced 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -59,7 +59,7 @@ std::tuple fused_score_for_moe_aux_loss_fwd( void fused_score_for_moe_aux_loss_bwd(at::Tensor intermediate_output, at::Tensor grad_scores, at::Tensor grad_logits, int topk, std::string score_function); -size_t get_moe_aux_loss_workspace_size(); +std::pair, DType> get_moe_aux_loss_workspace_config(); std::tuple fused_moe_aux_loss_fwd( at::Tensor probs, at::Tensor tokens_per_expert, int total_num_tokens, int num_experts, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 0f49ce80935..2195b203cd2 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -164,7 +164,7 @@ void init_router_bindings(pybind11::module &m) { m.def("fused_score_for_moe_aux_loss_bwd", &fused_score_for_moe_aux_loss_bwd, py::arg("intermediate_output"), py::arg("grad_scores"), py::arg("grad_logits"), py::arg("topk"), py::arg("score_function"), "Fused aux loss with score function bwd"); - m.def("get_moe_aux_loss_workspace_size", &get_moe_aux_loss_workspace_size); + m.def("get_moe_aux_loss_workspace_config", &get_moe_aux_loss_workspace_config); 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"), diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index fec71b422de..708776a35db 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -417,9 +417,14 @@ void fused_score_for_moe_aux_loss_bwd(at::Tensor intermediate_output, at::Tensor at::cuda::getCurrentCUDAStream()); } -size_t get_moe_aux_loss_workspace_size() { - // Python allocates FP32 elements; the common API returns bytes. - return nvte_get_moe_aux_loss_workspace_size() / sizeof(float); +std::pair, DType> get_moe_aux_loss_workspace_config() { + TensorWrapper workspace; + // Query the operation with an empty workspace, then let Python own the + // allocation and its per-stream or CUDA Graph lifetime. + nvte_fused_moe_aux_loss_forward_v2(nullptr, nullptr, 0, 0, 0, 0, 0, 0.0f, nullptr, nullptr, + workspace.data(), nullptr); + const auto shape = workspace.shape(); + return {std::vector(shape.data, shape.data + shape.ndim), workspace.dtype()}; } static TensorWrapper make_aux_loss_workspace(const std::optional &workspace, diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 59d35681221..cfe49322628 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -21,6 +21,8 @@ import torch import transformer_engine_torch as tex +from .constants import TE_DType_To_Torch + # Re-export the C++ enum NVTERoutingMapFormat under a friendlier Python name. # Members: # RoutingMapFormat.BYTEMAP — bool[num_tokens, num_experts] @@ -468,16 +470,18 @@ def fused_compute_score_for_moe_aux_loss( @functools.lru_cache(maxsize=None) -def _get_moe_aux_loss_workspace_size() -> int: - """Maximum FP32 scratch capacity, queried once from the CUDA implementation.""" - return tex.get_moe_aux_loss_workspace_size() +def _get_moe_aux_loss_workspace_config() -> tuple[tuple[int, ...], torch.dtype]: + """Scratch metadata, queried once through the common forward operation.""" + shape, dtype = tex.get_moe_aux_loss_workspace_config() + return tuple(shape), TE_DType_To_Torch[dtype] @functools.lru_cache(maxsize=None) -def _get_cached_moe_aux_loss_workspace(device_index: int, stream_handle: int) -> torch.Tensor: +def _get_cached_moe_aux_loss_workspace(device_index: int, _stream_handle: int) -> torch.Tensor: """Persistent scratch for eager calls, with independent storage for each stream.""" - # stream_handle is part of the cache key: concurrent streams must not share scratch. - return torch.empty(_get_moe_aux_loss_workspace_size(), dtype=torch.float32, device=device_index) + # _stream_handle is part of the cache key: concurrent streams must not share scratch. + shape, dtype = _get_moe_aux_loss_workspace_config() + return torch.empty(shape, dtype=dtype, device=device_index) def _get_moe_aux_loss_workspace(device: torch.device) -> torch.Tensor: @@ -486,9 +490,8 @@ def _get_moe_aux_loss_workspace(device: torch.device) -> torch.Tensor: if torch.cuda.is_current_stream_capturing(): # Allocate only during capture, never on replay. Do not cache graph-pool # storage: a graph may replay concurrently with eager work or another graph. - return torch.empty( - _get_moe_aux_loss_workspace_size(), dtype=torch.float32, device=device - ) + shape, dtype = _get_moe_aux_loss_workspace_config() + return torch.empty(shape, dtype=dtype, device=device) stream = torch.cuda.current_stream(device) return _get_cached_moe_aux_loss_workspace(stream.device.index, stream.cuda_stream)