[PyTorch] Fix mutable QB bounds in CUDA graphs

tests/pytorch/test_fused_router.py

@@ -1,18 +1,21 @@
 # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
-import torch
+from copy import deepcopy
 from typing import Optional
+
+import pytest
+import torch
+
 from transformer_engine.pytorch.router import (
     QBHistogramMode,
     RoutingMapFormat,
     fused_topk_with_score_function,
     fused_compute_score_for_moe_aux_loss,
     fused_moe_aux_loss,
+    mark_qb_bin_bounds_validated,
 )
 import transformer_engine_torch as tex
-import pytest
-from copy import deepcopy
 
 seed = 42
 torch.manual_seed(seed)
@@ -621,7 +624,8 @@ def test_qb_topk_rejects_invalid_bin_bounds(histogram_mode, invalid_bounds):
         )
 
 
-def test_qb_topk_revalidates_updated_bin_bounds():
+@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"])
+def test_qb_topk_revalidates_updated_bin_bounds(histogram_mode):
     logits = torch.randn(8, 16, device="cuda", dtype=torch.float32)
     expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32)
     histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32)
@@ -637,7 +641,7 @@ def test_qb_topk_revalidates_updated_bin_bounds():
         expert_bias,
         qb_histogram=histogram,
         qb_bin_bounds=bin_bounds,
-        qb_histogram_mode="fused_atomic",
+        qb_histogram_mode=histogram_mode,
     )
     bin_bounds.fill_(0.0)
     with pytest.raises(ValueError, match="finite with lower < upper"):
@@ -652,7 +656,7 @@ def test_qb_topk_revalidates_updated_bin_bounds():
             expert_bias,
             qb_histogram=histogram,
             qb_bin_bounds=bin_bounds,
-            qb_histogram_mode="fused_atomic",
+            qb_histogram_mode=histogram_mode,
         )
 
 
@@ -699,11 +703,15 @@ def test_qb_raw_binding_rejects_invalid_bin_bounds_recoverably(histogram_mode, u
 
 
 @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"])
-def test_qb_topk_cuda_graph_uses_prevalidated_bounds(histogram_mode):
+@pytest.mark.parametrize("use_dense_indices", [False, True])
+def test_qb_topk_cuda_graph_uses_mutable_bounds(histogram_mode, use_dense_indices):
     logits = torch.randn(8, 16, device="cuda", dtype=torch.float32)
     expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32)
     histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32)
     bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32)
+    topk_indices = (
+        torch.empty(8, 4, device="cuda", dtype=torch.int32) if use_dense_indices else None
+    )
 
     def run_router():
         return fused_topk_with_score_function(
@@ -715,19 +723,152 @@ def run_router():
             None,
             "sigmoid",
             expert_bias,
+            topk_indices=topk_indices,
             qb_histogram=histogram,
             qb_bin_bounds=bin_bounds,
             qb_histogram_mode=histogram_mode,
         )
 
     run_router()
+    bounds_data_ptr = bin_bounds.data_ptr()
+    validated_version = bin_bounds._version
+    bin_bounds.copy_(torch.tensor([-0.25, 0.75], device="cuda"))
+    assert bin_bounds.data_ptr() == bounds_data_ptr
+    assert bin_bounds._version != validated_version
+    mark_qb_bin_bounds_validated(bin_bounds)
+
+    histogram.zero_()
     graph = torch.cuda.CUDAGraph()
     with torch.cuda.graph(graph):
-        probs, routing_map = run_router()
+        probs, routing_output = run_router()
+    histogram.zero_()
+    graph.replay()
+    torch.cuda.synchronize()
+    reference = qb_topk_score_function_pytorch(
+        logits, 4, expert_bias, bin_bounds, histogram.shape[1]
+    )
+    torch.testing.assert_close(probs, reference["probs"])
+    if use_dense_indices:
+        torch.testing.assert_close(
+            topk_indices_to_routing_map(routing_output, logits.shape[1]),
+            reference["routing_map"],
+        )
+    else:
+        torch.testing.assert_close(routing_output, reference["routing_map"])
+    torch.testing.assert_close(histogram, reference["histogram"])
+    first_histogram = histogram.clone()
+
+    bin_bounds.copy_(torch.tensor([-0.75, 0.25], device="cuda"))
+    assert bin_bounds.data_ptr() == bounds_data_ptr
+    histogram.zero_()
     graph.replay()
     torch.cuda.synchronize()
