[Kernel] Support Nvfp4 Cutedsl Moe Swiglu-oai and Relu2(non-gated) Activation - #47106
Conversation
Enable FlashInfer CuteDSL NVFP4 MoE to use the activation-type API so SwiGLU-OAI variants and non-gated ReLU2 can share the same backend path. Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
Add focused coverage for FlashInfer CuteDSL NVFP4 SwiGLU-OAI layout handling and non-gated ReLU2 execution so activation-specific backend wiring is guarded. Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
flashinfer-ai/flashinfer#3737 was merged. We can test this with flashinfer main branch now. |
|
This pull request has merge conflicts that must be resolved before it can be |
…meter Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
Done, resolve merge conflict |
zyongye
left a comment
There was a problem hiding this comment.
Thanks for doing this. Just tiny changes on how to handle default value!
| self.gemm1_alpha = ( | ||
| float(quant_config.gemm1_alpha) | ||
| if quant_config.gemm1_alpha is not None | ||
| else 1.702 | ||
| ) | ||
| self.gemm1_beta = ( | ||
| float(quant_config.gemm1_beta) | ||
| if quant_config.gemm1_beta is not None | ||
| else 1.0 | ||
| ) | ||
| self.gemm1_clamp_limit = ( | ||
| float(quant_config.gemm1_clamp_limit) | ||
| if quant_config.gemm1_clamp_limit is not None | ||
| else 7.0 | ||
| ) |
There was a problem hiding this comment.
I think we want to set these args at None and don't apply them if the parameter in quant_config is None, otherwise there will be other model use this file which has these value in default, causing accuracy bug.
There was a problem hiding this comment.
Good point! It won't affect accuracy unless the activation type is swiglu-oai for flashinfer, but it's indeed a little confusing. Modify done.
Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
|
This pull request has merge conflicts that must be resolved before it can be |
Upstream vllm-project#49580 landed CuTeDSL NVFP4 MoE for ReLU2 and now passes activation_type unconditionally, so the version probe this branch added (has_flashinfer_cutedsl_moe_nvfp4_activation_type) is obsolete: an older FlashInfer already fails on upstream's own relu2 path. Drop the probe and keep only the SwiGLU-OAI work layered on upstream's implementation. - experts/flashinfer_cutedsl_moe.py: take upstream's unconditional activation_type, add SWIGLUOAI/SWIGLUOAI_UNINTERLEAVE to the supported activations, and forward only the SwiGLU params the model actually sets. The kernel defaults swiglu_{alpha,beta,limit} to floats (1.0/0.0/FLT_MAX), so None must not be passed through. - oracle/nvfp4.py: register FLASHINFER_CUTEDSL in NVFP4_BACKENDS_WITH_CLAMP unconditionally and list it in the clamp error message, which upstream had extended with cutlass/marlin/humming. - flashinfer_fp4_moe.py: keep the SWIGLUOAI interleaved-w13 handling as reorder_w13_to_w31_for_flashinfer_cutedsl() on top of upstream's is_gated guard. - flashinfer_utils.py: drop the duplicate SWIGLUOAI entries the auto-merge left in ACTIVATION_TO_FI_ACTIVATION; upstream already maps both. - tests: rebase onto upstream's test file, parametrized over SILU/RELU2_NO_MUL/SWIGLUOAI. Tested on B200 (SM100): test_flashinfer_cutedsl_nvfp4_moe.py 3 passed, test_flashinfer_cutedsl_layout.py 3 passed, test_cutedsl_moe.py 11 passed, test_nvfp4_moe.py 84 passed. Signed-off-by: Tiekai Bi <tiekaib@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The test inherited atol=rtol=2e-1 from the upstream relu2 test, but the outputs here are O(1e-2) while NVFP4 noise is O(1e-3): for SILU and RELU2_NO_MUL the tolerance exceeded the reference tensor's full dynamic range, so a kernel returning zeros would have passed. Two of the four activations the backend advertises had no coverage at all. - Parametrize over (activation, alpha, beta, limit) and add cases for SWIGLUOAI_UNINTERLEAVE and SILU-with-clamp. The latter exercises the NVFP4_BACKENDS_WITH_CLAMP entry this branch adds, which was untested. - Drop swiglu_limit from 7.0 to 0.3: the GEMM1 output is O(0.5) at these tensor scales, so the OAI default never clamped anything. - Build the reference from vLLM's own SiluAndMul / SiluAndMulWithClamp / SwigluOAIAndMul ops instead of torch_moe, which cannot parametrize the activation (it calls the registered op with default constructor args). - Tighten to atol=3e-2 (worst measured error 1.8e-2) and add a cosine check, which catches structural errors an absolute tolerance cannot. Mutation-tested: removing **swiglu_kwargs fails silu-clamp, swigluoai and swigluoai_uninterleave while leaving silu and relu2_no_mul passing; disabling the SWIGLUOAI interleaved-w13 branch fails swigluoai. Full run: 100 passed across the four cutedsl/nvfp4 MoE suites on B200. Signed-off-by: Tiekai Bi <tiekaib@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
/ci run |
…tivation (vllm-project#47106) Signed-off-by: Tiekai Bi <tiekaib@nvidia.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Purpose
FlashInfer modified the CuteDSL NVFP4 MoE kernel to support the SwiGLU-OAI activation (flashinfer-ai/flashinfer#3737). The CuteDSL NVFP4 MoE kernel now also supports ReLU² non-gated activation. This PR achieves compatibility with both new activations by passing the activation type and related parameters into the CuteDSL MoE kernel.
Test Plan
Test input layout prepare:
.venv/bin/python -m pytest tests/kernels/moe/test_flashinfer_cutedsl_layout.py -vTest swiglu oai and relu2 accuracy:
.venv/bin/python -m pytest tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py -vTest original model can still run successfully:
.venv/bin/python -m pytest tests/models/quantization/test_nvfp4.py -v -k "flashinfer_cutedsl"Test Result
Test input layout prepare:
.venv/bin/python -m pytest tests/kernels/moe/test_flashinfer_cutedsl_layout.py -vTest swiglu oai and relu2 accuracy:
.venv/bin/python -m pytest tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py -vTest original model can still run successfully:
.venv/bin/python -m pytest tests/models/quantization/test_nvfp4.py -v -k "flashinfer_cutedsl"Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.