diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 18efb6a190..e6990d600d 100644 --- a/tests/pytorch/test_fused_router.py +++ b/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( diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index ac38f7d1cf..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 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.