[seungrokj]

Overview

Since #59552, the gluon KDA prefill backend (AITER chunk_kda) is the default on CDNA4. It switches Kimi-K3's KDA prefill checkpoints off, because AITER chunk_kda cannot export the mid-prefill state. This PR keeps the checkpoints on with gluon: only prefill steps that carry checkpoint metadata go to the fused HIP kernel, and every other step stays on gluon.

Claims

• With kda_prefill_backend="gluon" and the fused kernel built, get_kv_cache_spec reserves the prefill checkpoint block (alignment KDA_CHECKPOINT_ALIGNMENT), as the fused backend does. Previously it reported 0 blocks.
• A step whose checkpoint metadata is set runs chunk_kda_prefill(backend="fused"). Steps without checkpoint metadata keep the configured backend.
• If the fused kernel is missing (is_fused_kda_chunk_supported() is False), gluon still opts out of checkpoints, so behavior is unchanged.

Validation

Kernel and spec tests, run on MI355X (1 GPU) in vllm/vllm-openai-rocm:nightly-rocm100-df6603f0583375f244d1ac6d36265462b7490392, with this diff applied to the image's vLLM (its kimi_k3/amd/kda.py matches the base commit):

Command	Without this PR	With this PR
pytest tests/models/kimi_k3/test_amd_kda_checkpoint.py -k gluon_backend	test_gluon_backend_keeps_the_checkpoint fails: assert 0 == 1 on num_prefill_checkpoint_blocks	2 passed
pytest tests/models/kimi_k3/test_amd_kda_checkpoint.py	n/a	28 passed
VLLM_ROCM_USE_AITER=1 pytest tests/models/kimi_k3/test_amd_kda_chunk.py	n/a	77 passed (includes the gluon parity tests and the new handoff test)
pytest tests/models/kimi_k3/test_amd_kda_checkpoint.py tests/models/kimi_k3/test_amd_kda_chunk.py (no AITER)	n/a	98 passed, 6 skipped (gluon needs AITER)

The new test_gluon_hands_checkpoint_steps_to_fused test prefills one 1024-token sequence in three steps: gluon [0, 256), then fused [256, 768) with a checkpoint at token 512, then gluon [768, 1024). It compares the result with a single Triton prefill of the whole sequence:

• Output: atol/rtol 2e-2.
• Final state: 5e-3.
• Checkpoint versus a Triton prefill of the 512-token prefix: 5e-3.

With VLLM_ROCM_USE_AITER=1, test_amd_kda_checkpoint.py has 11 metadata-builder failures (KimiK3ROCmKDAMetadataBuilder builds its own FLA chunk metadata ... 'aiter_flydsl'). They fail the same way on the unpatched image, and this PR does not touch them.

End-to-end check: Kimi-K3 TP8, DCP8, DSpark with 4 drafts and real rejection sampling, --enable-prefix-caching (mamba cache mode align), --max-num-batched-tokens 8192, on the same image. Both arms use the default auto KDA prefill backend, which resolves to gluon.

	Without this PR	With this PR
Server log	KDA prefill backend: gluon	KDA prefill backend: gluon, KDA prefill checkpoint enabled.
GSM8K 5-shot, 1319 questions (flexible / strict)	0.9560 / 0.9553 (±0.0056)	0.9659 / 0.9651 (±0.0050)
Prefix-cache check (below)	PASS	PASS
Tracebacks in server log	0	0

The GSM8K prompts are shorter than one mamba block, so they never reach a checkpoint (prefix_cache_hits_total grows only during the prefix check). GSM8K therefore shows that steps without a checkpoint still serve correctly on gluon. The prefix check is what exercises the checkpoint.

The prefix check sends four prompts of about 3.8k / 9.4k / 20.6k / 45.2k tokens. Each prompt is a list of facts plus a question about one of them, run with greedy decoding:

1. Request A, cold.
2. Request B: the same facts and a different question, served from the cache A left.
3. Request B again, after /reset_prefix_cache.

Prompt tokens	B cached tokens, without / with	Answers correct (A, B warm, B cold), without / with
3771	1536 / 1536	3/3 / 3/3
9371	7680 / 7680	3/3 / 3/3
20644	18432 / 18432	3/3 / 3/3
45151	43008 / 43008	3/3 / 3/3

Cached lengths are the same with and without the PR. Without it, the scheduler gets the same snapshot by ending the prefill step at the block boundary and running the tail as an extra step. With the PR, the snapshot comes out of the fused kernel inside the step, and the warm answers, which resume from it, are correct. This host is thermally throttled, so the time saved by dropping the extra step was measured on another node (below).

