benchmark: Add MXFP4/MXFP8 quantization mode support to FP4 MoE benchmark#2635
benchmark: Add MXFP4/MXFP8 quantization mode support to FP4 MoE benchmark#2635bkryu merged 5 commits intoflashinfer-ai:mainfrom
Conversation
bkryu
requested review from
Anerudhan,
cyx-6,
jiahanc and
nv-yunzheq
as code owners
Summary of ChangesHello @bkryu, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the FP4 MoE benchmarking capabilities by introducing support for MXFP4 and MXFP8 quantization modes, allowing for more diverse performance analysis. It also ensures broader compatibility with different versions of the FlashInfer library by adapting to API changes related to activation types. Additionally, a minor but important fix was applied to ensure accurate and complete data logging in the benchmark's CSV output. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
📝 WalkthroughWalkthroughAdds FP4 mode selection and ActivationType compatibility to MoE benchmark code: updates output columns (removes Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant CLI as CLI
participant Bench as Benchmark Runner
participant Q as Quantizer (moe_utils)
participant Kernel as Fused MoE Kernel
participant Store as Result Collector
CLI->>Bench: parse args (--fp4_mode, activation_type, other flags)
Bench->>Q: prepare/quantize tensors (fp4_mode, sf_vec_size)
Q-->>Bench: quantized tensors, scales, metadata
Bench->>Kernel: invoke kernel (passes activation kwarg via _activation_kwarg)
Kernel-->>Bench: runtime metrics, outputs
Bench->>Store: append result (includes fp4_mode, activation_type, formats)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Tip Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs). 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}
There was a problem hiding this comment.
Code Review
This pull request adds support for MXFP4/MXFP8 quantization modes to the FP4 MoE benchmark, including a new --fp4_mode CLI argument. It also introduces backward compatibility for older flashinfer versions by dynamically handling ActivationType and gated_act_type API differences, and fixes a CSV output column mismatch.
The changes are well-structured and the backward compatibility is handled cleanly. However, I've found a critical issue related to the quantization of hidden_states. The swizzled layout for scale factors requires the input tensor's row count to be a multiple of 128, but hidden_states can have an arbitrary number of rows (num_tokens). This will lead to runtime errors. I've provided suggestions to use a non-swizzled layout for activations to fix this.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
benchmarks/routines/moe.py (1)
1358-1383:⚠️ Potential issue | 🟡 Minor
activation_typenot written tocur_resintestTrtllmFp8BlockScaleMoe.The
"activation_type"column now exists inoutput_column_dict["moe"], andtestTrtllmFp8PerTensorScaleMoecorrectly populatescur_res["activation_type"], buttestTrtllmFp8BlockScaleMoeomits this field, leaving it empty in CSV output. Sinceactivation_typeis accepted by neither the FP8 block-scale kernel (not passed at all inrun_fp8_block_moe) nor written tocur_res, the omission is consistent but the CSV column will always be blank for this routine.Consider adding
cur_res["activation_type"] = args.activation_type.namefor completeness, or if the FP8 block-scale kernel genuinely ignores activation type, add a comment explaining why it's excluded.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@benchmarks/routines/moe.py` around lines 1358 - 1383, The CSV result for testTrtllmFp8BlockScaleMoe omits the activation_type field causing an empty column; update the code that builds cur_res in testTrtllmFp8BlockScaleMoe to include cur_res["activation_type"] = args.activation_type.name (or args.activation_type) so the activation_type is recorded like in testTrtllmFp8PerTensorScaleMoe, and if run_fp8_block_moe genuinely ignores activation type add a brief comment near run_fp8_block_moe explaining why activation_type is not used.
🧹 Nitpick comments (2)
benchmarks/routines/moe.py (2)
543-544: Redundant.to(torch.bfloat16)—hidden_statesis already BF16.
hidden_statesis created astorch.bfloat16at thecreate_trtllm_moe_test_datacall. The cast is a no-op and adds a small overhead in setup.♻️ Suggested fix
if fp4_mode == "mxfp4_bf16": - hidden_states_fp4 = hidden_states.to(torch.bfloat16) + hidden_states_fp4 = hidden_states hidden_states_scale_linear_fp4 = None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@benchmarks/routines/moe.py` around lines 543 - 544, Remove the redundant cast to bfloat16: in the block handling fp4_mode == "mxfp4_bf16" (where the variable hidden_states is set), delete the .to(torch.bfloat16) call or guard it so you only cast if hidden_states.dtype is not torch.bfloat16; this touches the fp4_mode check around hidden_states assignment (created by create_trtllm_moe_test_data) — simply rely on the input being BF16 or perform a dtype check before casting to avoid the no-op overhead.
61-68:inspect.signatureis called on every benchmark iteration — cache the result.
_activation_kwargis invoked insiderun_fp4_moeandrun_fp8_per_tensor_moe, which are called for each iteration bybench_gpu_time.inspect.signaturehas non-trivial overhead. Pre-compute it once before the benchmark loop.♻️ Suggested fix — compute kwargs once outside the closure
-def _activation_kwarg(fn, activation_type: ActivationType) -> dict: - """Return the correct activation keyword argument for *fn* in the installed version.""" - sig = inspect.signature(fn) - if "activation_type" in sig.parameters: - return {"activation_type": activation_type.value} - if "gated_act_type" in sig.parameters: - return {"gated_act_type": _ACTIVATION_TO_GATED_ACT.get(activation_type, 0)} - return {}Then at each call site, pre-compute outside the inner closure:
# Before defining run_fp4_moe / run_fp8_per_tensor_moe: activation_kwargs = _activation_kwarg(trtllm_fp4_block_scale_moe, activation_type) def run_fp4_moe(...): return trtllm_fp4_block_scale_moe( ..., **activation_kwargs, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@benchmarks/routines/moe.py` around lines 61 - 68, The helper _activation_kwarg currently calls inspect.signature(fn) on every benchmark iteration; compute and cache the activation kwargs once before the benchmark loop and pass them into the per-iteration functions instead of calling _activation_kwarg repeatedly. Concretely, call _activation_kwarg(trtllm_fp4_block_scale_moe, activation_type) (and any other target model functions) once and store the returned dict (e.g., activation_kwargs), then update run_fp4_moe and run_fp8_per_tensor_moe to accept/use the precomputed activation_kwargs and spread them into the model invocation instead of calling _activation_kwarg inside the per-iteration closure.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@benchmarks/routines/moe.py`:
- Around line 563-572: The code currently replaces
hidden_states_scale_linear_fp4 with a synthetic all-ones tensor when its element
count mismatches expected_scale_elems but only prints an info message if
args.verbose >= 1; change this to always emit a warning (unconditional) when the
fallback happens and include context (current and expected sizes and that an
all-ones tensor is being used) so users see the substitution even at default
verbosity. Locate the block using expected_scale_elems and
hidden_states_scale_linear_fp4 and replace the gated print with a call that
always logs or prints a clear warning message before creating the
torch.ones(...) fallback (include device and dtype details in the message if
helpful). Ensure behavior of assigning hidden_states_scale_linear_fp4 =
torch.ones(expected_scale_elems, device=device, dtype=torch.float8_e4m3fn)
remains unchanged.
- Around line 66-67: The code silently falls back to 0 for unknown
activation_type in the legacy gated_act_type path; update the block that checks
"gated_act_type" to detect when activation_type is not a key in
_ACTIVATION_TO_GATED_ACT and emit a warning (e.g., warnings.warn or
logger.warning) including the unexpected activation_type and that we're using
the legacy fallback value, then return the fallback {"gated_act_type":
_ACTIVATION_TO_GATED_ACT.get(activation_type, 0)} as before; reference the
activation_type variable and the _ACTIVATION_TO_GATED_ACT mapping to locate
where to add the warning.
- Around line 547-551: The code is reinterpreting uint8 scale bytes as float8
via hs_scale.view(torch.float8_e4m3fn) which is unsafe; replace that
reinterpretation by explicitly converting the scale values from the returned
uint8 representation into the float8 numeric type (or use an existing helper
like a mxfp8 scale decode/dequantize function) before reshaping. In other words,
locate mxfp8_quantize and the variables hs_scale and
hidden_states_scale_linear_fp4 and change the
hs_scale.view(torch.float8_e4m3fn).reshape(...) step to perform a proper
conversion/cast from torch.uint8 to torch.float8_e4m3fn (or call the mxfp8
scale-decoding routine) and then reshape to (num_tokens, -1), rather than
reinterpreting raw bytes.
---
Outside diff comments:
In `@benchmarks/routines/moe.py`:
- Around line 1358-1383: The CSV result for testTrtllmFp8BlockScaleMoe omits the
activation_type field causing an empty column; update the code that builds
cur_res in testTrtllmFp8BlockScaleMoe to include cur_res["activation_type"] =
args.activation_type.name (or args.activation_type) so the activation_type is
recorded like in testTrtllmFp8PerTensorScaleMoe, and if run_fp8_block_moe
genuinely ignores activation type add a brief comment near run_fp8_block_moe
explaining why activation_type is not used.
---
Nitpick comments:
In `@benchmarks/routines/moe.py`:
- Around line 543-544: Remove the redundant cast to bfloat16: in the block
handling fp4_mode == "mxfp4_bf16" (where the variable hidden_states is set),
delete the .to(torch.bfloat16) call or guard it so you only cast if
hidden_states.dtype is not torch.bfloat16; this touches the fp4_mode check
around hidden_states assignment (created by create_trtllm_moe_test_data) —
simply rely on the input being BF16 or perform a dtype check before casting to
avoid the no-op overhead.
- Around line 61-68: The helper _activation_kwarg currently calls
inspect.signature(fn) on every benchmark iteration; compute and cache the
activation kwargs once before the benchmark loop and pass them into the
per-iteration functions instead of calling _activation_kwarg repeatedly.
Concretely, call _activation_kwarg(trtllm_fp4_block_scale_moe, activation_type)
(and any other target model functions) once and store the returned dict (e.g.,
activation_kwargs), then update run_fp4_moe and run_fp8_per_tensor_moe to
accept/use the precomputed activation_kwargs and spread them into the model
invocation instead of calling _activation_kwarg inside the per-iteration
closure.
ℹ️ Review info
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
benchmarks/routines/flashinfer_benchmark_utils.pybenchmarks/routines/moe.pybenchmarks/routines/moe_utils.py
There was a problem hiding this comment.
🧹 Nitpick comments (2)
benchmarks/routines/moe.py (2)
549-551: Redundant.to(torch.bfloat16)—hidden_statesis alreadybfloat16.
create_trtllm_moe_test_dataconstructshidden_statesasdtype=torch.bfloat16, so the cast on line 550 is a no-op. Not harmful, but a minor clarity issue.♻️ Proposed fix
if fp4_mode == "mxfp4_bf16": - hidden_states_fp4 = hidden_states.to(torch.bfloat16) + hidden_states_fp4 = hidden_states # already bfloat16 hidden_states_scale_linear_fp4 = None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@benchmarks/routines/moe.py` around lines 549 - 551, Redundant cast: remove the unnecessary .to(torch.bfloat16) call and assign hidden_states_fp4 = hidden_states directly in the fp4_mode == "mxfp4_bf16" branch (leave hidden_states_scale_linear_fp4 = None); locate the branch using the fp4_mode variable and the hidden_states_fp4/hidden_states names in create_trtllm_moe_test_data to make the change.
68-72: Ruff TRY003 — consider shortening inline exception message or wrapping in a custom exception.Both this block (line 68–72) and the equivalent at lines 554–558 are flagged by Ruff
TRY003for long messages outside the exception class. For a benchmark file the impact is minimal, but aligning with the linter keeps CI clean.♻️ Suggested fix (example for lines 68–72)
- if activation_type not in _ACTIVATION_TO_GATED_ACT: - raise ValueError( - f"Activation type {activation_type.name} is not supported by the " - f"installed flashinfer version (pre-0.6.3 only supports " - f"{[k.name for k in _ACTIVATION_TO_GATED_ACT]})" - ) + if activation_type not in _ACTIVATION_TO_GATED_ACT: + supported = [k.name for k in _ACTIVATION_TO_GATED_ACT] + raise ValueError( + f"Activation {activation_type.name!r} unsupported pre-0.6.3; supported: {supported}" + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@benchmarks/routines/moe.py` around lines 68 - 72, The ValueError raised for unsupported activations has an overly long inline message; shorten it to a concise message like "Unsupported activation type: {activation_type.name}" and move the detailed list of supported activations (from _ACTIVATION_TO_GATED_ACT) into either a separate variable used for logging or into a small custom exception class (e.g., ActivationNotSupported) that formats the full detail in its __str__; update the raise sites (the ValueError instances around activation_type and the equivalent at the other location) to use the short message or raise the new ActivationNotSupported to satisfy Ruff TRY003 while preserving the detailed info elsewhere for debugging.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@benchmarks/routines/moe.py`:
- Around line 549-551: Redundant cast: remove the unnecessary
.to(torch.bfloat16) call and assign hidden_states_fp4 = hidden_states directly
in the fp4_mode == "mxfp4_bf16" branch (leave hidden_states_scale_linear_fp4 =
None); locate the branch using the fp4_mode variable and the
hidden_states_fp4/hidden_states names in create_trtllm_moe_test_data to make the
change.
- Around line 68-72: The ValueError raised for unsupported activations has an
overly long inline message; shorten it to a concise message like "Unsupported
activation type: {activation_type.name}" and move the detailed list of supported
activations (from _ACTIVATION_TO_GATED_ACT) into either a separate variable used
for logging or into a small custom exception class (e.g.,
ActivationNotSupported) that formats the full detail in its __str__; update the
raise sites (the ValueError instances around activation_type and the equivalent
at the other location) to use the short message or raise the new
ActivationNotSupported to satisfy Ruff TRY003 while preserving the detailed info
elsewhere for debugging.
| choices=["nvfp4", "mxfp4_mxfp8", "mxfp4_bf16"], | ||
| help=( | ||
| "FP4 quantization mode for trtllm_fp4_block_scale_moe: " | ||
| "nvfp4 (NvFP4 weights + NvFP4 hidden states, block_size=16), " |
There was a problem hiding this comment.
If these block_size are hardcoded numbers, we can create a utility function or a dict to store the block_size, e.g.:
sf_vec_size = {
"nvfp4": 16,
"mxfp4": 32,
}
benchmarks/routines/moe_utils.py
Outdated
| return 0.5 + 1 / 32 | ||
| elif fmt == "mxfp8": | ||
| # 1 e4m3 + 1 ue8m0 scale per 32-element block | ||
| return 1.0 + 1 / 32 |
There was a problem hiding this comment.
Ditto; addressed in latest commit d149911
There was a problem hiding this comment.
🧹 Nitpick comments (2)
benchmarks/routines/moe_utils.py (1)
752-752:sf_vec_size = 16should useSF_VEC_SIZE["nvfp4"]for consistency.This function is the only remaining site with a hardcoded block size now that
SF_VEC_SIZEexists.♻️ Proposed fix
- sf_vec_size = 16 + sf_vec_size = SF_VEC_SIZE["nvfp4"]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@benchmarks/routines/moe_utils.py` at line 752, Replace the hardcoded block size assignment "sf_vec_size = 16" with the centralized constant lookup SF_VEC_SIZE["nvfp4"] so the function uses the shared SF_VEC_SIZE map; locate the occurrence of the variable sf_vec_size in the file and change its value to be assigned from SF_VEC_SIZE["nvfp4"] (ensuring SF_VEC_SIZE is imported/available in the module).benchmarks/routines/moe.py (1)
62-75:ValueErrorfor unsupported activation on pre-0.6.3 APIs is a good improvement over a silent fallback.Minor: Ruff flags the long inline message at lines 69-73 (
TRY003). Consider moving it into a dedicated exception class or shortening it if you want to silence linting.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@benchmarks/routines/moe.py` around lines 62 - 75, The long inline ValueError message in _activation_kwarg triggers Ruff TRY003; replace the multi-line inline message with a concise raise using a dedicated exception class or a shorter message: define a custom exception (e.g., UnsupportedActivationError) near the top of the module that formats or stores the full multi-line guidance, then in _activation_kwarg raise UnsupportedActivationError(activation_type, list(_ACTIVATION_TO_GATED_ACT)) or raise ValueError with a single-line message that references the custom exception or points to docs; reference the symbols _activation_kwarg, ActivationType, _ACTIVATION_TO_GATED_ACT and the ValueError raise site when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@benchmarks/routines/moe_utils.py`:
- Line 752: Replace the hardcoded block size assignment "sf_vec_size = 16" with
the centralized constant lookup SF_VEC_SIZE["nvfp4"] so the function uses the
shared SF_VEC_SIZE map; locate the occurrence of the variable sf_vec_size in the
file and change its value to be assigned from SF_VEC_SIZE["nvfp4"] (ensuring
SF_VEC_SIZE is imported/available in the module).
In `@benchmarks/routines/moe.py`:
- Around line 62-75: The long inline ValueError message in _activation_kwarg
triggers Ruff TRY003; replace the multi-line inline message with a concise raise
using a dedicated exception class or a shorter message: define a custom
exception (e.g., UnsupportedActivationError) near the top of the module that
formats or stores the full multi-line guidance, then in _activation_kwarg raise
UnsupportedActivationError(activation_type, list(_ACTIVATION_TO_GATED_ACT)) or
raise ValueError with a single-line message that references the custom exception
or points to docs; reference the symbols _activation_kwarg, ActivationType,
_ACTIVATION_TO_GATED_ACT and the ValueError raise site when making the change.
…mark (flashinfer-ai#2635) <!-- .github/pull_request_template.md --> ## 📌 Description * Add --fp4_mode CLI argument to trtllm_fp4_block_scale_moe benchmark with three modes: * nvfp4 (default, existing behavior): NvFP4 weights + NvFP4 hidden states, block_size=16 * mxfp4_bf16: MXFP4 weights + BF16 hidden states, block_size=32 * mxfp4_mxfp8: MXFP4 weights + MXFP8 hidden states, block_size=32 * Add backward compatibility with flashinfer 0.6.0 (pre-0.6.3), where ActivationType was not yet exported from the top-level package and MoE APIs used gated_act_type instead of activation_type * Fix CSV output column mismatch: moe.py wrote to cur_res["activation_type"] but the CSV column was "gated_act", causing the field to be silently empty <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 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 * **New Features** * Added FP4 mode selection for benchmarks (nvfp4 default, mxfp4_mxfp8, mxfp4_bf16) with a CLI flag and updated benchmark output to include activation type and FP4 mode. * Benchmarks now report FP-format-aware size/scale behavior for new modes. * **Chores** * Backward-compatible handling of older activation argument names. * Adjusted quantization and vector-size logic to support mxfp4/mxfp8 modes. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>
3 participants
📌 Description
🔍 Related Issues
🚀 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
New Features
Chores