[ROCm][Kimi-K3][Perf] Fuse MLA decode KV-cache write and Q-prep via AITER by rbrugaro-amd · Pull Request #57640 · vllm-project/vllm

rbrugaro-amd · 2026-09-18T22:18:34Z

Purpose

On ROCm, Kimi-K3 MLA decode runs the per-layer query preparation and the KV-cache
write as separate launches:

vllm::concat_and_cache_mla_kernel — the fp8 KV-cache write
triton_poi_fused__to_copy_cat_clamp_mul_reciprocal_view_0 — the
[ql_nope | q_pe] concat and that query's static fp8 quant

This collapses both into a single AITER fused_qk_rope_concat_and_cache_mla
launch: 288 → 144 kernel launches and 1351.8 → 682.8 µs (1.98×) per rank over
6 decode steps, i.e. 9.39 → 4.74 µs per MLA layer per decode step on
MI355X at TP=8.

Kimi-K3 MLA is NoPE (mla_use_nope), while that kernel always applies RoPE and
takes the cos/sin caches as required arguments. It is fed an identity rotation
(cos = 1, sin = 0) so the rotation degenerates to a no-op — the same approach
ROCm/ATOM uses to route K3 through this kernel.

The identity cache is a single row. AITER's per-head and _opt decode kernels
clamp pos into [0, cos_cache.size(0)) before indexing, so real positions pass
straight through; a full-length constant cache would otherwise cost ~268 MB at
K3's max_position_embeddings of 1 048 576 to express a provable no-op. AITER's
general decode kernel does not clamp, so the layer refuses
the head counts that would select it (kv_lora_rank * num_heads < 2048, i.e.
≤3 heads/rank, reachable at TP=32).

Scope, and why this is small

This is a kernel-count / device-time optimization on the MLA decode path. The
two replaced kernels are a sub-1% slice of a decode step, on 24 MLA layers in a
model that is roughly 1:3 MLA:KDA. End-to-end serving throughput improvement is below 1%

Relationship to #47757 and #57452

This supersedes #57452, which I opened for the same optimization and will
close in favor of this one. That PR stacked on #47757: it carried #47757's six
commits so CI could run, and extended its hook in the shared
MLAAttention.forward_impl, also touching
vllm/model_executor/layers/mla.py. Two problems with that shape:

it could not merge until [ROCm][MLA][Perf] Fuse decode QK-RoPE + Q-concat + KV-concat + KV-cache write for sparse MLA #47757 did, and it is still under review
it modified shared MLA code, so every MLA model and backend was in the blast
radius of a Kimi-K3-only optimization

