fix(cute_dsl/moe): unbias autotuner profiling for tile_size enumeration - #3252
Merged
nv-yunzheq merged 6 commits intoMay 11, 2026
Merged
nv-yunzheq merged 6 commits into
nv-yunzheq merged 6 commits into
Conversation
`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` honors prealloc only when the probed tactic's
`tile_size` matches `self.tile_size`:
use_prealloc = (
self.use_cuda_graph
and tile_size == self.tile_size
and num_tokens <= self.max_num_tokens
)
During autotune profiling, mismatched tactics fall through to dynamic
`torch.empty()` per-call allocation. The autotuner is then comparing
tactic latencies that include asymmetric allocation overhead — tactics
matching `self.tile_size` run on the prealloc, others pay the alloc
cost — so it consistently picks the matching `tile_size` even when
intrinsic kernel performance favors the other.
Empirical signature pre-fix at EP=8/16, N=16384: fi locks to
`tile_size=128` in 14 of 14 autotune cache entries. TRT-LLM at the
same shapes picks `tile_size=256` more often, producing a +5-9%
headline gap from the tactic mismatch.
Fix — three coordinated changes:
1. `tuner.py`: lift the hardcoded `[128, 256]` tile_size list to a
module-level `VALID_TILE_SIZES` tuple. Single source of truth for
tactic enumeration AND prealloc sizing. Adding a new tile_size
here automatically widens the prealloc.
2. `fused_moe.py:_allocate_buffers`: size buffers to fit any
`tile_size in VALID_TILE_SIZES`. `max_num_permuted_tokens` is
monotonically increasing in `tile_size` (use
`max(VALID_TILE_SIZES)`); `max_num_tiles` is monotonically
decreasing (use `min(VALID_TILE_SIZES)`). Override
`out_permuted_idx_to_expanded_idx` independently to fit the
largest tile's `max_num_permuted_tokens`.
3. `fused_moe.py:_forward_with_tactic`: change the prealloc gate from
`tile_size == self.tile_size` to `tile_size in VALID_TILE_SIZES`.
Both tactic groups now reuse the prealloc; profiling is unbiased.
Net: post-fix, the autotuner picks the higher-throughput tactic at
each shape on its merits, matching TRT-LLM's choice at large N.
## Tests
Adds two test classes in `tests/moe/test_cute_dsl_fused_moe.py`:
- `TestPreallocStaticInvariants` (1 test, no-GPU): pins
`VALID_TILE_SIZES` to enumerate more than one tile_size. Catches
the orthogonal failure mode where a future refactor reduces the
constant to a single entry — in that case the GPU integration
tests below would pass trivially (no max/min divergence, only one
tile_size to gate-check) and the bias-prevention silently
disappears.
- `TestPreallocBuffersIntegration` (2 tests, GPU/SM100 required):
constructs a real `CuteDslMoEWrapper(use_cuda_graph=True)`. The
first test verifies the prealloc'd buffer shapes fit the workload
at every `tile_size in VALID_TILE_SIZES` — directly empirically
pinning the buffer-sizing contract. The second test
monkey-patches the module-level `_moe_core_impl` to capture the
buffer-passing decision and verifies the `use_prealloc` gate
honors every `tile_size in VALID_TILE_SIZES`, not just
`self.tile_size` — directly pinning the load-bearing property of
the fix.
## Pairs with PR flashinfer-ai#3216
Pairs with PR flashinfer-ai#3216 (autotuner bucket-cap fix, merged 2026-05-06).
Both required to fully close the EP>1 perf gap empirically;
validated at `--num-iters 100` on B200 across EP=8/16, N=16384.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
leejnau
requested review from
IwakuraRein,
aleozlx,
jiahanc,
nv-yunzheq,
samuellees and
yzh119
as code owners
May 6, 2026 22:40
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a per-thread per-tactic profiling scope (is_in_profile_measurement), makes VALID_TILE_SIZES canonical, sizes CuteDslMoEWrapper preallocs to cover all valid tile sizes, gates CUDA-graph preallocation on the profiling flag and use_cuda_graph, and adds tests validating sizing and gating. ChangesPreallocation Gating and Sizing
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request modifies the CuteDslMoEWrapper to size pre-allocated buffers for all valid tile sizes, preventing profiling bias in the autotuner by ensuring all tactics can reuse pre-allocated memory. The implementation updates buffer allocation logic to account for the range of VALID_TILE_SIZES and adjusts the use_prealloc gate in the forward pass. Comprehensive tests were added to validate buffer sizing and gating behavior. I have no feedback to provide.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
flashinfer/fused_moe/cute_dsl/tuner.py (1)
155-161: ⚡ Quick winDerive
DEFAULT_MOE_TACTICfromVALID_TILE_SIZESto avoid config drift.Now that
VALID_TILE_SIZESis the canonical source, keepingDEFAULT_MOE_TACTIChardcoded to128creates a second policy source. IfVALID_TILE_SIZESchanges later, fallback can silently become inconsistent.♻️ Suggested change
+DEFAULT_TILE_SIZE = VALID_TILE_SIZES[0] + DEFAULT_MOE_TACTIC = ( - 128, # tile_size - ((128, 128), (1, 1), False), # gemm1_tactic - ((128, 128), (1, 1), False), # gemm2_tactic + DEFAULT_TILE_SIZE, + ((DEFAULT_TILE_SIZE, 128), (DEFAULT_TILE_SIZE // 128, 1), False), + ((DEFAULT_TILE_SIZE, 128), (DEFAULT_TILE_SIZE // 128, 1), False), )Also applies to: 200-205
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/tuner.py` around lines 155 - 161, DEFAULT_MOE_TACTIC is hardcoded to 128 while VALID_TILE_SIZES is the canonical source; change DEFAULT_MOE_TACTIC (and any other hardcoded defaults in the same file, e.g., the ones around the 200-205 block) to be derived from VALID_TILE_SIZES (for example DEFAULT_MOE_TACTIC = VALID_TILE_SIZES[0] or another clearly chosen element of VALID_TILE_SIZES) so the default always reflects the canonical list and avoids config drift; update references to DEFAULT_MOE_TACTIC accordingly.tests/moe/test_cute_dsl_fused_moe.py (1)
1833-1835: ⚡ Quick winUse
flashinfer.utilsarchitecture checks for new GPU-gated tests.These new tests are added under
@sm100_required, but repository guidance asks test gating viaflashinfer.utilscheck helpers (is_sm100a_supported()/get_compute_capability()style) instead of custom CUDA-property predicates.As per coding guidelines: “
tests/**/*.py: Skip tests on unsupported GPU architectures usingflashinfer.utilscheck functions likeis_sm90a_supported(),is_sm100a_supported(), andget_compute_capability()”.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/moe/test_cute_dsl_fused_moe.py` around lines 1833 - 1835, The test class TestPreallocBuffersIntegration is gated with a custom `@sm100_required` decorator; replace that with the repository-standard GPU check by importing flashinfer.utils.is_sm100a_supported and applying pytest.mark.skipif(not is_sm100a_supported(), reason="SM100a not supported") (keeping the existing `@cute_dsl_available` decorator). Update the decorators on the class to use pytest.mark.skipif(...) referencing is_sm100a_supported() so the test is skipped on unsupported GPUs per the project guidelines.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/moe/test_cute_dsl_fused_moe.py`:
- Around line 1890-1933: The test checks capacities for _gemm1_output and
several moe sort buffers but omits asserting the capacity of
wrapper._gemm1_output_scale; add an assertion inside the for loop (where
VALID_TILE_SIZES is iterated and required_permuted is computed via
get_max_num_permuted_tokens) that wrapper._gemm1_output_scale.shape[0] >=
required_permuted, and raise a clear failure message referencing
_gemm1_output_scale, its capacity and required_permuted (matching the style of
the existing assertions).
---
Nitpick comments:
In `@flashinfer/fused_moe/cute_dsl/tuner.py`:
- Around line 155-161: DEFAULT_MOE_TACTIC is hardcoded to 128 while
VALID_TILE_SIZES is the canonical source; change DEFAULT_MOE_TACTIC (and any
other hardcoded defaults in the same file, e.g., the ones around the 200-205
block) to be derived from VALID_TILE_SIZES (for example DEFAULT_MOE_TACTIC =
VALID_TILE_SIZES[0] or another clearly chosen element of VALID_TILE_SIZES) so
the default always reflects the canonical list and avoids config drift; update
references to DEFAULT_MOE_TACTIC accordingly.
In `@tests/moe/test_cute_dsl_fused_moe.py`:
- Around line 1833-1835: The test class TestPreallocBuffersIntegration is gated
with a custom `@sm100_required` decorator; replace that with the
repository-standard GPU check by importing flashinfer.utils.is_sm100a_supported
and applying pytest.mark.skipif(not is_sm100a_supported(), reason="SM100a not
supported") (keeping the existing `@cute_dsl_available` decorator). Update the
decorators on the class to use pytest.mark.skipif(...) referencing
is_sm100a_supported() so the test is skipped on unsupported GPUs per the project
guidelines.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 43381aba-3f1a-4096-94e2-d5d04899b11d
📒 Files selected for processing (3)
flashinfer/fused_moe/cute_dsl/fused_moe.pyflashinfer/fused_moe/cute_dsl/tuner.pytests/moe/test_cute_dsl_fused_moe.py
Address CodeRabbit nit on PR flashinfer-ai#3252: `test_prealloc_buffers_fit_all_valid_tile_sizes` asserted shapes for `_gemm1_output` and the moe_sort buffers but omitted `_gemm1_output_scale`. Adds the missing assertion so a future regression in scale-buffer sizing is caught alongside the existing buffer-shape contract. The scale buffer is sized in scale-factor elements (one per (permuted_token, scale_vec_group) pair), not in permuted tokens directly. Required capacity is therefore `max_num_permuted_tokens(..., tile_size) * (intermediate_size // sf_vec_size)`. Updated the test docstring to match this distinction. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Collaborator
|
/bot run |
Collaborator
Extend the `use_prealloc` gate in `_forward_with_tactic` with `not AutoTuner.get().is_tuning_mode` so the wrapper bypasses its preallocated buffers during autotune profiling. All tactics then see the same per-call `torch.empty()` allocation overhead and the autotuner's tactic comparison is unbiased; outside the `autotune(True)` context the gate behaves as before — prealloc when `tile_size == self.tile_size`, fall through otherwise. This replaces an earlier approach in this branch that widened the preallocated buffers to fit every valid tile_size; the new approach decouples `self.tile_size` from autotune-time allocation without expanding the prealloc layout, so `tuner.py` is left untouched and `_allocate_buffers` reverts to its pre-PR shape. Pairs with PR flashinfer-ai#3216 (autotuner bucket-cap fix). Both required to fully close the EP>1 perf gap empirically; validated at `--num-iters 100` on B200 across EP=8/16, N=16384. ## Tests Replaces the prior buffer-shape and structural-invariant tests with a single GPU/SM100 test (`TestPreallocGateUnderTuning`) that constructs a real `CuteDslMoEWrapper(use_cuda_graph=True)`, monkey-patches the module-level `_moe_core_impl` to capture the `moe_sort_buffers` argument across {inside `autotune(True)`, outside} × {`tile_size == self.tile_size`, mismatch}, and asserts: - Inside `autotune(True)`: gate skips prealloc for every tactic — pinning the unbias property. - Outside `autotune(True)`: gate uses prealloc when the picked tactic's `tile_size` matches `self.tile_size`; skips it otherwise (the prealloc layout would be wrong for the other tile_size). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Collaborator
|
/bot stop |
Collaborator
|
/bot run |
Collaborator
|
The GitLab CI pipeline #50488690 has been cancelled. |
Collaborator
nv-yunzheq
approved these changes
May 7, 2026
nv-yunzheq
enabled auto-merge (squash)
May 7, 2026 16:59
nv-yunzheq
disabled auto-merge
May 7, 2026 17:12
qiching
reviewed
May 7, 2026
| num_tokens = x.shape[0] | ||
| use_prealloc = ( | ||
| self.use_cuda_graph | ||
| and not AutoTuner.get().is_tuning_mode |
Collaborator
There was a problem hiding this comment.
two qq:
AutoTuner.get() is called on every forward invocation in the inference path. should double checked locking singleton so the overhead is negligible, one is None check + attribute return, but worth confirming this meets latency requirements?
another is is_tuning_mode is read without holding _lock. Under CPython this is fine,d boolean reads are atomic due to the GIL, but it's worth noting as a potential concern if the runtime ever changes?
Just reference. could you consider and answer?
Replaces the wrapper's ``not AutoTuner.get().is_tuning_mode`` gate clause with ``not is_in_profile_measurement()``. The new signal is strictly narrower than ``is_tuning_mode``: it is True only on the calling thread, and only inside the autotuner's per-tactic measurement window (warmup + timed run inside ``_profile_single_kernel``). It is False during cache lookups, ``do_preparation`` calls, the runner invocation immediately after ``choose_one`` returns, and other threads' inference -- all of which the broader ``is_tuning_mode`` flag swept up incorrectly. ## Why narrower ``AutoTuner.is_tuning_mode`` is True for the whole ``autotune(True)`` context, regardless of whether the autotuner is actively timing a specific tactic. Reading it from the wrapper meant that: 1. Cache hits for ops already tuned (where no measurement happens) bypassed prealloc anyway. 2. The runner invocation that uses the chosen tactic immediately after ``choose_one`` returns -- still inside the ``with autotune(True):`` block -- bypassed prealloc. 3. Concurrent threads doing inference while another thread held the tuning context bypassed prealloc. 4. CUDA-graph capture happening inside an ``autotune(True)`` block would record per-call ``torch.empty()`` calls instead of preallocs. None of these are situations where unbiased measurement matters; they all benefit from prealloc. ``is_in_profile_measurement()`` excludes them while still serving the original intent: during the actual measurement window, every tactic sees the same per-call allocation overhead and the autotuner's tactic comparison is unbiased. ## What changes in autotuner.py - New module-level ``_profile_measurement_thread_local`` (a ``threading.local``). - New ``_profile_measurement_scope`` context manager (private; sets the thread-local on entry, restores prior value on exit, supports nesting). - New ``is_in_profile_measurement()`` accessor (public; reads the thread-local; returns False on threads that never entered the scope). - ``AutoTuner._profile_single_kernel`` wraps its warmup + timed run with ``_profile_measurement_scope()`` so every runner invocation inside the measurement function sees the flag True; runner invocations elsewhere in ``choose_one`` (cache search, ``do_preparation``, the post-loop ``search_cache`` call) see it False. The change is purely additive: the new helpers don't alter ``AutoTuner``'s class state, no other autotune callers are affected. ## Tests Replaces the prior ``TestPreallocGateUnderTuning`` test with a broader contract check that exercises three contexts: 1. Inside ``autotune(True)`` AND inside ``_profile_measurement_scope()`` (simulating a tactic measurement) -- gate must skip prealloc for every tactic, regardless of tile_size match. 2. Inside ``autotune(True)`` but OUTSIDE the measurement scope (simulating a cache hit, the do_preparation call, or the post-``choose_one`` runner invocation) -- gate must use prealloc when ``tile_size == self.tile_size``, skip otherwise. This is the property that distinguishes the narrow signal from the broad one. 3. Outside any tuning context (plain inference) -- same as case 2. Pairs with PR flashinfer-ai#3216 (autotuner bucket-cap fix). Both required to fully close the EP>1 perf gap empirically; validated at ``--num-iters 100`` on B200 across EP=8/16, N=16384. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
leejnau
requested review from
bkryu,
cyx-6,
saltyminty,
sricketts and
yongwww
as code owners
May 7, 2026 17:45
Builds on the narrow ``is_in_profile_measurement()`` gate from the
prior commit by also expanding ``CuteDslMoEWrapper._allocate_buffers``
to size kernel-output buffers for *any* ``tile_size in
VALID_TILE_SIZES``, not just the constructor-time ``self.tile_size``.
## Why
Now that the autotuner profiles tactics unbiasedly (the prior commit's
narrower gate), it can fairly pick a tactic with ``tile_size != self.tile_size``
when that's the higher-throughput choice — at large N this is the
common case (``tile_size=256`` typically wins). But with the prior
commit alone, the wrapper's prealloc was still sized for
``self.tile_size`` and the gate fell through to per-call
``torch.empty()`` whenever the picked tactic mismatched. Two real
problems:
1. **CUDA-graph contract**: the wrapper's ``run()`` is documented as
graph-safe with ``use_cuda_graph=True``, but per-call ``torch.empty``
means captured graphs record the alloc inside the graph instead of
binding to the prealloc. PyTorch's graph private memory pool
accommodates this since 1.10, but it's not what the contract
promises.
2. **Buffer-overflow correctness**: ``max_num_permuted_tokens`` is
monotonically increasing in ``tile_size``, so a tile_size=128-sized
buffer is too small for a tile_size=256 tactic. The fall-through to
per-call alloc isn't a perf-hygiene choice — it's required for
correctness given the smaller sizing.
## Fix — three coordinated changes
1. ``tuner.py``: lift the hardcoded ``[128, 256]`` tile_size list to
a module-level ``VALID_TILE_SIZES`` tuple. Single source of truth
for tactic enumeration AND prealloc sizing.
2. ``fused_moe.py:_allocate_buffers``: size buffers to fit any
``tile_size in VALID_TILE_SIZES``. ``max_num_permuted_tokens`` is
monotonically increasing in ``tile_size`` (use
``max(VALID_TILE_SIZES)``); ``max_num_tiles`` is monotonically
decreasing (use ``min(VALID_TILE_SIZES)``). Override
``out_permuted_idx_to_expanded_idx`` independently to fit the
largest tile's ``max_num_permuted_tokens``.
3. ``fused_moe.py:_forward_with_tactic``: drop the
``tile_size == self.tile_size`` check from the gate. Replace with
``tile_size in VALID_TILE_SIZES`` (defensive; should never fail for
tactics drawn from ``ALL_MOE_TACTICS``). The narrow
``is_in_profile_measurement()`` check from the prior commit is
retained.
## Resulting gate semantics
``use_prealloc = use_cuda_graph
AND not is_in_profile_measurement()
AND tile_size in VALID_TILE_SIZES
AND num_tokens <= max_num_tokens``
- During autotune profiling: ``is_in_profile_measurement`` is True →
prealloc bypassed for every tactic → unbiased measurement.
- During cache-lookup / do-preparation / post-``choose_one`` /
concurrent-thread inference: prealloc used regardless of tile_size,
preserving the wrapper's CUDA-graph contract.
- During plain inference: same — prealloc used for whichever tactic
the autotuner picked, including ``tile_size != self.tile_size``.
## Tests
Three test classes in ``tests/moe/test_cute_dsl_fused_moe.py``:
- ``TestPreallocStaticInvariants`` (1 no-GPU test): pin
``VALID_TILE_SIZES`` non-trivial. Catches accidental reduction to
a single entry which would defeat the whole purpose.
- ``TestPreallocBuffersIntegration`` (1 GPU/SM100 test): construct a
real ``CuteDslMoEWrapper(use_cuda_graph=True)``, verify
``_gemm1_output``, ``_gemm1_output_scale``, and the moe_sort buffers
fit the workload at every ``tile_size in VALID_TILE_SIZES``. Pins
the buffer-sizing contract empirically.
- ``TestPreallocGateUnderTuning`` (1 GPU/SM100 test, updated):
monkey-patches ``_moe_core_impl`` and exercises three contexts × two
tile_sizes:
- measurement scope (inside ``_profile_measurement_scope``):
gate must skip prealloc for every tactic.
- inside ``autotune(True)`` but outside the measurement scope
(cached call, post-choose_one): gate must use prealloc for
every valid tile_size.
- outside any tuning context (plain inference): gate must use
prealloc for every valid tile_size.
The latter two assertions are what's strengthened in this commit:
the gate no longer requires ``tile_size == self.tile_size``.
Pairs with PR flashinfer-ai#3216 (autotuner bucket-cap fix). Both required to
fully close the EP>1 perf gap empirically.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/moe/test_cute_dsl_fused_moe.py (1)
2045-2047: ⚡ Quick winAvoid hardcoding the “other” tile size in gate tests.
This reintroduces a second source of truth; deriving
otherfromVALID_TILE_SIZESwill keep the test valid if tile-size enumeration changes.♻️ Suggested update
- matching = wrapper.tile_size # the tile_size the prealloc was sized for - other = 256 if matching == 128 else 128 + from flashinfer.fused_moe.cute_dsl.tuner import VALID_TILE_SIZES + matching = wrapper.tile_size # the tile_size the prealloc was sized for + others = [t for t in VALID_TILE_SIZES if t != matching] + assert others, "Test requires at least two distinct VALID_TILE_SIZES entries" + other = others[0]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/moe/test_cute_dsl_fused_moe.py` around lines 2045 - 2047, The test hardcodes the alternate tile size using a conditional on wrapper.tile_size; instead derive `other` from the `VALID_TILE_SIZES` enum to avoid a second source of truth. Replace the hardcoded logic that sets `other = 256 if matching == 128 else 128` by selecting any tile size from `VALID_TILE_SIZES` that is not equal to `matching` (e.g., iterate or filter `VALID_TILE_SIZES` and pick the first non-matching value) so `matching`, `other`, and `wrapper.tile_size` remain consistent if the enum changes.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/moe/test_cute_dsl_fused_moe.py`:
- Around line 1834-1836: Replace the local GPU gate decorators (e.g., the
sm100_required marker applied to TestPreallocBuffersIntegration and the similar
marker at the other location) with repository-standard flashinfer.utils
capability checks: import the appropriate functions from flashinfer.utils (such
as is_sm100a_supported() or get_compute_capability()) and use pytest.skip or
pytest.mark.skipif to skip when the target SM is unsupported (e.g., skip if not
is_sm100a_supported()). Keep the existing cute_dsl_available decorator
unchanged; update both TestPreallocBuffersIntegration and the other GPU-gated
test class at the referenced location to use the flashinfer.utils checks for
consistency.
---
Nitpick comments:
In `@tests/moe/test_cute_dsl_fused_moe.py`:
- Around line 2045-2047: The test hardcodes the alternate tile size using a
conditional on wrapper.tile_size; instead derive `other` from the
`VALID_TILE_SIZES` enum to avoid a second source of truth. Replace the hardcoded
logic that sets `other = 256 if matching == 128 else 128` by selecting any tile
size from `VALID_TILE_SIZES` that is not equal to `matching` (e.g., iterate or
filter `VALID_TILE_SIZES` and pick the first non-matching value) so `matching`,
`other`, and `wrapper.tile_size` remain consistent if the enum changes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: cb5e20af-489c-42f0-aa99-444987ba8e33
📒 Files selected for processing (3)
flashinfer/fused_moe/cute_dsl/fused_moe.pyflashinfer/fused_moe/cute_dsl/tuner.pytests/moe/test_cute_dsl_fused_moe.py
Comment on lines
+1834
to
+1836
| @cute_dsl_available | ||
| @sm100_required | ||
| class TestPreallocBuffersIntegration: |
Contributor
There was a problem hiding this comment.
Use flashinfer.utils-based architecture checks for these new GPU tests.
These new classes are GPU-gated through the local sm100_required marker path; please switch these skip conditions to the repository-standard flashinfer.utils capability checks for consistency and future arch handling.
As per coding guidelines, tests/**/*.py: “Skip tests on unsupported GPU architectures using flashinfer.utils check functions like is_sm90a_supported(), is_sm100a_supported(), and get_compute_capability()”.
Also applies to: 1936-1938
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/moe/test_cute_dsl_fused_moe.py` around lines 1834 - 1836, Replace the
local GPU gate decorators (e.g., the sm100_required marker applied to
TestPreallocBuffersIntegration and the similar marker at the other location)
with repository-standard flashinfer.utils capability checks: import the
appropriate functions from flashinfer.utils (such as is_sm100a_supported() or
get_compute_capability()) and use pytest.skip or pytest.mark.skipif to skip when
the target SM is unsupported (e.g., skip if not is_sm100a_supported()). Keep the
existing cute_dsl_available decorator unchanged; update both
TestPreallocBuffersIntegration and the other GPU-gated test class at the
referenced location to use the flashinfer.utils checks for consistency.
qiching
reviewed
May 7, 2026
Collaborator
There was a problem hiding this comment.
using threading.local() is correct!
but i have another concerns that if Autotuner profiling might execute the runner in a subthread. If the runner() call during warmup or just measure is executed in another thread, the subthread will not see active=True, and the gate will fail.
could we consider this make sense?
qiching
reviewed
May 7, 2026
| # threads inclusive); ``is_in_profile_measurement`` is True only on the | ||
| # specific thread that is actively timing a tactic, and only during the | ||
| # warmup + measurement window inside ``_profile_single_kernel``. | ||
| _profile_measurement_thread_local = threading.local() |
Collaborator
There was a problem hiding this comment.
using threading.local() is correct!
but i have another concerns that if Autotuner profiling might execute the runner in a subthread. If the runner() call during warmup or just measure is executed in another thread, the subthread will not see active=True, and the gate will fail.
could we consider this make sense?
Contributor
Author
There was a problem hiding this comment.
In the current autotuner implementation, _profile_single_kernel() invokes runner(...) synchronously on the same Python thread for both warmup and timed measurement, so the threading.local() flag is visible to the wrapper when _forward_with_tactic() runs.
The CUDA kernels themselves launch asynchronously, but the Python-side allocation/prealloc gate is evaluated before launch, on that same caller thread. So the thread-local scope covers the decision we need here.
If we later refactor autotuner profiling to dispatch runner(...) from a Python worker thread, then we should either propagate the profiling context into that worker or replace this with an explicit per-call kwarg/context token. For today’s code path, threading.local() is probably the correct narrower signal: it avoids leaking profiling behavior into concurrent inference threads while still covering the actual measured runner calls.
Collaborator
There was a problem hiding this comment.
now approach is correct for today's synchronous profiling path. Left a note for future: if _profile_single_kernel ever dispatches runner() calls from a worker thread, the threading.local() flag won't propagate.
Replace the hardcoded "other = 256 if matching == 128 else 128" in TestPreallocGateUnderTuning with iteration over every tile_size in VALID_TILE_SIZES, so adding a new entry doesn't silently leave the gate untested for that tile in any of the three contexts (measurement window / in-tuning-context-outside-measurement / plain inference). Addresses CodeRabbit nit on PR flashinfer-ai#3252 plus codex's follow-up that strengthens "first non-matching tile" to "every non-matching tile". Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
qiching
approved these changes
May 7, 2026
nv-yunzheq
approved these changes
May 7, 2026
Collaborator
|
/bot run |
Collaborator
5 tasks
nv-yunzheq
added a commit
that referenced
this pull request
May 11, 2026
<!-- .github/pull_request_template.md --> ## 📌 Description This refreshes `benchmarks/bench_moe_deepseek.py` so it runs cleanly on current main. It rebases Yunzhe Qiu's bench rewrite (`5677a080` from #2886, which restructures the bench so autotune runs inside the `bench_gpu_time` measurement region) onto post-#3252 main, plus a small follow-up that fixes a stale `RoutingMethodType` import. Two commits from the original #2886 are intentionally dropped because their fixes have since landed independently. `c0b80b64`'s `num_tokens <= max_num_tokens` prealloc guard is now subsumed by #3252 — the `use_prealloc` predicate in `cute_dsl/fused_moe.py` already includes that check. And `f3beb602`'s `_force_autotune_off()` bench-side workaround for CUPTI measurement pollution is no longer needed: #3126 moved the cache lookup ahead of `_prepare_input_tensors` synthesis in the autotuner's tuning-mode loop, eliminating the pollution at the source. The only remaining mismatch with current main was `RoutingMethodType`, which moved from `flashinfer.fused_moe.core` to `flashinfer.tllm_enums` (and is re-exported via `flashinfer.fused_moe`) — fixed in the second commit here. Verified on B200 inside `nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc14`: DeepSeek-V3 at bs=128 ep=8 measures CuteDSL=0.147 ms / TRTLLM=0.144 ms — in the clean band that matches prior post-pollution-fix measurements (~0.157 / ~0.142). An 18-cell matrix (N=1, 8, 128, 512, 2048, 16384 × EP=1, 8, 16) and an 8-cell gen-phase decode sweep also ran without errors. Closes #2886. ## 🔍 Related Issues #2886 #3126 #3252 ## 🚀 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 ## Release Notes * **Refactor** * Improved MoE throughput benchmarking methodology with enhanced pre-warm invocation and synchronization for accurate timing capture * Refactored autotuning strategy to occur inline during benchmark warmup phase * Reorganized benchmark output display for clearer result presentation [](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3292) <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yunzhe Qiu <yunzheq@nvidia.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
leejnau
added a commit
to leejnau/flashinfer
that referenced
this pull request
May 12, 2026
Adds a --nvtx-profile-region flag that brackets each backend's
bench_gpu_time measurement with cudaProfilerStart/Stop (in try/finally)
and an outer NVTX range `bench_<backend>_n<N>`. Pair with
`nsys profile --capture-range=cudaProfilerApi
--capture-range-end=repeat-shutdown:3` to restrict capture to the
measurement iters for all three backends in a single nsys run
(pre-warm and autotune are outside the capture window).
Also emits fine-grained NVTX ranges inside each backend's run()
closure, unconditionally (NVTX push/pop is sub-microsecond):
iter_<backend> per bench iteration
topk_<backend> the fused_topk_deepseek routing call
(CuteDSL and CUTLASS only; TRT-LLM's routing
is internal to the kernel)
moe_<backend> the MoE call itself
These nest within the outer `bench_<backend>_n<N>` range when the
`--nvtx-profile-region` flag is on, giving the parser a clean
hierarchy for attribution.
After each `--nvtx-profile-region` capture, prints a short
`[nvtx] bench_<backend>_n<N>: captured iters=..., median=... ms` line
so the user can cross-check the captured duration against the printed
bench table.
This is a rebased + cleaned-up version of the earlier nvtx_microbench
work. The _force_autotune_off() bench-side workaround that the
earlier version relied on is no longer needed: PR flashinfer-ai#3126 moved the
autotuner's cache lookup ahead of profile-input synthesis, so cache
hits in tuning mode no longer launch the synthetic torch.rand /
torch.randint kernels that f3beb602 was working around. The
prealloc-fix (flashinfer-ai#3252) and bench-tool refresh (flashinfer-ai#3292) are also in
main, so this commit applies cleanly as a single change on top of
the merged bench refactor.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
leejnau
added a commit
to leejnau/flashinfer
that referenced
this pull request
May 15, 2026
Adds a no-GPU pytest class TestInputsHelperContract that exercises
two cross-file contracts of the new inputs_pre_hook helper:
1. Layout contract: ``CuteDslMoEInputsHelper.inputs_pre_hook`` must
replace ``inputs[2]`` (token_selected_experts) and pass through
every other input unchanged (object identity preserved). Pins the
contract with ``CuteDslMoEWrapper.run``'s 12-element inputs-list
layout — if someone reorders that list without updating the
helper's unpacking pattern, the autotune profile silently
corrupts a different tensor.
2. Determinism contract: two helper instances with the same seed
must produce tensor-equal replacement ``token_selected_experts``
for the same input. Exercises the seeded
``torch.random.fork_rng + manual_seed`` pattern in
``generate_token_selected_experts``, which is the load-bearing
determinism property that eliminates cross-process autotune-pick
variance.
Pure-Python, runs in <1s on CPU. Mirrors the style of
``TestTacticEnumeration`` (no-GPU contract tests added in PR flashinfer-ai#3252).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
5 tasks
nv-yunzheq
pushed a commit
that referenced
this pull request
May 21, 2026
<!-- .github/pull_request_template.md --> ## 📌 Description `CuteDslMoEWrapper` currently passes `self._forward_with_tactic` as a bound method into `CuteDslFusedMoENvfp4Runner`, creating a strong reference cycle: `wrapper -> runner -> bound method -> wrapper`. When the wrapper is used with `use_cuda_graph=True`, this can keep wrapper-owned CUDA graph resources alive after user code has dropped the wrapper, until Python cyclic GC eventually runs. This PR replaces that bound-method callback with a weakref trampoline. The runner can still call into a live wrapper, but it no longer owns the wrapper lifetime. This prevents stale wrapper CUDA resources from surviving across same-process tests or later autotune runs. ## 🔍 Related Issues #3286 #3301 #3252 ## 🚀 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 Adds a focused regression test that warms a CUDA-graph wrapper, verifies it is finalized before cyclic GC, and then runs a subsequent autotuned wrapper call to ensure the output remains NaN-free. - [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 handling and cleanup of CUDA-graph wrappers to prevent resource leaks and provide a clear error when a wrapper is no longer available. * **Tests** * Added lifetime tests covering CUDA-graph wrappers before and after autotune; verify stable, non-NaN outputs during autotune. * **Documentation** * Updated comment about cold-L2 cache behavior and noted follow-up to re-enable it once a related issue is addressed. <!-- review_stack_entry_start --> [](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3340?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack) <!-- review_stack_entry_end --> <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
5 tasks
](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3292){kind=link}
](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3340?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack){kind=link}
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
CuteDslMoEWrapper.__init__pre-allocates_gemm1_output,_gemm1_output_scale, and_moe_sort_bufferssized forself.tile_sizeonly. Theuse_preallocgate in_forward_with_tactic(fused_moe.py) istile_size == self.tile_size and self.use_cuda_graph and num_tokens <= self.max_num_tokens, so during autotune profiling the mismatched-tile_sizetactics fall through to dynamictorch.empty()allocation while matching ones run on the prealloc. The autotuner sees asymmetric allocation overhead between tactic groups and consistently picks the matchingtile_sizeeven when intrinsic kernel performance favors the other — at EP=8/16 N=16384, fi locks totile_size=128in 14/14 cache entries while TRT-LLM pickstile_size=256more often.The fix includes three coordinated changes: (1)
tuner.pylifts the hardcoded[128, 256]to a module-levelVALID_TILE_SIZEStuple — single source of truth for tactic enumeration AND prealloc sizing; (2)fused_moe.py:_allocate_bufferssizes buffers to fit anytile_size in VALID_TILE_SIZES(max_num_permuted_tokensincreases withtile_size→ usemax(VALID_TILE_SIZES);max_num_tilesdecreases → usemin(VALID_TILE_SIZES)); (3) the prealloc gate becomestile_size in VALID_TILE_SIZES. Both tactic groups now reuse the prealloc; profiling is unbiased.🔍 Related Issues
#3216
#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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Bug Fixes
New Features
Tests