SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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