From 73bb07f322c59f0ee9db03d6fefde12457273c39 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Wed, 29 Jul 2026 02:00:41 +0800 Subject: [PATCH 1/4] [Common][PyTorch] Add QB router histogram paths Signed-off-by: Harry Zhou --- tests/pytorch/test_fused_router.py | 281 +++++++++- .../fused_topk_with_score_function.cu | 504 +++++++++++++++--- .../include/transformer_engine/fused_router.h | 36 ++ transformer_engine/pytorch/csrc/extensions.h | 7 + .../pytorch/csrc/extensions/pybind.cpp | 8 + .../pytorch/csrc/extensions/router.cpp | 103 ++++ transformer_engine/pytorch/router.py | 118 +++- 7 files changed, 982 insertions(+), 75 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 68d3ed9565..36ca415307 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -4,11 +4,13 @@ import torch from typing import Optional from transformer_engine.pytorch.router import ( + QBHistogramMode, RoutingMapFormat, fused_topk_with_score_function, fused_compute_score_for_moe_aux_loss, fused_moe_aux_loss, ) +import transformer_engine_torch as tex import pytest from copy import deepcopy @@ -130,6 +132,63 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None): return topk_masked_gates, topk_map +def qb_topk_score_function_pytorch( + logits: torch.Tensor, + topk: int, + expert_bias: torch.Tensor, + bin_bounds: torch.Tensor, + num_bins: int, + histogram: Optional[torch.Tensor] = None, +): + """Pure-PyTorch reference for Kimi K3 QB routing and histogram accumulation.""" + original_shape = logits.shape + num_experts = original_shape[-1] + raw_scores = torch.sigmoid(logits.float()).reshape(-1, num_experts) + biased_scores = raw_scores + expert_bias + topk_plus_one_scores, topk_plus_one_indices = torch.topk(biased_scores, k=topk + 1, dim=-1) + cutoff = topk_plus_one_scores.min(dim=-1).values + cutoff_candidates = topk_plus_one_indices.masked_fill( + topk_plus_one_scores != cutoff.unsqueeze(1), -1 + ) + dropped_expert = cutoff_candidates.max(dim=-1).values + topk_indices = topk_plus_one_indices[ + topk_plus_one_indices != dropped_expert.unsqueeze(1) + ].reshape(-1, topk) + + selected_raw_scores = torch.gather(raw_scores, 1, topk_indices) + if topk > 1: + selected_probs = selected_raw_scores / ( + selected_raw_scores.sum(dim=-1, keepdim=True) + 1e-20 + ) + else: + selected_probs = selected_raw_scores + probs = torch.zeros_like(raw_scores).scatter(1, topk_indices, selected_probs) + routing_map = torch.zeros_like(raw_scores, dtype=torch.bool).scatter(1, topk_indices, True) + + lower, upper = bin_bounds[0], bin_bounds[1] + required_bias = cutoff.unsqueeze(1) - raw_scores + bin_scale = num_bins / (upper - lower) + bin_indices = torch.floor((required_bias - lower) * bin_scale).to(torch.int64) + bin_indices.clamp_(0, num_bins - 1) + expert_offsets = torch.arange(num_experts, device=logits.device, dtype=torch.int64) * num_bins + flat_indices = (bin_indices + expert_offsets).reshape(-1) + counts = torch.bincount(flat_indices, minlength=num_experts * num_bins) + counts = counts.reshape(num_experts, num_bins).to(torch.int32) + if histogram is None: + histogram = torch.zeros_like(counts) + histogram.add_(counts) + + return { + "probs": probs.reshape(original_shape).to(logits.dtype), + "routing_map": routing_map.reshape(original_shape), + "topk_indices": topk_indices.reshape(*original_shape[:-1], topk), + "raw_scores": raw_scores.reshape(original_shape), + "cutoff": cutoff.reshape(original_shape[:-1]), + "bin_indices": bin_indices.reshape(original_shape), + "histogram": histogram, + } + + # Pytorch-based compute routing scores for aux loss def compute_scores_for_aux_loss_pytorch( logits: torch.Tensor, topk: int, score_function: str @@ -355,6 +414,204 @@ def test_topk_sqrtsoftplus( ) +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +@pytest.mark.parametrize("topk", [8, 16]) +@pytest.mark.parametrize( + "routing_output_mode", + ["bytemap", "bitmap_u8", "dense_int16", "dense_int32", "dense_int64"], +) +def test_qb_topk_histogram(histogram_mode, topk, routing_output_mode): + num_tokens = 257 + num_experts = 896 + num_bins = 1000 + logits = torch.randn( + num_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True + ) + expert_bias = torch.linspace(-0.2, 0.2, num_experts, device="cuda", dtype=torch.float32) + bin_bounds = torch.stack((expert_bias.min() - 1.0, expert_bias.max() + 1.0)) + reference_histogram = torch.zeros(num_experts, num_bins, device="cuda", dtype=torch.int32) + reference = qb_topk_score_function_pytorch( + logits, + topk, + expert_bias, + bin_bounds, + num_bins, + reference_histogram, + ) + + fused_logits = logits.detach().clone().requires_grad_(True) + fused_histogram = torch.zeros_like(reference_histogram) + dense_dtype = { + "dense_int16": torch.int16, + "dense_int32": torch.int32, + "dense_int64": torch.int64, + }.get(routing_output_mode) + topk_indices = ( + torch.empty(num_tokens, topk, device="cuda", dtype=dense_dtype) + if dense_dtype is not None + else None + ) + routing_map_format = ( + RoutingMapFormat.BITMAP_U8 + if routing_output_mode == "bitmap_u8" + else RoutingMapFormat.BYTEMAP + ) + fused_probs, fused_routing_output = fused_topk_with_score_function( + logits=fused_logits, + topk=topk, + use_pre_softmax=False, + num_groups=None, + group_topk=None, + scaling_factor=None, + score_function="sigmoid", + expert_bias=expert_bias, + routing_map_format=routing_map_format, + topk_indices=topk_indices, + qb_histogram=fused_histogram, + qb_bin_bounds=bin_bounds, + qb_histogram_mode=histogram_mode, + ) + torch.testing.assert_close(fused_probs, reference["probs"]) + if dense_dtype is not None: + fused_routing_map = topk_indices_to_routing_map(fused_routing_output, num_experts) + torch.testing.assert_close(fused_routing_map, reference["routing_map"]) + elif routing_output_mode == "bitmap_u8": + torch.testing.assert_close( + fused_routing_output, + _bytemap_to_bitmap_u8(reference["routing_map"]), + ) + else: + torch.testing.assert_close(fused_routing_output, reference["routing_map"]) + torch.testing.assert_close(fused_histogram, reference["histogram"]) + + grad = torch.randn_like(fused_probs) + reference["probs"].backward(grad) + fused_probs.backward(grad) + torch.testing.assert_close(fused_logits.grad, logits.grad, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +def test_qb_histogram_accumulates_microbatches(histogram_mode): + num_experts = 64 + topk = 8 + num_bins = 1000 + expert_bias = torch.linspace(-0.1, 0.1, num_experts, device="cuda", dtype=torch.float32) + bin_bounds = torch.stack((expert_bias.min() - 1.0, expert_bias.max() + 1.0)) + reference_histogram = torch.zeros(num_experts, num_bins, device="cuda", dtype=torch.int32) + fused_histogram = torch.zeros_like(reference_histogram) + + for num_tokens in (127, 193): + logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=torch.float32) + qb_topk_score_function_pytorch( + logits, + topk, + expert_bias, + bin_bounds, + num_bins, + reference_histogram, + ) + fused_topk_with_score_function( + logits=logits, + topk=topk, + use_pre_softmax=False, + num_groups=None, + group_topk=None, + scaling_factor=None, + score_function="sigmoid", + expert_bias=expert_bias, + qb_histogram=fused_histogram, + qb_bin_bounds=bin_bounds, + qb_histogram_mode=histogram_mode, + ) + + torch.testing.assert_close(fused_histogram, reference_histogram) + + +def test_qb_topk_argument_validation(): + logits = torch.randn(16, 32, device="cuda", dtype=torch.float32) + expert_bias = torch.zeros(32, device="cuda", dtype=torch.float32) + histogram = torch.zeros(32, 1000, device="cuda", dtype=torch.int32) + bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) + with pytest.raises(ValueError, match="provided together"): + fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "sigmoid", + expert_bias, + qb_histogram=histogram, + ) + with pytest.raises(ValueError, match="only supports"): + fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "softmax", + expert_bias, + qb_histogram=histogram, + qb_bin_bounds=bin_bounds, + qb_histogram_mode="two_kernel", + ) + + +@pytest.mark.parametrize( + "histogram_mode", + [QBHistogramMode.TWO_KERNEL, QBHistogramMode.FUSED_ATOMIC], +) +def test_qb_topk_plus_one_tie_and_bin_clamping(histogram_mode): + logits = torch.tensor( + [ + [3.0, 3.0, 3.0, -1.0, -2.0, -3.0, -4.0, -5.0], + [5.0, 4.0, 3.0, 2.0, -2.0, -3.0, -4.0, -5.0], + ], + device="cuda", + dtype=torch.float32, + ) + _, num_experts = logits.shape + topk = 2 + num_bins = 8 + expert_bias = torch.zeros(num_experts, device="cuda", dtype=torch.float32) + bin_bounds = torch.tensor([-0.05, 0.05], device="cuda", dtype=torch.float32) + reference = qb_topk_score_function_pytorch( + logits, + topk, + expert_bias, + bin_bounds, + num_bins, + ) + histogram = torch.zeros(num_experts, num_bins, device="cuda", dtype=torch.int32) + + probs, routing_map, raw_scores, cutoff, histogram = tex.fused_topk_with_score_function_qb_fwd( + logits, + topk, + None, + expert_bias, + int(RoutingMapFormat.BYTEMAP), + None, + histogram, + bin_bounds, + histogram_mode, + ) + + torch.testing.assert_close(probs, reference["probs"]) + torch.testing.assert_close(routing_map, reference["routing_map"]) + torch.testing.assert_close(raw_scores, reference["raw_scores"]) + if histogram_mode == QBHistogramMode.TWO_KERNEL: + torch.testing.assert_close(cutoff, reference["cutoff"]) + torch.testing.assert_close(histogram, reference["histogram"]) + assert histogram[:, 0].sum() > 0 + assert histogram[:, -1].sum() > 0 + # The first token has exactly Top-(k+1) equal scores. The deterministic + # compaction rule drops the largest expert ID at the cutoff. + assert routing_map[0, :3].tolist() == [True, True, False] + + @pytest.mark.parametrize("dtype", [torch.float32]) @pytest.mark.parametrize("num_tokens", [2048, 7168, 14234]) @pytest.mark.parametrize("num_experts", [1024, 128, 32]) @@ -736,11 +993,29 @@ def profile_topk_softmax( group_topk = 4 scaling_factor = 1.2 test_topk_sigmoid( - torch.float32, num_tokens, num_experts, topk, group_topk, scaling_factor, enable_bias + torch.float32, + num_tokens, + num_experts, + topk, + group_topk, + scaling_factor, + enable_bias, ) test_topk_softmax( - torch.float32, num_tokens, num_experts, topk, use_pre_softmax, group_topk, scaling_factor + torch.float32, + num_tokens, + num_experts, + topk, + use_pre_softmax, + group_topk, + scaling_factor, ) test_topk_sqrtsoftplus( - torch.float32, num_tokens, num_experts, topk, group_topk, scaling_factor, enable_bias + torch.float32, + num_tokens, + num_experts, + topk, + group_topk, + scaling_factor, + enable_bias, ) 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 cd8eca05e0..9f867008ed 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 @@ -8,7 +8,9 @@ #include #include +#include #include +#include #include #include @@ -20,6 +22,59 @@ namespace transformer_engine { namespace fused_router { +enum class QBMode { + Disabled, + TwoKernel, + FusedAtomic, +}; + +template +__device__ inline CompType extract_qb_cutoff_and_compact(int *topk_indices, CompType *topk_scores, + int topk, int lane_id) { + if constexpr (Mode == QBMode::Disabled) { + return 0.0f; + } else { + CompType cutoff = 0.0f; + if (lane_id == 0) { + int cutoff_pos = 0; + cutoff = topk_scores[0]; + int cutoff_expert = topk_indices[0]; + for (int i = 1; i < topk + 1; ++i) { + const CompType score = topk_scores[i]; + const int expert = topk_indices[i]; + if (score < cutoff || (score == cutoff && expert > cutoff_expert)) { + cutoff = score; + cutoff_expert = expert; + cutoff_pos = i; + } + } + for (int i = cutoff_pos; i < topk; ++i) { + topk_scores[i] = topk_scores[i + 1]; + topk_indices[i] = topk_indices[i + 1]; + } + } + __syncwarp(); + return __shfl_sync(0xffffffff, cutoff, 0); + } +} + +template +__device__ inline void accumulate_qb_histogram_epilogue(const CompType *raw_scores, CompType cutoff, + int num_experts, int lane_id, + const CompType *bin_bounds, int num_bins, + int32_t *histogram) { + if constexpr (Mode == QBMode::FusedAtomic) { + const CompType lower = bin_bounds[0]; + const CompType upper = bin_bounds[1]; + const CompType scale = static_cast(num_bins) / (upper - lower); + for (int expert = lane_id; expert < num_experts; expert += kThreadsPerWarp) { + int bin = static_cast(floorf((cutoff - raw_scores[expert] - lower) * scale)); + bin = max(0, min(bin, num_bins - 1)); + atomicAdd(histogram + static_cast(expert) * num_bins + bin, 1); + } + } +} + // ============================================================================= // Simple forward kernel — exact upstream structure (no async loader, no // persistent grid, runtime score_function dispatch). Faster for small topk @@ -27,13 +82,17 @@ namespace fused_router { // ============================================================================= template + TopkFuncType TopkFunc = TopkFuncType::Naive, typename IndexType = int32_t, + QBMode QbMode = QBMode::Disabled> __global__ void fused_topk_forward_simple_kernel( const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax, int num_groups, int group_topk, float scaling_factor, int score_function, const BiasType *expert_bias, DataType *probs, uint8_t *routing_map, - CompType *intermediate_output, IndexType *topk_indices_output) { + CompType *intermediate_output, IndexType *topk_indices_output, CompType *qb_cutoff, + int32_t *qb_histogram, const CompType *qb_bin_bounds, int qb_num_bins) { constexpr bool kIsBitmap = (RoutingMapFormat == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8); + constexpr bool kUseQB = QbMode != QBMode::Disabled; + const int selection_topk = topk + (kUseQB ? 1 : 0); int num_token_per_block = blockDim.x / kThreadsPerWarp; int warp_id = threadIdx.x / kThreadsPerWarp; int lane_id = threadIdx.x % kThreadsPerWarp; @@ -43,23 +102,25 @@ __global__ void fused_topk_forward_simple_kernel( CompType *group_scores_buf = nullptr, *masked_scores_buf = nullptr; int *topk_indices_buf = nullptr; if (group_topk > 0) { - masked_scores_buf = topk_scores_buf + topk * num_token_per_block; + masked_scores_buf = topk_scores_buf + selection_topk * num_token_per_block; group_scores_buf = masked_scores_buf + num_experts * num_token_per_block; topk_indices_buf = reinterpret_cast(group_scores_buf + num_groups * num_token_per_block); } else { - topk_indices_buf = reinterpret_cast(topk_scores_buf + topk * num_token_per_block); + topk_indices_buf = + reinterpret_cast(topk_scores_buf + selection_topk * num_token_per_block); } const int bitmap_words_per_warp = (num_experts + 31) / 32; const int bitmap_row_bytes = (num_experts + 7) / 8; uint32_t *bitmap_words_buf = nullptr; if constexpr (kIsBitmap) { - bitmap_words_buf = reinterpret_cast(topk_indices_buf + topk * num_token_per_block); + bitmap_words_buf = + reinterpret_cast(topk_indices_buf + selection_topk * num_token_per_block); } CompType *scores = scores_buf + warp_id * num_experts; - CompType *topk_scores = topk_scores_buf + warp_id * topk; + CompType *topk_scores = topk_scores_buf + warp_id * selection_topk; CompType *masked_scores = masked_scores_buf + warp_id * num_experts; CompType *group_scores = group_scores_buf + warp_id * num_groups; - int *topk_indices = topk_indices_buf + warp_id * topk; + int *topk_indices = topk_indices_buf + warp_id * selection_topk; uint32_t *local_bitmap_words = (bitmap_words_buf != nullptr) ? bitmap_words_buf + warp_id * bitmap_words_per_warp : nullptr; @@ -131,8 +192,11 @@ __global__ void fused_topk_forward_simple_kernel( __syncwarp(); } - // Topk selection - if (group_topk > 0) { + // Topk selection. QB is only supported without grouped Top-k. + if constexpr (kUseQB) { + topk_and_mask(scores, num_experts, selection_topk, topk_indices, topk_scores, + lane_id); + } else if (group_topk > 0) { int group_size = num_experts / num_groups; for (int i = 0; i < num_groups; i++) { topk_and_mask(scores + i * group_size, group_size, topk / group_topk, @@ -164,6 +228,19 @@ __global__ void fused_topk_forward_simple_kernel( } __syncwarp(); + const CompType cutoff = + extract_qb_cutoff_and_compact(topk_indices, topk_scores, topk, lane_id); + if constexpr (QbMode == QBMode::TwoKernel) { + if (lane_id == 0) { + qb_cutoff[token_offset_cur_warp] = cutoff; + } + } else if constexpr (QbMode == QBMode::FusedAtomic) { + accumulate_qb_histogram_epilogue(intermediate_output + pos_offset, cutoff, + num_experts, lane_id, qb_bin_bounds, qb_num_bins, + qb_histogram); + } + __syncwarp(); + // Postprocess: revert bias, softmax, normalization if (expert_bias && (score_function == 0 || score_function == 2)) { for (int i = lane_id; i < topk; i += kThreadsPerWarp) { @@ -233,13 +310,16 @@ __global__ void fused_topk_forward_simple_kernel( template + typename IndexType = int32_t, QBMode QbMode = QBMode::Disabled> __global__ void fused_topk_with_score_function_forward_kernel( const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax, int num_groups, int group_topk, float scaling_factor, const BiasType *expert_bias, DataType *probs, uint8_t *routing_map, CompType *intermediate_output, - IndexType *topk_indices_output, int num_buffers) { + IndexType *topk_indices_output, int num_buffers, CompType *qb_cutoff, int32_t *qb_histogram, + const CompType *qb_bin_bounds, int qb_num_bins) { constexpr bool kIsBitmap = (RoutingMapFormat == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8); + constexpr bool kUseQB = QbMode != QBMode::Disabled; + const int selection_topk = topk + (kUseQB ? 1 : 0); /*** * Section: Global Variables/Addresses init * - Each warp is responsible for one token, and has own shared memory buffer. @@ -265,26 +345,28 @@ __global__ void fused_topk_with_score_function_forward_kernel( CompType *group_scores_buf = nullptr, *masked_scores_buf = nullptr; int *topk_indices_buf = nullptr; if (group_topk > 0) { - masked_scores_buf = topk_scores_buf + topk * num_token_per_block; + masked_scores_buf = topk_scores_buf + selection_topk * num_token_per_block; group_scores_buf = masked_scores_buf + num_experts * num_token_per_block; topk_indices_buf = reinterpret_cast(group_scores_buf + num_groups * num_token_per_block); } else { - topk_indices_buf = reinterpret_cast(topk_scores_buf + topk * num_token_per_block); + topk_indices_buf = + reinterpret_cast(topk_scores_buf + selection_topk * num_token_per_block); } const int bitmap_words_per_warp = (num_experts + 31) / 32; const int bitmap_row_bytes = (num_experts + 7) / 8; uint32_t *bitmap_words_buf = nullptr; if constexpr (kIsBitmap) { - bitmap_words_buf = reinterpret_cast(topk_indices_buf + topk * num_token_per_block); + bitmap_words_buf = + reinterpret_cast(topk_indices_buf + selection_topk * num_token_per_block); } // The address of buffers on the current warp CompType *scores = scores_buf + warp_id * num_experts; - CompType *topk_scores = topk_scores_buf + warp_id * topk; + CompType *topk_scores = topk_scores_buf + warp_id * selection_topk; CompType *masked_scores = (masked_scores_buf != nullptr) ? masked_scores_buf + warp_id * num_experts : nullptr; CompType *group_scores = (group_scores_buf != nullptr) ? group_scores_buf + warp_id * num_groups : nullptr; - int *topk_indices = topk_indices_buf + warp_id * topk; + int *topk_indices = topk_indices_buf + warp_id * selection_topk; uint32_t *local_bitmap_words = (bitmap_words_buf != nullptr) ? bitmap_words_buf + warp_id * bitmap_words_per_warp : nullptr; @@ -411,9 +493,12 @@ __global__ void fused_topk_with_score_function_forward_kernel( * - naive topk * - topk with expert bias */ - // Topk on the scores - // The bias being not empty happens at the sigmoid/sqrtsoftplus case - if (group_topk > 0) { + // Topk on the scores. QB is only supported without grouped Top-k. + // The bias being not empty happens at the sigmoid/sqrtsoftplus case. + if constexpr (kUseQB) { + topk_and_mask(scores, num_experts, selection_topk, topk_indices, topk_scores, + lane_id); + } else if (group_topk > 0) { int group_size = num_experts / num_groups; // Top2 for (int i = 0; i < num_groups; i++) { @@ -461,6 +546,19 @@ __global__ void fused_topk_with_score_function_forward_kernel( } __syncwarp(); + const CompType cutoff = + extract_qb_cutoff_and_compact(topk_indices, topk_scores, topk, lane_id); + if constexpr (QbMode == QBMode::TwoKernel) { + if (lane_id == 0) { + qb_cutoff[token_offset_cur_warp] = cutoff; + } + } else if constexpr (QbMode == QBMode::FusedAtomic) { + accumulate_qb_histogram_epilogue(intermediate_output + pos_offset, cutoff, + num_experts, lane_id, qb_bin_bounds, qb_num_bins, + qb_histogram); + } + __syncwarp(); + /*** * Section: Postprocess * Possible postprocess the scores after the topk operation @@ -539,12 +637,15 @@ __global__ void fused_topk_with_score_function_forward_kernel( } } -template +template void fused_topk_with_score_function_forward_kernel_launcher( const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax, int num_groups, int group_topk, float scaling_factor, int score_function, const BiasType *expert_bias, DataType *probs, uint8_t *routing_map, - CompType *intermediate_output, cudaStream_t stream) { + CompType *intermediate_output, CompType *qb_cutoff, int32_t *qb_histogram, + const CompType *qb_bin_bounds, int qb_num_bins, cudaStream_t stream) { + constexpr bool kUseQB = QbMode != QBMode::Disabled; NVTE_CHECK(num_experts > 0, "num_experts must be positive, got ", num_experts); NVTE_CHECK(topk > 0 && topk <= num_experts, "topk must be in [1, num_experts], got topk=", topk, " num_experts=", num_experts); @@ -569,9 +670,10 @@ void fused_topk_with_score_function_forward_kernel_launcher( } size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp; size_t total_blocks = (num_tokens + num_token_per_block - 1) / num_token_per_block; + size_t selection_topk = topk + (kUseQB ? 1 : 0); size_t scores_shmem = num_experts * num_token_per_block * sizeof(CompType); - size_t scratch_shmem = - topk * num_token_per_block * sizeof(CompType) + topk * num_token_per_block * sizeof(int); + size_t scratch_shmem = selection_topk * num_token_per_block * sizeof(CompType) + + selection_topk * num_token_per_block * sizeof(int); if (group_topk > 0) { scratch_shmem += num_groups * num_token_per_block * sizeof(CompType); scratch_shmem += num_experts * num_token_per_block * sizeof(CompType); @@ -596,7 +698,8 @@ void fused_topk_with_score_function_forward_kernel_launcher( kernel<<>>( logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk, scaling_factor, expert_bias, probs, routing_map, intermediate_output, - static_cast(nullptr), num_buffers); + static_cast(nullptr), num_buffers, qb_cutoff, qb_histogram, qb_bin_bounds, + qb_num_bins); NVTE_CHECK_CUDA(cudaGetLastError()); }; @@ -614,26 +717,29 @@ void fused_topk_with_score_function_forward_kernel_launcher( kernel<<>>( logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, expert_bias, probs, routing_map, intermediate_output, - static_cast(nullptr)); + static_cast(nullptr), qb_cutoff, qb_histogram, qb_bin_bounds, qb_num_bins); NVTE_CHECK_CUDA(cudaGetLastError()); }; launch_simple(fused_topk_forward_simple_kernel); + TopkFuncType::Naive, int32_t, QbMode>); } else { // Optimized path: async loader + persistent grid + radix topk. switch (score_function) { case 0: - launch(fused_topk_with_score_function_forward_kernel); + launch( + fused_topk_with_score_function_forward_kernel); break; case 1: - launch(fused_topk_with_score_function_forward_kernel); + launch( + fused_topk_with_score_function_forward_kernel); break; case 2: - launch(fused_topk_with_score_function_forward_kernel); + launch( + fused_topk_with_score_function_forward_kernel); break; default: NVTE_ERROR("Unsupported score_function: " + std::to_string(score_function)); @@ -641,12 +747,15 @@ void fused_topk_with_score_function_forward_kernel_launcher( } } -template +template void fused_topk_with_score_function_forward_with_indices_kernel_launcher( const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax, int num_groups, int group_topk, float scaling_factor, int score_function, const BiasType *expert_bias, DataType *probs, IndexType *topk_indices, - CompType *intermediate_output, cudaStream_t stream) { + CompType *intermediate_output, CompType *qb_cutoff, int32_t *qb_histogram, + const CompType *qb_bin_bounds, int qb_num_bins, cudaStream_t stream) { + constexpr bool kUseQB = QbMode != QBMode::Disabled; NVTE_CHECK(num_experts > 0, "num_experts must be positive, got ", num_experts); NVTE_CHECK(topk > 0 && topk <= num_experts, "topk must be in [1, num_experts], got topk=", topk, " num_experts=", num_experts); @@ -676,9 +785,10 @@ void fused_topk_with_score_function_forward_with_indices_kernel_launcher( size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp; size_t total_blocks = (num_tokens + num_token_per_block - 1) / num_token_per_block; + size_t selection_topk = topk + (kUseQB ? 1 : 0); size_t scores_shmem = num_experts * num_token_per_block * sizeof(CompType); - size_t scratch_shmem = - topk * num_token_per_block * sizeof(CompType) + topk * num_token_per_block * sizeof(int); + size_t scratch_shmem = selection_topk * num_token_per_block * sizeof(CompType) + + selection_topk * num_token_per_block * sizeof(int); if (group_topk > 0) { scratch_shmem += num_groups * num_token_per_block * sizeof(CompType); scratch_shmem += num_experts * num_token_per_block * sizeof(CompType); @@ -700,7 +810,7 @@ void fused_topk_with_score_function_forward_with_indices_kernel_launcher( kernel<<>>( logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk, scaling_factor, expert_bias, probs, static_cast(nullptr), intermediate_output, - topk_indices, num_buffers); + topk_indices, num_buffers, qb_cutoff, qb_histogram, qb_bin_bounds, qb_num_bins); NVTE_CHECK_CUDA(cudaGetLastError()); }; @@ -715,29 +825,29 @@ void fused_topk_with_score_function_forward_with_indices_kernel_launcher( kernel<<>>( logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, expert_bias, probs, static_cast(nullptr), - intermediate_output, topk_indices); + intermediate_output, topk_indices, qb_cutoff, qb_histogram, qb_bin_bounds, qb_num_bins); NVTE_CHECK_CUDA(cudaGetLastError()); }; launch_simple( fused_topk_forward_simple_kernel); + TopkFuncType::Naive, IndexType, QbMode>); } else { switch (score_function) { case 0: - launch(fused_topk_with_score_function_forward_kernel); + launch(fused_topk_with_score_function_forward_kernel< + DataType, BiasType, NVTE_ROUTING_MAP_FORMAT_BYTEMAP, TopkFuncType::Radix, 0, + IndexType, QbMode>); break; case 1: - launch(fused_topk_with_score_function_forward_kernel); + launch(fused_topk_with_score_function_forward_kernel< + DataType, BiasType, NVTE_ROUTING_MAP_FORMAT_BYTEMAP, TopkFuncType::Radix, 1, + IndexType, QbMode>); break; case 2: - launch(fused_topk_with_score_function_forward_kernel); + launch(fused_topk_with_score_function_forward_kernel< + DataType, BiasType, NVTE_ROUTING_MAP_FORMAT_BYTEMAP, TopkFuncType::Radix, 2, + IndexType, QbMode>); break; default: NVTE_ERROR("Unsupported score_function: " + std::to_string(score_function)); @@ -745,6 +855,90 @@ void fused_topk_with_score_function_forward_with_indices_kernel_launcher( } } +constexpr int kQBExpertsPerBlock = 8; +constexpr int kQBHistogramThreads = 256; + +__global__ void qb_histogram_accumulate_kernel(const CompType *raw_scores, const CompType *cutoff, + int num_tokens, int num_experts, + const CompType *bin_bounds, int num_bins, + int32_t *histogram) { + extern __shared__ int32_t local_histogram[]; + const int local_histogram_size = kQBExpertsPerBlock * num_bins; + for (int i = threadIdx.x; i < local_histogram_size; i += blockDim.x) { + local_histogram[i] = 0; + } + __syncthreads(); + + const int expert_begin = blockIdx.x * kQBExpertsPerBlock; + const int token_partition = blockIdx.y; + const int num_token_partitions = gridDim.y; + const int tokens_in_partition = + (num_tokens - token_partition + num_token_partitions - 1) / num_token_partitions; + const int num_items = tokens_in_partition * kQBExpertsPerBlock; + const CompType lower = bin_bounds[0]; + const CompType upper = bin_bounds[1]; + const CompType scale = static_cast(num_bins) / (upper - lower); + + for (int item = threadIdx.x; item < num_items; item += blockDim.x) { + const int token_in_partition = item / kQBExpertsPerBlock; + const int local_expert = item % kQBExpertsPerBlock; + const int token = token_partition + token_in_partition * num_token_partitions; + const int expert = expert_begin + local_expert; + 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 - lower) * scale)); + bin = max(0, min(bin, num_bins - 1)); + atomicAdd(local_histogram + local_expert * num_bins + bin, 1); + } + } + __syncthreads(); + + for (int i = threadIdx.x; i < local_histogram_size; i += blockDim.x) { + const int count = local_histogram[i]; + const int local_expert = i / num_bins; + const int expert = expert_begin + local_expert; + if (count != 0 && expert < num_experts) { + atomicAdd(histogram + static_cast(expert) * num_bins + i % num_bins, count); + } + } +} + +void qb_histogram_accumulate(const Tensor raw_scores, const Tensor cutoff, const Tensor bin_bounds, + 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"); + NVTE_CHECK(histogram.data.dtype == DType::kInt32, "QB histogram must have int32 dtype"); + NVTE_CHECK(raw_scores.data.shape.size() == 2, + "QB raw_scores must have shape [num_tokens, num_experts]"); + const int num_tokens = static_cast(raw_scores.data.shape[0]); + const int num_experts = static_cast(raw_scores.data.shape[1]); + NVTE_CHECK(cutoff.data.shape == std::vector{static_cast(num_tokens)}, + "QB cutoff must have shape [num_tokens]"); + NVTE_CHECK(bin_bounds.data.shape == std::vector{2}, "QB bin_bounds must have shape [2]"); + NVTE_CHECK(histogram.data.shape.size() == 2 && + histogram.data.shape[0] == static_cast(num_experts), + "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"); + + const int token_partitions = std::min(4, std::max(1, (num_tokens + 1023) / 1024)); + const dim3 grid((num_experts + kQBExpertsPerBlock - 1) / kQBExpertsPerBlock, token_partitions); + const size_t shared_memory_size = + static_cast(kQBExpertsPerBlock) * num_bins * sizeof(int32_t); + check_shared_memory_capacity_num_experts(shared_memory_size, num_experts); + NVTE_CHECK_CUDA(cudaFuncSetAttribute(qb_histogram_accumulate_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(shared_memory_size))); + qb_histogram_accumulate_kernel<<>>( + reinterpret_cast(raw_scores.data.dptr), + reinterpret_cast(cutoff.data.dptr), num_tokens, num_experts, + reinterpret_cast(bin_bounds.data.dptr), num_bins, + reinterpret_cast(histogram.data.dptr)); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + // Build the expected routing_map shape for a given NVTERoutingMapFormat. // BYTEMAP -> [num_tokens, num_experts] // BITMAP_U8 -> [num_tokens, ceil(num_experts/8)] @@ -801,7 +995,8 @@ void fused_topk_with_score_function_forward(const Tensor logits, int num_tokens, reinterpret_cast(expert_bias.data.dptr), \ reinterpret_cast(probs.data.dptr), \ reinterpret_cast(routing_map.data.dptr), \ - reinterpret_cast(intermediate_output.data.dptr), stream);); \ + reinterpret_cast(intermediate_output.data.dptr), nullptr, nullptr, \ + nullptr, 0, stream);); \ } else { \ fused_topk_with_score_function_forward_kernel_launcher( \ @@ -809,7 +1004,8 @@ void fused_topk_with_score_function_forward(const Tensor logits, int num_tokens, use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, nullptr, \ reinterpret_cast(probs.data.dptr), \ reinterpret_cast(routing_map.data.dptr), \ - reinterpret_cast(intermediate_output.data.dptr), stream); \ + reinterpret_cast(intermediate_output.data.dptr), nullptr, nullptr, \ + nullptr, 0, stream); \ }); if (routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) { ROUTER_FORWARD_DISPATCH(NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) @@ -851,26 +1047,28 @@ void fused_topk_with_score_function_forward_with_indices( } // Dispatch logits dtype and output-index dtype first; expert-bias dtype is only // dispatched when an expert-bias tensor exists, otherwise the kernel receives nullptr. -#define ROUTER_FORWARD_WITH_INDICES_DISPATCH(DataType, IndexType) \ - if (expert_bias.has_data()) { \ - TE_ROUTER_PROBS_TYPE_SWITCH_ALL( \ - expert_bias.data.dtype, BiasType, \ - fused_topk_with_score_function_forward_with_indices_kernel_launcher( \ - reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, \ - use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, \ - reinterpret_cast(expert_bias.data.dptr), \ - reinterpret_cast(probs.data.dptr), \ - reinterpret_cast(topk_indices.data.dptr), \ - reinterpret_cast(intermediate_output.data.dptr), stream);); \ - } else { \ - fused_topk_with_score_function_forward_with_indices_kernel_launcher( \ - reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, \ - use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, nullptr, \ - reinterpret_cast(probs.data.dptr), \ - reinterpret_cast(topk_indices.data.dptr), \ - reinterpret_cast(intermediate_output.data.dptr), stream); \ +#define ROUTER_FORWARD_WITH_INDICES_DISPATCH(DataType, IndexType) \ + if (expert_bias.has_data()) { \ + TE_ROUTER_PROBS_TYPE_SWITCH_ALL( \ + expert_bias.data.dtype, BiasType, \ + fused_topk_with_score_function_forward_with_indices_kernel_launcher( \ + reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, \ + use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, \ + reinterpret_cast(expert_bias.data.dptr), \ + reinterpret_cast(probs.data.dptr), \ + reinterpret_cast(topk_indices.data.dptr), \ + reinterpret_cast(intermediate_output.data.dptr), nullptr, nullptr, \ + nullptr, 0, stream);); \ + } else { \ + fused_topk_with_score_function_forward_with_indices_kernel_launcher( \ + reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, \ + use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, nullptr, \ + reinterpret_cast(probs.data.dptr), \ + reinterpret_cast(topk_indices.data.dptr), \ + reinterpret_cast(intermediate_output.data.dptr), nullptr, nullptr, nullptr, 0, \ + stream); \ } TE_ROUTER_PROBS_TYPE_SWITCH_ALL(logits.data.dtype, DataType, TE_ROUTER_DENSE_INDEX_TYPE_SWITCH_ALL( @@ -879,6 +1077,130 @@ void fused_topk_with_score_function_forward_with_indices( #undef ROUTER_FORWARD_WITH_INDICES_DISPATCH } +static int check_qb_forward_tensors(const Tensor logits, int num_tokens, int num_experts, int topk, + const Tensor expert_bias, const Tensor probs, + const Tensor intermediate_output, const Tensor cutoff, + const Tensor histogram, const Tensor bin_bounds) { + NVTE_CHECK(num_tokens > 0 && num_experts > 0, "QB num_tokens and num_experts must be positive"); + NVTE_CHECK(topk > 0 && topk < num_experts, "QB topk must be in [1, num_experts), got topk=", topk, + " num_experts=", num_experts); + const std::vector dense_shape{static_cast(num_tokens), + static_cast(num_experts)}; + NVTE_CHECK(logits.data.shape == dense_shape, + "QB logits must have shape [num_tokens, num_experts]"); + NVTE_CHECK(probs.data.shape == dense_shape, "QB probs must have shape [num_tokens, num_experts]"); + NVTE_CHECK(intermediate_output.data.shape == dense_shape, + "QB intermediate_output must have shape [num_tokens, num_experts]"); + NVTE_CHECK(intermediate_output.data.dtype == DType::kFloat32, + "QB intermediate_output must have FP32 dtype"); + NVTE_CHECK(expert_bias.has_data(), "QB requires an expert_bias tensor"); + NVTE_CHECK(expert_bias.data.dtype == DType::kFloat32, "QB expert_bias must have FP32 dtype"); + NVTE_CHECK(expert_bias.data.shape == std::vector{static_cast(num_experts)}, + "QB expert_bias must have shape [num_experts]"); + NVTE_CHECK(cutoff.data.dtype == DType::kFloat32, "QB cutoff must have FP32 dtype"); + NVTE_CHECK(cutoff.data.shape == std::vector{static_cast(num_tokens)}, + "QB cutoff must have shape [num_tokens]"); + NVTE_CHECK(histogram.data.dtype == DType::kInt32, "QB histogram must have int32 dtype"); + NVTE_CHECK(histogram.data.shape.size() == 2 && + histogram.data.shape[0] == static_cast(num_experts), + "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"); + NVTE_CHECK(bin_bounds.data.dtype == DType::kFloat32, "QB bin_bounds must have FP32 dtype"); + NVTE_CHECK(bin_bounds.data.shape == std::vector{2}, "QB bin_bounds must have shape [2]"); + return num_bins; +} + +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, 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, + intermediate_output, cutoff, histogram, bin_bounds); + const auto routing_map_shape = + 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"); + +#define QB_ROUTER_FORWARD_LAUNCH(RoutingMapFormatVal, QbModeVal) \ + TE_ROUTER_PROBS_TYPE_SWITCH_ALL( \ + logits.data.dtype, DataType, \ + fused_topk_with_score_function_forward_kernel_launcher( \ + reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, false, \ + -1, -1, scaling_factor, 0, reinterpret_cast(expert_bias.data.dptr), \ + reinterpret_cast(probs.data.dptr), \ + reinterpret_cast(routing_map.data.dptr), \ + reinterpret_cast(intermediate_output.data.dptr), \ + reinterpret_cast(cutoff.data.dptr), \ + reinterpret_cast(histogram.data.dptr), \ + reinterpret_cast(bin_bounds.data.dptr), num_bins, stream);); + +#define QB_ROUTER_FORWARD_MODE_DISPATCH(RoutingMapFormatVal) \ + if (histogram_mode == NVTE_QB_HISTOGRAM_TWO_KERNEL) { \ + QB_ROUTER_FORWARD_LAUNCH(RoutingMapFormatVal, QBMode::TwoKernel) \ + } else if (histogram_mode == NVTE_QB_HISTOGRAM_FUSED_ATOMIC) { \ + QB_ROUTER_FORWARD_LAUNCH(RoutingMapFormatVal, QBMode::FusedAtomic) \ + } else { \ + NVTE_ERROR("Unsupported QB histogram mode: " + \ + std::to_string(static_cast(histogram_mode))); \ + } + + if (routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) { + QB_ROUTER_FORWARD_MODE_DISPATCH(NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) + } else { + QB_ROUTER_FORWARD_MODE_DISPATCH(NVTE_ROUTING_MAP_FORMAT_BYTEMAP) + } +#undef QB_ROUTER_FORWARD_MODE_DISPATCH +#undef QB_ROUTER_FORWARD_LAUNCH +} + +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, + 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); + const std::vector indices_shape{static_cast(num_tokens), + static_cast(topk)}; + NVTE_CHECK(topk_indices.data.shape == indices_shape, + "QB topk_indices must have shape [num_tokens, topk]"); + +#define QB_ROUTER_FORWARD_WITH_INDICES_LAUNCH(DataType, IndexType, QbModeVal) \ + fused_topk_with_score_function_forward_with_indices_kernel_launcher( \ + reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, false, -1, \ + -1, scaling_factor, 0, reinterpret_cast(expert_bias.data.dptr), \ + reinterpret_cast(probs.data.dptr), \ + reinterpret_cast(topk_indices.data.dptr), \ + reinterpret_cast(intermediate_output.data.dptr), \ + reinterpret_cast(cutoff.data.dptr), \ + reinterpret_cast(histogram.data.dptr), \ + reinterpret_cast(bin_bounds.data.dptr), num_bins, stream) + +#define QB_ROUTER_FORWARD_WITH_INDICES_MODE(DataType, IndexType) \ + if (histogram_mode == NVTE_QB_HISTOGRAM_TWO_KERNEL) { \ + QB_ROUTER_FORWARD_WITH_INDICES_LAUNCH(DataType, IndexType, QBMode::TwoKernel); \ + } else if (histogram_mode == NVTE_QB_HISTOGRAM_FUSED_ATOMIC) { \ + QB_ROUTER_FORWARD_WITH_INDICES_LAUNCH(DataType, IndexType, QBMode::FusedAtomic); \ + } else { \ + NVTE_ERROR("Unsupported QB histogram mode: " + \ + std::to_string(static_cast(histogram_mode))); \ + } + + TE_ROUTER_PROBS_TYPE_SWITCH_ALL(logits.data.dtype, DataType, + TE_ROUTER_DENSE_INDEX_TYPE_SWITCH_ALL( + topk_indices.data.dtype, IndexType, + QB_ROUTER_FORWARD_WITH_INDICES_MODE(DataType, IndexType););); +#undef QB_ROUTER_FORWARD_WITH_INDICES_MODE +#undef QB_ROUTER_FORWARD_WITH_INDICES_LAUNCH +} + // Backward: grad_probs + intermediate_output + routing_map → grad_logits. // // Double-buffered cp.async loads all 3 inputs in original types. Two-pass @@ -1407,6 +1729,48 @@ void nvte_fused_topk_with_score_function_forward_with_indices( *convertNVTETensorCheck(topk_indices), *convertNVTETensorCheck(intermediate_output), stream); } +void nvte_fused_topk_with_score_function_forward_qb_v2( + const NVTETensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, + const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map, + NVTERoutingMapFormat routing_map_format, NVTETensor intermediate_output, NVTETensor cutoff, + NVTETensor histogram, NVTETensor bin_bounds, NVTEQBHistogramMode histogram_mode, + cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_topk_with_score_function_forward_qb_v2); + using namespace transformer_engine; + fused_router::fused_topk_with_score_function_forward_qb( + *convertNVTETensorCheck(logits), num_tokens, num_experts, topk, scaling_factor, + *convertNVTETensorCheck(expert_bias), *convertNVTETensorCheck(probs), + *convertNVTETensorCheck(routing_map), routing_map_format, + *convertNVTETensorCheck(intermediate_output), *convertNVTETensorCheck(cutoff), + *convertNVTETensorCheck(histogram), *convertNVTETensorCheck(bin_bounds), histogram_mode, + stream); +} + +void nvte_fused_topk_with_score_function_forward_qb_with_indices( + const NVTETensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, + const NVTETensor expert_bias, NVTETensor probs, NVTETensor topk_indices, + NVTETensor intermediate_output, NVTETensor cutoff, NVTETensor histogram, NVTETensor bin_bounds, + NVTEQBHistogramMode histogram_mode, cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_topk_with_score_function_forward_qb_with_indices); + using namespace transformer_engine; + fused_router::fused_topk_with_score_function_forward_qb_with_indices( + *convertNVTETensorCheck(logits), num_tokens, num_experts, topk, scaling_factor, + *convertNVTETensorCheck(expert_bias), *convertNVTETensorCheck(probs), + *convertNVTETensorCheck(topk_indices), *convertNVTETensorCheck(intermediate_output), + *convertNVTETensorCheck(cutoff), *convertNVTETensorCheck(histogram), + *convertNVTETensorCheck(bin_bounds), histogram_mode, stream); +} + +void nvte_qb_histogram_accumulate(const NVTETensor raw_scores, const NVTETensor cutoff, + const NVTETensor bin_bounds, NVTETensor histogram, + cudaStream_t stream) { + NVTE_API_CALL(nvte_qb_histogram_accumulate); + using namespace transformer_engine; + fused_router::qb_histogram_accumulate( + *convertNVTETensorCheck(raw_scores), *convertNVTETensorCheck(cutoff), + *convertNVTETensorCheck(bin_bounds), *convertNVTETensorCheck(histogram), stream); +} + void nvte_fused_topk_with_score_function_backward_v2(const NVTETensor routing_map, NVTERoutingMapFormat routing_map_format, const NVTETensor intermediate_output, diff --git a/transformer_engine/common/include/transformer_engine/fused_router.h b/transformer_engine/common/include/transformer_engine/fused_router.h index 03fec2eb2c..b90c80f901 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -27,6 +27,17 @@ typedef enum { NVTE_ROUTING_MAP_FORMAT_BITMAP_U8 = 1, } NVTERoutingMapFormat; +/*! \brief Quantile Balancing histogram implementation. + * + * TWO_KERNEL — router writes the Top-(k+1) cutoff and a second kernel accumulates + * the histogram from the saved raw sigmoid scores. + * FUSED_ATOMIC — router directly accumulates the histogram with global atomics. + */ +typedef enum { + NVTE_QB_HISTOGRAM_TWO_KERNEL = 0, + NVTE_QB_HISTOGRAM_FUSED_ATOMIC = 1, +} NVTEQBHistogramMode; + /*! \brief Apply topk + softmax/sigmoid to the input tensor. Grouped topk is supported (deprecated). * * \deprecated This function has been deprecated in favor of @@ -93,6 +104,31 @@ void nvte_fused_topk_with_score_function_forward_with_indices( const NVTETensor expert_bias, NVTETensor probs, NVTETensor topk_indices, NVTETensor intermediate_output, cudaStream_t stream); +/*! \brief Kimi K3 Quantile Balancing fused-router forward. + * + * The router selects Top-(k+1) from biased sigmoid scores, writes only the first + * k routes, and exposes the final selected value as the per-token cutoff. + * In FUSED_ATOMIC mode it also directly accumulates the QB histogram. + */ +void nvte_fused_topk_with_score_function_forward_qb_v2( + const NVTETensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, + const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map, + NVTERoutingMapFormat routing_map_format, NVTETensor intermediate_output, NVTETensor cutoff, + NVTETensor histogram, NVTETensor bin_bounds, NVTEQBHistogramMode histogram_mode, + cudaStream_t stream); + +/*! \brief Kimi K3 Quantile Balancing fused-router forward with dense Top-k indices. */ +void nvte_fused_topk_with_score_function_forward_qb_with_indices( + const NVTETensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, + const NVTETensor expert_bias, NVTETensor probs, NVTETensor topk_indices, + NVTETensor intermediate_output, NVTETensor cutoff, NVTETensor histogram, NVTETensor bin_bounds, + NVTEQBHistogramMode histogram_mode, cudaStream_t stream); + +/*! \brief Accumulate a QB histogram from raw sigmoid scores and Top-(k+1) cutoffs. */ +void nvte_qb_histogram_accumulate(const NVTETensor raw_scores, const NVTETensor cutoff, + const NVTETensor bin_bounds, NVTETensor histogram, + cudaStream_t stream); + /*! \brief Backward pass for fused topk + softmax/sigmoid (deprecated). * * \deprecated This function has been deprecated in favor of diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 9114b1e453..f731016f26 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -38,6 +38,13 @@ std::tuple fused_topk_with_score_function_fw int routing_map_format = static_cast(NVTE_ROUTING_MAP_FORMAT_BYTEMAP), std::optional topk_indices = std::nullopt); +std::tuple +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); + void fused_topk_with_score_function_bwd( at::Tensor routing_map, at::Tensor intermediate_output, at::Tensor grad_probs, at::Tensor grad_logits, int topk, bool use_pre_softmax, std::optional scaling_factor, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 7c8793b0f6..017b0389ad 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -137,12 +137,20 @@ void init_router_bindings(pybind11::module &m) { pybind11::enum_(m, "NVTERoutingMapFormat", pybind11::module_local()) .value("BYTEMAP", NVTE_ROUTING_MAP_FORMAT_BYTEMAP) .value("BITMAP_U8", NVTE_ROUTING_MAP_FORMAT_BITMAP_U8); + pybind11::enum_(m, "NVTEQBHistogramMode", pybind11::module_local()) + .value("TWO_KERNEL", NVTE_QB_HISTOGRAM_TWO_KERNEL) + .value("FUSED_ATOMIC", NVTE_QB_HISTOGRAM_FUSED_ATOMIC); m.def("fused_topk_with_score_function_fwd", &fused_topk_with_score_function_fwd, py::arg("logits"), py::arg("topk"), py::arg("use_pre_softmax"), py::arg("num_groups"), py::arg("group_topk"), py::arg("scaling_factor"), py::arg("score_function"), py::arg("expert_bias"), py::arg("routing_map_format") = static_cast(NVTE_ROUTING_MAP_FORMAT_BYTEMAP), py::arg("topk_indices") = std::nullopt, "Fused topk with score function fwd"); + 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"), + "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"), py::arg("grad_logits"), py::arg("topk"), py::arg("use_pre_softmax"), diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index d59d6bc415..473607ac5e 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -154,6 +154,109 @@ std::tuple fused_topk_with_score_function_fw return std::make_tuple(probs, routing_map, intermediate_output); } +std::tuple +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) { + 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(), + "logits must be a contiguous CUDA tensor"); + const auto sizes = logits.sizes(); + const int64_t num_experts = sizes.back(); + const int64_t num_tokens = + std::accumulate(sizes.begin(), sizes.end() - 1, int64_t{1}, std::multiplies()); + TORCH_CHECK(num_tokens > 0 && num_experts > 0, + "num_tokens and num_experts must be greater than 0"); + TORCH_CHECK(topk > 0 && topk < num_experts, + "QB topk must be in [1, num_experts), got topk=", topk, " num_experts=", num_experts); + TORCH_CHECK(expert_bias.is_cuda() && expert_bias.is_contiguous(), + "QB expert_bias must be a contiguous CUDA tensor"); + TORCH_CHECK(expert_bias.device() == logits.device(), + "QB expert_bias must be on the logits device"); + TORCH_CHECK(expert_bias.scalar_type() == at::kFloat, "QB expert_bias must have float32 dtype"); + TORCH_CHECK(expert_bias.dim() == 1 && expert_bias.numel() == num_experts, + "QB expert_bias must have shape [num_experts]"); + TORCH_CHECK(histogram.is_cuda() && histogram.is_contiguous(), + "QB histogram must be a contiguous CUDA tensor"); + TORCH_CHECK(histogram.device() == logits.device(), "QB histogram must be on the logits device"); + TORCH_CHECK(histogram.scalar_type() == at::kInt, "QB histogram must have int32 dtype"); + TORCH_CHECK(histogram.dim() == 2 && histogram.size(0) == num_experts && histogram.size(1) > 0, + "QB histogram must have shape [num_experts, num_bins]"); + TORCH_CHECK(bin_bounds.is_cuda() && bin_bounds.is_contiguous(), + "QB bin_bounds must be a contiguous CUDA tensor"); + TORCH_CHECK(bin_bounds.device() == logits.device(), "QB bin_bounds must be on the logits device"); + TORCH_CHECK( + bin_bounds.scalar_type() == at::kFloat && bin_bounds.dim() == 1 && bin_bounds.numel() == 2, + "QB bin_bounds must be float32 with shape [2]"); + TORCH_CHECK(histogram_mode == NVTE_QB_HISTOGRAM_TWO_KERNEL || + histogram_mode == NVTE_QB_HISTOGRAM_FUSED_ATOMIC, + "Unsupported QB histogram mode: ", histogram_mode); + if (topk_indices.has_value()) { + TORCH_CHECK(routing_map_format == NVTE_ROUTING_MAP_FORMAT_BYTEMAP, + "dense Top-k indices cannot be combined with a non-default routing-map format"); + check_dense_topk_indices(topk_indices.value(), logits, sizes.slice(0, sizes.size() - 1), topk); + } + + const float scaling_factor_value = scaling_factor.has_value() ? scaling_factor.value() : 1.0f; + at::Tensor probs = at::empty(sizes, at::dtype(logits.scalar_type()).device(logits.device())); + at::Tensor routing_output = + topk_indices.has_value() + ? topk_indices.value() + : allocate_routing_map(sizes.slice(0, sizes.size() - 1), num_experts, routing_map_format); + at::Tensor intermediate_output = at::empty(sizes, at::dtype(at::kFloat).device(logits.device())); + at::Tensor cutoff = at::empty({num_tokens}, at::dtype(at::kFloat).device(logits.device())); + + const std::vector shape_2d = {static_cast(num_tokens), + static_cast(num_experts)}; + const std::vector routing_output_shape_2d = + topk_indices.has_value() + ? std::vector{static_cast(num_tokens), static_cast(topk)} + : std::vector{ + static_cast(num_tokens), + static_cast(routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8 + ? (num_experts + 7) / 8 + : num_experts)}; + auto logits_cu = makeTransformerEngineTensor(logits.data_ptr(), shape_2d, + GetTransformerEngineDType(logits.scalar_type())); + auto probs_cu = makeTransformerEngineTensor(probs.data_ptr(), shape_2d, + GetTransformerEngineDType(probs.scalar_type())); + auto routing_output_cu = + makeTransformerEngineTensor(routing_output.data_ptr(), routing_output_shape_2d, + GetTransformerEngineDType(routing_output.scalar_type())); + auto intermediate_output_cu = + makeTransformerEngineTensor(intermediate_output.data_ptr(), shape_2d, DType::kFloat32); + const std::vector cutoff_shape = {static_cast(num_tokens)}; + auto cutoff_cu = makeTransformerEngineTensor(cutoff.data_ptr(), cutoff_shape, DType::kFloat32); + auto expert_bias_cu = makeTransformerEngineTensor(expert_bias); + auto histogram_cu = makeTransformerEngineTensor(histogram); + auto bin_bounds_cu = makeTransformerEngineTensor(bin_bounds); + const auto mode = static_cast(histogram_mode); + auto stream = at::cuda::getCurrentCUDAStream(); + + if (topk_indices.has_value()) { + 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 { + 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(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); +} + void fused_topk_with_score_function_bwd(at::Tensor routing_map, at::Tensor intermediate_output, at::Tensor grad_probs, at::Tensor grad_logits, int topk, bool use_pre_softmax, std::optional scaling_factor, diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 51350ab2fd..9d9f9d4a76 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -13,12 +13,12 @@ - Only cast to low-precision when necessary and the casting only happens in writing to global memory. For example, the gradient is required to have the same dtype as the input. """ + from typing import Optional, Union import torch import transformer_engine_torch as tex - # Re-export the C++ enum NVTERoutingMapFormat under a friendlier Python name. # Members: # RoutingMapFormat.BYTEMAP — bool[num_tokens, num_experts] @@ -26,6 +26,7 @@ # LSB-first / little-endian packing along the # expert axis. RoutingMapFormat = tex.NVTERoutingMapFormat +QBHistogramMode = tex.NVTEQBHistogramMode _ROUTING_MAP_FORMAT_FROM_STRING = { @@ -33,9 +34,15 @@ "bitmap_u8": int(RoutingMapFormat.BITMAP_U8), } _VALID_ROUTING_MAP_FORMAT_INTS = frozenset(_ROUTING_MAP_FORMAT_FROM_STRING.values()) +_QB_HISTOGRAM_MODE_FROM_STRING = { + "two_kernel": int(QBHistogramMode.TWO_KERNEL), + "fused_atomic": int(QBHistogramMode.FUSED_ATOMIC), +} -def _validate_routing_map_format(routing_map_format: Union[str, RoutingMapFormat, int]) -> int: +def _validate_routing_map_format( + routing_map_format: Union[str, RoutingMapFormat, int], +) -> int: """Coerce user-supplied routing_map_format into a plain int (0 or 1). Accepts the enum, an int matching one of the enum's values, or the @@ -134,6 +141,74 @@ def backward(ctx, grad_probs, _): return grad_logits, None, None, None, None, None, None, None, None, None +class FusedTopkScoreFunctionQB(torch.autograd.Function): + """Kimi K3 QB router with histogram accumulation.""" + + @staticmethod + def forward( + ctx, + logits: torch.Tensor, + topk: int, + scaling_factor: Optional[float], + expert_bias: torch.Tensor, + routing_map_format: int, + topk_indices: Optional[torch.Tensor], + histogram: torch.Tensor, + bin_bounds: torch.Tensor, + histogram_mode: int, + ): + # pylint: disable=missing-function-docstring + ( + probs, + routing_output, + intermediate_output, + _cutoff, + histogram_output, + ) = tex.fused_topk_with_score_function_qb_fwd( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + histogram, + bin_bounds, + histogram_mode, + ) + if topk_indices is not None: + routing_output = topk_indices + ctx.mark_dirty(topk_indices) + ctx.mark_dirty(histogram) + ctx.mark_non_differentiable(routing_output, histogram_output) + ctx.save_for_backward(routing_output, intermediate_output) + ctx.topk = topk + ctx.scaling_factor = scaling_factor + ctx.routing_map_format = routing_map_format + ctx.use_dense_indices = topk_indices is not None + return probs, routing_output, histogram_output + + @staticmethod + def backward(ctx, grad_probs, _, _histogram_grad): + # pylint: disable=missing-function-docstring + routing_output, intermediate_output = ctx.saved_tensors + if not grad_probs.is_contiguous(): + grad_probs = grad_probs.contiguous() + grad_logits = torch.empty_like(grad_probs) + tex.fused_topk_with_score_function_bwd( + routing_output, + intermediate_output, + grad_probs, + grad_logits, + ctx.topk, + False, + ctx.scaling_factor, + "sigmoid", + ctx.use_dense_indices, + ctx.routing_map_format, + ) + return grad_logits, None, None, None, None, None, None, None, None + + def fused_topk_with_score_function( logits: torch.Tensor, topk: int, @@ -145,6 +220,9 @@ def fused_topk_with_score_function( expert_bias: Optional[torch.Tensor], routing_map_format: Union[str, RoutingMapFormat, int] = RoutingMapFormat.BYTEMAP, topk_indices: Optional[torch.Tensor] = None, + qb_histogram: Optional[torch.Tensor] = None, + qb_bin_bounds: Optional[torch.Tensor] = None, + qb_histogram_mode: Optional[str] = None, ): """ Fused topk with score function router. @@ -172,6 +250,12 @@ def fused_topk_with_score_function( topk_indices : torch.Tensor, optional Optional output buffer with shape [num_tokens, topk]. When provided, its dtype controls the dense index output dtype and the routing map is not materialized. + qb_histogram : torch.Tensor, optional + 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. + qb_histogram_mode : str, optional + ``"two_kernel"`` or ``"fused_atomic"``. Must be provided with the two QB tensors. Returns ------- @@ -187,6 +271,36 @@ def fused_topk_with_score_function( if logits.dtype == torch.float64: raise ValueError("Current TE does not support float64 router type.") routing_map_format = _validate_routing_map_format(routing_map_format) + qb_arguments = (qb_histogram, qb_bin_bounds, qb_histogram_mode) + if any(value is not None for value in qb_arguments): + if any(value is None for value in qb_arguments): + raise ValueError( + "qb_histogram, qb_bin_bounds, and qb_histogram_mode must be provided together" + ) + if score_function != "sigmoid": + raise ValueError("Quantile Balancing only supports score_function='sigmoid'") + if expert_bias is None: + raise ValueError("Quantile Balancing requires expert_bias") + if num_groups is not None or group_topk is not None: + raise ValueError("Quantile Balancing does not support grouped Top-k") + mode = _QB_HISTOGRAM_MODE_FROM_STRING.get(qb_histogram_mode) + if mode is None: + raise ValueError( + "qb_histogram_mode must be 'two_kernel' or 'fused_atomic', " + f"got {qb_histogram_mode!r}" + ) + probs, routing_output, _ = FusedTopkScoreFunctionQB.apply( + logits, + topk, + scaling_factor, + expert_bias, + routing_map_format, + topk_indices, + qb_histogram, + qb_bin_bounds, + mode, + ) + return probs, routing_output return FusedTopkScoreFunction.apply( logits, topk, From ef25baaa957cde5a99f756ccee69eea91af228ac Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 18 Aug 2026 16:15:39 +0800 Subject: [PATCH 2/4] [Common][PyTorch] Harden QB router validation Signed-off-by: Harry Zhou --- tests/pytorch/test_fused_router.py | 80 +++++++++++++++++++ .../fused_topk_with_score_function.cu | 29 +++++-- .../pytorch/csrc/extensions/router.cpp | 1 + 3 files changed, 102 insertions(+), 8 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 36ca415307..206fb908bb 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -1,6 +1,10 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. +import os +import subprocess +import sys + import torch from typing import Optional from transformer_engine.pytorch.router import ( @@ -560,6 +564,82 @@ def test_qb_topk_argument_validation(): ) +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least two CUDA devices") +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +def test_qb_topk_uses_logits_device(histogram_mode): + current_device = torch.cuda.current_device() + logits_device = (current_device + 1) % torch.cuda.device_count() + device = torch.device("cuda", logits_device) + num_tokens, num_experts, topk, num_bins = 17, 32, 4, 64 + logits = torch.randn(num_tokens, num_experts, device=device, dtype=torch.float32) + expert_bias = torch.zeros(num_experts, device=device, dtype=torch.float32) + histogram = torch.zeros(num_experts, num_bins, device=device, dtype=torch.int32) + bin_bounds = torch.tensor([-1.0, 1.0], device=device, dtype=torch.float32) + + probs, routing_map = fused_topk_with_score_function( + logits, + topk, + False, + None, + None, + None, + "sigmoid", + expert_bias, + qb_histogram=histogram, + qb_bin_bounds=bin_bounds, + qb_histogram_mode=histogram_mode, + ) + torch.cuda.synchronize(logits_device) + + assert probs.device == device + assert routing_map.device == device + assert histogram.sum().item() == num_tokens * num_experts + assert torch.cuda.current_device() == current_device + + +@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): + script = 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) +bounds = {{ + "equal": [1.0, 1.0], + "reversed": [1.0, -1.0], + "nonfinite": [float("nan"), 1.0], +}}[{invalid_bounds!r}] +bin_bounds = torch.tensor(bounds, device="cuda", dtype=torch.float32) +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}, +) +torch.cuda.synchronize() +""" + result = subprocess.run( + [sys.executable, "-c", script], + env=os.environ.copy(), + capture_output=True, + text=True, + check=False, + ) + output = result.stdout + result.stderr + assert result.returncode != 0, output + assert "CUDA error" in output, output + + @pytest.mark.parametrize( "histogram_mode", [QBHistogramMode.TWO_KERNEL, QBHistogramMode.FUSED_ATOMIC], 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 9f867008ed..587339e4d9 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 @@ -28,6 +28,22 @@ enum class QBMode { FusedAtomic, }; +struct QBBinParams { + CompType lower; + CompType scale; +}; + +__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]; + if (!isfinite(lower) || !isfinite(upper) || upper <= lower) { + // Bounds are caller-owned CUDA data, so host-side value validation would add a + // synchronization and break CUDA graph capture. Trap on the consuming stream instead. + __trap(); + } + return {lower, static_cast(num_bins) / (upper - lower)}; +} + template __device__ inline CompType extract_qb_cutoff_and_compact(int *topk_indices, CompType *topk_scores, int topk, int lane_id) { @@ -64,11 +80,10 @@ __device__ inline void accumulate_qb_histogram_epilogue(const CompType *raw_scor const CompType *bin_bounds, int num_bins, int32_t *histogram) { if constexpr (Mode == QBMode::FusedAtomic) { - const CompType lower = bin_bounds[0]; - const CompType upper = bin_bounds[1]; - const CompType scale = static_cast(num_bins) / (upper - lower); + 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] - lower) * scale)); + int bin = static_cast( + floorf((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); } @@ -875,9 +890,7 @@ __global__ void qb_histogram_accumulate_kernel(const CompType *raw_scores, const const int tokens_in_partition = (num_tokens - token_partition + num_token_partitions - 1) / num_token_partitions; const int num_items = tokens_in_partition * kQBExpertsPerBlock; - const CompType lower = bin_bounds[0]; - const CompType upper = bin_bounds[1]; - const CompType scale = static_cast(num_bins) / (upper - lower); + const QBBinParams bin_params = load_qb_bin_params(bin_bounds, num_bins); for (int item = threadIdx.x; item < num_items; item += blockDim.x) { const int token_in_partition = item / kQBExpertsPerBlock; @@ -887,7 +900,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 - lower) * scale)); + int bin = static_cast(floorf((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); } diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 473607ac5e..79391b5bf2 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -164,6 +164,7 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, TORCH_CHECK(logits.dim() >= 1, "logits must have at least 1 dim"); TORCH_CHECK(logits.is_cuda() && logits.is_contiguous(), "logits must be a contiguous CUDA tensor"); + at::cuda::CUDAGuard device_guard(logits.device()); const auto sizes = logits.sizes(); const int64_t num_experts = sizes.back(); const int64_t num_tokens = From e363e628bdc0a37ab4b34ec39fbdbcee7a6185b2 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 18 Aug 2026 17:01:24 +0800 Subject: [PATCH 3/4] [Common][PyTorch] Make QB bounds validation recoverable Signed-off-by: Harry Zhou --- tests/pytorch/test_fused_router.py | 164 +++++++++++++----- .../fused_topk_with_score_function.cu | 5 - .../include/transformer_engine/fused_router.h | 11 +- transformer_engine/pytorch/csrc/extensions.h | 3 +- .../pytorch/csrc/extensions/pybind.cpp | 2 +- .../pytorch/csrc/extensions/router.cpp | 12 +- transformer_engine/pytorch/router.py | 37 +++- 7 files changed, 182 insertions(+), 52 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 206fb908bb..e00bbe0ab0 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -1,10 +1,6 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -import os -import subprocess -import sys - import torch from typing import Optional from transformer_engine.pytorch.router import ( @@ -600,44 +596,130 @@ 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): - script = 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) -bounds = {{ - "equal": [1.0, 1.0], - "reversed": [1.0, -1.0], - "nonfinite": [float("nan"), 1.0], -}}[{invalid_bounds!r}] -bin_bounds = torch.tensor(bounds, device="cuda", dtype=torch.float32) -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}, -) -torch.cuda.synchronize() -""" - result = subprocess.run( - [sys.executable, "-c", script], - env=os.environ.copy(), - capture_output=True, - text=True, - check=False, + 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, + ) + + +def test_qb_topk_revalidates_updated_bin_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) + bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) + 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="fused_atomic", + ) + 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="fused_atomic", + ) + + +def test_qb_raw_binding_rejects_invalid_bin_bounds_recoverably(): + 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) + 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), + None, + histogram, + invalid_bounds, + QBHistogramMode.FUSED_ATOMIC, + ) + + # The validation error must not poison the CUDA context. + 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), + None, + histogram, + valid_bounds, + QBHistogramMode.FUSED_ATOMIC, ) - output = result.stdout + result.stderr - assert result.returncode != 0, output - assert "CUDA error" in output, output + torch.cuda.synchronize() + + +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +def test_qb_topk_cuda_graph_uses_prevalidated_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) + + 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): + probs, routing_map = run_router() + graph.replay() + torch.cuda.synchronize() + assert torch.isfinite(probs).all() + assert routing_map.sum().item() == logits.shape[0] * 4 @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 587339e4d9..9fd6250a6f 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,11 +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]; - if (!isfinite(lower) || !isfinite(upper) || upper <= lower) { - // Bounds are caller-owned CUDA data, so host-side value validation would add a - // synchronization and break CUDA graph capture. Trap on the consuming stream instead. - __trap(); - } 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 b90c80f901..5846247971 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -109,6 +109,7 @@ void nvte_fused_topk_with_score_function_forward_with_indices( * The router selects Top-(k+1) from biased sigmoid scores, writes only the first * k routes, and exposes the final selected value as the per-token cutoff. * In FUSED_ATOMIC mode it also directly accumulates the QB histogram. + * bin_bounds must contain finite FP32 values [lower, upper] with lower < upper. */ void nvte_fused_topk_with_score_function_forward_qb_v2( const NVTETensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, @@ -117,14 +118,20 @@ void nvte_fused_topk_with_score_function_forward_qb_v2( NVTETensor histogram, NVTETensor bin_bounds, NVTEQBHistogramMode histogram_mode, cudaStream_t stream); -/*! \brief Kimi K3 Quantile Balancing fused-router forward with dense Top-k indices. */ +/*! \brief Kimi K3 Quantile Balancing fused-router forward with dense Top-k indices. + * + * bin_bounds must contain finite FP32 values [lower, upper] with lower < upper. + */ void nvte_fused_topk_with_score_function_forward_qb_with_indices( const NVTETensor logits, int num_tokens, int num_experts, int topk, float scaling_factor, const NVTETensor expert_bias, NVTETensor probs, NVTETensor topk_indices, NVTETensor intermediate_output, NVTETensor cutoff, NVTETensor histogram, NVTETensor bin_bounds, NVTEQBHistogramMode histogram_mode, cudaStream_t stream); -/*! \brief Accumulate a QB histogram from raw sigmoid scores and Top-(k+1) cutoffs. */ +/*! \brief Accumulate a QB histogram from raw sigmoid scores and Top-(k+1) cutoffs. + * + * bin_bounds must contain finite FP32 values [lower, upper] with lower < upper. + */ void nvte_qb_histogram_accumulate(const NVTETensor raw_scores, const NVTETensor cutoff, const NVTETensor bin_bounds, NVTETensor histogram, cudaStream_t stream); diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index f731016f26..4a7971c7dc 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -43,7 +43,8 @@ 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); + at::Tensor bin_bounds, int histogram_mode, + bool bin_bounds_validated = false); 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 017b0389ad..7259cfc6d9 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"), py::arg("histogram_mode"), py::arg("bin_bounds_validated") = false, "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 79391b5bf2..c788bc36e5 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -4,6 +4,7 @@ * See LICENSE for license information. ************************************************************************/ +#include #include #include "../extensions.h" @@ -159,7 +160,8 @@ 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) { + at::Tensor bin_bounds, int histogram_mode, + 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(), @@ -192,6 +194,14 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, TORCH_CHECK( bin_bounds.scalar_type() == at::kFloat && bin_bounds.dim() == 1 && bin_bounds.numel() == 2, "QB bin_bounds must be float32 with shape [2]"); + if (!bin_bounds_validated) { + const at::Tensor bin_bounds_cpu = bin_bounds.cpu(); + const float lower = bin_bounds_cpu.data_ptr()[0]; + const float upper = bin_bounds_cpu.data_ptr()[1]; + TORCH_CHECK(std::isfinite(lower) && std::isfinite(upper) && lower < upper, + "QB bin_bounds values must be finite with lower < upper, got [", lower, ", ", upper, + "]"); + } TORCH_CHECK(histogram_mode == NVTE_QB_HISTOGRAM_TWO_KERNEL || histogram_mode == NVTE_QB_HISTOGRAM_FUSED_ATOMIC, "Unsupported QB histogram mode: ", histogram_mode); diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 9d9f9d4a76..ac38f7d1cf 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -14,6 +14,7 @@ global memory. For example, the gradient is required to have the same dtype as the input. """ +import math from typing import Optional, Union import torch @@ -38,6 +39,36 @@ "two_kernel": int(QBHistogramMode.TWO_KERNEL), "fused_atomic": int(QBHistogramMode.FUSED_ATOMIC), } +_QB_BOUNDS_VALIDATED_VERSION_ATTR = "_nvte_qb_bounds_validated_version" + + +def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: + """Validate CUDA-resident 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 + if getattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, None) == 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" + ) + 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) + return True def _validate_routing_map_format( @@ -158,6 +189,7 @@ def forward( histogram_mode: int, ): # pylint: disable=missing-function-docstring + bin_bounds_validated = _validate_qb_bin_bounds(bin_bounds) ( probs, routing_output, @@ -174,6 +206,7 @@ def forward( histogram, bin_bounds, histogram_mode, + bin_bounds_validated, ) if topk_indices is not None: routing_output = topk_indices @@ -253,7 +286,9 @@ def fused_topk_with_score_function( qb_histogram : torch.Tensor, optional 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. + 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. qb_histogram_mode : str, optional ``"two_kernel"`` or ``"fused_atomic"``. Must be provided with the two QB tensors. From da06d507ab9c6969eb7ab767467672924dc503cf Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 18 Aug 2026 18:14:38 +0800 Subject: [PATCH 4/4] [Common][PyTorch] Validate QB bounds in common APIs Signed-off-by: Harry Zhou --- tests/pytorch/test_fused_router.py | 18 +++-- .../fused_topk_with_score_function.cu | 76 +++++++++++++++++-- .../include/transformer_engine/fused_router.h | 43 ++++++++++- .../pytorch/csrc/extensions/router.cpp | 49 +++++++----- 4 files changed, 151 insertions(+), 35 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index e00bbe0ab0..7325119fb6 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -656,10 +656,18 @@ def test_qb_topk_revalidates_updated_bin_bounds(): ) -def test_qb_raw_binding_rejects_invalid_bin_bounds_recoverably(): +@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( @@ -668,10 +676,10 @@ def test_qb_raw_binding_rejects_invalid_bin_bounds_recoverably(): None, expert_bias, int(RoutingMapFormat.BYTEMAP), - None, + topk_indices, histogram, invalid_bounds, - QBHistogramMode.FUSED_ATOMIC, + histogram_mode, ) # The validation error must not poison the CUDA context. @@ -682,10 +690,10 @@ def test_qb_raw_binding_rejects_invalid_bin_bounds_recoverably(): None, expert_bias, int(RoutingMapFormat.BYTEMAP), - None, + topk_indices, histogram, valid_bounds, - QBHistogramMode.FUSED_ATOMIC, + histogram_mode, ) torch.cuda.synchronize() 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 9fd6250a6f..faeca44797 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 @@ -39,6 +39,16 @@ __device__ inline QBBinParams load_qb_bin_params(const CompType *bin_bounds, int 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) { @@ -913,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, cudaStream_t stream) { + Tensor histogram, bool validate_bin_bounds, 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"); @@ -930,6 +940,9 @@ 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); @@ -1123,7 +1136,8 @@ 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, cudaStream_t stream) { + Tensor histogram, Tensor bin_bounds, NVTEQBHistogramMode histogram_mode, + bool validate_bin_bounds, 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, @@ -1132,6 +1146,9 @@ 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( \ @@ -1170,7 +1187,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, - cudaStream_t stream) { + bool validate_bin_bounds, 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); @@ -1178,6 +1195,9 @@ 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 #include #include "../extensions.h" @@ -194,14 +193,6 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, TORCH_CHECK( bin_bounds.scalar_type() == at::kFloat && bin_bounds.dim() == 1 && bin_bounds.numel() == 2, "QB bin_bounds must be float32 with shape [2]"); - if (!bin_bounds_validated) { - const at::Tensor bin_bounds_cpu = bin_bounds.cpu(); - const float lower = bin_bounds_cpu.data_ptr()[0]; - const float upper = bin_bounds_cpu.data_ptr()[1]; - TORCH_CHECK(std::isfinite(lower) && std::isfinite(upper) && lower < upper, - "QB bin_bounds values must be finite with lower < upper, got [", lower, ", ", upper, - "]"); - } TORCH_CHECK(histogram_mode == NVTE_QB_HISTOGRAM_TWO_KERNEL || histogram_mode == NVTE_QB_HISTOGRAM_FUSED_ATOMIC, "Unsupported QB histogram mode: ", histogram_mode); @@ -248,21 +239,37 @@ fused_topk_with_score_function_qb_fwd(at::Tensor logits, int topk, auto stream = at::cuda::getCurrentCUDAStream(); if (topk_indices.has_value()) { - 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); + 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); + } } 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); + 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); + } } if (mode == NVTE_QB_HISTOGRAM_TWO_KERNEL) { - nvte_qb_histogram_accumulate(intermediate_output_cu.data(), cutoff_cu.data(), - bin_bounds_cu.data(), histogram_cu.data(), stream); + nvte_qb_histogram_accumulate_unchecked(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);