diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index c6c82fa3d78..ac2f46826df 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -661,109 +661,55 @@ def test_qb_topk_uses_logits_device(histogram_mode): @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -@pytest.mark.parametrize("invalid_bounds", ["equal", "reversed", "nonfinite"]) -def test_qb_topk_rejects_invalid_bin_bounds(histogram_mode, invalid_bounds): - 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) - bounds = { - "equal": [1.0, 1.0], - "reversed": [1.0, -1.0], - "nonfinite": [float("nan"), 1.0], - }[invalid_bounds] - bin_bounds = torch.tensor(bounds, device="cuda", dtype=torch.float32) - with pytest.raises(ValueError, match="finite with lower < upper"): - 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("use_dense_indices", [False, True]) +@pytest.mark.parametrize( + "invalid_bounds", + [[1.0, 1.0], [1.0, -1.0], [float("nan"), 1.0], [float("-inf"), float("inf")]], + ids=["equal", "reversed", "nan", "infinite"], +) +def test_qb_topk_invalid_bin_bounds_stay_in_histogram( + histogram_mode, use_dense_indices, invalid_bounds +): + num_tokens, num_experts, num_bins, topk = 8, 16, 32, 4 + logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=torch.float32) + expert_bias = torch.zeros(num_experts, device="cuda", dtype=torch.float32) + histogram = torch.zeros(num_experts, num_bins, device="cuda", dtype=torch.int32) + bin_bounds = torch.tensor(invalid_bounds, device="cuda", dtype=torch.float32) + topk_indices = ( + torch.empty(num_tokens, topk, device="cuda", dtype=torch.int32) + if use_dense_indices + else None + ) -@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) - bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) - fused_topk_with_score_function( + # bin_bounds values are not validated, so invalid bounds must neither raise nor write outside + # the histogram. + probs, routing_output = fused_topk_with_score_function( logits, - 4, + topk, False, None, None, None, "sigmoid", expert_bias, + topk_indices=topk_indices, qb_histogram=histogram, qb_bin_bounds=bin_bounds, qb_histogram_mode=histogram_mode, ) - bin_bounds.fill_(0.0) - with pytest.raises(ValueError, match="finite with lower < upper"): - 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", - [QBHistogramMode.TWO_KERNEL, QBHistogramMode.FUSED_ATOMIC], -) -@pytest.mark.parametrize("use_dense_indices", [False, True]) -def test_qb_raw_binding_rejects_invalid_bin_bounds_recoverably(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) - topk_indices = ( - torch.empty(8, 4, device="cuda", dtype=torch.int32) if use_dense_indices else None - ) - invalid_bounds = torch.tensor([1.0, 1.0], device="cuda", dtype=torch.float32) - with pytest.raises(RuntimeError, match="finite with lower < upper"): - tex.fused_topk_with_score_function_qb_fwd( - logits, - 4, - None, - expert_bias, - int(RoutingMapFormat.BYTEMAP), - topk_indices, - histogram, - invalid_bounds, - histogram_mode, - ) + torch.cuda.synchronize() - # The validation error must not poison the CUDA context. + # Routing does not read the bounds, and every (token, expert) pair is counted in some bin. valid_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) - tex.fused_topk_with_score_function_qb_fwd( - logits, - 4, - None, - expert_bias, - int(RoutingMapFormat.BYTEMAP), - topk_indices, - histogram, - valid_bounds, - histogram_mode, - ) - torch.cuda.synchronize() + reference = qb_topk_score_function_pytorch(logits, topk, expert_bias, valid_bounds, num_bins) + torch.testing.assert_close(probs, reference["probs"]) + if use_dense_indices: + torch.testing.assert_close( + topk_indices_to_routing_map(routing_output, num_experts), reference["routing_map"] + ) + else: + torch.testing.assert_close(routing_output, reference["routing_map"]) + assert histogram.sum().item() == num_tokens * num_experts @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) @@ -848,7 +794,6 @@ 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 @@ -871,98 +816,11 @@ def run_iteration(): torch.testing.assert_close(bin_bounds, next_bounds) -@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -def test_qb_topk_eager_call_revalidates_after_captured_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) +def test_mark_qb_bin_bounds_validated_is_deprecated(): 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_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() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - bin_bounds.copy_(next_bounds) + with pytest.warns(DeprecationWarning, match="mark_qb_bin_bounds_validated is deprecated"): mark_qb_bin_bounds_validated(bin_bounds) - # Replayed device writes do not advance bin_bounds._version. The capture-time marker must not - # let a subsequent eager router call inherit trust in the replay-produced values. - next_bounds.zero_() - graph.replay() - torch.cuda.synchronize() - with pytest.raises(ValueError, match="finite with lower < upper"): - run_router() - - -@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( "histogram_mode", 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 32d525bb027..e9d6626a2b0 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 @@ -10,7 +10,6 @@ #include #include -#include #include #include @@ -33,22 +32,16 @@ struct QBBinParams { CompType scale; }; +// bin_bounds values are not validated, since that would need a host sync. Users of these params +// compute bin indices with __float2int_rd, which saturates instead of invoking undefined behavior, +// and clamp them to [0, num_bins), so invalid bounds give a wrong histogram but never an +// out-of-bounds access. __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]; return {lower, static_cast(num_bins) / (upper - lower)}; } -void validate_qb_bin_bounds(const Tensor bin_bounds, cudaStream_t stream) { - float bounds[2]; - NVTE_CHECK_CUDA(cudaMemcpyAsync(bounds, bin_bounds.data.dptr, sizeof(bounds), - cudaMemcpyDeviceToHost, stream)); - NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - NVTE_CHECK(std::isfinite(bounds[0]) && std::isfinite(bounds[1]) && bounds[0] < bounds[1], - "QB bin_bounds values must be finite with lower < upper, got [", bounds[0], ", ", - bounds[1], "]"); -} - template __device__ inline CompType extract_qb_cutoff_and_compact(int *topk_indices, CompType *topk_scores, int topk, int lane_id) { @@ -87,8 +80,7 @@ __device__ inline void accumulate_qb_histogram_epilogue(const CompType *raw_scor if constexpr (Mode == QBMode::FusedAtomic) { const QBBinParams bin_params = load_qb_bin_params(bin_bounds, num_bins); for (int expert = lane_id; expert < num_experts; expert += kThreadsPerWarp) { - int bin = static_cast( - floorf((cutoff - raw_scores[expert] - bin_params.lower) * bin_params.scale)); + int bin = __float2int_rd((cutoff - raw_scores[expert] - bin_params.lower) * bin_params.scale); bin = max(0, min(bin, num_bins - 1)); atomicAdd(histogram + static_cast(expert) * num_bins + bin, 1); } @@ -913,7 +905,7 @@ __global__ void qb_histogram_accumulate_kernel(const CompType *raw_scores, const if (token < num_tokens && expert < num_experts) { const CompType required_bias = cutoff[token] - raw_scores[static_cast(token) * num_experts + expert]; - int bin = static_cast(floorf((required_bias - bin_params.lower) * bin_params.scale)); + int bin = __float2int_rd((required_bias - bin_params.lower) * bin_params.scale); bin = max(0, min(bin, num_bins - 1)); atomicAdd(local_histogram + local_expert * num_bins + bin, 1); } @@ -931,7 +923,7 @@ __global__ void qb_histogram_accumulate_kernel(const CompType *raw_scores, const } void qb_histogram_accumulate(const Tensor raw_scores, const Tensor cutoff, const Tensor bin_bounds, - Tensor histogram, bool validate_bin_bounds, cudaStream_t stream) { + Tensor histogram, cudaStream_t stream) { NVTE_CHECK(raw_scores.data.dtype == DType::kFloat32, "QB raw_scores must have FP32 dtype"); NVTE_CHECK(cutoff.data.dtype == DType::kFloat32, "QB cutoff must have FP32 dtype"); NVTE_CHECK(bin_bounds.data.dtype == DType::kFloat32, "QB bin_bounds must have FP32 dtype"); @@ -948,9 +940,6 @@ void qb_histogram_accumulate(const Tensor raw_scores, const Tensor cutoff, const "QB histogram must have shape [num_experts, num_bins]"); const int num_bins = static_cast(histogram.data.shape[1]); NVTE_CHECK(num_bins > 0, "QB num_bins must be positive"); - if (validate_bin_bounds) { - validate_qb_bin_bounds(bin_bounds, stream); - } const int token_partitions = std::min(4, std::max(1, (num_tokens + 1023) / 1024)); const dim3 grid((num_experts + kQBExpertsPerBlock - 1) / kQBExpertsPerBlock, token_partitions); @@ -1144,8 +1133,7 @@ void fused_topk_with_score_function_forward_qb( const Tensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, const Tensor expert_bias, Tensor probs, Tensor routing_map, NVTERoutingMapFormat routing_map_format, Tensor intermediate_output, Tensor cutoff, - Tensor histogram, Tensor bin_bounds, NVTEQBHistogramMode histogram_mode, - bool validate_bin_bounds, cudaStream_t stream) { + Tensor histogram, Tensor bin_bounds, NVTEQBHistogramMode histogram_mode, cudaStream_t stream) { check_routing_map_format(routing_map_format); const int num_bins = check_qb_forward_tensors(logits, num_tokens, num_experts, topk, expert_bias, probs, @@ -1154,9 +1142,6 @@ void fused_topk_with_score_function_forward_qb( expected_routing_map_shape(num_tokens, num_experts, routing_map_format); NVTE_CHECK(routing_map.data.shape == routing_map_shape, "QB routing_map shape does not match routing_map_format"); - if (validate_bin_bounds) { - validate_qb_bin_bounds(bin_bounds, stream); - } #define QB_ROUTER_FORWARD_LAUNCH(RoutingMapFormatVal, QbModeVal) \ TE_ROUTER_PROBS_TYPE_SWITCH_ALL( \ @@ -1195,7 +1180,7 @@ void fused_topk_with_score_function_forward_qb_with_indices( const Tensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, const Tensor expert_bias, Tensor probs, Tensor topk_indices, Tensor intermediate_output, Tensor cutoff, Tensor histogram, Tensor bin_bounds, NVTEQBHistogramMode histogram_mode, - bool validate_bin_bounds, cudaStream_t stream) { + cudaStream_t stream) { const int num_bins = check_qb_forward_tensors(logits, num_tokens, num_experts, topk, expert_bias, probs, intermediate_output, cutoff, histogram, bin_bounds); @@ -1203,9 +1188,6 @@ void fused_topk_with_score_function_forward_qb_with_indices( static_cast(topk)}; NVTE_CHECK(topk_indices.data.shape == indices_shape, "QB topk_indices must have shape [num_tokens, topk]"); - if (validate_bin_bounds) { - validate_qb_bin_bounds(bin_bounds, stream); - } #define QB_ROUTER_FORWARD_WITH_INDICES_LAUNCH(DataType, IndexType, QbModeVal) \ fused_topk_with_score_function_forward_with_indices_kernel_launcher scaling_factor, at::Tensor expert_bias, int routing_map_format, std::optional topk_indices, at::Tensor histogram, - at::Tensor bin_bounds, int histogram_mode, - bool bin_bounds_validated = false); + at::Tensor bin_bounds, int histogram_mode); void fused_topk_with_score_function_bwd( at::Tensor routing_map, at::Tensor intermediate_output, at::Tensor grad_probs, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 89e7d92aca9..ebcdd17b07c 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -149,7 +149,7 @@ void init_router_bindings(pybind11::module &m) { m.def("fused_topk_with_score_function_qb_fwd", &fused_topk_with_score_function_qb_fwd, py::arg("logits"), py::arg("topk"), py::arg("scaling_factor"), py::arg("expert_bias"), py::arg("routing_map_format"), py::arg("topk_indices"), py::arg("histogram"), - py::arg("bin_bounds"), py::arg("histogram_mode"), py::arg("bin_bounds_validated") = false, + py::arg("bin_bounds"), py::arg("histogram_mode"), "Kimi K3 QB fused topk with histogram accumulation"); m.def("fused_topk_with_score_function_bwd", &fused_topk_with_score_function_bwd, py::arg("routing_map"), py::arg("intermediate_output"), py::arg("grad_probs"), diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 77ee80c0995..79391b5bf2f 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -159,8 +159,7 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, std::optional scaling_factor, at::Tensor expert_bias, int routing_map_format, std::optional topk_indices, at::Tensor histogram, - at::Tensor bin_bounds, int histogram_mode, - bool bin_bounds_validated) { + at::Tensor bin_bounds, int histogram_mode) { 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,37 +238,21 @@ 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) { - 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(), - intermediate_output_cu.data(), cutoff_cu.data(), histogram_cu.data(), - bin_bounds_cu.data(), mode, stream); - } else { - nvte_fused_topk_with_score_function_forward_qb_with_indices( - 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(), - intermediate_output_cu.data(), cutoff_cu.data(), histogram_cu.data(), - bin_bounds_cu.data(), mode, stream); - } + nvte_fused_topk_with_score_function_forward_qb_with_indices( + 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(), + intermediate_output_cu.data(), cutoff_cu.data(), histogram_cu.data(), bin_bounds_cu.data(), + mode, stream); } else { - 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(), - static_cast(routing_map_format), intermediate_output_cu.data(), - cutoff_cu.data(), histogram_cu.data(), bin_bounds_cu.data(), mode, stream); - } else { - nvte_fused_topk_with_score_function_forward_qb_v2( - 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(), - static_cast(routing_map_format), intermediate_output_cu.data(), - cutoff_cu.data(), histogram_cu.data(), bin_bounds_cu.data(), mode, stream); - } + nvte_fused_topk_with_score_function_forward_qb_v2( + 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(), + static_cast(routing_map_format), intermediate_output_cu.data(), + cutoff_cu.data(), histogram_cu.data(), bin_bounds_cu.data(), mode, stream); } if (mode == NVTE_QB_HISTOGRAM_TWO_KERNEL) { - nvte_qb_histogram_accumulate_unchecked(intermediate_output_cu.data(), cutoff_cu.data(), - bin_bounds_cu.data(), histogram_cu.data(), stream); + nvte_qb_histogram_accumulate(intermediate_output_cu.data(), cutoff_cu.data(), + bin_bounds_cu.data(), histogram_cu.data(), stream); } return std::make_tuple(probs, routing_output, intermediate_output, cutoff, histogram); diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 898310b0d61..508c2cba769 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -14,7 +14,7 @@ global memory. For example, the gradient is required to have the same dtype as the input. """ -import math +import warnings from typing import Optional, Union import torch @@ -39,72 +39,24 @@ "two_kernel": int(QBHistogramMode.TWO_KERNEL), "fused_atomic": int(QBHistogramMode.FUSED_ATOMIC), } -_QB_BOUNDS_VALIDATED_VERSION_ATTR = "_nvte_qb_bounds_validated_version" -_QB_BOUNDS_CAPTURE_VALIDATED_VERSION_ATTR = "_nvte_qb_bounds_capture_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. +def mark_qb_bin_bounds_validated( + bin_bounds: torch.Tensor, # pylint: disable=unused-argument +) -> None: + """Deprecated no-op. - This function validates metadata only; it does not inspect or schedule validation of tensor - values. The caller must guarantee finite FP32 bounds with ``lower < upper``. When called during - CUDA graph capture, call it immediately after dispatching the trusted in-place update on the - same stream. The Python marker itself is not replayed, so every value produced on replay must - satisfy the bounds contract for every possible input and control-flow path. + QB ``bin_bounds`` values are no longer validated, so there is nothing to mark. This function + will be removed in a future release. """ - 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]") - version = bin_bounds._version - setattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, version) - with torch.cuda.device(bin_bounds.device): - capture_validated_version = version if torch.cuda.is_current_stream_capturing() else None - setattr( - bin_bounds, - _QB_BOUNDS_CAPTURE_VALIDATED_VERSION_ATTR, - capture_validated_version, + warnings.warn( + "mark_qb_bin_bounds_validated is deprecated and has no effect because QB bin_bounds " + "values are no longer validated. It will be removed in a future release.", + category=DeprecationWarning, + stacklevel=2, ) -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 - and bin_bounds.is_contiguous() - and bin_bounds.dtype == torch.float32 - and bin_bounds.shape == (2,) - ): - # 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) - capture_validated_version = getattr(bin_bounds, _QB_BOUNDS_CAPTURE_VALIDATED_VERSION_ATTR, None) - if validated_version == version and capture_validated_version != version: - return True - with torch.cuda.device(bin_bounds.device): - is_capturing = torch.cuda.is_current_stream_capturing() - if validated_version == version and is_capturing: - return True - if is_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}]" - ) - mark_qb_bin_bounds_validated(bin_bounds) - return True - - def _validate_routing_map_format( routing_map_format: Union[str, RoutingMapFormat, int], ) -> int: @@ -223,7 +175,6 @@ def forward( histogram_mode: int, ): # pylint: disable=missing-function-docstring - bin_bounds_validated = _validate_qb_bin_bounds(bin_bounds) ( probs, routing_output, @@ -240,7 +191,6 @@ def forward( histogram, bin_bounds, histogram_mode, - bin_bounds_validated, ) if topk_indices is not None: routing_output = topk_indices @@ -321,13 +271,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``. 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. CUDA graph replay - does not execute Python validation or advance PyTorch tensor-version counters. Any bounds - update visible to replay, whether issued outside the graph or captured inside it, is - therefore trusted and must produce finite values with ``lower < upper`` on every replay. - Invalid replay-time bounds may silently produce an incorrect histogram. + finite with ``lower < upper``. The values are only read on the device and are not + validated, so the tensor may be updated in place at any time, including inside a CUDA + graph. Invalid bounds silently produce an incorrect histogram. qb_histogram_mode : str, optional ``"two_kernel"`` or ``"fused_atomic"``. Must be provided with the two QB tensors.