Skip to content

Prevent MoE autotuner buffer overflow on large token buckets #3025

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nv-yunzheq merged 1 commit into from
Apr 10, 2026
Merged

Prevent MoE autotuner buffer overflow on large token buckets #3025

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 12 additions & 6 deletions flashinfer/fused_moe/cute_dsl/fused_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -607,7 +613,7 @@ def run(
_, best_tactic = tuner.choose_one(
"CuteDslMoEWrapper::run",
[self._runner],
CuteDslFusedMoENvfp4Runner.tuning_config,
self._runner.tuning_config,
inputs,
)

Expand Down Expand Up @@ -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,
)
Expand Down
82 changes: 41 additions & 41 deletions flashinfer/fused_moe/cute_dsl/tuner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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(

@yzh119 yzh119 Apr 9, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we need to fixed random see to guarantee consistent here?

@leejnau leejnau Apr 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we need to fixed random see to guarantee consistent here?

No, since these are throwaway dummy tensors for timing kernel execution during profiling. The values don't affect tactic selection.

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(
(
Expand Down
Loading