[ROCm][Perf] Use AITER tgemm for DeepSeek V4 compressors - #48727
[ROCm][Perf] Use AITER tgemm for DeepSeek V4 compressors#48727Fangzhou-Ai wants to merge 10 commits into
Conversation
mergify
Bot
added
deepseek
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Can you explain AITER's behaviour in more details? What are initialized? Why can't we expect it is just a regular function where you call the tgemm like regular function during graph capture? |
|
Another thing, there is a thread on slack discussing whether it will cause accuracy issue by swapping from FP32 accumulation to BF16 accumulation Original torch.mm(..., out_dtype=torch.float32)This PR tgemm.mm(..., otype=torch.bfloat16) |
I personally don't think we can turn this on by default. But we can keep both path and let user choose if needed. |
|
This pull request has merge conflicts that must be resolved before it can be |
7e2509c to
546d394
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
|
This pull request has merge conflicts that must be resolved before it can be |
b7e1551 to
76cbbc9
Compare
|
Hi @tjtanaa @zyongye I have made a couple of changes to make this PR much simpler, the CUDA graph changes are removed and some per-forward functional calls are cached for better performance. Currently the only concern would be the fp32-->bf16 conversion, although from my test I didn't see any precision regression, do you guys have any suggestions on it? |
Use tuned BF16 GEMMs for the two compressed-attention projections in full CUDA graphs. Scope the warmup signal to a ROCm platform hook and compute indexer weights in the AMD implementation, leaving the shared model and NVIDIA behavior unchanged. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: fai <fangzhouai@gmail.com>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Assisted-by: OpenAI Codex Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
Call AITER tgemm directly during full graph capture and remove the platform warmup-context plumbing. Preserve ROCm preshuffled FP8 projections through the polymorphic projection helper. Assisted-by: OpenAI Codex Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
Select tgemm directly from cached static eligibility, BF16 input dtype, and the V1 FULL graph mode. This removes the per-layer capture-state query while leaving V1 uBatch and V2 on the existing fallback. Assisted-by: OpenAI Codex Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
Route all eligible BF16 compressor projections through AITER tgemm without inspecting CUDA graph runtime mode. This applies consistently to V1, V2, eager, FULL, and PIECEWISE execution. Assisted-by: OpenAI Codex Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
Keep metadata-dependent DeepSeek V4 attention work in an eager segment for PIECEWISE breakable graphs. FULL captures, V2 outer captures, and eager execution continue to run the tgemm path directly. Assisted-by: OpenAI Codex Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
Route eligible compressor projections through AITER tgemm at the existing input-GEMM hook, preserving the shared metadata eager boundary without a model-specific CUDA graph decorator. Assisted-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
Keep AITER tgemm inside FULL graph capture while routing PIECEWISE and eager execution through the shared input-GEMM path. This avoids replaying per-layer tgemm dispatch eagerly and keeps fresh tgemm outputs from crossing the breakable attention boundary. Assisted-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
4f70b30 to
4c13bdd
Compare
|
Rebased and force-pushed the simplified implementation on the refreshed base.
These are intermediate measurements. The final refreshed-base side-by-side Validation: 47 focused tests passed, and all pre-commit hooks passed on the two |
Use AITER tgemm only for FULL graph batches up to four tokens, where both DeepSeek V4 compressor projections have tuned entries. Larger batches retain the original FP32-output GEMM, avoiding serving regressions and the M=128 ASM graph failure. Assisted-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
|
Updated the PR description with the final same-base side-by-side sweep. Positive
Output-throughput geometric-mean gain: +1.78%. The Final-head GSM8K, 5-shot greedy, 1,319 examples:
Both scores match the prior baseline exactly. |
|
I will close this PR for now as a conversion from fp32 to bf16 can cause some precision issue under some corner cases. |
3 participants
Summary
Use AITER tuned BF16 GEMMs for the DeepSeek V4 main and optional indexer
fused_wkv_wgatecompressor projections ongfx950.The optimized path is intentionally limited to FULL CUDA-graph execution with
M <= 4, whereMis the compressor projection row count. Larger shapes,PIECEWISE graphs, and eager execution retain the original FP32-output
torch.mmpath.This changes the two compressor score outputs from FP32 to BF16 only on that
optimized path. AITER remains opt-in through
VLLM_ROCM_USE_AITER=1; users canretain the original compressor path with
VLLM_ROCM_USE_AITER_LINEAR=0.Implementation
initialization so non-dynamic predicates are not repeated per layer and
forward pass.
capture/replay when
M <= 4.batches also skip that repeated context lookup.
attn_gemm_parallel_executeinput-GEMM hook.Reuse its inherited path and existing attention boundary for
M > 4,PIECEWISE graphs, and eager execution.
decorator, capture hook, or alternate attention forward is added.
_fused_wqa_wkv_gemm()helper so FP8 preshuffledprojections remain correct.
dtype, both compressor weights, and the inherited tuple contract.
M <= 4is the range where the current DSV4 AITER table has tuned entries forboth compressor projections. In the same-base serving sweep, unrestricted
tgemm improved CONC=4, was neutral at CONC=8, and regressed throughput by
8.23%, 10.91%, and 7.61% at CONC=16, 32, and 64 respectively. The bounded path
preserves the low-concurrency optimization without those larger-batch
regressions.
An earlier all-runtime-mode version was also evaluated and rejected. At
CONC=4 it reached 154.462 output tok/s, 8.99% below its supplied 169.714 tok/s
baseline, because PIECEWISE execution exposed 132 per-forward tgemm dispatches
to the CPU. Keeping the optimization inside FULL graphs removes that repeated
dispatch overhead.
The unrestricted CONC=128 run also encountered an HSA memory-access fault when
the indexer projection selected its M=128 ASM split-K kernel. Separately,
ROCm/aiter#4494 tracks an ASM split-K
semaphore deadlock under graph capture. That upstream report supports avoiding
the larger captured ASM path for now, but it does not prove that its reported
deadlock caused the HSA fault observed here.
Duplicate-work check
No issue number was supplied, so an issue-specific lookup was not applicable.
Open-PR searches for
DeepSeek V4 AITER tgemm,compressor tgemm, andDeepSeek V4 compressor ROCmfound no other PR implementing this change.The closest open work is materially different:
tgemm into generic linear layers; these compressor projections use direct
torch.mmcalls and are not routed through that integration.path that fuses the broader compressor, normalization, RoPE, quantization,
and cache-storage pipeline; this PR only replaces two BF16 projection GEMMs.
Tests
.venv/bin/python -m pytest tests/kernels/attention/test_rocm_triton_attn_dsv4.py -q— 49 passed.
.venv/bin/pre-commit run --files vllm/models/deepseek_v4/amd/rocm.py tests/kernels/attention/test_rocm_triton_attn_dsv4.py— all hooks passed, including Ruff, mypy, and DCO sign-off checks.
Performance and model evaluation
The side-by-side comparison uses the branch's exact parent,
821717118fc26667dd474b9b0ab81d29259dfc5c, as the baseline andaaa3fee24014237015e6b9578f28c7f21b8034bcas the PR revision. Both sides used fresh servers in the samerefreshed container, editable environment, model weights, and server
configuration.
Settings: DeepSeek-V4-Pro, TP=8,
gpu-memory-utilization=0.8, random8K-input/1K-output requests,
random-range-ratio=0.8,seed=0,ignore-eos,2 * CONCwarmups, and10 * CONCmeasured requests. Each fresh serverreceived one discarded CONC=4 preconditioning pass before its recorded sweep.
The baseline and PR used pairwise-identical input/output token-length arrays.
The output-throughput geometric-mean gain across the six concurrency points is
+1.78%. CONC=128 completed all 1,280 requests without the HSA fault seen
with unrestricted tgemm.
For throughput, delta is
(PR / baseline - 1) * 100; for latency, delta is(baseline - PR) / baseline * 100. Positive values therefore indicateimprovement.
Final-head GSM8K, 5-shot greedy, 1,319 examples:
Both scores match the prior baseline exactly.
AI assistance disclosure
OpenAI Codex assisted with review-comment analysis, simplifying the
implementation, running tests and model evaluations, benchmarking, and drafting
this description. The human submitter reviewed every changed line and can
understand and defend the change end-to-end.