Repository navigation
[AMD] Select Quark MXFP4 dense GEMM via --fp4-gemm-backend aiter|triton - #43360
Draft
vorapolsiloai wants to merge 2 commits into
Draft
vorapolsiloai wants to merge 2 commits into
vorapolsiloai wants to merge 2 commits into
Conversation
Quark W4A4 MXFP4 dense linear layers on ROCm always call AITER's Triton gemm_afp4wfp4, so they cannot use the a4w4_blockscale_tuned_gemm CSV tables that AITER gemm_a4w4 and vLLM share. Add an opt-in path that does. - Add SGLANG_ROCM_USE_AITER_FP4_ASM_GEMM (Envs, default off) for gfx95 - Preshuffle FP4 weights and E8M0 weight scales once after loading - Quantize activations with AITER's shuffled per-1x32 FP4 quantizer - Dispatch through gemm_a4w4 so model-specific tuned CSV entries apply - Keep the Triton path for layers whose scale shape cannot be swizzled - Reject Quark tuple-input fusion projections while the flag is on
Replace the SGLANG_ROCM_USE_AITER_FP4_ASM_GEMM switch with two ROCm choices on the existing --fp4-gemm-backend flag, as requested in review. - Add 'aiter' (gfx95 only; AITER CK/ASM gemm_a4w4 with preshuffled weights and the a4w4 tuned CSVs) and 'triton' (AITER Triton gemm_afp4wfp4) to FP4_GEMM_RUNNER_BACKEND_CHOICES and the Fp4GemmRunnerBackend enum - Resolve the flag for ROCm in initialize_fp4_gemm_config: 'auto' maps to 'triton', so the AITER path stays opt-in; 'aiter' off gfx95 and 'aiter'/'triton' off ROCm raise at startup - QuarkW4A4MXFP4 reads the resolved backend at construction instead of the env var at import - Deprecate SGLANG_ROCM_USE_AITER_FP4_ASM_GEMM: still honored with a DeprecationWarning while the flag is 'auto', ignored with a warning when the flag is set explicitly - Document the new choices in server_arguments.mdx and quantization.mdx
github-actions
Bot
added
documentation
3 of 5 tasks
This branch has not been deployed
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
1 participant
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.
Summary
Follow-up to #39598 as requested in review: extends the existing
--fp4-gemm-backendCLI argument to supportaiterandtritonon ROCm, removing the need for a separate environment variable (SGLANG_ROCM_USE_AITER_FP4_ASM_GEMM).'aiter'(AITER CK/ASMgemm_a4w4with preshuffled weights on gfx95) and'triton'(AITER Tritongemm_afp4wfp4) toFp4GemmRunnerBackendand valid choices.--fp4-gemm-backend automaps totritonto keep existing baseline behavior unchanged.QuarkW4A4MXFP4reads the resolved backend at initialization instead of inspecting environment variables at import time.SGLANG_ROCM_USE_AITER_FP4_ASM_GEMMwith a clear deprecation warning when--fp4-gemm-backendis unset (auto).server_arguments.mdxandquantization.mdx.test_quark_mxfp4_asm.py.Stack: Depends on #39598.
CI States
Latest PR Test (Base): ❌ Run #37920891223
Latest PR Test (Extra): ❌ Run #37920890866
Latest PR Test (AMD ROCm 10): ❌ Run #37920891326