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 <gtest/gtest.h>
+#include <transformer_engine/fused_router.h>
+
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+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<bool> {};
+
+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<size_t>{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<size_t>{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<size_t>{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 c6c82fa3d78..9c18bd4a5db 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,289 @@ 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())
+    shape, dtype = te_router._get_moe_aux_loss_workspace_config()
+    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 == 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
+    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 +1482,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 +1497,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 +1526,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..814f9cdccdb 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 <cuda_runtime.h>
 #include <transformer_engine/fused_router.h>
 
+#include <algorithm>
+
 #include "../common.h"
 #include "../util/logging.h"
 #include "../utils.cuh"
@@ -17,11 +19,39 @@
 namespace transformer_engine {
 namespace fused_router {
 
-template <typename DataType, typename IndexType>
+// 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 populate_aux_loss_workspace_config(Tensor& workspace, int num_rows) {
+  const size_t num_partials = aux_loss_reduction_blocks(num_rows);
+  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());
+}
+
+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 <typename DataType, typename IndexType, bool kDeterministic>
 __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 +96,13 @@ __global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs,
     CompType block_sum = warp_reduce_on_shmem<CompType, ReduceFuncType::SUM>(
         shmem_block, static_cast<int>(blockDim.x), lane_id);
     if (lane_id == 0) {
-      atomicAdd(&Coeff_buf[1], static_cast<float>(block_sum * coeff));
+      if constexpr (kDeterministic) {
+        // Every CTA writes its own slot, including empty row partitions.
+        partial_sums[blockIdx.x] = static_cast<float>(block_sum * coeff);
+      } else {
+        // Retained only for callers of the original C API.
+        atomicAdd(&Coeff_buf[1], static_cast<float>(block_sum * coeff));
+      }
     }
   }
 }
@@ -77,16 +113,34 @@ __global__ void convert_accum_to_output(const float* Coeff_buf, DataType* aux_lo
   aux_loss[0] = static_cast<DataType>(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 <typename DataType>
+__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<DataType>(sum);
+  }
+}
+
 /* -------------------------------------------------------------------------
  *  Kernel launcher -- simplified (no cluster launch).
  * ------------------------------------------------------------------------- */
