[KV Cache] GLM-5.3-Flash: use the generic packed KV layout; CircularBufferSpec tail with speculative ring slots

[ivanium]

Purpose

Follow-up to #53906 (GLM-5.3-Flash). Remove the model-specific KV cache grouping / tensor layout and the matching NIXL and Mooncake special cases by routing GLM through the generic packed (block-outermost) layout that DeepSeek-V4 and Qwen3.8-Next already use, and fix two issues in the kpool tail ring.

• KpoolTailSpec is dropped; the tail is a plain CircularBufferSpec (its manager already was CircularBufferManager).
• The ring is sized for speculative tokens, kpool * cdiv(kpool + num_spec, kpool) as the QSA ring does. With exactly kpool slots, a rejected pool-completing draft left the drafts behind it overwriting the committed keys the redo needs (covered by new CUDA kernel and index-level tests). The seed/decode tail kernels (NVIDIA and AMD) take the ring size and the tail view's strides.
• GLM's MLA and indexer pages have incompatible sizes, which _get_kv_cache_groups_glm5_next / _glm5_next_tensor_layout handled with a hand-rolled layer-outer layout, storage_block_size, worker-side virtual block splitting and NIXL "dense virtual transfer pages". The generic packed path already handles mixed page sizes, but a block-outermost block cannot be split by the worker into the 32/64-row pages DeepGEMM paged-MQA and the TRT-LLM sparse MLA kernel take. This PR re-pages inside the kernels' callers instead: kpool_page_geometry / _kpool_flat_page_view / _indexer_page_table for the indexer (fed by a block_stride_bytes builder hook), _kernel_paged_view for the FlashInfer sparse MLA impls (TRT-LLM on SM100 now MultipleOf(32), 32-row pages; SM120 64-row pages; SM90 uses page size 1 so it only needs the row stride and no longer forces the LBHNC layout), and the packed block stride is aligned to the lcm of the kernel pages (_packed_block_alignment_bytes, only when a layer needs re-paging; DeepSeek-V4 and Qwen are unaffected). The page size a backend re-pages to is declared by a new AttentionBackend.get_kernel_page_rows() and recorded on MLAAttentionSpec.kernel_page_rows, so the layout code does not hard-code kernel constants. GLM keeps block 1152 and 3 Mamba groups; the alignment costs ~4.7% KV capacity on GB200 TP4.
• Net: 34 files, +750/-1303. kv_cache_utils.py and the NIXL worker carry no GLM-specific code path any more. The NIXL wire layout is unchanged (the removed virtual pages were never part of a release), so no NIXL_CONNECTOR_VERSION bump.

Duplicate check: no open PR touches the GLM KV layout, KpoolTailSpec, storage_block_size or the NIXL virtual transfer pages.

Test Plan

pytest tests/v1/core/test_kv_cache_utils.py tests/v1/core/test_contiguous_kv_packing.py \
  tests/v1/attention/test_kpool_page_geometry.py tests/v1/attention/test_kpool_tail_slot_mapping.py \
  tests/v1/kv_connector/unit/test_nixl_desc_geometry.py tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py \
  tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py tests/v1/worker/test_attn_utils.py \
  tests/v1/attention/test_sparse_indexer_decode_seq_lens.py tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py \
  tests/kernels/test_kpool_decode_update_batched.py tests/v1/attention/test_flashinfer_mla_sparse_sm90.py \
  tests/v1/attention/test_sparse_mla_backends.py

GLM-5.3-Flash, TP4-to-TP4 NIXL PD on GB200, GSM8K 5-shot x2 (mqa), no MTP:

Test Result

• 751 passed, 548 skipped (GPU/ROCm-only) locally on the listed files; pre-commit incl. mypy clean.
• GSM8K PD: 0.9727 flexible / 0.9742 strict with this PR (BLHNC, block 1152, 3 Mamba groups, NIXL transfers of the tail on every request), vs 0.9731 / 0.9738 on the layer-outer layout it replaces. 0 request errors.
• MTP decode: the kpool indexer launches DeepGEMM's paged MQA logits without the indices its schedule was built with ([Bugfix] GLM-5.3-Flash: launch the kpool paged MQA logits in the varlen mode its schedule was built with #55270, independent of this PR). On DeepGEMM 2.5.0 that deadlocks GLM MTP decode within minutes; on the pinned 2.6.1 it is tolerated. With [Bugfix] GLM-5.3-Flash: launch the kpool paged MQA logits in the varlen mode its schedule was built with #55270 applied on top of this branch: single-node TP4 MTP3 GSM8K 0.9723 / 0.9738 and TP4-to-TP4 NIXL PD + MTP3 0.9708 / 0.9719 on DeepGEMM 2.5.0; on the pinned 2.6.1 single-node MTP3 passes in eager (0.9704 / 0.9719) and with CUDA graphs (0.97). All with 0 request errors. ROCm kernels mirror the CUDA change but were not run. SM90 and SM120 are covered by unit tests only (no H100 / RTX Pro available): the sparse MLA wiring and the kpool block alignment (32-state indexer pages on SM90/SM100, 64 on SM120, matching the pinned DeepGEMM's get_paged_mqa_logits_metadata constraints). GLM on SM120 additionally needs the rope-free sparse MLA path tracked in GLM-5.3-Flash (glm5_next): no SM120 sparse-MLA path for rope-free MLA (qk_rope_head_dim=0) — three failure modes on RTX PRO 6000 #53963.

AI assistance (Claude Code) was used for analysis, implementation and review; the submitter reviewed every line.

🤖 Generated with Claude Code

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

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

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

[200lz]

I'm preparing a focused fix for ordinary 64-row V3.2 sparse MLA/indexer packing. At the current head, #55219's packed alignment is tied to kernel re-paging, while ordinary heterogeneous row widths still require a common packed-block row alignment. Would it make sense for a shared alignment helper to carry both requirements?

[200lz]

Opened #55528 with the focused ordinary-row alignment change. Happy to coordinate the shared allocator helper if useful.