[Bugfix][KV Connector] Mooncake: heterogeneous-TP support for hybrid GDN/Mamba models

[lucifer1004]

What

Fix MooncakeConnector for heterogeneous-TP disaggregated prefill/decode (e.g. PP2TP2 prefill → TP4 decode) on hybrid GDN/Mamba + full-attention models (Qwen3.5 / Qwen3-Next family).

Without these changes, het-TP transfer on hybrid models either fails region validation or silently corrupts the mamba/GDN state on the consumer side.

Changes

1. Register every mamba state with its real, unpadded byte length (register_kv_caches). The block page packs conv + ssm states contiguously before padding; registering the padded page length lets the TP-ratio slice cut across padding and misalign states. Each state now registers as its own region with its true payload size.

2. Split the GDN conv state into per-sub-projection regions (Q/K/V), reusing derive_mamba_conv_split from the NIXL path. Under the DS (dim, state_len) layout each sub-projection is contiguous and shards across TP independently; treating the conv page as one region lets the het-TP byte-split cut across the Q/K/V boundary and scrambles the state (consumer ranks receive Q-segment bytes in K position, etc.). Because the conv state only influences the first kernel-1 decoded tokens, the corruption shows up as a moderate accuracy drop with occasional empty generations rather than a hard failure. Falls back to a single conv region under the SD layout (homogeneous-TP whole-region copies are unaffected by layout).

3. Whole-region copy when the consumer TP exceeds the KV-head count (consumer_kv_replicated). With 2 KV heads and consumer TP4, each consumer rank replicates its head group's region rather than slicing the producer's region; the generic ratio slicing would otherwise deliver only half the heads to each consumer rank.

4. Tolerate peer-PP-stage layers in _get_transfer_regions: remote layers from another PP stage are absent from the local spec table, so use .get() (their split flag is unused — region alignment only matches locally registered layers).

5. Group-aware per-region length validation in _validate_asymmetric_region_lengths: attention regions follow the replicated-consumer rule while mamba/GDN regions keep the TP-ratio rule.

Tests

• test_register_kv_caches_splits_gdn_conv_sub_projections: DS layout registers Q/K/V/ssm as separate regions with correct offsets and real byte lengths.
• test_gdn_conv_sub_projection_regions_align_across_het_tp: P TP2 → D TP4 per-region transfer plan keeps every slice inside its sub-projection.
• Existing hybrid-mamba cases updated (per-state registration, replicated consumer KV, PP peer layers) — the module's 12 tests pass.

Notes

• Homogeneous-TP paths (e.g. TP4→TP4) are unaffected: per-region copies degenerate to whole-region copies.
• Heterogeneous-TP GDN requires VLLM_SSM_CONV_STATE_LAYOUT=DS, the same requirement as the NIXL connector.

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