Sf only per tok quant - #4042
IwakuraRein wants to merge 13 commits into
Conversation
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds scale-only NVFP4 quantization kernels for FP32 block scales, extends MoE runner dtype contracts, and selects between E2m1 scale-only and BF16 NVFP4 quantization paths. FP4 launcher allocation, environment handling, kernel filtering, artifact references, benchmark inputs, and test fixtures are updated. ChangesNVFP4 MoE quantization
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant FP4BlockScaleLauncher
participant MoERunner
participant PermuteGemm1
participant ScaleOnlyQuantization
participant NVFP4Quantization
FP4BlockScaleLauncher->>MoERunner: configure dtypeGemm1Output
MoERunner->>PermuteGemm1: run FC1 with selected output dtype
alt E2m1 output
MoERunner->>ScaleOnlyQuantization: process FP32 block scales
else BF16 output
MoERunner->>NVFP4Quantization: quantize BF16 activation
end
Possibly related PRs
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 |
![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 introduces support for fast-math-free FP4 quantization in MoE kernels, adding the scaleOnlyQuantAndPerTokenScaleKernel to quantize FP32 NVFP4 dequantization scales to E4M3 and find row-wise amax when raw NVFP4 input is unavailable. It also updates the MoE runner and launcher to handle BF16 or E2M1 GEMM1 output datatypes based on the fast math configuration. Feedback on these changes includes: removing the redundant and potentially undefined tg:: namespace prefix in trtllm_batched_gemm_runner.cu, adding a null check for workspace.activation_output_scale in trtllm_fused_moe_runner.cu to prevent dereferencing issues, and handling potential shared memory limit exhaustion in quantization.cu if smem_size exceeds 48KB.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
|
/bot run |
|
[FAILED] Pipeline #58862832: 7/20 passed |
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
|
/bot run |
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
csrc/trtllm_fused_moe_kernel_launcher.cu (1)
538-553: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffConsider including
mDtypeGemm1OutputinrunnerKey.The runner cache key omits
mDtypeGemm1Output, yet the runner is constructed withmDtypeGemm1Output.value_or(mDtypeAct)(Line 570). This is only safe today because that dtype is fully derivable fromusePerTokenScalingGemm2(already in the key) plus the process-constantdisableFP4QuantFastMath. If the selection ever gains a runtime-varying input, cached runners could be returned with a stale FC1 output dtype. Adding the effectivedtypeGemm1Outputto the key makes this robust to future changes.🤖 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 `@csrc/trtllm_fused_moe_kernel_launcher.cu` around lines 538 - 553, Update the runnerKey construction in the launcher to include the effective mDtypeGemm1Output value, matching the value passed to runner construction via mDtypeGemm1Output.value_or(mDtypeAct). Ensure distinct FC1 output dtypes produce distinct cached runners while preserving the existing key fields and ordering.
🤖 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 `@csrc/trtllm_fused_moe_kernel_launcher.cu`:
- Around line 538-553: Update the runnerKey construction in the launcher to
include the effective mDtypeGemm1Output value, matching the value passed to
runner construction via mDtypeGemm1Output.value_or(mDtypeAct). Ensure distinct
FC1 output dtypes produce distinct cached runners while preserving the existing
key fields and ordering.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9f79e499-11d9-4ac4-81d9-5885369eb398
📒 Files selected for processing (9)
csrc/nv_internal/cpp/kernels/quantization.cucsrc/nv_internal/tensorrt_llm/kernels/quantization.cuhcsrc/nv_internal/tensorrt_llm/kernels/quantization.hcsrc/trtllm_batched_gemm_runner.cucsrc/trtllm_fused_moe_kernel_launcher.cucsrc/trtllm_fused_moe_runner.cuflashinfer/artifacts.pyinclude/flashinfer/trtllm/fused_moe/runner.htests/moe/utils.py
|
[FAILED] Pipeline #59145927: 6/20 passed |
|
/bot run tests/moe |
|
[FAILED] Pipeline #59274664 — 16/18 executed test jobs passed Compared with nightly #59054944. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsNew relative to nightly (attribution uncertain)
|
| Optional<Tensor> activation_output_scale; | ||
| TensorView topk_ids; // [num_tokens, top_k] - pre-computed or output top-k expert indices | ||
| TensorView topk_weights; // [num_tokens, top_k] - pre-computed or output top-k routing weights | ||
| static bool const disableFP4QuantFastMath; |
There was a problem hiding this comment.
Is caching an env var in a static variable common practice in the codebase? I can imagine it gets initialized once and remains unchanged across test fixtures that enable/disable FP4 quantization fast math.
There was a problem hiding this comment.
Right. This variable will be fixed after .so is loaded. Thanks for point out.
|
LGTM overall. Left a minor comment. |
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
|
/bot run tests/moe |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
benchmarks/bench_nvfp4_quantize_backend_comparison.py (1)
92-107: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the backend-specific scale placement.
The CPU/GPU split at Line [104] is intentional but easy to regress. Add a short comment explaining that CUDA extracts the scalar from the host copy, while CuTe-DSL consumes the CUDA tensor.
As per coding guidelines, intentional departures from established efficiency or defensiveness must be documented with rationale.
🤖 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 `@benchmarks/bench_nvfp4_quantize_backend_comparison.py` around lines 92 - 107, Add a concise inline comment in _run_nvfp4_quantize immediately before the global_sf_cpu if backend == "cuda" selection, documenting that CUDA extracts the scalar from the host copy while CuTe-DSL consumes the CUDA tensor. Preserve the existing backend-specific argument behavior.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.
Inline comments:
In `@benchmarks/bench_nvfp4_quantize_backend_comparison.py`:
- Around line 149-151: Move the global_sf.cpu() call inside the try block in
verify_nvfp4_correctness, keeping the CUDA-to-host copy covered by the existing
exception handling so deferred CUDA errors return the failure tuple instead of
aborting the sweep.
---
Nitpick comments:
In `@benchmarks/bench_nvfp4_quantize_backend_comparison.py`:
- Around line 92-107: Add a concise inline comment in _run_nvfp4_quantize
immediately before the global_sf_cpu if backend == "cuda" selection, documenting
that CUDA extracts the scalar from the host copy while CuTe-DSL consumes the
CUDA tensor. Preserve the existing backend-specific argument behavior.
🪄 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: a3887fc2-0878-4074-8bfa-884b8c03489c
📒 Files selected for processing (2)
benchmarks/bench_nvfp4_quantize_backend_comparison.pycsrc/trtllm_fused_moe_kernel_launcher.cu
🚧 Files skipped from review as they are similar to previous changes (1)
- csrc/trtllm_fused_moe_kernel_launcher.cu
| global_sf_cpu = global_sf.cpu() | ||
|
|
||
| try: |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate and inspect the referenced function around the reported lines.
file="benchmarks/bench_nvfp4_quantize_backend_comparison.py"
if [ ! -f "$file" ]; then
echo "Referenced file not found: $file"
fd -a 'bench_nvfp4_quantize_backend_comparison.py' .
exit 0
fi
wc -l "$file"
sed -n '110,190p' "$file" | cat -n -v
echo '--- all verify_nvfp4_correctness definitions/usages ---'
rg -n "def verify_nvfp4_correctness|verify_nvfp4_correctness|_split_quantize_result" "$file" benchmarks || trueRepository: flashinfer-ai/flashinfer
Length of output: 4881
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="benchmarks/bench_nvfp4_quantize_backend_comparison.py"
echo '--- verify_nvfp4_correctness full body ---'
sed -n '119,320p' "$file" | cat -n -v
echo '--- try/except structure search in file ---'
python3 - <<'PY'
from pathlib import Path
import ast
path = Path("benchmarks/bench_nvfp4_quantize_backend_comparison.py")
tree = ast.parse(path.read_text())
for node, label in (
(tree.body[0], "module level"),
(tree.body[1], "verify_nvfp4_correctness"),
):
if isinstance(node, ast.FunctionDef):
print(f"{label}: lines {node.lineno}-{node.end_lineno}")
for c in node.body:
print("- " + ast.unparse(c).splitlines()[0][:120])
PY
echo '--- call site surrounding try/except ---'
sed -n '380,445p' "$file" | cat -n -vRepository: flashinfer-ai/flashinfer
Length of output: 11108
Keep the CUDA-to-host copy inside verify_nvfp4_correctness.
global_sf.cpu() still runs before the try block that returns verify_nvfp4_correctness’s failure tuple. Since .cpu() can synchronize and surface deferred CUDA errors, wrap it inside the same error boundary so failures are reported instead of aborting the sweep.
🤖 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 `@benchmarks/bench_nvfp4_quantize_backend_comparison.py` around lines 149 -
151, Move the global_sf.cpu() call inside the try block in
verify_nvfp4_correctness, keeping the CUDA-to-host copy covered by the existing
exception handling so deferred CUDA errors return the failure tuple instead of
aborting the sweep.
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
|
/bot run tests/moe |
|
[FAILED] Pipeline #60004868 — 16/18 executed test jobs passed Compared with nightly #59757517 (different CI configuration). Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 4/6 passed
No individual test or infrastructure failures could be extracted. |
|
[FAILED] Pipeline #60004868 — 16/18 executed test jobs passed No usable JUnit artifact was available; individual tests and nightly comparison could not be recovered. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 4/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
1 similar comment
|
[FAILED] Pipeline #60004868 — 16/18 executed test jobs passed No usable JUnit artifact was available; individual tests and nightly comparison could not be recovered. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 4/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
3 participants
📌 Description
Add a faster nvfp4 per-token quantization path for trtllm-gen Nvfp4 MoE. The new quantization kernel achieves speedup for large batch size.
scaleOnlyQuantAndPerTokenScaleKernelPerformance
python benchmarks/bench_moe_deepseek.py --use-per-token-activation --ep 8🔍 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
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Chores