[qescccczmr]

Motivation

FLA 0.5.2's chunk_kda adapter does not forward internal recurrent checkpoint arguments. The validated vLLM FlashKDA fork exposes them through its native operator, while its public Python wrapper omits them.

Provide a shared recurrent prefill entry point for ordinary prefill and checkpoint export, while keeping ordinary prefill usable through FLA when the optional native checkpoint operator is unavailable.

Modification

Kda -> CudaKdaImpl._prefill_recurrent
    -> torch.ops.flash_kda.fwd  (compatible native operator available)
    -> FLA chunk_kda           (ordinary prefill fallback)

• Add optional, keyword-only checkpoint_state and checkpoint_offsets to the KDA module/backend interfaces. A native checkpoint call processes the full supplied interval and also exports the recurrent state at the requested position.
• Select the native path using tensor/device requirements: NVIDIA SM90, BF16 inputs, K/V dimensions of 128, FP32 recurrent state and a gate lower bound in [-5, 0]. Cache hardware support by device index with lru_cache(maxsize=16); no model identifier is passed to the backend.
• Cache optional native availability per backend instance. Ordinary prefill retains FLA dispatch if the extension cannot load, the native operator is absent, or its checkpoint schema is incompatible. Explicit checkpoint requests require native support and fail before cache mutation. Native execution errors are propagated without retrying computation through FLA.
• Before convolution or MTP cache access, validate paired checkpoint arguments, state dtype/shape/device/contiguity, offset shape/device/contiguity/dtype, and corresponding cu_seqlens metadata. Offset values remain the caller's responsibility so the hot path does not synchronize GPU tensors to the CPU.
• Allocate native output, final-state and workspace buffers in the CUDA backend; checkpoint buffers remain caller-owned. Preserve FLA convolution/gate operations, AR/MTP decode kernels, state-ring behavior, prefix-checkpoint lifecycle, and existing paths for other models/configurations.

Scope

This PR provides the recurrent checkpoint interface. It does not integrate IC scheduling, supply checkpoint arguments from the normal model forward, or remove checkpoint-related prefill splitting. Scheduler/model wiring, matching convolution snapshots, and checkpoint reservation/publication/restoration remain separate work.

A multi-request view such as ring[:, slot] is generally non-contiguous and is rejected before cache mutation. Batched IC integration must provide contiguous scratch storage and scatter the exported states into reserved request/ring slots. That lifecycle remains outside this API PR.

Compatibility and dependency

Official flash-linear-attention==0.5.2 remains the common dependency. Ordinary prefill can use its existing dispatch without installing a checkpoint-capable FlashKDA fork. No new backend-selection flag or global dependency upgrade is introduced. FLA_FLASH_KDA keeps its meaning inside FLA; it does not control the direct native path.

Internal checkpoint export requires the separate vllm-project/FlashKDA fork, validated at commit ee0be888cd0e972f9409bf53756f8c38c6652173, merged into that fork's dev by PR #7. The tested checkout has no local source modifications. MoonshotAI/FlashKDA master at 7afb9f4 does not provide this checkpoint ABI. Official FLA plus MoonshotAI FlashKDA therefore does not enable IC export.

The fork's PyTorch 2.10+/CUDA Toolkit 12.9+ build requirements apply to the optional native path. See docs/en/advance/flash_kda.md for the pinned installation and buffer contract.

Validation

Local test environment: PyTorch 2.13.0+cu129, CUDA 12.9, official FLA 0.5.2. This PyTorch version is newer than the repository's current torch<=2.12.1 pin; the results below are kernel/API validation, not verification of the default installation matrix.

• tests/pytorch/test_kda_prefill.py: 91 passed. Includes missing package/extension/operator, old checkpoint ABI, simulated extension-loading errors, ordinary/MTP fallback, cache preservation on invalid buffers, and rejection of strided multi-request ring views.
• tests/pytorch/kernel/test_kda.py: 17 passed on an idle H200 with the specified fork and official FLA 0.5.2. Includes existing IC/prefix/resume/ring numerical checks and exact output/final-state agreement between fallback and direct FLA at 129 and 2048 tokens. The absence of FlashKDA is simulated for the fallback test; the FLA GPU kernels execute normally.
• pre-commit run --all-files: all hooks passed under Python 3.12.

Native-versus-FLA differences are recorded as JUnit properties, separately from the exact fallback regression check:

Tokens	Output max absolute error	Output relative RMS error	Final-state max absolute error	Final-state relative RMS error
129	6.10e-5	0.512%	7.36e-4	0.450%
2048	1.22e-4	0.553%	9.20e-4	0.417%

These implementations are not elementwise equivalent: final-state comparison at rtol=1e-2, atol=1e-4 has outliers on these inputs. The metrics do not establish full-model accuracy. CUDA 12.8/PyTorch 2.9 installations, long-context workspace peaks, TTFT/TPS, DP/EP deployment and PCG performance have not been validated. This PR does not claim the end-to-end prefill speedup of a complete IC scheduler integration.

Checklist

• Run the repository's full pre-commit checks.
• Cover optional dependency fallback and cache preservation before invalid checkpoint requests.
• Validate recurrent checkpoint behavior and exact FLA fallback behavior on H200.
• Document the fork dependency and remaining IC integration/performance scope.

[RunningLeon]

pls remove model_type, this argument is a model specific argument, not a good design

[qescccczmr]

Removed in 7311bae2. Kda, KdaBuildSpec, and CudaKdaImpl no longer accept model_type, and the model call site is back to Kda(). Dispatch now depends only on the tensor/device requirements. GLM-5.3 is currently the only model using this shared KDA module, so existing paths for other models are unchanged.

Validation: 49 CPU contract tests, 15 H200 numerical tests with vllm-project/FlashKDA ee0be88 without local source modifications, and all pre-commit hooks passed.

Correction to the dependency attribution in my original reply: the tested FlashKDA revision belongs to the vLLM fork, not MoonshotAI's default branch. Official FLA 0.5.2 alone, or together with MoonshotAI/FlashKDA master at 7afb9f4, does not provide the native checkpoint operator required by this PR. The PR description and installation documentation have been corrected accordingly.

[RunningLeon]

can we use lru cache for this repeated check? such as

@lru_cache(maxsize=16)
def _supports_flash_kda_device(device_index: int) -> bool:
    return (
        torch.version.hip is None
        and torch.cuda.get_device_capability(device_index) == (9, 0)
    )

[qescccczmr]

Added in 7311bae2 as suggested: a module-level _supports_flash_kda_device(device_index) with @lru_cache(maxsize=16), called with q.device.index. The hardware check is shared across backend instances, while tensor/dtype/gate checks remain per call.

Tests cover caching both supported and unsupported capabilities, reuse across instances, and separate cache entries for different CUDA devices.