[Liang-jianhao97]

Motivation

FlashKDA still spoke a dense ABI: initial_state[N, H, V, K] plus a freshly allocated out. vLLM Kimi-K3 serving keeps recurrent state in a paged pool [slots, H, V, K] with state_indices / has_initial_state.

Wiring that through AITER meant:

1. gather_initial_states (or state_cache[indices]) into a dense tensor

2. run chunk_kimi_delta_attn

3. scatter final_state back

4. out.copy_(kernel_out)

Those copies sit in the serving timer. On MI350, a K3-shaped FlashKDA prefill paid ~40–60 µs for them (4K: 171 → 130 µs, 16K: 499 → 440 µs with the same segment pin). The fused vLLM path already reads and writes the pool in-kernel (vLLM #54038); this PR gives FlashKDA the same cache contract so a later vLLM dispatch (#56036) does not have to gather around AITER.

Technical Details

Public kwargs on chunk_kimi_delta_attn / flash_kda_fwd:

• state_cache: fp32 V-first [slots, H, V, K]. Each slot plane is dense; stride(0) may pad.

• state_indices: int32 [N]

• has_initial_state: bool [N]; false starts that sequence from zero and does not read a dirty slot

• out: required with state_cache, same shape as v (K2 still loads v, so out must not alias it)

K2 and the scan load/store slot * stride(0) + head * V * K + …. The last segment writes the final state back to the same row. Segmented pass A does not bind state_cache (keeps the existing K2 binary / autotune; binding it here previously ~2× slowed K2).

@input_guard(skip=("state_cache",)) leaves the pool’s storage and stride(0) alone. Other operands are still packed. The skip is an AITER decorator default, not a vLLM flag; the caller just passes state_cache= as a keyword.

Dense initial_state is unchanged and mutually exclusive with state_cache. final_state is None when the cache is used.

The default five-kernel path indexes a dense h0 and is GVA / C=64 / K-first. If state_cache is set and FlashKDA is not selected,

chunk_delta_attn_fwd raises instead of ignoring the argument.

Test Plan

pytest op_tests/triton_tests/chunk_delta_attn/test_flash_kda.py \

  -k 'paged_cache or flash_kda' -q

Added cases:

• test_paged_cache_matches_dense — paged vs dense initial_state, bit-identical out and selected slots; empty slots untouched; unsegmented and segmented

• test_paged_cache_has_initial_state_false — false flag matches a zero initial_state

• test_paged_cache_padded_stride_writes_live_pool — non-contiguous stride(0) updates the original storage pointer

• test_paged_cache_rejected_on_default_pipeline — non-FlashKDA call with state_cache raises

Serving-shaped A/B (K3 H=12, D=128, lower_bound=-5, varlen cu_seqlens)

vs gather and vs vLLM fused prefill on gfx950.

Test Result

Paged vs dense FlashKDA: torch.equal on output and selected cache rows (rel err 0).

Paged FlashKDA vs vLLM fused prefill (same op, not e2e): output rel L2 ~4.8e-3, state ~2.5e-3 at 4K/16K.

MI350X, official vLLM ROCm image, serving-style call (µs, smaller is better):

shape	gather FlashKDA	paged FlashKDA	fused
1×4096, H=12	171	130	125
1×16384, H=12, --segment 8	499	440	389
1×16384, H=12, auto segment	—	380	388

Default-pipeline reject and padded-stride tests pass.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5754 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[Copilot]

Copilot review overview

🟡 Changes recommended

Unresolved safety, output-semantics, performance, and test-coverage issues remain.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1 · 4

Open (5)

• Reject outer strides smaller than the cache plane · New
• Reject or forward out when FlashKDA is disabled · New
• Mask zero-start cache loads to avoid reading dirty slots · New
• Avoid synchronous state-index reductions on the hot path · New
• Preserve caller-owned output for non-contiguous buffers · New

What changed in this PR

Adds in-kernel fp32 paged state-cache I/O to FlashKDA, avoiding dense gather/scatter operations.

Changes:

• Adds paged cache, slot indices, initial-state flags, and caller-owned output support.
• Implements cache reads/writes for segmented and unsegmented paths.
• Adds cache equivalence, stride, initialization, fallback, and dispatch tests.

File	Summary
op_tests/​triton_tests/​chunk_delta_attn/​test_flash_kda.py	Adds paged-cache tests; coverage should use the public wrapper and include multi-sequence, permuted-slot cases.
aiter/​ops/​triton/​kimi_delta_attn/​chunk_delta_attn.py	Extends the public FlashKDA API for paged state-cache arguments.
aiter/​ops/​triton/​_triton_kernels/​chunk_delta_attn/​flash_kda.py	Implements cache I/O; unresolved issues include outer-stride validation, masked zero-state loads, synchronous index validation, and non-contiguous output handling.
aiter/​ops/​triton/​_triton_kernels/​chunk_delta_attn/​chunk_fwd.py	Dispatches FlashKDA cache requests, but supplied out is ignored when the default pipeline is selected.
aiter/​ops/​triton/​_triton_kernels/​chunk_delta_attn/​chunk_delta_attn_utils.py	Preserves cache storage and strides through input guarding.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot review overview

🟡 Changes recommended

Positional cache/output handling and missing multi-sequence coverage leave moderate correctness issues unresolved.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1

Open (1)

• Positional out and state_cache bypass skip handling · New

Resolved since last review (5)

• Reject outer strides smaller than the cache plane
• Preserve caller-owned output for non-contiguous buffers
• Avoid synchronous state-index reductions on the hot path
• Mask zero-start cache loads to avoid reading dirty slots
• Reject or forward out when FlashKDA is disabled