[ROCm][Bugfix] Kimi-K3 Fix KDA NaN on mixed batches and racy autotune config

[kliuae]

Purpose

For Kimi-K3, there are two correctness bugs on the ROCm Kimi-K3 KDA path.

• KimiGatedDeltaNetAttention._forward sends the decode sequences of a mixed decode+prefill non-speculative batch through the KDA chunk kernel chunk_kda_with_fused_gate. The chunk kernel treats each decode as length=1 sequence, and returns NaN for them. This PR separates the decode out of the mixed batch and pass them to fused_recurrent_kda.
• recompute_w_u_fwd_kernel autotunes the configuration of num_stages=4. This configuration, particularly in combination with num_warps=4, is racy for this kernel, and can silently produce large (e.g. ~1e38) garbage values that corrupts the downstream. This PR addresses this by leaving the num_stages=4 out in autotuning so that it doesn't land on the erroneous configuration.

Test Plan

moonshotai/Kimi-K3 with gsm8k on MI355X

Server command

AITER_BF16_FP8_MOE_BOUND=0 \
VLLM_ROCM_USE_AITER=1 \
SAFETENSORS_FAST_GPU=1 \
AITER_SITUV2_A8W4=1 \
VLLM_USE_BREAKABLE_CUDAGRAPH=0 \
vllm serve 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 \
  --enable-auto-tool-choice \
  --tool-call-parser kimi_k3 \
  --reasoning-parser kimi_k3

lm_eval

lm_eval --model local-completions --model_args model=moonshotai/Kimi-K3,base_url=http://0.0.0.0:8000/v1/completions,tokenizer_backend=None,tokenized_requests=False,num_concurrent=128,max_retries=10,max_gen_toks=2048,timeout=60000 --batch_size auto --tasks gsm8k --num_fewshot $FEWSHOTS

Test Result

Before, 5-shot

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.8901	±	0.0086
		strict-match	5	exact_match	↑	0.8901	±	0.0086

PR, 5-shot

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9689	±	0.0048
		strict-match	5	exact_match	↑	0.9697	±	0.0047

PR, 20-shot

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	20	exact_match	↑	0.9697	±	0.0047
		strict-match	20	exact_match	↑	0.9704	±	0.0047

---

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.

[ZJY0516]

@kliuae Could you copy vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py into kimi_k3/amd/ and then modify? like #50364

[kliuae]

@kliuae Could you copy vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py into kimi_k3/amd/ and then modify? like #50364

Sure, copied the kimi-k3 dedicated kda to kimi_k3/amd/kda.py and applied the fix there.

[tjtanaa]

LGTM

[ZJY0516]

LGTM now