feat(moe): enable MxInt4 in unified API - #4320
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📥 CommitsReviewing files that changed from the base of the PR and between c4415b286c1bac30a91c9d08a9b3d39d18418a29 and 66ea1df. 📒 Files selected for processing (8)
🚧 Files skipped from review as they are similar to previous changes (8)
📝 WalkthroughWalkthroughThe PR adds unified TRTLLM MxInt4 MoE support. It prepares packed INT4 weights, registers a routed runner, supports packed and ChangesUnified TRTLLM MxInt4 MoE
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant MoELayer
participant TrtllmMxInt4RoutedRunner
participant TRTLLMMoERunner
MoELayer->>TrtllmMxInt4RoutedRunner: pack activations, routing, and weights
TrtllmMxInt4RoutedRunner->>TRTLLMMoERunner: construct inputs and launch parameters
TRTLLMMoERunner-->>TrtllmMxInt4RoutedRunner: return MoE output
TrtllmMxInt4RoutedRunner-->>MoELayer: return output buffer
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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 |
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
/bot run tests/moe |
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
flashinfer/fused_moe/prepare.py (1)
986-995: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPack the nibbles in an unsigned dtype.
quantizedisint8.(quantized[..., 1] & 0x0F) << 4produces values up to 240, which does not fit inint8. The final.to(torch.uint8)restores the intended byte, so the current result is correct, but the code depends on signed-overflow wrap behavior in the shift. Cast touint8before the shift to make the packing well-defined.♻️ Proposed refactor
quantized = ( (blocks * scales.reciprocal()) .round() .clamp(-8, 7) .to(torch.int8) .reshape(-1, sf_vec_size // 2, 2) ) - packed = (quantized[..., 0] & 0x0F) | ((quantized[..., 1] & 0x0F) << 4) + nibbles = (quantized & 0x0F).to(torch.uint8) + packed = nibbles[..., 0] | (nibbles[..., 1] << 4) return ( - packed.reshape(*weights.shape[:-1], weights.shape[-1] // 2).to(torch.uint8), + packed.reshape(*weights.shape[:-1], weights.shape[-1] // 2), scales.to(torch.bfloat16), )🤖 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/prepare.py` around lines 986 - 995, Update the nibble-packing expression in the quantization flow to cast the masked quantized values to an unsigned 8-bit dtype before applying the left shift and combining them. Preserve the existing low/high nibble layout and returned packed tensor shape.
🤖 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/prepare.py`:
- Around line 1050-1075: Separate payload and scale permutation caches in the
preparation flow around prepare_trtllm_mxint4_weights to prevent cache-key
collisions between tensors using different num_elts_per_sf values. Create
distinct cache dictionaries for payload and scale permutations, use the payload
cache for w1_q and w2_q calls, and pass scale_permute_cache to the w1_sf and
w2_sf calls that specify num_elts_per_sf.
---
Nitpick comments:
In `@flashinfer/fused_moe/prepare.py`:
- Around line 986-995: Update the nibble-packing expression in the quantization
flow to cast the masked quantized values to an unsigned 8-bit dtype before
applying the left shift and combining them. Preserve the existing low/high
nibble layout and returned packed tensor shape.
🪄 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 Plus
Run ID: 4c18b9aa-d455-4575-910a-d019b73537e4
📥 Commits
Reviewing files that changed from the base of the PR and between 1d2fc15 and e324e4b4a71780d69512b0dace85a89103bd6fdc.
📒 Files selected for processing (7)
flashinfer/fused_moe/__init__.pyflashinfer/fused_moe/api.pyflashinfer/fused_moe/layer.pyflashinfer/fused_moe/prepare.pyflashinfer/fused_moe/runners.pytests/moe/test_unified_moe_fuzz.pytests/moe/test_unified_moe_mxint4.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/core.py`:
- Around line 173-180: Scope every TRTLLM permutation cache key by the input
tensor’s CUDA device so cached index tensors are never reused across GPUs.
Update the key construction for _TRTLLM_PERMUTE_CACHE,
_TRTLLM_FP8_PERMUTE_CACHE, _TRTLLM_FP8_PER_TENSOR_PERMUTE_CACHE, and
_TRTLLM_MXINT4_PERMUTE_CACHE, including the cache path around the second
referenced block, using the relevant tensor device while preserving all existing
key parameters.
🪄 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 Plus
Run ID: 15cf60bc-52f9-4b6c-b9ad-89afa1eceb46
📥 Commits
Reviewing files that changed from the base of the PR and between e324e4b4a71780d69512b0dace85a89103bd6fdc and c4415b286c1bac30a91c9d08a9b3d39d18418a29.
📒 Files selected for processing (3)
flashinfer/fused_moe/core.pyflashinfer/fused_moe/prepare.pytests/moe/test_unified_moe_mxint4.py
🚧 Files skipped from review as they are similar to previous changes (1)
- flashinfer/fused_moe/prepare.py
|
[SUCCESS] Pipeline #60838824: 18/18 executed test jobs passed |
Prevent CUDA-graph stale-output false passes, add explicit autotune coverage, and tighten dedicated and fuzz numerical thresholds using measured error bounds.
c4415b2 to
66ea1df
Compare
3 participants
📌 Description
Enable TensorRT-LLM MxInt4 execution through FlashInfer’s unified MoE API.
This change adds production BF16-to-MxInt4 weight preparation through
TrtllmMxInt4Config.prepare_weights(), including signed INT4 packing, BF16 block scales, gated-activation row permutation, scale interleaving, and BlockMajorK conversion.It also adds and registers
TrtllmMxInt4RoutedRunner, supporting:FromLogitssupport with BF16 router logitsThe unified MoE fuzzer now includes an MxInt4 handler, a quantization-aware reference, and curated packed and
FromLogitsseeds.🚀 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
SM100
SM107
Summary by CodeRabbit
New Features
Bug Fixes
Tests