fix(cute_dsl/moe): unbias autotuner profiling for tile_size enumeration (#3252)
<!-- .github/pull_request_template.md --> ## ๐ Description `CuteDslMoEWrapper.__init__` pre-allocates `_gemm1_output`, `_gemm1_output_scale`, and `_moe_sort_buffers` sized for `self.tile_size` only. The `use_prealloc` gate in `_forward_with_tactic` (`fused_moe.py`) is `tile_size == self.tile_size and self.use_cuda_graph and num_tokens <= self.max_num_tokens`, so during autotune profiling the mismatched-`tile_size` tactics fall through to dynamic `torch.empty()` allocation while matching ones run on the prealloc. The autotuner sees **asymmetric allocation overhead** between tactic groups and consistently picks the matching `tile_size` even when intrinsic kernel performance favors the other โ at EP=8/16 N=16384, fi locks to `tile_size=128` in 14/14 cache entries while TRT-LLM picks `tile_size=256` more often. The fix includes three coordinated changes: (1) `tuner.py` lifts the hardcoded `[128, 256]` to a module-level `VALID_TILE_SIZES` tuple โ single source of truth for tactic enumeration AND prealloc sizing; (2) `fused_moe.py:_allocate_buffers` sizes buffers to fit any `tile_size in VALID_TILE_SIZES` (`max_num_permuted_tokens` increases with `tile_size` โ use `max(VALID_TILE_SIZES)`; `max_num_tiles` decreases โ use `min(VALID_TILE_SIZES)`); (3) the prealloc gate becomes `tile_size in VALID_TILE_SIZES`. Both tactic groups now reuse the prealloc; profiling is unbiased. ## ๐ Related Issues https://github.com/flashinfer-ai/flashinfer/pull/3216 https://github.com/flashinfer-ai/flashinfer/pull/3171 ## ๐ Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### โ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## ๐งช Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved CUDA preallocation gating so preallocated buffers are reused safely when CUDA graphs are enabled, avoiding allocation mismatches across supported tile sizes and improving memory efficiency. * **New Features** * Autotuner now scopes per-tactic timing so preallocation is skipped during measurement windows but used outside them; buffer sizing now supports all valid tile sizes. * **Tests** * Added CPU/GPU tests validating preallocation capacity and correct gating behavior during tuning and inference. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
L
Lee Nau committed
e744910fc3bda965542862e5fcff98f2d3edda2b
Parent: 4f072c4
Committed by GitHub <noreply@github.com>
on 5/11/2026, 4:32:57 PM