Time to first token. Kimi-K3 TP8 on one unthrottled MI355X node, same image, default auto (gluon) backend, --enable-prefix-caching, DSpark with synthetic acceptance. Two server configurations from the InferenceX Kimi-K3 recipe:

• c1: DCP1, --max-num-batched-tokens 16384.
• c16: DCP8, --max-num-batched-tokens 8192.

The client sends one random prompt at a time with 1 output token, after a warmup, using the same seeds in both arms. No prompt hit the prefix cache. Each cell is the median of 10 prompts (6 at 32768 and 65536), in ms.

Prompt tokens	c1 without	c1 with	Δ	c16 without	c16 with	Δ
1024	203.4	199.7	−1.9%	198.6	197.2	−0.7%
4096	413.0	237.2	−42.6%	398.6	244.9	−38.5%
8192	484.5	373.8	−22.9%	489.5	383.2	−21.7%
16384	793.3	686.3	−13.5%	905.1	905.5	0.0%
32768	1537.8	1540.0	+0.1%	1733.7	1737.3	+0.2%
65536	3263.1	3143.7	−3.7%	3670.8	3544.1	−3.5%

Where the PR removes the split, it saves 107–176 ms per prompt. The unchanged rows are expected:

• With DSpark's block drop, the checkpoint sits one 1536-token block before the last full block.
• 1024 is shorter than one block, so there is no checkpoint.
• At 32768 the checkpoint (30720) is already a chunk boundary at both chunk sizes, so neither arm splits.
• At 16384 with c16, the checkpoint (13824) falls in a chunk that is not the prompt's last, where both arms stop anyway.

prefix_check.py

"""Prefix-cache check for Kimi-K3 KDA prefill checkpoints.

Per prompt length: a cold request A (facts + question about fact i), then B
(same facts + question about fact j) which should hit the prefix cache, then
B again after a prefix-cache reset (cold). B's answer must be right both ways.
"""
import json, random, sys, time, urllib.request

URL = "http://localhost:8000"
TOK = sys.argv[1] if len(sys.argv) > 1 else "/model"
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained(TOK, trust_remote_code=True)

def post(path, body=None):
    req = urllib.request.Request(URL + path, data=json.dumps(body or {}).encode(),
                                 headers={"Content-Type": "application/json"}, method="POST")
    t = time.time()
    with urllib.request.urlopen(req, timeout=3600) as r:
        raw = r.read()
    return (json.loads(raw) if raw.strip().startswith(b"{") else raw.decode()), time.time() - t

def metrics():
    with urllib.request.urlopen(URL + "/metrics") as r:
        txt = r.read().decode()
    out = {}
    for ln in txt.splitlines():
        for k in ("vllm:prefix_cache_hits_total", "vllm:prefix_cache_queries_total"):
            if ln.startswith(k):
                out[k] = out.get(k, 0) + float(ln.split()[-1])
    return out

def build(n_tokens, seed):
    rnd = random.Random(seed)
    lines, facts = [], []
    while True:
        name = "".join(rnd.choice("BCDFGHJKLMNPRSTVZ") + rnd.choice("aeiou") for _ in range(3)).capitalize()
        code = rnd.randint(1000, 9999)
        facts.append((name, code))
        lines.append(f"Record {len(facts)}: the access code for the city of {name} is {code}.")
        if len(lines) % 50 == 0 and len(tok.encode("\n".join(lines))) >= n_tokens:
            break
    return "\n".join(lines), facts

def ask(doc, name):
    p = (f"Below is a list of records.\n\n{doc}\n\nQuestion: What is the access code for "
         f"the city of {name}? Answer with the four-digit number only.\nAnswer: The access code for the city of {name} is")
    r, dt = post("/v1/completions", {"model": "moonshotai/Kimi-K3", "prompt": p, "max_tokens": 8,
                                      "temperature": 0, "seed": 0})
    c = r["choices"][0]["text"]
    cached = (r.get("usage", {}).get("prompt_tokens_details") or {}).get("cached_tokens")
    return c, r["usage"]["prompt_tokens"], cached, dt

