[KV Connector] Support NIXL P/D for hybrid MLA+SSM models - #49762
Merged
Conversation
JaredforReal
approved these changes
Jul 27, 2026
JaredforReal
left a comment
JaredforReal
left a comment
Contributor
There was a problem hiding this comment.
look good to me, cc @NickLucche @ZhanqiuHu
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
njhill
marked this pull request as ready for review
July 28, 2026 16:23
njhill
requested review from
ApostaC,
NickLucche,
ivanium,
orozery and
xuechendi
as code owners
July 28, 2026 16:23
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Re-implements the intent of vllm-project#44848 on top of the reworked NIXL connector, from the current architecture rather than the original patch. KimiLinear pools its KDA (GDN-typed MambaSpec) and MLA layers into shared HMA tensors, making every region dual-purpose. Since the mamba-page unification raises the attention block size until the MLA page equals the unified page, and FlashMLA fixes the kernel block at 64 tokens, both sides always exchange kernel-granularity geometry that is TP-independent for MLA: block_size_ratio stays 1 and TP variance is absorbed by the existing per-engine physical_blocks_per_logical handling. What remains are three real gaps: - register_kv_caches marked a shared region's FA view by whichever layer registered it first; a KDA-first ordering left dual-purpose regions flagged as head-sharded. Merge the MLA flag on the HMA dedup path so split handles replicate their FA descriptors. - The push worker's replicated-MLA WRITE branch asserted a single source rank, which hybrid MLA+SSM violates (sharded SSM targets every covered D rank). Replicate only the attention groups across the write ranks and route hybrid writes through the per-rank split handles, mirroring the pull side. - Handshake validation skipped all FA block_len checks for mamba models; hybrid MLA can and now does enforce the TP-independent invariant that kernel-granularity block_lens match exactly. The KDA conv state decomposes as GDN (equal Q/K/V sub-projections), so the existing 3-read conv transfer applies unchanged. Heterogeneous physical_blocks_per_logical still requires disabling prefix caching, as already enforced at handshake time. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Jared Wen <w13431838023@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
…el pages The MLA per-token page is TP-independent, so kernel block lengths that differ by anything other than the block-size ratio must fail the handshake loudly rather than transfer at mismatched geometry. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QVJLHN2a1DP6iBJjayotrs Signed-off-by: Nick Hill <nickhill123@gmail.com>
| # KDA+MLA in KimiLinear); the region's FA view is MLA whichever | ||
| # layer registered it first. | ||
| idx = seen_base_addresses.index(base_addr) | ||
| self._region_is_mla[idx] |= is_mla_region |
Contributor
There was a problem hiding this comment.
nit:
self._region_is_mla[idx] = self._region_is_mla[idx] or is_mla_region for readibility
ZJY0516
approved these changes
Jul 29, 2026
This was referenced Aug 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Re-implements the intent of #44848 on top of the reworked NIXL connector, from the current architecture rather than the original patch. KimiLinear pools its KDA (GDN-typed MambaSpec) and MLA layers into shared HMA tensors, making every region dual-purpose. Since the mamba-page unification raises the attention block size until the MLA page equals the unified page, and FlashMLA fixes the kernel block at 64 tokens, both sides always exchange kernel-granularity geometry that is TP-independent for MLA:
block_size_ratiostays 1 and TP variance is absorbed by the existing per-enginephysical_blocks_per_logicalhandling. What remains are three real gaps:register_kv_cachesmarked a shared region's FA view by whichever layer registered it first; a KDA-first ordering left dual-purpose regions flagged as head-sharded. Merge the MLA flag on the HMA dedup path so split handles replicate their FA descriptors.The KDA conv state decomposes as GDN (equal Q/K/V sub-projections), so the existing 3-read conv transfer applies unchanged. Heterogeneous
physical_blocks_per_logicalstill requires disabling prefix caching, as already enforced at handshake time.Claude was used for this.