This PR drops both. It touches no shared code at all — only
vllm/models/kimi_k3/amd/ — mirroring the vendor-local q/kv RMSNorm fusion that
the same wrapper already owns (#52494, merged). Consequences:

no dependency on [ROCm][MLA][Perf] Fuse decode QK-RoPE + Q-concat + KV-concat + KV-cache write for sparse MLA #47757, so CI runs standalone and this can merge in any order
zero blast radius: no other model, backend or platform can be affected
it is a single file — the AITER call lives directly in the K3 wrapper

#47757 remains complementary: it fuses the sparse MLA path (deepseek_v32)
through the shared hook, while this fuses the dense Kimi-K3 path inside the
vendor wrapper. Different call sites, no conflict in either merge order.

Updated after review (thanks @Rohan138) — the fused call is no longer a
registered custom op, the torch.compiler.is_compiling() guard is gone, and the
branch now lives directly in forward() instead of a helper returning None.
That folded the separate ops module away, so this is now one file. All numbers
below were re-measured on the refactored code.

On is_compiling(): the baseline's triton_poi_fused… kernel comes from
mla_decode_concat_quant_fp8 compiling its own forward_native (class docstring
at mla_attention.py:1541), not from the model being compiled — Kimi-K3 is not
@support_torch_compile. The guard was dead code.

Why not a duplicate

[ROCm][MLA][Perf] Fuse decode QK-RoPE + Q-concat + KV-concat + KV-cache write for sparse MLA #47757 — same kernel, sparse MLA on deepseek_v32; this is the dense K3 path
and the NoPE handling is new.
[ROCm] Fuse decode rope + MLA KV-cache write + fp8 query assembly via AITER #55230 — same kernel family, DeepSeek-R1, edits shared layers/mla.py.
[ROCm][MLA] Fuse RoPE + MLA KV-cache write via AITER Triton kernel #45798 / [ROCm] Add AITER RoPE + KV cache fusion for MLA prefill and decode #38313 / [ROCm] Add AITER RoPE+KV cache and dual RMSNorm fusion for MLA #41099 / [ROCm] Add unified AITER RoPE+KV cache and dual RMSNorm fusion for MLA #41101 — the same fusion as a torch.compile
graph-rewrite pass. Kimi-K3 is not @support_torch_compile, so those never fire.
[Perf][ROCm][Attention] Fuse Kimi-K3 MLA chunked-context gather on AITER #53166 — K3 ROCm MLA, but chunked-context gather, a different op.
[AMD][kimik3][ROCm][Perf] Fuse MLA q/kv RMSNorm in AMD Kimi-K3 MLA wrapper #52494 (merged) — the sibling q/kv RMSNorm fusion in this same wrapper.

Implementation

One file, vllm/models/kimi_k3/amd/mla.py, on the existing K3 wrapper:

_fused_qk_prep_supported() resolves everything fixed at load time —
platform, KV-cache dtype, shapes, DCP/PCP/HiSparse, head padding, and which
W_UK variant is live. Evaluated once in __init__.
_fused_decode() mirrors MLAAttention.forward_impl's W_UK dispatch
including its precedence (fp4 before fp8; both flags can be set at once) so
ql_nope is bit-identical to what the unfused path feeds its concat+quant,
then issues the fused call and continues through impl.forward_mqa and
_v_up_proj.

With the static gating hoisted into __init__, the only per-batch condition left
in forward() is that the batch is decode-only. Everything else takes the
stock MLAAttention path, unchanged:

any batch that is not purely decode (the fused call covers only the MQA slice,
so a mixed batch's prefill rows would never reach the KV cache)
DCP (the kernel skips slot < 0 rows, which would then be garbage after the
query all-gather), PCP and HiSparse (both rewrite the KV inputs/slots
before the standalone write, which this path bypasses)
the fp8-bmm and AMX W_UK variants, and any q_pad_num_heads head padding
profile/dummy runs, an unallocated KV cache, or a missing slot mapping
non-ROCm, non-fp8 KV cache, or head counts that would select AITER's
non-clamping general decode kernel

Test Plan

All numbers below are from the stock upstream nightly
vllm/vllm-openai-rocm:nightly-rocm100-0bfc7a15d095fe83ecc82b50561a93c177fece2d
(2026-09-16), 8× MI355X, Kimi-K3 fp4, TP=8, with no benchmark harness and no
vendor patches — plain vllm serve + vllm bench serve:

Server (identical for both arms)

export VLLM_ROCM_USE_AITER=1

vllm serve $MODEL \
  --served-model-name moonshotai/Kimi-K3 \
  --trust-remote-code \
  --load-format fastsafetensors \
  --language-model-only \
  --tensor-parallel-size 8 \
  --gpu-memory-utilization 0.9 \
  --kv-cache-dtype fp8 \
  --max-model-len 32768 \
  --no-enable-prefix-caching \
  --host 0.0.0.0 --port 8010

VLLM_ROCM_USE_AITER=1 is the only environment variable set.

Perf

vllm bench serve --backend openai-chat --endpoint /v1/chat/completions \
  --model moonshotai/Kimi-K3 --tokenizer $MODEL --trust-remote-code \
  --host 127.0.0.1 --port 8010 \
  --dataset-name random --random-input-len 1024 --random-output-len 8192 \
  --random-range-ratio 0 --ignore-eos \
  --num-prompts $CONC --max-concurrency $CONC \
  --seed 42 --percentile-metrics ttft,tpot,itl,e2el --metric-percentiles 50,90,99

Kernel census

Add --profiler-config to the same serve command, then bracket a decode phase:

# appended to the serve command above
--profiler-config '{"profiler":"torch","torch_profiler_dir":"/out/traces",
                    "torch_profiler_use_gzip":true,"ignore_frontend":true,
                    "delay_iterations":0,"max_iterations":6}'

# with a bench in flight and past prefill:
curl -X POST http://127.0.0.1:8010/start_profile
sleep 15
curl -X POST http://127.0.0.1:8010/stop_profile

Accuracy

lm_eval --model local-chat-completions --apply_chat_template \
  --tasks gsm8k --output_path out --log_samples \
  --model_args model=moonshotai/Kimi-K3,base_url=http://127.0.0.1:8010/v1/chat/completions,\
api_key=EMPTY,eos_string='</s>',num_concurrent=32,tokenized_requests=False,max_length=32768 \
  --gen_kwargs max_tokens=4096,temperature=0,top_p=1

Results

1. Decode kernel census — two kernels become one

Rank-0 torch profile, 6 engine steps, 24 MLA layers, 144 = 24 × 6:

kernel	baseline	fused
`vllm::concat_and_cache_mla_kernel`	144 (668.0 µs)	0
`triton_poi_fused__to_copy_cat_clamp_mul_reciprocal_view_0`	144 (683.8 µs)	0
`aiter::fuse_qk_rope_concat_and_cache_mla_kernel`	0	144 (682.8 µs)

288 → 144 launches, 1351.8 → 682.8 µs = 1.98×.
Per MLA layer per decode step: 9.39 → 4.74 µs.

Dispatch lands on AITER's _opt decode kernel rather than per_head —
per_head requires kv_cache.size(1) == 1 and vLLM derives 1536 here. The _opt
gate is kv_lora_rank × num_heads ≥ 2048; at TP=8 that is 512 × 12 = 6144.
_opt is also the position-clamping kernel the single-row identity cache needs.

2. GSM8K — no accuracy regression

Full 1319 questions, 5-shot, temperature=0, top_p=1, max_tokens=4096,
chat template applied:

arm	strict-match	flexible-extract	stderr
fused	0.9545	0.9553	0.0057

3. E2E serving — unchanged within noise

ISL 1024 / OSL 8192 C=1|8|32 one repeat showed TPOT p50 (ms), ITL p50 (ms), ITL p90 (ms), out tok/s for Conc 1 and 8 are flat to within 0.2%. Conc 32 moves ~0.7% in the baseline's favor, having moved ~1% the other way before the refactor — same essential code, opposite sign, so nothing is claimed in either direction.

This is the expected outcome. The fusion removes a sub-1% slice of a decode
step's device time; for scale, cross_device_reduce_2stage alone is ~46% of
rank-0 decode device time in the same trace.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

rbrugaro-amd · 2026-09-21T16:00:46Z

/ci run

Rohan138 · 2026-09-21T16:47:36Z

+            and q_dcp_replicated is None
+            and not torch.compiler.is_compiling()
+        ):
+            attn_out = self._try_fused_decode(


so instead of a try_fused_decode, let's just move the logic here to the main block so if the requirements are satisfied we always do the fused kv_write+q_prep, else we do mla_attn. You also don't need to check torch.compiler.is_compiling since that'll always be False.

@Rohan138 Done in 4e333f070b. Apologies that this landed as a force-push over the commit you reviewed... the later fixes are stacked as their own commit.

All three parts:

is_compiling gone.

No more _try_fused_decode returning None. It's now _fused_decode -> torch.Tensor, with the eight return None bail-outs split by when they're actually knowable:

load-time (DCP, PCP, HiSparse, q_pad_num_heads, which W_UK variant is live, shapes, kv-cache dtype) -> _fused_qk_prep_supported(), evaluated once in __init__

per-batch -> inline in forward()

Logic moved into the main block:

attn_metadata = layer = kv_cache = slot_mapping = None fuse = self._fused_qk_prep and q_dcp_replicated is None if fuse: attn_metadata, layer, kv_cache, slot_mapping = get_attention_context( self.mla_attn.layer_name ) fuse = ( attn_metadata is not None and slot_mapping is not None and kv_cache.numel() > 0 and attn_metadata.num_actual_tokens == attn_metadata.num_decode_tokens ) if fuse: attn_out = self._fused_decode(...) else: attn_out = self.mla_attn(...)

So if the requirements hold we always take the fused kv_write+q_prep, else mla_attn — the only per-batch condition left is that the batch is decode-only, plus the degenerate cases (profile/dummy run, unallocated KV cache, missing slot mapping).

One consequence: _fused_decode now takes attn_metadata, layer, kv_cache, slot_mapping as arguments, since get_attention_context() has to be called in forward() to evaluate the gate. That's why the signature went from 5 args to 9

Rohan138 · 2026-09-21T16:48:41Z

+    return cos, sin
+
+
+def _fused_mla_qk_prep_impl(


I don't think this needs to be a custom op anymore, since again, torch.compile isn't active. Can you try just folding this into mla.py directly?

@Rohan138 Done in 4e333f070b — folded into mla.py and the file is deleted.

Removed: direct_register_custom_op, _fused_mla_qk_prep_impl, _fused_mla_qk_prep_fake, the torch.ops.vllm.kimi_k3_fused_mla_qk_prep indirection and the fused_mla_qk_prep wrapper. supports_fused_qk_prep() and make_identity_rope_cache() became _fused_qk_prep_supported() and an inline _identity_rope construction on the class.

_fused_decode now calls the kernel directly:

from aiter.ops.cache import fused_qk_rope_concat_and_cache_mla ... fused_qk_rope_concat_and_cache_mla( ql_nope.contiguous(), q_pe, kv_c_normed, k_pe.squeeze(1), kv_cache.view(kv_cache.shape[0], -1, head_size), q_out, slot_mapping.flatten(), layer._k_scale, layer._q_scale, positions, cos_cache, sin_cache, is_neox=True, is_nope_first=True, )

Net across both comments: -249/+91, and the PR is now a single file.

Re-validated on the stock nightly after the refactor — GSM8K 0.9545 strict / 0.9553 flexible, and the kernel census still shows 144x fuse_qk_rope_concat_and_cache_mla with zero concat_and_cache_mla and zero triton_poi_fused.

…ITER On a stock build, Kimi-K3 MLA decode runs two kernels per layer per step: the fp8 KV-cache write (`vllm::concat_and_cache_mla_kernel`) and the `[ql_nope | q_pe]` concat plus that query's static fp8 quant, which the `mla_decode_concat_quant_fp8` custom op emits as a single `triton_poi_fused__to_copy_cat_clamp_mul_reciprocal_view` kernel. Collapse both into one AITER `fused_qk_rope_concat_and_cache_mla` launch. Measured on MI355X at TP=8, plain `vllm serve` on the stock upstream nightly with no speculative decoding and no compilation-config override: 288 -> 144 launches and 1351.8 -> 682.8 us per rank over 6 decode steps, i.e. 9.39 -> 4.74 us per MLA layer per decode step. End-to-end serving throughput is unchanged within measurement noise and this change does not claim a throughput win. K3 MLA is NoPE, while the kernel always applies RoPE and takes the cos/sin caches as required arguments. Feed it an identity rotation (cos=1, sin=0) so the rotation is a no-op -- the approach ROCm/ATOM uses for K3. The cache is a single row: the `_opt` decode kernel clamps `pos` into `[0, cos_cache.size(0))`, so real positions pass through untouched and K3's 1M max_position_embeddings does not cost ~268 MB to express a no-op. AITER's general decode kernel does not clamp, so the head counts that would select it are excluded. Kept entirely in `vllm/models/kimi_k3/amd/mla.py`, mirroring the q/kv RMSNorm fusion this wrapper already owns (vllm-project#52494); no shared MLA code is touched, so this carries no dependency on vllm-project#47757 and cannot affect any other model or backend. Everything that is fixed at load time -- platform, KV dtype, shapes, DCP/PCP/HiSparse, head padding and the W_UK variant -- is resolved once in `_fused_qk_prep_supported()`, so the only per-batch condition is that the batch is decode-only; anything else falls through to the stock `MLAAttention` path. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>

rbrugaro-amd · 2026-09-22T06:10:28Z

Thanks @Rohan138 for the comments. I have made the edits and updated the PR description

Rohan138 · 2026-09-23T14:12:45Z

 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.mla import MultiHeadLatentAttentionWrapper
+from vllm.platforms import current_platform
+
+# AITER dispatch constants, mirrored from csrc/kernels/cache_kernels.cu. With all


remove comment

Rohan138 · 2026-09-23T14:13:17Z

+    """Kimi-K3 MLA wrapper with eager AITER q/kv RMSNorm fusion and a fused
+    decode Q-prep path.
+
+    Both fusions are kept here rather than in the shared MLA layer so only AMD


remove second comment

Rohan138 · 2026-09-23T14:13:28Z

+        self._identity_rope: tuple[torch.Tensor, torch.Tensor] | None = None
+
+    def _fused_qk_prep_supported(self) -> bool:
+        """Whether this layer can ever take the fused decode path.


remove comment

Rohan138 · 2026-09-23T14:18:50Z

+        """
+        attn = self.mla_attn
+        if not (
+            current_platform.is_rocm()


platform will always be rocm under this path, similarly you can also skip checking fp8_ds_mla since that's NV deepgemm-only.

Rohan138 · 2026-09-23T14:19:58Z

+        # Only the W_UK variants reproduced in _fused_decode are supported.
+        return attn.is_aiter_triton_fp4_bmm_enabled or (
+            not attn.is_aiter_triton_fp8_bmm_enabled
+            and not attn.is_amx_bmm_enabled


amx_bmm is CPU-only

amd-mghanimi

nice work. I just highlighted potential improvements over aiter_ops imports.

amd-mghanimi · 2026-09-23T16:22:46Z

+        the concat and quant already arrive as one inductor-fused kernel, so
+        this collapses two launches into one.
+        """
+        from aiter.ops.cache import fused_qk_rope_concat_and_cache_mla


it seems @xaguilar-amd had prior attempts to import this kernel in aiter_ops:

https://github.com/vllm-project/vllm/pull/47757/changes#diff-6df46339473cb226edf1075bc70cd5fe9fab4e4dd67558a472146352f512eaa0R804

maybe we can refactor to something like that?

answered below

amd-mghanimi · 2026-09-23T16:25:51Z

+        # path feeds its concat+quant, because this only replaces what happens
+        # after the bmm. Both variants return bf16, which the kernel consumes.
+        if attn.is_aiter_triton_fp4_bmm_enabled:
+            from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4


we have a better import option already:

vllm/vllm/_aiter_ops.py

Line 3405 in 0549e8d

def batched_gemm_a16wfp4(

Thanks @amd-mghanimi
batched_gemm_a16wfp4: the wrapper takes the output tensor as a required
argument, so it fits _v_up_proj (which already has out to hand) but not the
decode-absorb site I'm mirroring, which lets the kernel allocate:

vllm/vllm/model_executor/layers/attention/mla_attention.py

Lines 1027 to 1037 in 92b40f5

if self.is_aiter_triton_fp4_bmm_enabled:

from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4

mqa_ql_nope = batched_gemm_a16wfp4(

mqa_q_nope,

self.W_K,

self.W_K_scale,

transpose_bm=True,

prequant=True,

y_scale=self._q_scale if fp8_attention else None,

)

This PR's correctness argument is that ql_nope is bit-identical to what the
unfused path feeds its concat+quant, so I'd rather keep this call than add a caller-allocated buffer now.

fused_qk_rope_concat_and_cache_mla: agreed it belongs in _aiter_ops, but
it isn't on main yet — the binding you linked is one #47757 adds. Using it
now means either editing _aiter_ops.py here (shared code, conflicting with
#47757 adding the same method) or depending on #47757, and this PR was
restructured to have neither so it can merge standalone.

My proposal would be once #47757 lands and we decide on the refactor @Rohan138 is proposing below I'll do both in one follow-up. How does that sound @amd-mghanimi @Rohan138 ?

Rohan138 · 2026-09-23T16:32:49Z

+                slot_mapping,
+            )
+        else:
+            attn_out = self.mla_attn(


Honestly I wonder if we should just do a clean aiter mla impl for Kimi similar to https://github.com/vllm-project/vllm/blob/main/vllm/models/kimi_k3/nvidia/mla.py

Then you don't need a separate self._fused_decode that does GEMM+RoPE+kv_cache+q_prep+MLA+out_proj, you can just do something like:

if supported: fused_qk_rope_concat_and_cache_mla else: rope() kv_cache_write() q_prep()

Within your mla function. This may be out of scope for this PR, cc @dllehr-amd

@Rohan138 @dllehr-amd I thought about that too, my 2 cents:

the mla_attention.py has lots of commits over time(sparse, DCP,PCP…) and doing mla.py full on amd side would require track those by hand, right?. With the fused_path weight prep, the backend impl behind forward_mqa, _v_up_proj, the KV cache spec and every non-fused batch (prefill, mixed,..) goes through unchanged mla_attention.py. Current solution might be lower maintenance at the expense of little less readability.

-For nvidia might be more justified because they diverge from the shared implementation more

Thoughts? please correct me if I am misunderstanding

Per review on vllm-project#57640: - remove the three explanatory comment blocks - drop `current_platform.is_rocm()`: kimi_k3/__init__.py only imports the amd/ wrapper under is_rocm(), so it cannot be instantiated elsewhere - drop the `fp8_ds_mla` check: only FLASHMLA_SPARSE and FLASHINFER_MLA_SPARSE_SM120 promote to that dtype, both NVIDIA-only - drop `is_amx_bmm_enabled`: set only by AMXMLAImpl, which is CPU-only All four gates are unreachable on this path, so the fused path's behavior is unchanged. Verified by kernel census on 8x MI355X TP=8: the fused kernel still fires 144x (24 MLA layers x 6 decode steps) with zero concat_and_cache_mla and zero triton_poi_fused, matching the pre-change build. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>

Rohan138

LGTM! Can do further refactoring in follow-ups per discussion in comments

github-actions · 2026-09-28T19:13:59Z

✅ @rbrugaro-amd, CI is now available for this PR.

/ci run starts upstream CI; /amd-ci run starts AMD CI only.
Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
/ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
/amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
/ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

rbrugaro-amd · 2026-09-29T03:40:35Z

/ci run

okorzh-amd · 2026-09-29T04:45:01Z

/ci run

github-actions · 2026-09-29T04:45:17Z

❌ @okorzh-amd, Only reviewers with write access can use CI commands before CI is delegated to the PR author.

AndreasKaratzas · 2026-09-29T04:48:56Z

/ci run

AndreasKaratzas · 2026-09-29T04:49:03Z

/ci run

github-actions · 2026-09-29T04:49:15Z

✅ Triggered Buildkite CI #91746 for commit 991734b1693d.

github-actions · 2026-09-29T04:49:34Z

✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/91746

dllehr-amd

Looks good! THanks @rbrugaro-amd

rbrugaro-amd requested review from dllehr-amd, hongxiayang and tjtanaa as code owners September 18, 2026 22:18

claude Bot reviewed Sep 18, 2026

View reviewed changes

mergify Bot added kimi k3 rocm Related to AMD ROCm labels Sep 18, 2026

github-project-automation Bot added this to AMD Sep 18, 2026

github-project-automation Bot moved this to Todo in AMD Sep 18, 2026

rbrugaro-amd mentioned this pull request Sep 18, 2026

[ROCm][Kimi-K3][Perf] Fuse MLA decode KV-cache write and Q-prep via AITER #57452

Closed

rbrugaro-amd force-pushed the rbrugaro/kimi-k3-fused-mla-qkprep-vendor branch from 4ec495c to 30573bf Compare September 21, 2026 16:11

Rohan138 reviewed Sep 21, 2026

View reviewed changes

rbrugaro-amd force-pushed the rbrugaro/kimi-k3-fused-mla-qkprep-vendor branch from 30573bf to 4e333f0 Compare September 22, 2026 04:51

Rohan138 reviewed Sep 23, 2026

View reviewed changes

amd-mghanimi reviewed Sep 23, 2026

View reviewed changes

Rohan138 reviewed Sep 23, 2026

View reviewed changes

Rohan138 approved these changes Sep 28, 2026

View reviewed changes

Rohan138 added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 28, 2026

rbrugaro-amd added 2 commits September 28, 2026 14:14

Merge branch 'main' into rbrugaro/kimi-k3-fused-mla-qkprep-vendor

45fbf87

Merge branch 'main' into rbrugaro/kimi-k3-fused-mla-qkprep-vendor

8f53630

Merge branch 'main' into rbrugaro/kimi-k3-fused-mla-qkprep-vendor

991734b

Monishver11 mentioned this pull request Sep 29, 2026

[Perf][Kernel] Vectorized MLA KV-cache write for large launches #59289

Draft

3 of 4 tasks

dllehr-amd approved these changes Sep 30, 2026

View reviewed changes

dllehr-amd merged commit 8f24ab3 into vllm-project:main Sep 30, 2026
92 checks passed

okorzh-amd mentioned this pull request Oct 10, 2026

[ROCm] ROCm binding on common Kimi-K3 Dspark MLA module #60857

Draft

6 tasks

	if self.is_aiter_triton_fp4_bmm_enabled:
	from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4

	mqa_ql_nope = batched_gemm_a16wfp4(
	mqa_q_nope,
	self.W_K,
	self.W_K_scale,
	transpose_bm=True,
	prequant=True,
	y_scale=self._q_scale if fp8_attention else None,
	)

Conversation

rbrugaro-amd commented Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Scope, and why this is small

Relationship to #47757 and #57452

Why not a duplicate

Implementation

Test Plan

Server (identical for both arms)

Perf

Kernel census

Accuracy

Results

1. Decode kernel census — two kernels become one

2. GSM8K — no accuracy regression

3. E2E serving — unchanged within noise

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

rbrugaro-amd commented Sep 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rbrugaro-amd commented Sep 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Rohan138 Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

amd-mghanimi left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Rohan138 Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Rohan138 left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Sep 28, 2026

Uh oh!

rbrugaro-amd commented Sep 29, 2026

Uh oh!

okorzh-amd commented Sep 29, 2026

Uh oh!

github-actions Bot commented Sep 29, 2026

Uh oh!

AndreasKaratzas commented Sep 29, 2026

Uh oh!

AndreasKaratzas commented Sep 29, 2026

Uh oh!

github-actions Bot commented Sep 29, 2026

rbrugaro-amd commented Sep 18, 2026 •

edited

Loading

rbrugaro-amd commented Sep 21, 2026 •

edited

Loading

Rohan138 Sep 23, 2026 •

edited

Loading

Rohan138 Sep 23, 2026 •

edited

Loading