res = []
ok_all = True
for n, seed in ((3000, 1), (9000, 2), (20000, 3), (45000, 4)):
    doc, facts = build(n, seed)
    rnd = random.Random(seed + 100)
    (na, ca), (nb, cb) = facts[rnd.randrange(len(facts))], facts[rnd.randrange(len(facts) // 4)]
    post("/reset_prefix_cache")
    a = ask(doc, na)
    m0 = metrics()
    b_warm = ask(doc, nb)
    m1 = metrics()
    post("/reset_prefix_cache")
    b_cold = ask(doc, nb)
    row = dict(target_tokens=n, prompt_tokens=a[1],
               A=dict(expect=ca, text=a[0], ok=str(ca) in a[0], cached=a[2], s=round(a[3], 2)),
               B_warm=dict(expect=cb, text=b_warm[0], ok=str(cb) in b_warm[0], cached=b_warm[2], s=round(b_warm[3], 2),
                           hit_tokens=m1.get("vllm:prefix_cache_hits_total", 0) - m0.get("vllm:prefix_cache_hits_total", 0)),
               B_cold=dict(text=b_cold[0], ok=str(cb) in b_cold[0], cached=b_cold[2], s=round(b_cold[3], 2)),
               warm_eq_cold=b_warm[0] == b_cold[0])
    ok_all &= row["A"]["ok"] and row["B_warm"]["ok"] and row["B_cold"]["ok"]
    res.append(row)
    print(json.dumps(row), flush=True)
print("PREFIX_CHECK", "PASS" if ok_all else "FAIL")

Details

• Root cause: can_checkpoint required kda_prefill_backend == "fused". When auto resolved to gluon, the spec reserved no checkpoint block. The scheduler then fell back to splitting the final prefill chunk at the block boundary (scheduler.py), which costs an extra step per prompt and drops the interior checkpoints.
• Why route to fused, not export from gluon: AITER chunk_kda (amd-aiter 0.1.25 / 0.1.26 and AITER main) has no checkpoint output. The proper fix is to export the state from the chunk walk inside the kernel; that needs an AITER change first. This PR is the vLLM-only bridge. When AITER grows that output, the routing can be replaced by a signature check.
• Trade-off: checkpoint metadata is per step, so a batched prefill step in which any request crosses a snapshot position runs entirely on the fused kernel. Steps without a checkpoint keep gluon's speed. Decode throughput was not part of this measurement.
• Related PRs (not duplicates):
  • [ROCm][Perf] Integrate AITER KDA prefill for Kimi-K3 #56036 adds a separate flashkda prefill backend and does not touch checkpoint gating.
  • [ROCm][GLM-5.3-Flash] Enable KDA prefill checkpoints on the Triton backend #61951 adds Triton prefill checkpoints for GLM-5.3-Flash KDA through a host-side split.
  • [Mamba] Declare prefill checkpoint alignment once per layer #61466 moves the prefill checkpoint alignment into MambaBase.get_prefill_checkpoint_alignment() and edits the same get_kv_cache_spec. Whichever lands second needs a small rebase; the gating condition here carries over unchanged.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

AI assistance: Claude Opus 5.5 (claude-opus-5-5, Claude Code) drafted the change, the tests and this description, and ran the validation above. The human submitter reviews every changed line before marking the PR ready.

🤖 Generated with Claude Code

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[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.

[ChuanLi1101]

This can land. Gluon cannot write the mid-prefill state, so the scheduler splits the last chunk into an extra step. This PR keeps the checkpoint and runs only that step on the fused kernel. The other steps stay on gluon. That is why TTFT goes from 413 ms to 237 ms at 4K and from 484 ms to 374 ms at 8K. 32K does not change, because the checkpoint is already on a chunk boundary.

[github-actions]

✅ @seungrokj, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, 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.

[seungrokj]

/ci run

[github-actions]

✅ Triggered Buildkite CI #96999 for commit 33837647068b.

[seungrokj]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97581 for commit 8fbdbe24cb89.

ℹ️ This PR is 1 commit behind upstream main. Merged upstream main into this branch as 8fbdbe24cb89 before starting CI. Pull the branch before pushing more changes.

[fxmarty-amd]

LGTM

@omuhamma PTAL

[dllehr-amd]

Looks good! Thanks @seungrokj!

[omuhamma]

lgtm

[fxmarty-amd]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97984 for commit 62a4ee6b9d87.

ℹ️ This PR is 15 commits behind upstream main. Merged upstream main into this branch as 62a4ee6b9d87 before starting CI. Pull the branch before pushing more changes.

[seungrokj]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98396 for commit f235224c6647.

[seungrokj]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98449 for commit 9d4f545c9f9c.

ℹ️ This PR is 12 commits behind upstream main. Merged upstream main into this branch as 9d4f545c9f9c before starting CI. Pull the branch before pushing more changes.

[seungrokj]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98790 for commit 2909be057d22.