[ROCm][Perf] Fused MoE W4A16 HIP kernel for AMD RDNA3 (gfx1100) - #44075
Conversation
|
Hi @JartX, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, Tip Is
|
d061257 to
e283487
Compare
|
cc @dllehr-amd Who should review this PR? Is there anyone from the you know familiar with gfx1100? |
ff3d9e9 to
29d6afd
Compare
29d6afd to
dd0a25c
Compare
![depthfirst-app[bot]](https://avatars.githubusercontent.com/in/1021002?s=60&v=4){kind=small}
dd0a25c to
a948f5d
Compare
DarkLight1337
added
the
verified
Last time similar PR was reviewed by @mgehre-amd . It would be great if we have someone from AMD who can review this. I would like to also see the Radeon support being brought up in vLLM as there are in fact many Radeon users based on the reddit local inferencing community. |
|
@JartX Please provide the lmeval score for the model and make sure to validate with large concurrency to test the batching logic of the kernels. In e.g. MODEL=deepseek-ai/DeepSeek-V4-Pro
lm_eval --model local-completions --model_args model=$MODEL,base_url=http://0.0.0.0:8001/v1/completions,num_concurrent=256,max_retries=10,max_gen_toks=2048,max_length=1048576,timeout=60000 --batch_size auto --tasks gsm8k --num_fewshot 5 \
--output_path ./results_deepseekv4pro_validatepr2_numshot5 \
--log_samples \
| tee lmeval_deepseekv4pro_validatepr2_numshot5.log |
The AMD CI image (e.g. mi325) exposes csrc/ + CMakeLists.txt under /vllm-workspace for building but installs the vllm python package as a wheel, so vllm/_custom_ops.py and compressed_tensors_moe.py are not in the source tree. _find_repo_root() still matched /vllm-workspace via CMakeLists.txt + csrc/, so the python static-analysis tests raised FileNotFoundError instead of skipping. Add _read_source_or_skip() which skips the test when the target file is absent. The C++/CMake guard checks still run (those files are present); only the python-source checks skip where the python tree isn't shipped. Signed-off-by: JartX <sagformas@epdcenter.es>
Instead of skipping the _custom_ops.py / compressed_tensors_moe.py guard checks when the python source tree is absent, resolve them from the installed vllm package via vllm.__file__. This makes the checks actually run in CI (e.g. AMD mi325, which ships the wheel and only checks out csrc/ + CMakeLists.txt for building) and verifies the code that is truly imported at runtime. C/CMake guard checks still read from the repo tree (csrc/, CMakeLists.txt exist only in a source checkout); only those skip when the tree is unavailable. Signed-off-by: JartX <sagformas@epdcenter.es>
There was a problem hiding this comment.
I'd suggest following the modular kernel / oracle style and create a new expert class for your rdna3 w4a16 kernel.
oracle and is_supported_config will handle the platform / config triage, and other quantization methods that support w4a16 can also benefit from your kernel if applicable.
There was a problem hiding this comment.
Hi @BowenBao thanks for answer :)
I Appreciate the suggestion. Keeping the dispatch in compressed_tensors_moe was a deliberate call, and I'd lean toward keeping it scoped for this PR unless you feel strongly.
The intent is to keep the RDNA path self-contained and easy to follow. This kernel targets the RDNA consumer line (RDNA3 today, very likely RDNA4 and onward), and the plan is for each arch to have its own simple, dedicated trace. This also follows earlier guidance we got to keep these paths easily traceable and cleanly separated. A contained dispatch lets me do that without threading RDNA-specific gating into the shared wNa16 oracle, which is otherwise CUDA/CDNA-oriented. I'd rather keep RDNA and CDNA (MI2xx/MI3xx) as cleanly separated worlds, since an RDNA fused kernel has very different tradeoffs from the CDNA parts.
It's also fully gated and hermetic. RDNA3 code can't leak into non-gfx1100 builds (I added tests to enforce that), so there's zero blast radius on the CUDA/CDNA side.
So I'd suggest landing it as-is with the contained dispatch. That said, if you'd like, I'm very happy to work through together how to make the RDNA separation even cleaner. Either way works for me. What do you think?
There was a problem hiding this comment.
Synced offline with @JartX, I'd recommend following the oracle and expert class convention that vllm is heading, which keeps RDNA logic contained and managed in the respective expert class. #43693 is related and is merging CompressedTensorsWNA16MarlinMoEMethod and CompressedTensorsWNA16MoEMethod. That said, if you are to refactor the PR it will depend on #43693 to land first.
Another option is to merge this as is and refactor later after #43693 lands.
I'll let other folks chime in as for preferences, cc @AndreasKaratzas , @tjtanaa , @mgoin , @bnellnm .
There was a problem hiding this comment.
Thanks @BowenBao for summing it up. Both work for me, but I'd lean slightly toward merging as-is now and doing the oracle/expert refactor as a follow-up once #43693 lands. The kernel is already validated, tested and benchmarked, so this gets it into users' hands now. RDNA stays its own separate path either way, and I'm committed to the expert-class refactor right after #43693 merges. Happy to open a tracking issue for that follow-up so it doesn't fall through the cracks. That said, I'll defer to whatever you all prefer :)
There was a problem hiding this comment.
I'm fine with either approach. #43693 already depends on another PR to land first so it might be awhile before it lands.
BowenBao
left a comment
BowenBao
left a comment
There was a problem hiding this comment.
LGTM on quant side per follow-up refactors.
| if quant_method_name in ( | ||
| "CompressedTensorsWNA16MarlinMoEMethod", | ||
| "CompressedTensorsWNA16MoEMethod", | ||
| "CompressedTensorsWNA16RDNA3MoEMethod", |
There was a problem hiding this comment.
There's another check for CompressedTensorsWNA16MoEMethod above (at around ~490). Does that need to be updated with the new class?
There was a problem hiding this comment.
@bnellnm many thanks for the review!
It is a good question, thanks for flagging it. I traced again it and I don't think that block needs the new class. Here's why:
CompressedTensorsWNA16RDNA3MoEMethod subclasses the non-Marlin CompressedTensorsWNA16MoEMethod and inherits its create_weights unchanged, which hardcodes load_full_w2=False on the w2 scale. intermediate_size_full is only ever consumed inside the Marlin/auto-GPTQ create_weights, gated on act-order (load_full_w2 = actorder and group_size != -1). On the non-Marlin path load_full_w2 is always False, so the w2 gets sharded normally by intermediate_size_per_partition and intermediate_size_full is never read.
|
@JartX HI! I have some similar work in progress which serves all of RDNA family: #43389. Quickly benchmarking shows the following on gfx1100(vllm serve, 200 prompts, sharegpt, gpu-utilization=0.9): Model 1. cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit (compressed-tensors, symmetric)
50 prompts, max_num_seqs=1
Model 2: Qwen/Qwen3-30B-A3B-AWQ (AWQ, asymmetric zeros). This PR doesn't affect this at all, changing only Compressed-tensor pass.
Since your native HIP kernel beats, I think it makes sense to extends this kernel to asymmetric awq and symetric GPTQ models, which are currently dispatched differently through vllm/model_executor/layers/quantization/moe_wna16.py. I understand there is ongoing refactoring #44570 and you want to push this first. When you come back later to refactor, can you make sure this extends to awq/gptq models as well? |
|
Many thanks @amd-xavierwang Thanks a lot for the detailed benchmarks and the comparison, really appreciate you taking the time to run both configurations. Great data to have. And thanks for the kind words about the kernel, means a lot coming from someone at AMD. It has been validated by 2 AMD members, both dispatch and performance. It is getting closer :) |
I am happy that I can help! This PR is great work and I understand the toughness developing native HIP kernels.. Everyone knows native kernel is the ultimate efficient solution, and yet I chose a much easier path as I am not proficient developing HIP kernels right now. Can't wait to collaborate in the future!! |
…-project#44075) Signed-off-by: JartX <sagformas@epdcenter.es> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
…-project#44075) Signed-off-by: JartX <sagformas@epdcenter.es> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
7 participants
Summary
Native HIP kernel for W4A16 MoE on RDNA3 (gfx1100), replacing the Triton
fused_moe_kernel_gptq_awqpath. It uses the same dequant + dot primitives asthe dense W4A16 kernel:
v_dot2_f32_f16/v_dot2_f32_bf16, exllama bit-trickdequant, and 64-bit CAS atomic output.
What it does
csrc/rocm/moe_q_gemm_rdna3.cu): expert routing + W4A16GEMM in a single kernel launch. Templated on
BLOCK_SIZE_M = {1, 2, 4, 8}.The bf16
M=1fast path skips LDS staging (direct global read, with av_dot2_f32_bf16opacity trick to defeat InstCombine).output_topkparameter: fusesmoe_suminto the w2 kernel — writesdirectly to
out[token_id / top_k]via atomics, eliminating a separate kernellaunch and the intermediate buffer.
BLOCK_SIZE_M = 1for smallM(eliminates ~75% ofpadding waste); pre-allocated w1/act buffers (eliminates per-layer
torch.zeros)._try_get_rocm_moe_method()): checks architecture andop availability in priority order. New architectures (RDNA4, CDNA) can be added
as branches. Falls through to Triton WNA16 if no native kernel is available.
.cuis only compiled whenVLLM_GPU_ARCHEScontainsgfx1100; the torch op is registered under#ifdef VLLM_ROCM_GFX1100; thePython side checks
hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3").Benchmarks
Measured on 1× RX 7900 XTX (gfx1100) via EvalScope
perf,openqadataset,50 requests. Native HIP kernel vs. upstream Triton
fused_moe_kernel_gptq_awqpath.
Model:
cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit(MoE, 128 experts / 8 active, W4A16 AWQ).
Throughput (completion tok/s)
Per-request latency (max_num_seqs = 32)
The speedup scales with concurrency: ~1.2× single-stream, growing to ~2.2× under
batch (8–32 seqs), consistent with a kernel that improves batched decode
efficiency rather than single-request latency.
ACC EVAL
vllm serve cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit --gpu-memory-utilization 0.90 --max-model-len 16384 --tensor-parallel-size 1 --port 8000 --enable-log-requests --tool-call-parser hermes --enable-auto-tool-choice --attention-backend TRITON_ATTN --served-model-name QWEN3
lm_eval --model local-chat-completions --model_args model=QWEN3,base_url=http://localhost:8000/v1/chat/completions,num_concurrent=35,max_retries=3,tokenized_requests=False --tasks gsm8k --batch_size auto --apply_chat_template --output_path ./results/qwen_gsm8k_eval.json
HIP
TRITON
Testing
Model:
cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit, GPU: 1× RX 7900 XTX(gfx1100). Output parity checked via EvalScope on
openqa, HIP kernel vs. theTriton WNA16 reference path:
```bash
evalscope eval
--model cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit
--api-url http://localhost:8000/v1
--eval-type openai_api
--datasets openqa
```