[Bugfix][KV Connector] Mooncake: transfer GLM-5.3-Flash's kpool indexer pages and tail like NIXL

[lucifer1004]

Overview

GLM-5.3-Flash P/D over Mooncake fails: its kpool indexer stores several kernel blocks in one tensor row and keeps its kpool tail inside the same allocation, which #53906 taught NIXL but not Mooncake. This PR gives Mooncake the same handling, sharing NIXL's helpers.

Claims

• GLM-5.3-Flash prefill/decode over Mooncake works: chat GSM8K 0.970 (0.974 with MTP), against 0.972 for a single server, with no failed transfers. Before this PR most transfers failed and requests errored.
• Other models keep their registration: the change only applies to a compressed cache whose rows hold several kernel blocks, the same predicate NIXL uses.

Validation

GLM-5.3-Flash-NVFP4, TP4 prefill -> TP4 decode on one host (4+4 RTX PRO 6000 Blackwell, two NUMA halves), Mooncake over RDMA, --kv-cache-dtype fp8 --max-model-len 65536 --max-num-seqs 32 --max-num-batched-tokens 8192, one run each:

	Before	After
Mooncake transfers (prefill, per 5 s window)	~4 succeeded, ~300 failed	~200 succeeded, 0 failed
Descriptors per transfer	837	65
Chat GSM8K (1,319, flexible-extract, 32 in flight)	requests failed	0.970
Same, with {"method":"mtp","num_speculative_tokens":1}	—	0.974 (94-95% draft acceptance)

Before the fix, the prefill logged Memory region not registered by any active device(s) and Failed to get segment descriptor ... 0x...6c00--0x...7000 (1 KiB, the kpool tail's half page).

DeepSeek-V4.1-Flash, same topology, as a regression check of another compressed-cache model: its caches do not match the predicate (the same 18 packed layer regions are registered), chat GSM8K 0.944 against 0.947 without this PR, no failed transfers.

Tests:

pytest tests/v1/kv_connector/unit/test_mooncake_connector.py \
  tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py \
  tests/v1/kv_connector/unit/test_nixl_desc_geometry.py   # 345 passed
pytest tests/v1/kv_connector/unit/test_nixl_connector_hma.py  # 96 passed

Validation environment

The serving runs used vLLM 5f30fc7031 with this branch merged, alongside unrelated changes for other models (among them #53078, which reworks Mooncake's Mamba regions in the same function, and #58305). The unit tests above ran on this branch alone.

Details

Root cause. The kpool indexer is an MLAAttentionSpec with tokens_per_state > 1 whose tensor row holds several kernel blocks, and the KpoolTailSpec cache is a view at the start of the same allocation. Mooncake registered the indexer per tensor row and the tail as a region of its own, then _logical_to_kernel_block_ids expanded both groups' block ids to kernel blocks (x36 here: 2304-token blocks over 64-token kernel blocks). Kernel block j was addressed as row j instead of page j, so the indexer and tail descriptors landed in the wrong bytes or past the registered memory.

Change.

• register_kv_caches registers a cache matching NIXL's uses_dense_virtual_transfer_pages page by page (stride and length = the kernel page), so kernel block j is page j.
• A KpoolTailSpec view inside such an allocation is not registered; its group is added to the indexer region's shared groups (_SHARED_REGION_GROUP_ID), so the tail group's blocks are read through the pages that hold them, as NIXL reads every attention group through each region. A tail that is not fully covered raises instead of transferring partially.
• tensor_byte_span_end and uses_dense_virtual_transfer_pages move from NIXL's base_worker.py to kv_connector/utils.py unchanged, so both connectors use one predicate.

Related. #57169 (and #55219 before it) and #59297 move GLM-5.3-Flash to the generic packed layout with a CircularBufferSpec tail and drop this layout from both connectors. This PR fixes Mooncake for the layout on main until one of them lands; if one does, this handling goes with it. No other open PR changes Mooncake's handling of this layout (searched Mooncake kpool, Mooncake KpoolTail, Mooncake GLM-5.3-Flash, Mooncake compressed indexer, Mooncake tokens_per_state, mooncake register_kv_caches).

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR. (Claude Code wrote the change and this description; I reviewed every changed line and ran the tests and evaluations above.)

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[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.

[snadampal]

I thought the current PR was removing this constraint, why is it still added here?

[lucifer1004]

The PR doesn't remove a constraint. _physical_blocks_per_logical_kv_block already exists on main: it is the number of kernel blocks per logical block. The new code divides by it to get the kernel page size (_physical_page_size). The assert only pinned this test's fixture to a ratio of 1. You're right that this left the real case uncovered: GLM-5.3-Flash runs 2304-token blocks over 64-token kernel blocks (a ratio of 36), with one tensor row per logical block. I've parametrized the test over the ratio (1 and 4). The ratio-4 case lays out one row per logical block, split into kernel pages. The test now also checks that every kernel block of a logical block addresses its own page in the indexer tensor. Both cases fail with the dense-page registration disabled and pass with it.

[snadampal]

since this is a layout change, it requires connector version upgrade to match on both P and D.

[snadampal]

But looks like mooncake connector doesn't maintain version and compat hash similar to nixl connector. It's required to add version support either in this PR or a different one.

[lucifer1004]

Agreed that Mooncake has nothing like NIXL's version and compatibility hash, and it should. This PR leaves the wire format alone: MooncakeXferMetadata and the handshake are unchanged. What changes is the region geometry each side registers, and only for caches that match uses_dense_virtual_transfer_pages, i.e. GLM-5.3-Flash's kpool indexer and its tail. Every other model registers the same regions as before. A P/D pair that mixes builds from before and after this PR would therefore only disagree for GLM-5.3-Flash, and that model's Mooncake P/D doesn't work on main today anyway.

I'd prefer to add a Mooncake connector version and a compatibility hash, modeled on NIXL's (NIXL_CONNECTOR_VERSION and compute_nixl_compatibility_hash, checked in the handshake), as a separate PR. It would cover every future layout change, not just this one, and I can open it right after this one. If you'd rather have it here, I can fold it in.

[mergify]

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

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