Conversation
Signed-off-by: Rebecca Lee <rebecca.lee@amd.com>
Signed-off-by: Rebecca Lee <rebecca.lee@amd.com>
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
rebklee
marked this pull request as ready for review
September 9, 2026 17:07
rebklee
requested review from
AndreasKaratzas,
DarkLight1337,
dllehr-amd,
hongxiayang,
tjtanaa and
ywang96
as code owners
September 9, 2026 17:07
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
dllehr-amd
reviewed
Sep 9, 2026
|
|
||
| return self._gated_o_proj(attn_out, hidden_states) | ||
|
|
||
| def _gated_o_proj( |
Collaborator
There was a problem hiding this comment.
I don't think we need this here. the unit test won't call this ever, and it'll just confuse other callers
| ) | ||
| t, k = x.shape | ||
| _, fp8_max = get_fp8_min_max() | ||
| if not HAS_TRITON or t == 0 or k > _MAX_K: |
Contributor
There was a problem hiding this comment.
- Question: not sure how common k > 8192 is, but in that case is it still better to run the fused torch version rather than the unfused variant? Some perf numbers here would be great
- Suggestion: if we do point 1, can we drop the triton check here (also as its already in the maybe_fused_mla_oproj_ptpc). In that case we could also move _sigmoid_mul_fp8_torch to the test file and just have it as a reference implementation
rebklee
added a commit
to rebklee/vllm
that referenced
this pull request
Sep 16, 2026
This branch carried its own copy of vllm-project#54248: the three aiter.py input_quant_key() methods, the online/fp8.py QuantizedActivation plumbing, and the contract test's AITER kernel registrations. Those files now come from vllm-project#54248 alone, and this branch keeps only the KDA gated-RMSNorm fusion, inert until vllm-project#54248 lands -- the same dependency that vllm-project#55043 already has. Signed-off-by: Rebecca Lee <rebecca.lee@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
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
3 participants
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.
Purpose
Depends on #54248 (K3 o_proj becomes Fp8PtpcOnlineLinearMethod). Until this lands, this path is a no-op.
Fuse MLA
attn_out * sigmoid(g_proj)with PTPCo_proj's per-token FP8 quant into one Triton kernel (replaces ATen sigmoid, ATen mul, anddynamic_per_token_scaled_quant).Wrap
(fp8, scale)asQuantizedActivation(kFp8DynamicTokenSym)so PTPCo_projskips in-kernel quant.g_projGEMM stays separate.Not a duplicate. No other PR produces per-token FP8 from the MLA sigmoid gate. #50664 stops at BF16. #54254 is KDA.
Profile
128k/1k, conc=1, 8×MI355X, TP8, Kimi-K3 + DSpark, 50 decode steps. Image
vllm/vllm-openai-rocm:nightly-7c5dc571cbd1064ecc8a9b1045637ff647aa22cb. Overlay is--quantization-configPTPCo_projonly.Unfused
Sigmoid, mul, and per-token quant are still three launches.
vectorized_elementwise_kernel<8, sigmoid_kernel_cuda, BFloat16>vectorized_elementwise_kernel<8, MulFunctor<float>, BFloat16>aiter::dynamic_per_token_scaled_quant_kernel<fp8e4m3, 8>(The 4650 quant launches are 1200 MLA + 3450 KDA.)
MLA-only epilogue ≈ 12.4 µs (4.0 + 4.3 + 4.1).
Fused (this PR)
Sigmoid, mul, and per-token quant are one launch. Kineto names the Triton kernel
_kernel.kd._sigmoid_mul_fp8_per_token_kernel(_kernel.kd)vectorized_elementwise_kernel<8, sigmoid_kernel_cuda, BFloat16>vectorized_elementwise_kernel<8, MulFunctor<float>, BFloat16>aiter::dynamic_per_token_scaled_quant_kernel<fp8e4m3, 8>MLA epilogue 12.4 → 4.1 µs (−8.3 µs / MLA layer, ~0.20 ms/step × 24 MLA).
Accuracy
GSM8K
lm_eval5-shot, speculative decoding off, n=1319, temperature 0.Tests
tests/models/kimi_k3/test_amd_mla_oproj_ptpc.py— fused vs sequential; decline withoutinput_quant_keyor on mismatched shapes;_gated_o_projfalls through to BF16 when the layer is not PTPC, and passesQuantizedActivation(kFp8DynamicTokenSym)when it is.sigmoid_mul_fp8_per_tokenmatches the torch reference (x * sigmoid(g)then per-token amax). Scale rtol/atol1e-4; dequant rtol/atol1e-2. Shapes includeT=0and TP8K=1536.