flatten_kv_cache zero padding

[grimoire]

fa3 might read dirty data in padding block.

[Copilot]

Pull request overview

This PR addresses a correctness issue where consumers (notably FA3 paths) may read beyond the “real” flattened KV length and encounter uninitialized/dirty values in the padded tail. It updates the Triton flatten kernels to explicitly zero-fill the padded tail region and adds a regression test to ensure the padding is zeros (not NaNs/garbage).

Changes:

• Add a regression test asserting flatten_kv_cache(..., out_size=padded) produces zeroed values in the padded tail for both hsd and shd output layouts.
• Modify Triton kernels to avoid copying “invalid” KV positions into the padded tail and instead zero-fill that tail via an extra “pseudo batch” in the launch grid.
• Apply the same padded-tail zero-fill approach to the quantized and MLA-FP8 flatten kernels.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
tests/pytorch/kernel/test_flatten_kv_cache.py	Adds a new regression test verifying padded tail is zeroed (not NaN/dirty) for both output layouts.
lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py	Updates Triton flatten kernels to zero-fill padded tail by launching an extra batch and writing zeros into the tail region.

Comments suppressed due to low confidence (2)

lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py:172

• Same concern as above: making BATCH_SIZE a tl.constexpr will specialize/compile the quantized kernel per batch size. Prefer passing it as a runtime scalar or computing the tail batch id from tl.num_programs(1) so compilation is not tied to batch size.

    stride_boff,
    quant_policy: tl.constexpr,
    OUT_SIZE,
    BATCH_SIZE: tl.constexpr,
    HEAD_DIM_K: tl.constexpr,
    HEAD_DIM_V: tl.constexpr,

lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py:503

• Same concern as above: BATCH_SIZE being tl.constexpr will generate separate MLA-FP8 kernel variants per batch size. Consider using a runtime BATCH_SIZE (no tl.constexpr) or computing the tail batch index via tl.num_programs(1) to avoid recompilation across batches.

    stride_kod: tl.constexpr,
    stride_boff,
    OUT_SIZE,
    BATCH_SIZE: tl.constexpr,
    BLOCK_BS: tl.constexpr,
    BLOCK_NOPE: tl.constexpr,

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[RunningLeon]

LGTM