Repository navigation
Conversation
… that share its KV heads A paged pool with `context_parallel: true` stripes each page's token slots over the ranks that would otherwise hold the same KV heads: logical slot s lives on partner s % P at physical slot s // P, and page ids stay logical, so allocation, prefix keys, offload and admission are unchanged. One key, two modes: over the TP ranks of a KV head when tp > kv_heads (vLLM's DCP; MLA's one latent head included), or over a node group's SP ranks, where the node runs replicated across SP with its heads sharded over TP alone. A spec says whether it places positions (`NodeResourceSpec.places_positions`), and one predicate decides a node's SP role for the engine, the engine manager and validation. Attention over a placed pool (FlashInfer GQA and MLA) attends each rank's slots, returns the partial with its log-sum-exp, exchanges it (an all-gather in SP mode; in TP mode the queries and partials over the TP group) and merges exactly. A step that isn't one token a row attends its own tokens apart, from the K/V the layer passes; in a captured bucket the own-token kernel never sees the padding rows' keys (FA3 reads past its plan). One shell on AttentionManager builds and runs the wrappers; a backend supplies its context and own-token wrappers. Every rank of the lockstep instance issues the same collectives with the same shapes, captured or eager; MLA's FA planner cap is a symmetric admission refusal the worker retries smaller, sized by the step's real rows. LTX-2.5's SP attention over sharded keys uses the same gather-and-merge. The engine narrows `EngineResourceInfo.head_shards` to TP alone for a node under SP placement; a resource over a placed pool must declare `supports_context_parallel`. Every write-addressing path, the host-staged captured plans included, maps through the placement, and a dropped pre-plan forgets the own-token wrapper's held plan. Shared scheduling, for every model: a too-big refusal from a resource that evicts nothing (MLA's work-item cap, or a KV pool without host offload) now keys its batch cap on no resource, so any freeing teardown lifts it; that is safe for other models because only KV pools reclaim, so with one paged pool on a worker it lifts at the same teardowns as before, and with more it costs at most one extra refusal. The `glm` extra pins FlashInfer to ==0.7.2 (most extras already do), because MLA over a context-parallel pool mirrors that release's FA planner and refuses any other version at load. Along the way: `KVConfig.shard` refuses a head split it used to round to one KV head a rank; capture configs are sized by `head_shards`; MLA `copy_slots` indexes the token axis; MLA `context_only` plans come from the plan's indptrs (`PagedIndptrs.context_only`). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…acement Conflicts: - engine.py: main's Engine.evictable(node, resources) replaces this branch's identical-purpose evictable(node, affected_resources); the CP caller passes the set positionally. - distributed/linear.py: keep both imports (kv_head_replicas, LoRA). - engine_manager.py: main's serving_resources replaces apply_yaml_overrides; keep the positions_over_sp import. - test_admit_failure_handling.py: this branch's KV-only evictable stub, with main's keyword name. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The pool refused any node on SP ranks, as each SP rank holds a shard of the tokens. A node whose pools place positions over its SP ranks runs replicated over them instead, heads split over TP alone, so every rank sees every row and the pool's TP shard and joint-group agreement already fit. Refuse only a node whose heads split over SP (head_shards beyond TP). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
LoraTarget.slice_rows kept its own copy of the rule for which TP ranks hold a replicated KV head; use the helper the K/V projections and context parallelism share, so the three cannot drift apart. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ding KV heads that don't split evenly over the TP ranks are a config error, and LoRA's slicing now reaches the helper too, so raise ValueError rather than a bare assert (which python -O drops). Say in the docs and the pool's refusal that only a node running its own sequence parallelism can't serve LoRA, and that placing its KV positions over the SP ranks is the way to keep SP. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A resource is built once, sized by its first node's head-shard count. A node placing its KV positions over its SP ranks (heads over TP) and one running its own SP (heads over TP x SP) can't share one, so refuse that at load. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…fferent ways Mirror the engine's load-time refusal before spawn, from the positions_over_sp answers _sequence_parallel already computes. The LoRA SP refusal goes back to its two working remedies: a node in sp_enabled_nodes can't take context_parallel. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- _sequence_parallel refuses a LoRA spec over a node that runs its own sequence parallelism, as LoraAdapterPool.build does on every worker. - When the LoRA specs fail to build, the placement checks still run over the other specs, so their problems are reported beside the build failure. - kv_head_replicas names "heads", since LoRA slicing reaches it too. - The CP docs give the reason chunked prefill stays off on SP ranks. - The gloo LoRA agreement test also runs over the SP ranks of a node that places its KV positions over them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… specs - Report a node group whose nodes split heads different ways once, naming every resource they share, rather than once per resource. - When building the LoRA specs fails, keep the ones with targets, as the build would have, so their placement checks still run. - Docstring and the chunked-prefill docs line say what they cover. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
dghosef
changed the title
… into sp/1-strategies Conflicts resolved to this branch's design: no sp_enabled_nodes or positions_over_sp. _sequence_parallel keeps sequence_parallel_refusals and gains #1124's refusal of a resource over nodes that split heads different ways, from whether each node places its positions (heads over TP) or splits its steps (heads over TP x SP). LoRA over a node splitting its steps is refused by LoraSpec.split_step_refusals, which the engine and validate_deployment both run, in place of the pool's build-time check. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…-pcp Conflicts: docs/adding_models.rst (SP mode: this branch's split-step text plus the parent's "chunked prefill over a placed pool is not supported yet"), docs/models.rst (both Qwen3-Omni notes), and test_qwen3omni_vision_mixed.py (both sides' tests). Under this branch a node storing KV by position may split its steps, so the head-split refusal (engine and validate_deployment) now names the other side "keeps KV by head", and the LoRA docs and comment say LoRA is refused wherever the steps split (SeqShardSpec.splits), placed pool or not. A validate_deployment case pins it: placed + all_gather/ring refused, placed + sp_exchange unset served. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
NSagan271
left a comment
NSagan271
left a comment
There was a problem hiding this comment.
One blocking comment on a merge-related issue, and I'm also requesting a relatively minor abstraction-level change.
Live tests are running.
| supports_context_parallel: bool = False | ||
|
|
||
| @property | ||
| def placement(self) -> "ContextParallel | None": |
There was a problem hiding this comment.
ContextParallel mixes two concerns, and this hook exposes both of them on every Resource. The placement itself (size, index, group, over_sp, partners, and the slot math in slot_range / write_slots) is what the KV pool needs, and I think whether tokens are divided across ranks is a reasonable thing for a resource to expose, so the hook can stay. But attend, exchange_queries, exchange_partials, local_indptrs and local_geometry are attention execution and plan shaping, and their only callers are ContextParallelWrapper and the backends' _bind_placement. Nothing outside attention calls them later in the stack either.
Could those methods move to the attention side, into something that takes a ContextParallel (as ContextParallelWrapper already does), so this hook returns plain placement? It's cheapest to do here, since #1125 onward adds more to these methods (merged_share, for example).
| if self.placement is not None: | ||
| pages, offsets = self.placement.write_slots(pages, offsets) |
There was a problem hiding this comment.
Blocking: after rebasing onto main, resident planning (#1127) bypasses this conversion, so decode KV writes go to the wrong slots.
On a resident step _setup_plan_states doesn't run. The replay fills the addressing from _gather_write_slots_kernel (kv/resident.py), which takes n // PAGE and n % PAGE at the logical page size. Nothing keeps a placed pool off that path: resident_planning defaults to auto, Qwen3.5 names decode in resident_walks, and neither KVManager.resident_refusal nor FlashInfer's looks at placement. So Qwen3.5 with context_parallel, in TP or SP mode, would have every rank write every token, at offsets up to page_size - 1 in a cache with page_size // P slots a page.
A placement refusal in KVManager.resident_refusal looks like the smallest fix.
| Chunked prefill stays off on SP ranks, as for the model's own SP, so SP mode prefills a prompt | ||
| in one step: chunked prefill over a placed pool is not supported yet. |
There was a problem hiding this comment.
Nit: this reads as if chunked prefill is off for any placed pool, and the description says the same, but only SP ranks turn it off (_unchunkable_nodes). TP mode still chunks prompts over 2048 tokens. Can you scope this sentence to SP mode?
| # a captured prefill's bucket holds 8 tokens past the step | ||
| n = case.q.shape[0] | ||
| pad = 8 if captured and rows == PREFILL_ROWS else 0 | ||
| bucket = BucketKey("w", bs=len(rows), num_tokens=n + pad) if captured else None |
There was a problem hiding this comment.
This bucket never has a padding row, so the description's "NaN in the padding rows" doesn't cover them (to be tested live)
With bs=len(rows) there are only 8 padding tokens on the last real row. That misses a span-0 row, and a padded decode bucket where one rank attends SINK_PAGE and masks it. For GQA the served rows probably reach those. MLA has no served run, so it has no GPU evidence for them. Could this also parametrize bs > len(rows)?
There was a problem hiding this comment.
I ran this: with padding rows in the bucket, NaN in their keys turns every real row NaN on captured GQA prefill.
That was this test's capture and replay on 2×H200 with bs > len(rows), with padding rows on the sink as KVManager builds them. For GQA prefill, in TP and SP mode, at three padding shapes:
- NaN in the padding rows' K/V: all real rows NaN
- NaN only in their queries, or finite padding: real rows match the reference
With NaN padding at one shape (2 rows of 2 tokens), MLA prefill and decode, and GQA decode, keep real rows correct. Ranks issued the same collectives in every case, so nothing hangs.
The planned mask in FlashInferRaggedPrefillWrapper.run only zeroes keys past the plan's end, and padding rows' keys are inside the plan. I haven't traced how the NaN crosses rows in the kernel. I also haven't shown a served model's padding keys can go NaN, since they start from dummy zeros, so this may only be hardening. But the docstring on test_a_captured_buckets_padding_merges_to_0_and_reaches_no_real_row does say they can be NaN. Could the own-token path mask padding rows' keys too, with this test parametrized on bs > len(rows)?
Live test: Context / Sequence Parallel 1/5: split the KV cache by position across GPUs (DCP)Works on every served path tested: 2B TP4 and 27B TP2×SP2 with context parallelism give correct answers, KV per rank halves, and offload and disconnects hold up. Non-CP models show no regression. NaN in padding keys does reach real rows on captured GQA prefill, and resident planning and served MLA have no run. 4×H200 (2×H200 for GPU tests and LTX), TP4 and TP2×SP2 · PR head Claims
Risks
Performance
c=32 comes from 320-request runs; c=8 and c=128 from 480 and 7,680. Client commandsc=32 on 0.8B used Other
|
Context / Sequence Parallel stack: 1/5 #1124 → 2/5 #1125 → 3/5 #1138 → 4/5 #1141 → 5/5 #1147
This adds context parallelism for paged KV pools. Today the ranks that hold the same KV heads keep identical copies of the cache: the TP ranks past a model's KV-head count, and the replicas of a node that can't shard further. With
context_parallel: trueon a paged pool, those ranks stripe each page's slots between them instead, so each rank stores1/Pof every request's KV and a deployment can holdPtimes the context at the same memory. It works in two modes from one key: over the TP ranks that share a KV head whentp > kv_heads(vLLM's DCP; MLA's one latent head included), or over a node group's SP ranks, where the node runs replicated across SP with its heads sharded over TP alone. The model declares nothing. Related to #356, which proposes a token-axis design; this is the placement half of that space.Logical slot
slives on partners % Pat physical slots // P, and page ids stay logical, so allocation, prefix keys, offload and admission are unchanged. A token is written by its owner and to the sink page everywhere else. Attention (FlashInfer GQA and MLA) runs on each rank over its own slots, returns the partial with its log-sum-exp, exchanges it (an all-gather in SP mode; queries and partials over the TP group in TP mode) and merges exactly. A step that isn't one token a row attends its own tokens apart, from the K/V the layer passesrun. One shell onAttentionManagerbuilds and runs these wrappers, and a backend supplies its context and own-token wrappers. Every rank of a lockstep instance issues the same collectives with the same shapes, eager or captured. MLA's FlashInfer FA planner has a work-item cap that TP mode reachesPtimes sooner and that counts each rank's own KV lengths, so admission refuses a step past it from its query lengths alone, the same on every rank, and the worker retries it with the real requests that fit. The engine narrowsEngineResourceInfo.head_shards(TP×SP) to TP alone for a node under SP placement, a spec says whether it places positions (NodeResourceSpec.places_positions), and any resource over a placed pool must declaresupports_context_parallelor the engine refuses to load. LTX-2.5's SP attention over sharded keys now uses the same gather-and-merge (gather_merge).Changes outside CP itself that a reviewer should know about:
glmextra pinsflashinfer-python==0.7.2(most extras already do): MLA over a placed pool mirrors that release's FA planner rules and refuses other versions at load. GLM without CP loads on any version.KVConfig.shardraises aValueErroron a second, different shard count, and on a head split it used to round down to one KV head a rank (a TP×SP count that neither divides nor is a multiple of the KV heads). Capture configs are sized byhead_shardstoo, which only differs from before under SP placement.flashinfer_backend: autoruns asfa2from SM100 (the rule main already applies to window pools), anddecode_fmhaisn't used, because the merge needs each partial's log-sum-exp.KVCache.read_tokens,KVManager.read_kv,write_kv(return_tensor=),context_only_views;attention_over_sharded_keysbecamecontext_parallel.gather_merge. MLAcopy_slotsnow slices slots rather than latent columns.validate_deploymentlets a node outsidesp_enabled_nodestakesp_size > 1when its paged pools setcontext_parallel.validate_deploymentnow catches that before spawn. The engine andvalidate_deploymentalso refuse one resource over nodes that split heads different ways.On this head (ae8ff84, with main eceea1b merged in):
--core --compare main, with the Rust extension), ruffGPU, run before the rebase (at f1f6bc9). A rerun on this head is pending: main has since changed paths these runs take (#367 TP prefill chunking at 2048 tokens, #852 direct planners, #1076 plan skipping, #988 prefill buckets), and the rerun also covers the LTX fp32-merge check. I'll update this table when it lands.
test_two_gpus_attend_a_placed_pool_like_one(tp/sp × prefill/decode × eager/captured × GQA/MLA; captured on a full-bucket plan, NaN in the padding rows)owns_padding)test_kv_context_parallel,test_gpu_mla_attention,test_attn_decode_tensor_cores,test_flashinfer_mla_gates,test_context_only_attention,test_sequence_parallel,test_glm52_capture_gpucuGetProcAddress, NCCL's no-kernel-image probes); initcheck: 0 errors. initcheck on captured SP-mode MLA reported only FlashInfer's planner workspace copy, then hung under the sanitizer and timed outflashinfer_backend: fa2, needles at 1k/32kmax_num_pagesHardware and builds: Eras H200 (sm90, driver 580.178), torch 2.12.1+cu132 (CUDA 13.2), FlashInfer 0.7.2,
flashinfer_backend: auto(FA3) unless noted; the served rows ran with async scheduling on (main's default since #1075) unless marked; the earlier serial runs gave the same verdicts. Not run: TP8 arms and GLM-5.2 (MLA over a placed pool in serving), because no 8-GPU node was free; MLA is covered by the 2-GPU kernel tests above.🤖 Generated with Claude Code