From 127868db2110d74ea898fa4ba5c9a7527b3329bf Mon Sep 17 00:00:00 2001 From: Lee Nau Date: Mon, 6 Apr 2026 15:59:27 -0700 Subject: [PATCH] prevent buffer overflow when autotuner probes large token buckets --- flashinfer/fused_moe/cute_dsl/fused_moe.py | 18 +++-- flashinfer/fused_moe/cute_dsl/tuner.py | 82 +++++++++++----------- 2 files changed, 53 insertions(+), 47 deletions(-) diff --git a/flashinfer/fused_moe/cute_dsl/fused_moe.py b/flashinfer/fused_moe/cute_dsl/fused_moe.py index 8558bf92c20..8ed6a8ba727 100644 --- a/flashinfer/fused_moe/cute_dsl/fused_moe.py +++ b/flashinfer/fused_moe/cute_dsl/fused_moe.py @@ -485,10 +485,16 @@ def _forward_with_tactic( **kwargs, ) -> torch.Tensor: """Forward implementation called by auto-tuner.""" - # Pre-allocated buffers are sized for self.tile_size. When the tactic - # uses a different tile_size (e.g. during autotune), fall back to - # dynamic allocation to avoid buffer overflow in moe_sort. - use_prealloc = self.use_cuda_graph and tile_size == self.tile_size + # Pre-allocated buffers are sized for self.tile_size and + # self.max_num_tokens. Fall back to dynamic allocation when the + # tactic uses a different tile_size or the batch exceeds what the + # buffers were sized for (e.g. autotuner probing larger buckets). + num_tokens = x.shape[0] + use_prealloc = ( + self.use_cuda_graph + and tile_size == self.tile_size + and num_tokens <= self.max_num_tokens + ) return _moe_core_impl( x=x, x_sf=x_sf, @@ -607,7 +613,7 @@ def run( _, best_tactic = tuner.choose_one( "CuteDslMoEWrapper::run", [self._runner], - CuteDslFusedMoENvfp4Runner.tuning_config, + self._runner.tuning_config, inputs, ) @@ -784,7 +790,7 @@ def cute_dsl_fused_moe_nvfp4( _, best_tactic = tuner.choose_one( "CuteDslFusedMoE::run_moe_nvfp4", [runner], - CuteDslFusedMoENvfp4Runner.tuning_config, + runner.tuning_config, inputs, aux_stream=aux_stream, ) diff --git a/flashinfer/fused_moe/cute_dsl/tuner.py b/flashinfer/fused_moe/cute_dsl/tuner.py index 8dd0c26af09..0cc8628ed97 100644 --- a/flashinfer/fused_moe/cute_dsl/tuner.py +++ b/flashinfer/fused_moe/cute_dsl/tuner.py @@ -246,47 +246,6 @@ class CuteDslFusedMoENvfp4Runner(TunableRunner): output_dtype: Output data type (default: torch.bfloat16). """ - # Tensor initializers for dynamic tensors (indices 0, 1, 2, 3, 11) - # These create valid dummy tensors for profiling with different num_tokens - dynamic_tensor_initializers = [ - # 0: x - FP4 quantized input (uint8 packed) - lambda shapes, dtype, device: torch.randint( - 0, 256, shapes, dtype=torch.uint8, device=device - ), - # 1: x_sf - FP8 scale factors (uint8) - lambda shapes, dtype, device: torch.randint( - 1, 128, shapes, dtype=torch.uint8, device=device - ), - # 2: token_selected_experts - expert indices (int32, 0 to num_experts-1) - lambda shapes, dtype, device: torch.randint( - 0, - 8, - shapes, - dtype=torch.int32, - device=device, # num_experts=8 typical - ), - # 3: token_final_scales - routing weights (float32, softmax normalized) - lambda shapes, dtype, device: torch.softmax( - torch.randn(shapes, device=device), dim=-1 - ).to(torch.float32), - # 11: moe_output - output buffer (bfloat16) - lambda shapes, dtype, device: torch.empty(shapes, dtype=dtype, device=device), - ] - - # Tuning config with dynamic tensor specs for num_tokens dimension - # Indices 0, 1, 2, 3, 11 all have num_tokens as their first dimension - tuning_config = TuningConfig( - dynamic_tensor_specs=( - DynamicTensorSpec( - input_idx=(0, 1, 2, 3, 11), # x, x_sf, experts, scales, moe_output - dim_idx=(0, 0, 0, 0, 0), # First dimension is num_tokens for all - gen_tuning_buckets=get_last_power_of_2_num_tokens_buckets(8192), - map_to_tuning_buckets=lambda x: min(last_positive_power_of_2(x), 8192), - tensor_initializers=dynamic_tensor_initializers, - ), - ), - ) - def __init__( self, forward_impl: Callable, @@ -307,6 +266,47 @@ def __init__( self.output_dtype = output_dtype self.enable_pdl = enable_pdl + # Instance-level so dummy expert IDs span all local experts + # (randint(0, num_experts)) for realistic profiling. + self.tuning_config = TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + input_idx=(0, 1, 2, 3, 11), + dim_idx=(0, 0, 0, 0, 0), + gen_tuning_buckets=get_last_power_of_2_num_tokens_buckets(8192), + map_to_tuning_buckets=lambda x: min( + last_positive_power_of_2(x), 8192 + ), + tensor_initializers=[ + # 0: x — FP4 quantized input (uint8 packed) + lambda shapes, dtype, device: torch.randint( + 0, 256, shapes, dtype=torch.uint8, device=device + ), + # 1: x_sf — FP8 scale factors (uint8) + lambda shapes, dtype, device: torch.randint( + 1, 128, shapes, dtype=torch.uint8, device=device + ), + # 2: token_selected_experts — expert indices [0, num_experts) + lambda shapes, dtype, device: torch.randint( + 0, + max(num_experts, 1), + shapes, + dtype=torch.int32, + device=device, + ), + # 3: token_final_scales — routing weights (softmax normalized) + lambda shapes, dtype, device: torch.softmax( + torch.randn(shapes, device=device), dim=-1 + ).to(torch.float32), + # 11: moe_output — output buffer + lambda shapes, dtype, device: torch.empty( + shapes, dtype=dtype, device=device + ), + ], + ), + ), + ) + def __hash__(self): return hash( (