-    assert torch.isfinite(probs).all()
-    assert routing_map.sum().item() == logits.shape[0] * 4
+    reference = qb_topk_score_function_pytorch(
+        logits, 4, expert_bias, bin_bounds, histogram.shape[1]
+    )
+    torch.testing.assert_close(histogram, reference["histogram"])
+    assert not torch.equal(histogram, first_histogram)
+
+
+@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"])
+def test_qb_topk_cuda_graph_captures_bounds_update(histogram_mode):
+    logits = torch.randn(8, 16, device="cuda", dtype=torch.float32)
+    expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32)
+    histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32)
+    bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32)
+    next_bounds = torch.tensor([-0.75, 0.25], device="cuda", dtype=torch.float32)
+
+    def run_iteration():
+        probs, routing_map = fused_topk_with_score_function(
+            logits,
+            4,
+            False,
+            None,
+            None,
+            None,
+            "sigmoid",
+            expert_bias,
+            qb_histogram=histogram,
+            qb_bin_bounds=bin_bounds,
+            qb_histogram_mode=histogram_mode,
+        )
+        bin_bounds.copy_(next_bounds)
+        mark_qb_bin_bounds_validated(bin_bounds)
+        return probs, routing_map
+
+    # Match full-iteration capture: an eager warmup ends with a bounds update, and the same
+    # update is part of the captured iteration.
+    run_iteration()
+    histogram.zero_()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        probs, routing_map = run_iteration()
+    histogram.zero_()
+    graph.replay()
+    torch.cuda.synchronize()
+
+    reference = qb_topk_score_function_pytorch(
+        logits, 4, expert_bias, next_bounds, histogram.shape[1]
+    )
+    torch.testing.assert_close(probs, reference["probs"])
+    torch.testing.assert_close(routing_map, reference["routing_map"])
+    torch.testing.assert_close(histogram, reference["histogram"])
+    torch.testing.assert_close(bin_bounds, next_bounds)
+
+
+@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"])
+def test_qb_topk_cuda_graph_rejects_unvalidated_bounds(histogram_mode):
+    logits = torch.randn(8, 16, device="cuda", dtype=torch.float32)
+    expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32)
+    histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32)
+    bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32)
+
+    with pytest.raises(RuntimeError, match="current version must be validated"):
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            fused_topk_with_score_function(
+                logits,
+                4,
+                False,
+                None,
+                None,
+                None,
+                "sigmoid",
+                expert_bias,
+                qb_histogram=histogram,
+                qb_bin_bounds=bin_bounds,
+                qb_histogram_mode=histogram_mode,
+            )
+
+
+@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"])
+def test_qb_topk_cuda_graph_rejects_stale_bounds_validation(histogram_mode):
+    logits = torch.randn(8, 16, device="cuda", dtype=torch.float32)
+    expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32)
+    histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32)
+    bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32)
+
+    def run_router():
+        return fused_topk_with_score_function(
+            logits,
+            4,
+            False,
+            None,
+            None,
+            None,
+            "sigmoid",
+            expert_bias,
+            qb_histogram=histogram,
+            qb_bin_bounds=bin_bounds,
+            qb_histogram_mode=histogram_mode,
+        )
+
+    run_router()
+    bin_bounds.zero_()
+    with pytest.raises(RuntimeError, match="current version must be validated"):
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            run_router()
 
 
 @pytest.mark.parametrize(

transformer_engine/pytorch/router.py

@@ -42,8 +42,26 @@
 _QB_BOUNDS_VALIDATED_VERSION_ATTR = "_nvte_qb_bounds_validated_version"
 
 
+def mark_qb_bin_bounds_validated(bin_bounds: torch.Tensor) -> None:
+    """Mark the current QB bounds version as valid after a trusted device-side update.
+
+    This function does not inspect tensor values. The caller must guarantee finite FP32 bounds
+    with ``lower < upper``. During CUDA graph capture, call it immediately after dispatching the
+    trusted in-place update on the same stream.
+    """
+    if not (
+        isinstance(bin_bounds, torch.Tensor)
+        and bin_bounds.is_cuda
+        and bin_bounds.is_contiguous()
+        and bin_bounds.dtype == torch.float32
+        and bin_bounds.shape == (2,)
+    ):
+        raise ValueError("QB bin_bounds must be a contiguous FP32 CUDA tensor with shape [2]")
+    setattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, bin_bounds._version)
+
+
 def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool:
-    """Validate CUDA-resident QB bounds once per PyTorch tensor version."""
+    """Validate QB bounds once per PyTorch tensor version."""
     if not (
         isinstance(bin_bounds, torch.Tensor)
         and bin_bounds.is_cuda
@@ -55,19 +73,20 @@ def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool:
         return False
 
     version = bin_bounds._version
-    if getattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, None) == version:
+    validated_version = getattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, None)
+    if validated_version == version:
         return True
     with torch.cuda.device(bin_bounds.device):
         if torch.cuda.is_current_stream_capturing():
             raise RuntimeError(
-                "QB bin_bounds must be validated by an eager router call before CUDA graph capture"
+                "QB bin_bounds current version must be validated before CUDA graph capture"
             )
     lower, upper = bin_bounds.detach().cpu().tolist()
     if not (math.isfinite(lower) and math.isfinite(upper) and lower < upper):
         raise ValueError(
             f"QB bin_bounds values must be finite with lower < upper, got [{lower}, {upper}]"
         )
-    setattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, version)
+    mark_qb_bin_bounds_validated(bin_bounds)
     return True
 
 
@@ -287,8 +306,12 @@ def fused_topk_with_score_function(
         Caller-owned int32 ``[num_experts, num_bins]`` histogram accumulated in place.
     qb_bin_bounds : torch.Tensor, optional
         FP32 CUDA tensor ``[lower, upper]`` defining uniform QB histogram bins. Values must be
-        finite with ``lower < upper``. Bounds are revalidated after PyTorch-tracked in-place
-        updates; validate once with an eager call before CUDA graph capture.
+        finite with ``lower < upper``. Eager calls revalidate PyTorch-tracked in-place updates.
+        Before CUDA graph capture, validate the current version eagerly or use
+        :func:`mark_qb_bin_bounds_validated` after a trusted device-side update. The marker may be
+        called immediately after an in-place update dispatched during capture. Graph replays read
+        later trusted updates through the captured device pointer, so callers must preserve
+        validity.
     qb_histogram_mode : str, optional
         ``"two_kernel"`` or ``"fused_atomic"``. Must be provided with the two QB tensors.