feat(moe): add unified CUTLASS BF16 and W4A16 runners - #4328
Merged
feih-nv merged 19 commits intoAug 10, 2026
Merged
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
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:
📝 WalkthroughWalkthroughThe PR adds architecture-specific CUTLASS BF16 and W4A16 MoE configurations, weight preparation, runners, dispatch registration, autotuning, workspace caching, and CUDA test coverage. ChangesCUTLASS MoE integration
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant MoELayer
participant CutlassBf16Runner
participant CUTLASS fused-MoE backend
MoELayer->>CutlassBf16Runner: validate inputs and prepare routing
CutlassBf16Runner->>CutlassBf16Runner: select tactics and activate workspace bucket
CutlassBf16Runner->>CUTLASS fused-MoE backend: launch fused MoE execution
CUTLASS fused-MoE backend-->>MoELayer: return output tensor
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 |
feih-nv
changed the title
feih-nv
changed the title
feih-nv
marked this pull request as ready for review
August 4, 2026 09:50
feih-nv
requested review from
Aneureka,
IwakuraRein,
aleozlx,
jiahanc,
nv-yunzheq,
samuellees and
yzh119
as code owners
August 4, 2026 09:50
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Collaborator
Author
|
/bot run /tests/moe |
Collaborator
![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: 3
🧹 Nitpick comments (3)
flashinfer/fused_moe/runners.py (1)
339-369: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffConsider a single max-sized workspace instead of a per-bucket cache.
_workspace_cachenever evicts, and it keeps one allocation per(bucket, hidden_size)pair. The hybrid bucket ladder up totune_max_num_tokensproduces many buckets, so a long-lived layer can retain many multi-MiB/GiB buffers at once.
cutlass_fused_moe_workspace_sizedocuments monotonic sizing: "a buffer allocated for the maximum shape is valid for all smaller shapes on the same call". One allocation sized attune_max_num_tokenstherefore satisfies every bucket, keeps the pointer stable for captured CUDA graphs, and removes the cache. Keep the current design only if a per-bucket size reduction is a measured requirement; then document the peak-memory tradeoff here.Note that the existing tests assert per-bucket cache identity (
tests/moe/test_unified_moe_cutlass.pylines 363-386, 604-639), so this change requires test updates.🤖 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/runners.py` around lines 339 - 369, Replace the per-(num_tokens, hidden_size) allocation and _workspace_cache logic in _ensure_workspace with one workspace sized for the maximum configured token count, tune_max_num_tokens, while preserving the current hidden-size and dtype/configuration inputs to cutlass_fused_moe_workspace_size. Reuse this stable allocation for all smaller buckets and remove cache-specific state and assignments. Update the affected tests to assert workspace reuse and stable identity across buckets rather than per-bucket cache entries.tests/moe/test_unified_moe_cutlass.py (2)
456-515: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the shared case builder.
_make_w4a16_caserepeats the seed, device, geometry, weight, routing, and activation setup of_make_case. Only the quant variant, backend candidate, and reference-weight construction differ. Extract the common part into one helper that takes the quant variant and backend config. This prevents the two fixtures from diverging as the CUTLASS contracts evolve.🤖 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_unified_moe_cutlass.py` around lines 456 - 515, Refactor _make_w4a16_case and _make_case to use a shared case-builder helper for seed/device setup, tensor geometry, weights, routing, and activation construction. Parameterize that helper with the quant variant and backend configuration, while keeping each caller’s distinct reference-weight preparation and quantized-weight construction in its respective function.
389-398: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGate the CUTLASS SM90 tests with
flashinfer.utils.is_sm90a_supported.Use the shared architecture helper instead of the local compute-capability check, so the skip also covers the CUDA toolkit requirement used by CUTLASS SM90 support.
♻️ Proposed refactor
-def _is_cutlass_sm90_arch() -> bool: - if not torch.cuda.is_available(): - return False - major, minor = torch.cuda.get_device_capability() - return major * 10 + minor == 90 - - cutlass_sm90_required = pytest.mark.skipif( - not _is_cutlass_sm90_arch(), reason="requires an SM90 CUTLASS GPU" + not is_sm90a_supported(torch.device("cuda")), + reason="requires an SM90 CUTLASS GPU", )Add the import:
from flashinfer.utils import is_sm90a_supported🤖 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_unified_moe_cutlass.py` around lines 389 - 398, Replace the local _is_cutlass_sm90_arch compute-capability check with flashinfer.utils.is_sm90a_supported, importing the shared helper and using it directly in the cutlass_sm90_required pytest skip condition. Remove the now-unused local helper while preserving the existing skip reason.Sources: Coding guidelines, Learnings
🤖 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 `@flashinfer/fused_moe/api.py`:
- Around line 488-492: Update the CUTLASS BF16 and W4A16 public configuration
docstrings near the unified MoE API declarations to document the concrete
execution contract: packed precomputed routing, SwiGLU, do_finalize=True, no
expert parallelism or shared experts, and for W4A16 require hidden_size and
intermediate_size to be divisible by 128. Keep the existing hardware coverage
documentation intact.
In `@tests/moe/test_unified_moe_cutlass.py`:
- Around line 406-426: The CUTLASS case builders generate outputs too small for
their fixed tolerances, allowing incorrect results to pass. In
tests/moe/test_unified_moe_cutlass.py lines 406-426 and 461-481, rescale
activations and weights so reference outputs are O(1); then update the BF16
assertions at lines 551, 626, 639, and 684 to compare relative to the reference
norm, and tighten the W4A16 assertions at lines 568, 707, and 719 using bounds
derived from MXFP4 quantization error instead of fixed atol=1e-1.
- Line 1: Update the module docstring in tests/moe/test_unified_moe_cutlass.py
to mention both BF16 and W4A16 coverage, including the MXFP4 quantization tests
represented by CutlassW4A16Config and CutlassW4A16Runner.
---
Nitpick comments:
In `@flashinfer/fused_moe/runners.py`:
- Around line 339-369: Replace the per-(num_tokens, hidden_size) allocation and
_workspace_cache logic in _ensure_workspace with one workspace sized for the
maximum configured token count, tune_max_num_tokens, while preserving the
current hidden-size and dtype/configuration inputs to
cutlass_fused_moe_workspace_size. Reuse this stable allocation for all smaller
buckets and remove cache-specific state and assignments. Update the affected
tests to assert workspace reuse and stable identity across buckets rather than
per-bucket cache entries.
In `@tests/moe/test_unified_moe_cutlass.py`:
- Around line 456-515: Refactor _make_w4a16_case and _make_case to use a shared
case-builder helper for seed/device setup, tensor geometry, weights, routing,
and activation construction. Parameterize that helper with the quant variant and
backend configuration, while keeping each caller’s distinct reference-weight
preparation and quantized-weight construction in its respective function.
- Around line 389-398: Replace the local _is_cutlass_sm90_arch
compute-capability check with flashinfer.utils.is_sm90a_supported, importing the
shared helper and using it directly in the cutlass_sm90_required pytest skip
condition. Remove the now-unused local helper while preserving the existing skip
reason.
🪄 Autofix
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 Plus
Run ID: 0387beea-3f58-42c8-86a4-5095e112c7f3
📥 Commits
Reviewing files that changed from the base of the PR and between d7e390c and b8834a6e1ad934405fb7bf2dfb082a5265b60f93.
📒 Files selected for processing (8)
flashinfer/fused_moe/__init__.pyflashinfer/fused_moe/api.pyflashinfer/fused_moe/core.pyflashinfer/fused_moe/layer.pyflashinfer/fused_moe/prepare.pyflashinfer/fused_moe/runners.pytests/moe/test_unified_moe.pytests/moe/test_unified_moe_cutlass.py
Expose the existing CUTLASS fused-MoE implementation through the unified API with bucket-aware independent GEMM tuning, reusable workspace management, and persistent cache key coverage. Add canonical BF16 weight preparation, backend registration, documentation, and focused correctness and CUDA graph tests. Developed with AI assistance.
feih-nv
force-pushed
the
feih/unified-moe-cutlass-sm90
branch
from
August 5, 2026 01:16
b8834a6 to
8f8f4d8
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
3 tasks
Collaborator
Author
|
/bot run /tests/moe |
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (3)
flashinfer/fused_moe/runners.py (1)
554-558: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider adding
_enable_pdlto the cache-key extras.
get_cache_key_extrasreturns only_device_arch. The inner CUTLASSMoERunner.get_cache_key_extrasinflashinfer/fused_moe/core.pyreturns the full fixed launch configuration, includingenable_pdl,top_k, andactivation_type.At the outer level,
top_kis encoded through thetopk_ids/topk_weightsprofile shapes, andcheck_supportpins the activation to Swiglu._enable_pdlis not encoded anywhere. Two runners that differ only in_enable_pdlshare one outer tactic cache entry.
__hash__already includesself.config, so the in-memory runner identity is distinct. The concern is limited to the persisted tactic cache.♻️ Proposed extras extension
def get_cache_key_extras(self, _inputs: List[torch.Tensor]) -> tuple: - return (self._device_arch,) + return (self._device_arch, self._enable_pdl)🤖 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/runners.py` around lines 554 - 558, Update the outer runner’s get_cache_key_extras method to include _enable_pdl alongside _device_arch, ensuring persisted tactic cache entries remain distinct for different PDL settings. Leave __hash__ and the existing profile-shape and activation handling unchanged.tests/moe/test_unified_moe_cutlass.py (2)
293-306: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd a case that exercises the
local_num_expertsbranch alone.The third parametrization sets both
local_expert_offset=2andlocal_num_experts=2._CutlassRunnerBase.check_supportrejects the config on thelocal_expert_offset != 0term, so thelocal_num_experts != num_expertsterm is never the sole cause of rejection. Add a case withlocal_expert_offset=0and alocal_num_expertsvalue belowrouting.num_expertsto cover that branch.🤖 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_unified_moe_cutlass.py` around lines 293 - 306, Update the parametrized configurations in the unified MoE support test to add a case with local_expert_offset=0 and local_num_experts below routing.num_experts, so the local_num_experts != num_experts rejection branch is exercised independently. Keep the existing combined offset-and-count case unless it is redundant with the test’s intended coverage.
712-719: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winMark CUTLASS BF16 MoE tests with an architecture skip.
test_cutlass_autotuned_compound_tactic_numerics_and_cuda_graph,test_cutlass_autotune_override_reuses_max_workspace, andtest_cutlass_forward_reuses_max_workspace_after_smaller_overridecreateMoELayer(config)with the CUTLASS BF16 backend. Add apytest.mark.skipif/module-level skip that uses the appropriate CUDA compute-capability or backend capability API for unsupported devices, or add a helper decorator and apply it to these tests.🤖 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_unified_moe_cutlass.py` around lines 712 - 719, Add an architecture guard for the CUTLASS BF16 MoE tests so they are skipped on unsupported devices before constructing MoELayer(config). Apply a shared pytest.mark.skipif or helper decorator to test_cutlass_autotuned_compound_tactic_numerics_and_cuda_graph, test_cutlass_autotune_override_reuses_max_workspace, and test_cutlass_forward_reuses_max_workspace_after_smaller_override, using the existing CUDA compute-capability or backend capability check available in this test module.Source: Coding 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.
Nitpick comments:
In `@flashinfer/fused_moe/runners.py`:
- Around line 554-558: Update the outer runner’s get_cache_key_extras method to
include _enable_pdl alongside _device_arch, ensuring persisted tactic cache
entries remain distinct for different PDL settings. Leave __hash__ and the
existing profile-shape and activation handling unchanged.
In `@tests/moe/test_unified_moe_cutlass.py`:
- Around line 293-306: Update the parametrized configurations in the unified MoE
support test to add a case with local_expert_offset=0 and local_num_experts
below routing.num_experts, so the local_num_experts != num_experts rejection
branch is exercised independently. Keep the existing combined offset-and-count
case unless it is redundant with the test’s intended coverage.
- Around line 712-719: Add an architecture guard for the CUTLASS BF16 MoE tests
so they are skipped on unsupported devices before constructing MoELayer(config).
Apply a shared pytest.mark.skipif or helper decorator to
test_cutlass_autotuned_compound_tactic_numerics_and_cuda_graph,
test_cutlass_autotune_override_reuses_max_workspace, and
test_cutlass_forward_reuses_max_workspace_after_smaller_override, using the
existing CUDA compute-capability or backend capability check available in this
test module.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: fd413730-059d-4e86-911e-d68645b3320e
📒 Files selected for processing (3)
flashinfer/fused_moe/layer.pyflashinfer/fused_moe/runners.pytests/moe/test_unified_moe_cutlass.py
🚧 Files skipped from review as they are similar to previous changes (1)
- flashinfer/fused_moe/layer.py
Collaborator
Author
|
/bot run /tests/moe |
Collaborator
Collaborator
|
[FAILED] Pipeline #61209428 — 15/18 executed test jobs passed Compared with nightly #61004434. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
3 tasks
5 tasks
Aneureka
approved these changes
Aug 7, 2026
Collaborator
|
Approved. @aleozlx may take a look as well. |
feih-nv
enabled auto-merge (squash)
August 10, 2026 01:06
feih-nv
disabled auto-merge
August 10, 2026 01:07
Leave a FIXME that the method currently nests choose_one() because the autotuner cannot express factorized GEMM tactics yet.
feih-nv
enabled auto-merge (squash)
August 10, 2026 03:26
feih-nv
added a commit
that referenced
this pull request
Aug 12, 2026
…otuning in unified MoE API (#4376) ## 📌 Description Enforce an explicit `check_support() -> build() -> execute` lifecycle for all registered unified MoE runners, and refine CUTLASS compound-tactic autotuning. ### Runner lifecycle - Centralize check_support() → build() → execute enforcement in MoERunner. - Make builds idempotent and defer backend initialization until support validation succeeds. - Keep shape-dependent TRTLLM and b12x inner runners lazy. - Update direct-runner callers and add lifecycle regression tests. ### CUTLASS staged autotuning - Add AutoTuner.rank_tactics() for ranked stage candidates. - Retain the top two tactics for GEMM1 and GEMM2, then profile their four combinations end to end. - Cache complete ranked shortlists in-process while preserving the existing persistent winner format. - Rebuild shortlists from persisted winners when tuning resumes. - Select the correct dynamic optimization profile before ranking. This reduces the search cost from O(n1 × n2) to O(n1 + n2 + k²) while improving compound tactic selection. ## 🔍 Related Issues - Follow up tasks for CUTLASS BF16/W4A16 runners merged in #4328. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used my 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. ## 🧪 Tests Focused validation: - 496 passed - 31 skipped - All changed autotuner and unified MoE test files passed ## Reviewer Notes - `MoELayer` now validates each runner with `check_support()` before calling its idempotent `build()` method. Direct-runner users must perform the same lifecycle explicitly. - TRTLLM and b12x inner runners remain lazily constructed because their configuration depends on runtime shapes. - CUTLASS staged ranking is currently orchestrated by get_valid_tactics(); a future autotuner abstraction could model multi-stage tuning declaratively.
10 tasks
jefby
pushed a commit
to jefby/flashinfer
that referenced
this pull request
Aug 19, 2026
…otuning in unified MoE API (flashinfer-ai#4376) ## 📌 Description Enforce an explicit `check_support() -> build() -> execute` lifecycle for all registered unified MoE runners, and refine CUTLASS compound-tactic autotuning. ### Runner lifecycle - Centralize check_support() → build() → execute enforcement in MoERunner. - Make builds idempotent and defer backend initialization until support validation succeeds. - Keep shape-dependent TRTLLM and b12x inner runners lazy. - Update direct-runner callers and add lifecycle regression tests. ### CUTLASS staged autotuning - Add AutoTuner.rank_tactics() for ranked stage candidates. - Retain the top two tactics for GEMM1 and GEMM2, then profile their four combinations end to end. - Cache complete ranked shortlists in-process while preserving the existing persistent winner format. - Rebuild shortlists from persisted winners when tuning resumes. - Select the correct dynamic optimization profile before ranking. This reduces the search cost from O(n1 × n2) to O(n1 + n2 + k²) while improving compound tactic selection. ## 🔍 Related Issues - Follow up tasks for CUTLASS BF16/W4A16 runners merged in flashinfer-ai#4328. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used my 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. ## 🧪 Tests Focused validation: - 496 passed - 31 skipped - All changed autotuner and unified MoE test files passed ## Reviewer Notes - `MoELayer` now validates each runner with `check_support()` before calling its idempotent `build()` method. Direct-runner users must perform the same lifecycle explicitly. - TRTLLM and b12x inner runners remain lazily constructed because their configuration depends on runtime shapes. - CUTLASS staged ranking is currently orchestrated by get_valid_tactics(); a future autotuner abstraction could model multi-stage tuning declaratively.
5 tasks
YangXu1990uiuc
added a commit
to YangXu1990uiuc/flashinfer
that referenced
this pull request
Aug 28, 2026
The trtllm SM100 MoERunner (_get_trtllm_moe_sm100_module_impl) did not override get_cache_key_extras, so its persisted v2 key was (op, class, profile, ()). runner_hash is dropped from file keys by design, leaving extras the only channel for constructor-fixed config -- so two layers with identical dims/dtypes but different activation (or weight layout, quantization, expert structure) aliased to one stored entry: the second publish() clobbered the first and serving mis-matched tactics under a normal "cache hit" log. In-process tuning was unaffected (runner_hash distinguishes there); only the persistent store collided. Mirror the CUTLASS MoERunner: return the non-shape-derived config as a tuple of int/bool (enums -> int keeps it JSON-round-trippable for preload). flashinfer-ai#4328 fixed only the CUTLASS MoERunner, not this trtllm one. Found by Vincent Tombari during 5-GPU autotuner-v2 validation. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
YangXu1990uiuc
added a commit
to YangXu1990uiuc/flashinfer
that referenced
this pull request
Sep 1, 2026
The trtllm SM100 MoERunner (_get_trtllm_moe_sm100_module_impl) did not override get_cache_key_extras, so its persisted v2 key was (op, class, profile, ()). runner_hash is dropped from file keys by design, leaving extras the only channel for constructor-fixed config -- so two layers with identical dims/dtypes but different activation (or weight layout, quantization, expert structure) aliased to one stored entry: the second publish() clobbered the first and serving mis-matched tactics under a normal "cache hit" log. In-process tuning was unaffected (runner_hash distinguishes there); only the persistent store collided. Mirror the CUTLASS MoERunner: return the non-shape-derived config as a tuple of int/bool (enums -> int keeps it JSON-round-trippable for preload). flashinfer-ai#4328 fixed only the CUTLASS MoERunner, not this trtllm one. Found by Vincent Tombari during 5-GPU autotuner-v2 validation. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
YangXu1990uiuc
added a commit
to YangXu1990uiuc/flashinfer
that referenced
this pull request
Sep 3, 2026
The trtllm SM100 MoERunner (_get_trtllm_moe_sm100_module_impl) did not override get_cache_key_extras, so its persisted v2 key was (op, class, profile, ()). runner_hash is dropped from file keys by design, leaving extras the only channel for constructor-fixed config -- so two layers with identical dims/dtypes but different activation (or weight layout, quantization, expert structure) aliased to one stored entry: the second publish() clobbered the first and serving mis-matched tactics under a normal "cache hit" log. In-process tuning was unaffected (runner_hash distinguishes there); only the persistent store collided. Mirror the CUTLASS MoERunner: return the non-shape-derived config as a tuple of int/bool (enums -> int keeps it JSON-round-trippable for preload). flashinfer-ai#4328 fixed only the CUTLASS MoERunner, not this trtllm one. Found by Vincent Tombari during 5-GPU autotuner-v2 validation. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
YangXu1990uiuc
added a commit
to YangXu1990uiuc/flashinfer
that referenced
this pull request
Sep 4, 2026
The trtllm SM100 MoERunner (_get_trtllm_moe_sm100_module_impl) did not override get_cache_key_extras, so its persisted v2 key was (op, class, profile, ()). runner_hash is dropped from file keys by design, leaving extras the only channel for constructor-fixed config -- so two layers with identical dims/dtypes but different activation (or weight layout, quantization, expert structure) aliased to one stored entry: the second publish() clobbered the first and serving mis-matched tactics under a normal "cache hit" log. In-process tuning was unaffected (runner_hash distinguishes there); only the persistent store collided. Mirror the CUTLASS MoERunner: return the non-shape-derived config as a tuple of int/bool (enums -> int keeps it JSON-round-trippable for preload). flashinfer-ai#4328 fixed only the CUTLASS MoERunner, not this trtllm one. Found by Vincent Tombari during 5-GPU autotuner-v2 validation. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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
3 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.
Expose the existing CUTLASS fused-MoE implementation through the unified API for BF16 and W4A16, with independent GEMM tuning, safe workspace reuse, and CUDA graph support.
Introduce a
build()method toMoERunnerto skip the JIT overhead for unsupported cases.📌 Description
Add unified CUTLASS runners for:
CutlassBf16Config→CutlassBf16RunnerCutlassW4A16Config→CutlassW4A16RunnerCutlassConfigremains as a compatibility placeholder, but it is not registered withMoELayerand is not directly runnable. We can remove it later after we have all the CUTLASS variants supported (e.g. nvfp4).CUTLASS BF16 follows the flat API architecture coverage and is currently GPU-validated here on SM90. CUTLASS W4A16 remains SM90-only because it uses Hopper-specific mixed-input layouts.
Key changes for the runners:
MoELayerregistration.do_finalize=True.top_kin persistent tuning keys.tune_max_num_tokens.For the API change:
__init__()to the idempotentbuild()method.MoELayerenforcecheck_support()beforebuild(), preventing unsupported configurations from triggering JIT compilation.🚀 Pull Request Checklist
✅ Pre-commit Checks
pre-commit.pre-commit run --all-filesand fixed reported issues.🧪 Tests
python3 -m pytest -q -p no:cacheprovider \ tests/moe/test_unified_moe_cutlass.py \ tests/moe/test_unified_moe.py::TestBackendOptions 31 passed, 2 warnings in 5.60s