Repository navigation
[Review only][NIXL] Simplify draft head slicing from #58471 - #60368
Draft
NickLucche wants to merge 3 commits into
Draft
NickLucche wants to merge 3 commits into
NickLucche wants to merge 3 commits into
Conversation
Adapt the core logic from vllm-project#58471 for review, reusing layer specs and a shared head-offset calculation. Keep draft-region classification at registration and preserve equal-TP push transfers. Co-authored-by: Stu Cao <stu@primeintellect.ai> Co-authored-by: Codex <noreply@openai.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Run shared handshake validation before applying the push-only TP restriction. Describe the shared TP constraint using local and remote ranks. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Keep the draft mapping and descriptor slicing from vllm-project#58471 while retaining the shared offset helper, registration-time draft classification, existing layer specs, and push-worker validation. Co-authored-by: Stu Cao <stu@primeintellect.ai> Co-authored-by: Codex <noreply@openai.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
3 of 4 tasks
stu-cao
added a commit
to stu-cao/vllm
that referenced
this pull request
Oct 14, 2026
… draft head slicing Take NickLucche's review branch (vllm-project#60368) as the implementation: draft regions are classified once at the end of register_kv_caches from the recorded layer specs (`_head_sharded_draft_regions`) instead of a per-call check backed by a parallel `_region_num_kv_heads` list; the head-offset arithmetic moves into `compute_head_offset()` so the draft mapping and the target mapping share it; and the push-mode rule lives in NixlPushConnectorWorker, rejecting only heterogeneous TP. The source files are byte-identical to vllm-project#60368 cherry-picked onto main. Tests follow the new surface: the region-identity test drives `_get_head_sharded_draft_regions()` through `_layer_specs`, the push rejection is tested against the push worker's override, and the worker test doubles seed `_head_sharded_draft_kv_heads` / `_head_sharded_draft_regions` instead of `_region_num_kv_heads`. Co-authored-by: NickLucche <nicolo.lucchesi@mistral.ai> Signed-off-by: Stu Cao <stu@primeintellect.ai>
stu-cao
added a commit
to stu-cao/vllm
that referenced
this pull request
Oct 15, 2026
Squash of the three commits from vllm-project#60368 ("[Review only][NIXL] Simplify draft head slicing from vllm-project#58471"), applied on top of this PR's history: - [KV Connector][NIXL] Simplify draft head slicing under MLA - [NIXL] Keep push draft validation in the push worker - [NIXL] Reuse draft TP mappings with shared head-offset logic Draft-region identity is computed once at register_kv_caches into _head_sharded_draft_regions (via _layer_specs); the push-mode guard lives in NixlPushConnectorWorker and rejects only heterogeneous TP; the rank-offset arithmetic moves into tp_mapping.compute_head_offset and is shared by the target and draft mappings. Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Stu Cao <stu@primeintellect.ai> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Signed-off-by: Stu Cao <stu@primeintellect.ai>
stu-cao
added a commit
to stu-cao/vllm
that referenced
this pull request
Oct 15, 2026
…ucture Follow the restructuring from vllm-project#60368: the handshake fixture seeds _head_sharded_draft_regions instead of the removed per-region head-count list, the layer-identity test drives _get_head_sharded_draft_regions through _layer_specs, the push-mode rejection test moves to the push worker (equal TP passes, heterogeneous TP raises), and the geometry doubles seed the new attribute register_kv_caches now reads. Signed-off-by: Stu Cao <stu@primeintellect.ai>
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
This branch has not been deployed
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
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.
No description provided.