-template <typename DataType, typename IndexType>
-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 <typename DataType, typename IndexType, bool kDeterministic = false>
+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 +151,47 @@ void fused_moe_aux_loss_forward_kernel_launcher(const DataType* probs,
   const int block_size = ((std::min(1024, num_cols) + static_cast<int>(kThreadsPerWarp) - 1) /
                           static_cast<int>(kThreadsPerWarp)) *
                          static_cast<int>(kThreadsPerWarp);
-  const int grid_size = cuda::sm_count() * 2;
+  const 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<DataType, IndexType>
+  if constexpr (!kDeterministic) {
+    NVTE_CHECK_CUDA(cudaMemsetAsync(Coeff_buf + 1, 0, sizeof(float), stream));
+  }
+  fused_moe_aux_loss_forward_kernel<DataType, IndexType, kDeterministic>
       <<<grid_size, block_size, smem_size, stream>>>(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<DataType><<<1, 1, 0, stream>>>(Coeff_buf, aux_loss);
+  if constexpr (kDeterministic) {
+    reduce_aux_loss_partials<DataType>
+        <<<1, kThreadsPerWarp, 0, stream>>>(partial_sums, grid_size, Coeff_buf, aux_loss);
+  } else {
+    convert_accum_to_output<DataType><<<1, 1, 0, stream>>>(Coeff_buf, aux_loss);
+  }
   NVTE_CHECK_CUDA(cudaGetLastError());
 }
 
+template <bool kDeterministic = false>
 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<DataType, IndexType>(
+          fused_moe_aux_loss_forward_kernel_launcher<DataType, IndexType, kDeterministic>(
               reinterpret_cast<DataType*>(probs.data.dptr),
               reinterpret_cast<IndexType*>(tokens_per_expert.data.dptr), total_num_tokens,
               num_experts, num_rows, num_cols, topk, coeff,
               reinterpret_cast<DataType*>(aux_loss.data.dptr),
-              reinterpret_cast<float*>(Coeff_buf.data.dptr), stream);););
+              reinterpret_cast<float*>(Coeff_buf.data.dptr), stream, partial_sums);););
 }
 
 /* -------------------------------------------------------------------------
@@ -138,10 +200,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 <typename DataType, typename IndexType>
+template <typename DataType, typename IndexType, bool kDeterministic>
 __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 +233,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<float>(block_sum * C_coeff));
+      if constexpr (kDeterministic) {
+        partial_sums[blockIdx.x] = static_cast<float>(block_sum * C_coeff);
+      } else {
+        atomicAdd(&Coeff_buf[1], static_cast<float>(block_sum * C_coeff));
+      }
     }
   }
 }
 
-template <typename DataType, typename IndexType>
+template <typename DataType, typename IndexType, bool kDeterministic = false>
 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 +255,35 @@ void fused_moe_aux_loss_forward_kernel_launcher_graph_safe(
   const int block_size = ((std::min(1024, num_cols) + static_cast<int>(kThreadsPerWarp) - 1) /
                           static_cast<int>(kThreadsPerWarp)) *
                          static_cast<int>(kThreadsPerWarp);
-  const int grid_size = cuda::sm_count() * 2;
+  const 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<DataType, IndexType>
+  // 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<DataType, IndexType, kDeterministic>
       <<<grid_size, block_size, smem_size, stream>>>(probs, tokens_per_expert, total_num_tokens_dev,
                                                      num_experts, num_rows, num_cols, topk, coeff,
-                                                     Coeff_buf);
+                                                     Coeff_buf, partial_sums);
   NVTE_CHECK_CUDA(cudaGetLastError());
 
-  convert_accum_to_output<DataType><<<1, 1, 0, stream>>>(Coeff_buf, aux_loss);
+  if constexpr (kDeterministic) {
+    reduce_aux_loss_partials<DataType>
+        <<<1, kThreadsPerWarp, 0, stream>>>(partial_sums, grid_size, Coeff_buf, aux_loss);
+  } else {
+    convert_accum_to_output<DataType><<<1, 1, 0, stream>>>(Coeff_buf, aux_loss);
+  }
   NVTE_CHECK_CUDA(cudaGetLastError());
 }
 
+template <bool kDeterministic = false>
 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<int>(total_num_tokens.data.dtype));
@@ -218,12 +293,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<DataType, IndexType>(
+          fused_moe_aux_loss_forward_kernel_launcher_graph_safe<DataType, IndexType,
+                                                                kDeterministic>(
               reinterpret_cast<DataType*>(probs.data.dptr),
               reinterpret_cast<IndexType*>(tokens_per_expert.data.dptr),
               reinterpret_cast<const int64_t*>(total_num_tokens.data.dptr), num_experts, num_rows,
               num_cols, topk, coeff, reinterpret_cast<DataType*>(aux_loss.data.dptr),
-              reinterpret_cast<float*>(Coeff_buf.data.dptr), stream);););
+              reinterpret_cast<float*>(Coeff_buf.data.dptr), stream, partial_sums);););
 }
 
 template <typename DataType, typename IndexType>
@@ -303,6 +379,42 @@ void nvte_fused_moe_aux_loss_forward_graph_safe(const NVTETensor probs,
       *convertNVTETensorCheck(aux_loss), *convertNVTETensorCheck(Coeff_buf), stream);
 }
 
+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::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<true>(
+      *convertNVTETensorCheck(probs), *convertNVTETensorCheck(tokens_per_expert), total_num_tokens,
+      num_experts, num_rows, num_cols, topk, coeff, *convertNVTETensorCheck(aux_loss), constants,
+      stream, reinterpret_cast<float*>(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::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<true>(
+      *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<float*>(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..921e6bd38a9 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,43 @@ void nvte_fused_moe_aux_loss_forward_graph_safe(const NVTETensor probs,
                                                 NVTETensor aux_loss, NVTETensor Const_buf,
                                                 cudaStream_t stream);
 
+/*! \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.
+ *
+ *  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  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,
+                                        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/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 5a4378ea328..89c9c263ced 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -59,15 +59,17 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> 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<at::Tensor, at::Tensor> 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);
+std::pair<std::vector<size_t>, DType> get_moe_aux_loss_workspace_config();
+
+std::tuple<at::Tensor, at::Tensor> 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<at::Tensor> workspace = std::nullopt);
 
 std::tuple<at::Tensor, at::Tensor> 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<at::Tensor> 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..2195b203cd2 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_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"),
+        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..708776a35db 100644
--- a/transformer_engine/pytorch/csrc/extensions/router.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/router.cpp
@@ -417,11 +417,34 @@ void fused_score_for_moe_aux_loss_bwd(at::Tensor intermediate_output, at::Tensor
       at::cuda::getCurrentCUDAStream());
 }
 
+std::pair<std::vector<size_t>, 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<size_t>(shape.data, shape.data + shape.ndim), workspace.dtype()};
+}
+
+static TensorWrapper make_aux_loss_workspace(const std::optional<at::Tensor> &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<at::Tensor, at::Tensor> 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<at::Tensor> 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 +458,27 @@ std::tuple<at::Tensor, at::Tensor> 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<at::Tensor, at::Tensor> 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<at::Tensor> 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 +499,19 @@ std::tuple<at::Tensor, at::Tensor> 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..cfe49322628 100644
--- a/transformer_engine/pytorch/router.py
+++ b/transformer_engine/pytorch/router.py
@@ -14,12 +14,15 @@
   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
 
 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]
@@ -466,6 +469,33 @@ 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_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:
+    """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.
+    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:
+    """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.
+            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)
+
+
 class FusedAuxLoss(torch.autograd.Function):
     """
     Fused MoE aux loss. ``total_num_tokens`` may be either a Python int
@@ -482,10 +512,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 +528,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 +541,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 +560,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 +570,7 @@ def fused_moe_aux_loss(
     num_experts: int,
     topk: int,
     coeff: float,
+    deterministic: bool = False,
 ) -> torch.Tensor:
     """
     Fused MoE aux loss.
@@ -552,10 +589,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
+    )