From 77653a56a73d895d5bd7f5b9c428f9e8c87d57c4 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Wed, 26 Aug 2026 22:08:23 +0800 Subject: [PATCH 1/5] [PyTorch] Fix mutable QB bounds in CUDA graphs Signed-off-by: Harry Zhou --- tests/pytorch/test_fused_router.py | 75 +++++++++++++++++++++++++--- transformer_engine/pytorch/router.py | 9 +++- 2 files changed, 75 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 18efb6a190..32c857ad48 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -621,7 +621,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 +638,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 +653,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 +700,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_prevalidated_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 +720,75 @@ 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 + + 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_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="validated by an eager router call"): + 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( diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index ac38f7d1cf..aeb67b1ed5 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -43,7 +43,7 @@ def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: - """Validate CUDA-resident QB bounds once per PyTorch tensor version.""" + """Validate QB bounds eagerly while allowing prevalidated mutable CUDA-graph inputs.""" if not ( isinstance(bin_bounds, torch.Tensor) and bin_bounds.is_cuda @@ -55,10 +55,15 @@ 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(): + if validated_version is not None: + # CUDA graphs support persistent input tensors whose values are updated in place + # between capture and replay. The kernel reads the current bounds by device pointer. + return True raise RuntimeError( "QB bin_bounds must be validated by an eager router call before CUDA graph capture" ) From 2c69379c82570864d11d68155fd262ef4635ec75 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Wed, 26 Aug 2026 22:22:16 +0800 Subject: [PATCH 2/5] [PyTorch] Document mutable QB graph contract Signed-off-by: Harry Zhou --- transformer_engine/pytorch/router.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index aeb67b1ed5..3b040ce4a4 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -62,7 +62,8 @@ def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: if torch.cuda.is_current_stream_capturing(): if validated_version is not None: # CUDA graphs support persistent input tensors whose values are updated in place - # between capture and replay. The kernel reads the current bounds by device pointer. + # between capture and replay. Python validation is not re-entered by replay, so + # keeping those mutable values valid is necessarily the caller's contract. return True raise RuntimeError( "QB bin_bounds must be validated by an eager router call before CUDA graph capture" @@ -292,8 +293,9 @@ 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. + CUDA graphs read updates through the captured device pointer, so callers must preserve + valid values across capture and replay; validate once eagerly before capture. qb_histogram_mode : str, optional ``"two_kernel"`` or ``"fused_atomic"``. Must be provided with the two QB tensors. From 8841e0481a7c6a77a376db8d2ec91f105ca8e0d0 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Wed, 26 Aug 2026 22:38:37 +0800 Subject: [PATCH 3/5] [PyTorch] Require explicit QB bounds validation Signed-off-by: Harry Zhou --- tests/pytorch/test_fused_router.py | 34 +++++++++++++++++++++++++- transformer_engine/pytorch/router.py | 36 ++++++++++++++++++++-------- 2 files changed, 59 insertions(+), 11 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 32c857ad48..df097026c1 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -9,6 +9,7 @@ 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 @@ -732,6 +733,7 @@ def run_router(): 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() @@ -773,7 +775,7 @@ def test_qb_topk_cuda_graph_rejects_unvalidated_bounds(histogram_mode): 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="validated by an eager router call"): + with pytest.raises(RuntimeError, match="current version must be validated"): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): fused_topk_with_score_function( @@ -791,6 +793,36 @@ def test_qb_topk_cuda_graph_rejects_unvalidated_bounds(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( "histogram_mode", [QBHistogramMode.TWO_KERNEL, QBHistogramMode.FUSED_ATOMIC], diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 3b040ce4a4..0b37134c77 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -42,8 +42,28 @@ _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`` and call this function outside CUDA graph capture. + """ + 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]") + with torch.cuda.device(bin_bounds.device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("QB bin_bounds must be marked valid outside CUDA graph capture") + setattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, bin_bounds._version) + + def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: - """Validate QB bounds eagerly while allowing prevalidated mutable CUDA-graph inputs.""" + """Validate QB bounds once per PyTorch tensor version.""" if not ( isinstance(bin_bounds, torch.Tensor) and bin_bounds.is_cuda @@ -60,20 +80,15 @@ def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: return True with torch.cuda.device(bin_bounds.device): if torch.cuda.is_current_stream_capturing(): - if validated_version is not None: - # CUDA graphs support persistent input tensors whose values are updated in place - # between capture and replay. Python validation is not re-entered by replay, so - # keeping those mutable values valid is necessarily the caller's contract. - return True 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 @@ -294,8 +309,9 @@ def fused_topk_with_score_function( qb_bin_bounds : torch.Tensor, optional FP32 CUDA tensor ``[lower, upper]`` defining uniform QB histogram bins. Values must be finite with ``lower < upper``. Eager calls revalidate PyTorch-tracked in-place updates. - CUDA graphs read updates through the captured device pointer, so callers must preserve - valid values across capture and replay; validate once eagerly before capture. + Before CUDA graph capture, validate the current version eagerly or use + :func:`mark_qb_bin_bounds_validated` after a trusted device-side update. Graph replays read + subsequent 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. From 5b2d1e12b2a70ef83147c845b687a4aad2bc2591 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Thu, 27 Aug 2026 00:19:02 +0800 Subject: [PATCH 4/5] [Common][PyTorch] Validate QB bounds in CUDA graphs Signed-off-by: Harry Zhou --- tests/pytorch/test_fused_router.py | 115 +++++++++++------- .../fused_topk_with_score_function.cu | 10 ++ .../include/transformer_engine/fused_router.h | 20 +-- .../pytorch/csrc/extensions/router.cpp | 6 +- transformer_engine/pytorch/router.py | 47 ++----- 5 files changed, 109 insertions(+), 89 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index df097026c1..d54429d13a 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -1,19 +1,23 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -import torch +from copy import deepcopy +import subprocess +import sys +from textwrap import dedent 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) @@ -702,7 +706,7 @@ def test_qb_raw_binding_rejects_invalid_bin_bounds_recoverably(histogram_mode, u @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) @pytest.mark.parametrize("use_dense_indices", [False, True]) -def test_qb_topk_cuda_graph_uses_mutable_prevalidated_bounds(histogram_mode, use_dense_indices): +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) @@ -727,13 +731,9 @@ def run_router(): 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() @@ -769,39 +769,15 @@ def run_router(): @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -def test_qb_topk_cuda_graph_rejects_unvalidated_bounds(histogram_mode): +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) - 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( + def run_iteration(): + probs, routing_map = fused_topk_with_score_function( logits, 4, False, @@ -814,13 +790,70 @@ def run_router(): qb_bin_bounds=bin_bounds, qb_histogram_mode=histogram_mode, ) + bin_bounds.copy_(next_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() - run_router() - bin_bounds.zero_() - with pytest.raises(RuntimeError, match="current version must be validated"): + 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_invalid_bounds_on_device(histogram_mode): + script = dedent( + f""" + import torch + from transformer_engine.pytorch.router import fused_topk_with_score_function + + 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) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - run_router() + 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!r}, + ) + graph.replay() + torch.cuda.synchronize() + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + timeout=60, + ) + output = result.stdout + result.stderr + assert result.returncode != 0 + assert "QB bin_bounds values must be finite with lower < upper." in output + assert "CUDA error" in output @pytest.mark.parametrize( diff --git a/transformer_engine/common/fused_router/fused_topk_with_score_function.cu b/transformer_engine/common/fused_router/fused_topk_with_score_function.cu index 4ea11f05dc..1b609ba77c 100644 --- a/transformer_engine/common/fused_router/fused_topk_with_score_function.cu +++ b/transformer_engine/common/fused_router/fused_topk_with_score_function.cu @@ -36,6 +36,16 @@ struct QBBinParams { __device__ inline QBBinParams load_qb_bin_params(const CompType *bin_bounds, int num_bins) { const CompType lower = bin_bounds[0]; const CompType upper = bin_bounds[1]; + const bool valid = isfinite(lower) && isfinite(upper) && lower < upper; + if (!valid) { + if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) { + printf("QB bin_bounds values must be finite with lower < upper.\n"); + __trap(); + } + // Keep non-reporting threads away from invalid floating-point-to-integer conversions while + // the device-side assertion is propagated to the host. + return {0.0f, 0.0f}; + } return {lower, static_cast(num_bins) / (upper - lower)}; } diff --git a/transformer_engine/common/include/transformer_engine/fused_router.h b/transformer_engine/common/include/transformer_engine/fused_router.h index 323f04905e..431a976565 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -120,11 +120,11 @@ void nvte_fused_topk_with_score_function_forward_qb_v2( NVTETensor histogram, NVTETensor bin_bounds, NVTEQBHistogramMode histogram_mode, cudaStream_t stream); -/*! \brief Same as nvte_fused_topk_with_score_function_forward_qb_v2, but skips bin_bounds value - * validation. +/*! \brief Same as nvte_fused_topk_with_score_function_forward_qb_v2, but skips host-side + * bin_bounds value validation. * - * The caller must have validated that bin_bounds contains finite FP32 values [lower, upper] with - * lower < upper. Use this variant to avoid host synchronization in a hot path or CUDA graph. + * The histogram kernel raises a device-side error if bin_bounds is not finite and ordered. Use + * this variant to avoid host synchronization in a hot path or CUDA graph. */ void nvte_fused_topk_with_score_function_forward_qb_v2_unchecked( const NVTETensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, @@ -145,10 +145,10 @@ void nvte_fused_topk_with_score_function_forward_qb_with_indices( NVTEQBHistogramMode histogram_mode, cudaStream_t stream); /*! \brief Same as nvte_fused_topk_with_score_function_forward_qb_with_indices, but skips - * bin_bounds value validation. + * host-side bin_bounds value validation. * - * The caller must have validated that bin_bounds contains finite FP32 values [lower, upper] with - * lower < upper. Use this variant to avoid host synchronization in a hot path or CUDA graph. + * The histogram kernel raises a device-side error if bin_bounds is not finite and ordered. Use + * this variant to avoid host synchronization in a hot path or CUDA graph. */ void nvte_fused_topk_with_score_function_forward_qb_with_indices_unchecked( const NVTETensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, @@ -165,10 +165,10 @@ void nvte_qb_histogram_accumulate(const NVTETensor raw_scores, const NVTETensor const NVTETensor bin_bounds, NVTETensor histogram, cudaStream_t stream); -/*! \brief Same as nvte_qb_histogram_accumulate, but skips bin_bounds value validation. +/*! \brief Same as nvte_qb_histogram_accumulate, but skips host-side bin_bounds value validation. * - * The caller must have validated that bin_bounds contains finite FP32 values [lower, upper] with - * lower < upper. Use this variant to avoid host synchronization in a hot path or CUDA graph. + * The histogram kernel raises a device-side error if bin_bounds is not finite and ordered. Use + * this variant to avoid host synchronization in a hot path or CUDA graph. */ void nvte_qb_histogram_accumulate_unchecked(const NVTETensor raw_scores, const NVTETensor cutoff, const NVTETensor bin_bounds, NVTETensor histogram, diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 77ee80c099..713406a276 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -160,7 +160,7 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, int routing_map_format, std::optional topk_indices, at::Tensor histogram, at::Tensor bin_bounds, int histogram_mode, - bool bin_bounds_validated) { + bool skip_bin_bounds_host_validation) { check_routing_map_format(routing_map_format); TORCH_CHECK(logits.dim() >= 1, "logits must have at least 1 dim"); TORCH_CHECK(logits.is_cuda() && logits.is_contiguous(), @@ -239,7 +239,7 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, auto stream = at::cuda::getCurrentCUDAStream(); if (topk_indices.has_value()) { - if (bin_bounds_validated) { + if (skip_bin_bounds_host_validation) { nvte_fused_topk_with_score_function_forward_qb_with_indices_unchecked( logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, scaling_factor_value, expert_bias_cu.data(), probs_cu.data(), routing_output_cu.data(), @@ -253,7 +253,7 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, bin_bounds_cu.data(), mode, stream); } } else { - if (bin_bounds_validated) { + if (skip_bin_bounds_host_validation) { nvte_fused_topk_with_score_function_forward_qb_v2_unchecked( logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, scaling_factor_value, expert_bias_cu.data(), probs_cu.data(), routing_output_cu.data(), diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 0b37134c77..69ca23276d 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -42,28 +42,8 @@ _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`` and call this function outside CUDA graph capture. - """ - 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]") - with torch.cuda.device(bin_bounds.device): - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("QB bin_bounds must be marked valid outside CUDA graph capture") - setattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, bin_bounds._version) - - -def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: - """Validate QB bounds once per PyTorch tensor version.""" +def _prepare_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: + """Select recoverable host or graph-safe device validation for QB bounds.""" if not ( isinstance(bin_bounds, torch.Tensor) and bin_bounds.is_cuda @@ -74,21 +54,20 @@ def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: # The C++ binding owns metadata validation and its detailed error messages. return False - version = bin_bounds._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 current version must be validated before CUDA graph capture" - ) + # The no-host-sync common path validates the device values in the histogram kernel. + # This check is replayed, unlike Python tensor-version bookkeeping. + return True + version = bin_bounds._version + if getattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, None) == version: + return True 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}]" ) - mark_qb_bin_bounds_validated(bin_bounds) + setattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, version) return True @@ -210,7 +189,7 @@ def forward( histogram_mode: int, ): # pylint: disable=missing-function-docstring - bin_bounds_validated = _validate_qb_bin_bounds(bin_bounds) + skip_bin_bounds_host_validation = _prepare_qb_bin_bounds(bin_bounds) ( probs, routing_output, @@ -227,7 +206,7 @@ def forward( histogram, bin_bounds, histogram_mode, - bin_bounds_validated, + skip_bin_bounds_host_validation, ) if topk_indices is not None: routing_output = topk_indices @@ -309,9 +288,7 @@ def fused_topk_with_score_function( qb_bin_bounds : torch.Tensor, optional FP32 CUDA tensor ``[lower, upper]`` defining uniform QB histogram bins. Values must be 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. Graph replays read - subsequent updates through the captured device pointer, so callers must preserve validity. + CUDA graphs validate the device values on every execution without a host synchronization. qb_histogram_mode : str, optional ``"two_kernel"`` or ``"fused_atomic"``. Must be provided with the two QB tensors. From 583702dde8aaec82ded1df9932a4e3b7d26b341f Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Thu, 27 Aug 2026 00:32:26 +0800 Subject: [PATCH 5/5] [PyTorch] Allow trusted QB updates during graph capture Signed-off-by: Harry Zhou --- tests/pytorch/test_fused_router.py | 73 +++++++++++-------- .../fused_topk_with_score_function.cu | 10 --- .../include/transformer_engine/fused_router.h | 20 ++--- .../pytorch/csrc/extensions/router.cpp | 6 +- transformer_engine/pytorch/router.py | 47 +++++++++--- 5 files changed, 92 insertions(+), 64 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index d54429d13a..e6990d600d 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -2,9 +2,6 @@ # # See LICENSE for license information. from copy import deepcopy -import subprocess -import sys -from textwrap import dedent from typing import Optional import pytest @@ -16,6 +13,7 @@ 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 @@ -731,9 +729,13 @@ def run_router(): 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() @@ -791,6 +793,7 @@ def run_iteration(): 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 @@ -814,16 +817,13 @@ def run_iteration(): @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -def test_qb_topk_cuda_graph_rejects_invalid_bounds_on_device(histogram_mode): - script = dedent( - f""" - import torch - from transformer_engine.pytorch.router import fused_topk_with_score_function - - 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 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( @@ -837,23 +837,38 @@ def test_qb_topk_cuda_graph_rejects_invalid_bounds_on_device(histogram_mode): expert_bias, qb_histogram=histogram, qb_bin_bounds=bin_bounds, - qb_histogram_mode={histogram_mode!r}, + qb_histogram_mode=histogram_mode, ) - graph.replay() - torch.cuda.synchronize() - """ - ) - result = subprocess.run( - [sys.executable, "-c", script], - capture_output=True, - text=True, - check=False, - timeout=60, - ) - output = result.stdout + result.stderr - assert result.returncode != 0 - assert "QB bin_bounds values must be finite with lower < upper." in output - assert "CUDA error" in output + + +@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( diff --git a/transformer_engine/common/fused_router/fused_topk_with_score_function.cu b/transformer_engine/common/fused_router/fused_topk_with_score_function.cu index 1b609ba77c..4ea11f05dc 100644 --- a/transformer_engine/common/fused_router/fused_topk_with_score_function.cu +++ b/transformer_engine/common/fused_router/fused_topk_with_score_function.cu @@ -36,16 +36,6 @@ struct QBBinParams { __device__ inline QBBinParams load_qb_bin_params(const CompType *bin_bounds, int num_bins) { const CompType lower = bin_bounds[0]; const CompType upper = bin_bounds[1]; - const bool valid = isfinite(lower) && isfinite(upper) && lower < upper; - if (!valid) { - if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) { - printf("QB bin_bounds values must be finite with lower < upper.\n"); - __trap(); - } - // Keep non-reporting threads away from invalid floating-point-to-integer conversions while - // the device-side assertion is propagated to the host. - return {0.0f, 0.0f}; - } return {lower, static_cast(num_bins) / (upper - lower)}; } diff --git a/transformer_engine/common/include/transformer_engine/fused_router.h b/transformer_engine/common/include/transformer_engine/fused_router.h index 431a976565..323f04905e 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -120,11 +120,11 @@ void nvte_fused_topk_with_score_function_forward_qb_v2( NVTETensor histogram, NVTETensor bin_bounds, NVTEQBHistogramMode histogram_mode, cudaStream_t stream); -/*! \brief Same as nvte_fused_topk_with_score_function_forward_qb_v2, but skips host-side - * bin_bounds value validation. +/*! \brief Same as nvte_fused_topk_with_score_function_forward_qb_v2, but skips bin_bounds value + * validation. * - * The histogram kernel raises a device-side error if bin_bounds is not finite and ordered. Use - * this variant to avoid host synchronization in a hot path or CUDA graph. + * The caller must have validated that bin_bounds contains finite FP32 values [lower, upper] with + * lower < upper. Use this variant to avoid host synchronization in a hot path or CUDA graph. */ void nvte_fused_topk_with_score_function_forward_qb_v2_unchecked( const NVTETensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, @@ -145,10 +145,10 @@ void nvte_fused_topk_with_score_function_forward_qb_with_indices( NVTEQBHistogramMode histogram_mode, cudaStream_t stream); /*! \brief Same as nvte_fused_topk_with_score_function_forward_qb_with_indices, but skips - * host-side bin_bounds value validation. + * bin_bounds value validation. * - * The histogram kernel raises a device-side error if bin_bounds is not finite and ordered. Use - * this variant to avoid host synchronization in a hot path or CUDA graph. + * The caller must have validated that bin_bounds contains finite FP32 values [lower, upper] with + * lower < upper. Use this variant to avoid host synchronization in a hot path or CUDA graph. */ void nvte_fused_topk_with_score_function_forward_qb_with_indices_unchecked( const NVTETensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, @@ -165,10 +165,10 @@ void nvte_qb_histogram_accumulate(const NVTETensor raw_scores, const NVTETensor const NVTETensor bin_bounds, NVTETensor histogram, cudaStream_t stream); -/*! \brief Same as nvte_qb_histogram_accumulate, but skips host-side bin_bounds value validation. +/*! \brief Same as nvte_qb_histogram_accumulate, but skips bin_bounds value validation. * - * The histogram kernel raises a device-side error if bin_bounds is not finite and ordered. Use - * this variant to avoid host synchronization in a hot path or CUDA graph. + * The caller must have validated that bin_bounds contains finite FP32 values [lower, upper] with + * lower < upper. Use this variant to avoid host synchronization in a hot path or CUDA graph. */ void nvte_qb_histogram_accumulate_unchecked(const NVTETensor raw_scores, const NVTETensor cutoff, const NVTETensor bin_bounds, NVTETensor histogram, diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 713406a276..77ee80c099 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -160,7 +160,7 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, int routing_map_format, std::optional topk_indices, at::Tensor histogram, at::Tensor bin_bounds, int histogram_mode, - bool skip_bin_bounds_host_validation) { + bool bin_bounds_validated) { check_routing_map_format(routing_map_format); TORCH_CHECK(logits.dim() >= 1, "logits must have at least 1 dim"); TORCH_CHECK(logits.is_cuda() && logits.is_contiguous(), @@ -239,7 +239,7 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, auto stream = at::cuda::getCurrentCUDAStream(); if (topk_indices.has_value()) { - if (skip_bin_bounds_host_validation) { + if (bin_bounds_validated) { nvte_fused_topk_with_score_function_forward_qb_with_indices_unchecked( logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, scaling_factor_value, expert_bias_cu.data(), probs_cu.data(), routing_output_cu.data(), @@ -253,7 +253,7 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, bin_bounds_cu.data(), mode, stream); } } else { - if (skip_bin_bounds_host_validation) { + if (bin_bounds_validated) { nvte_fused_topk_with_score_function_forward_qb_v2_unchecked( logits_cu.data(), static_cast(num_tokens), static_cast(num_experts), topk, scaling_factor_value, expert_bias_cu.data(), probs_cu.data(), routing_output_cu.data(), diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 69ca23276d..d4bbb42e4b 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -42,8 +42,26 @@ _QB_BOUNDS_VALIDATED_VERSION_ATTR = "_nvte_qb_bounds_validated_version" -def _prepare_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: - """Select recoverable host or graph-safe device validation for QB bounds.""" +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 QB bounds once per PyTorch tensor version.""" if not ( isinstance(bin_bounds, torch.Tensor) and bin_bounds.is_cuda @@ -54,20 +72,21 @@ def _prepare_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: # The C++ binding owns metadata validation and its detailed error messages. return False - with torch.cuda.device(bin_bounds.device): - if torch.cuda.is_current_stream_capturing(): - # The no-host-sync common path validates the device values in the histogram kernel. - # This check is replayed, unlike Python tensor-version bookkeeping. - return True 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 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 @@ -189,7 +208,7 @@ def forward( histogram_mode: int, ): # pylint: disable=missing-function-docstring - skip_bin_bounds_host_validation = _prepare_qb_bin_bounds(bin_bounds) + bin_bounds_validated = _validate_qb_bin_bounds(bin_bounds) ( probs, routing_output, @@ -206,7 +225,7 @@ def forward( histogram, bin_bounds, histogram_mode, - skip_bin_bounds_host_validation, + bin_bounds_validated, ) if topk_indices is not None: routing_output = topk_indices @@ -288,7 +307,11 @@ def fused_topk_with_score_function( qb_bin_bounds : torch.Tensor, optional FP32 CUDA tensor ``[lower, upper]`` defining uniform QB histogram bins. Values must be finite with ``lower < upper``. Eager calls revalidate PyTorch-tracked in-place updates. - CUDA graphs validate the device values on every execution without a host synchronization. + 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.