[kliuae]

Purpose

On ROCm's Kimi-K3 path, each prefill/mixed step has a stall in prepare_chunk_indices, caused by two interacting factors:

def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Tensor:
    indices = torch.cat(
        [
            torch.arange(n)
            for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()
        ]
    )
    return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens)

• tolist() D2H copy makes the host wait until the device queue is emptied
• Host ops executing before the next H2D copy, while the device is sitting idle

The data that the host waits on is already sitting on the host. At each step, GDNAttentionMetadataBuilder builds chunk_indices and chunk_offsets from query_start_loc_cpu and saves them in attention metadata. However, the ROCm Kimi-K3 path does not use it, and the stall waiting on and rebuilding the data from the device is a duplication of the same work.

This PR addresses this and removes the stall by

• Making the chunk kda kernels optionally accept chunk_indices and chunk_offsets. When they're provided, they don't need to be rebuilt through prepare_chunk_indices.
• Passing these metadata to the kernels to actually remove the stall.
• For the prepare_chunk_offsets call in metadata builder, building them directly on the device, instead of having a synchronization point waiting on the host-side python loop and host-device copy before the next device workload can be queued.

These remove the stall from prepare_chunk_indices, and achieves TTFT and performance gain especially in low concurrency scenarios where the stall cannot be hidden by the heavy device workloads.

Test Plan

E2E benchmark and lm_eval with gsm8k on MI355X

Server command

vllm serve moonshotai/Kimi-K3 \
  --served-model-name moonshotai/Kimi-K3 \
  --trust-remote-code \
  --moe-backend auto \
  --tensor-parallel-size 8 \
  --load-format auto \
  --gpu-memory-utilization 0.95 \
  --mm-encoder-tp-mode data \
  --max-num-seqs 128 \
  --max-num-batched-tokens 4096 \
  --max-model-len 10240 \
  --no-enable-prefix-caching \
  --enable-auto-tool-choice \
  --tool-call-parser kimi_k3 \
  --reasoning-parser kimi_k3 \
  --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY","custom_ops":["+fused_rms_norm_gated"]}' \
  --attention-backend TRITON_MLA

Benchmark

vllm bench serve \
  --model moonshotai/Kimi-K3 --backend vllm --trust-remote-code \
  --dataset-name random \
  --random-input-len 8192 --random-output-len 128 \
  --max-concurrency "$C" --num-prompts $((C * 4)) \
  --temperature 0 --ignore-eos --seed 1234

lm_eval

lm_eval --model local-chat-completions \
  --model_args "model=moonshotai/Kimi-K3,base_url=http://localhost:8000/v1/chat/completions,num_concurrent=64,max_retries=3,tokenized_requests=False,max_gen_toks=2048,timeout=3600" \
  --tasks gsm8k --num_fewshot 20 \
  --apply_chat_template --fewshot_as_multiturn \
  --seed 1234

Test Result

gsm8k on MI355X

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	20	exact_match	↑	0.9757	±	0.0042
		strict-match	20	exact_match	↑	0.9765	±	0.0042

Benchmark results

ISL/OSL	Conc	Config	Total tokens/s	Delta (%)	Total reqs/s	Delta (%)	Output tokens/s	Delta (%)	mean TTFT (ms)	Delta (%)	mean TPOT (ms)	Delta (%)	mean ITL (ms)	Delta (%)
8192/128	8	main	8052.42	-	0.97	-	123.88	-	1813.30	-	50.47	-	50.47	-
8192/128	8	PR	8316.70	+3.28%	1.00	+3.09%	127.95	+3.29%	1766.05	-2.61%	48.78	-3.35%	48.78	-3.35%
8192/128	16	main	10067.08	-	1.21	-	154.88	-	2481.99	-	83.84	-	83.84	-
8192/128	16	PR	10358.07	+2.89%	1.24	+2.48%	159.35	+2.89%	2447.98	-1.37%	81.17	-3.18%	81.17	-3.18%
8192/128	32	main	11674.19	-	1.40	-	179.60	-	3729.02	-	148.50	-	148.50	-
8192/128	32	PR	11766.98	+0.79%	1.41	+0.71%	181.03	+0.80%	3713.48	-0.42%	147.20	-0.88%	147.20	-0.88%
8192/128	64	main	12760.26	-	1.53	-	196.31	-	6267.18	-	275.10	-	275.10	-
8192/128	64	PR	12822.07	+0.48%	1.54	+0.65%	197.26	+0.48%	6255.47	-0.19%	273.57	-0.56%	273.57	-0.56%

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

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

[njhill]

Thanks @kliuae, I think this is a duplicate of #51540. Would you mind reviewing / trying out that one?

[njhill]

I guess they are not completely duplicate but overlapping

[kliuae]

Thanks for pointing this to the #51540 PR @njhill. I believe this PR is not a duplicate, but it extends the coverage and enhances sync/stall removal for Kimi-K3 on ROCm.

For ROCm KDA, #51540 only handles the Kimi-Linear path, and Kimi-K3 still uses prepare_chunk_indices and thus the sync when it is called. This PR passes the existing field in metadata (for correctness, ROCm Kimi-K3 requires a different chunk_indices field that doesn't contain decode, but it has already been built) to the consumer sites on the ROCm Kimi-K3 path so the sync/stall is removed.

This PR additionally wires chunk_offsets, which is also already built in metadata, to ROCm Kimi-K3 call sites to bypass prepare_chunk_offsets's host ops and H2D copy.

This PR also builds chunk_indices and chunk_offsets directly on device rather than on the host to avoid the host ops and H2D copy for the consumer in metadata builder.

[tjtanaa]

@kliuae @njhill how would you guys like consolidate them and ship them quickly?

[tjtanaa]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83506 for commit 4aea8cbcf91f.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @kliuae.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[tjtanaa]

LGTM. For this PR it does not affect NVIDIA, and only affects AMD. This PR contains full necessary feature to reduce D2H and H2D on ROCm. We can ship this PR then only ship PR #51540 which is more general and affecting both NVIDIA and ROCm path.

NOTE: #51540 does not cover kimi k3 on ROCm so both PRs are orthogonal.

[github-actions]

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

• /ci run starts a CI build.
• /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.
• /ci cancel cancels scheduled or running CI builds for this PR branch.

[kliuae]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83526 for commit ea5df9efbe1c.