Repository navigation
Conversation
Kimi-K3 under pipeline parallelism (verified as pp2 x tp4), including DSpark speculative decoding and prefix caching of the KDA state. PP still needs --enforce-eager: decode has no CUDA graphs yet. Model and transport - Carry K3's block_residual across the stage boundary with the hidden states, and load the DSpark draft's own embedding on the last stage, where the target's embedding is a PPMissingLayer. - Relay DSpark's aux hidden states across stages (PPAuxRelay). The drafter lives on the last stage, so a hook on a tap owned by an earlier stage lands on a PPMissingLayer and never fires. Each stage captures its own taps and forwards them, concatenated in layer_ids order, in the intermediate-tensor payload. The embedding tap (-1) is resolved as the drafter resolves it: the embedding module, identity extract. Speculative decoding under PP - Let a PP engine verify its own drafts. PP output is undeferred, so: the sampler reports this step's rejection and bonus counts on the non-deferred path; they are published to every stage's attention backends, so non-sampling stages roll back linear-attention state; postprocess keeps ONE run of exactly mtp_k + 1 placeholders after every step, whatever was rejected. - Rebuild the verify window from that run. All four backends read `verify_context_shift()`: `num_rejected` under deferred output, 1 under undeferred. The arithmetic lives in token_layout.decode so it is testable without aiter. The window's positions and KV slots are anchored where the scheduler staged them. - Gate and size the verify window on the run's speculative config, not on whether this stage drafts: only the last stage has a drafter, but every stage verifies. Per-stage buffers are sized the same way. - A P/D producer still never speculates, under PP or not. - Round MLA's query-head padding to a width aiter dispatches (24 -> 32 at TP4; TP8 already pads 12 -> 16). KDA state and prefix caching under PP - Progress advances at schedule time, so a chunk ending on a checkpoint rung would let the next chunk overwrite the state slot before the checkpoint is filed. Hold the next chunk (both prefill scans skip it) until the batch that carries the rung files its hashes; the hold is released with that batch even if the seq was preempted, aborted or retired meanwhile. - The checkpoint is filed after the prompt, so its state fork lands on the first decode. Honour it there: publish the source slot and move conv and recurrent state onto the write slot before the kernels run. With DSpark that decode is a verify step, so the spec path does the same. Only a backend whose model does that copy may take a fork on the spec path (`GDNStateMixin.spec_honours_state_fork`, set by K3); every other backend keeps the assert. K3's mono plan refuses a forked step, since a mono step replaces the branch that does the copy. CI - Add a Kimi-K3 pp2 x tp4 DSpark accuracy case with prefix caching and ptpc_fp8 online quant on, as the TP8 Kimi-K3 case runs it, the only accuracy job that reaches any of the paths above. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- tokenIDProcessor.verify_ctx_lens(context_lens) replaces token_layout.decode.verify_context_shift: the attention builders now take their verify-window context lengths straight from the runner. - Rename the runner's `spec_enabled` -> `is_spec_run` (speculative decoding configured, every PP stage runs verify windows) and `use_spec` -> `has_drafter` (this stage owns the drafter and rejection sampler; under PP the last stage only). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…unner tests The non-GPU runner has no aiter, so importing model_runner aborted collection. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…heck it Every stage now orders its columns by sorted(layer_ids), and a stage's incoming taps are those below its first layer, so the layout no longer depends on the order the drafter lists its taps in. pack/absorb reject a payload whose width does not match the taps they expect. Round-trip test over pp 2/4/8 with unsorted taps and stages that own none. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A filing that raised mid-batch (logged and swallowed by the PP head loop) left the holds of every later seq in the batch set, keeping them out of the prefill scan forever. Drop the whole batch's holds up front in both register_prefill_hashes and postprocess. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The non-first PP stage overrode skip_weight_prefixes on the instance while the class declared it a ClassVar list. It already rebound rather than mutated, but the annotation said the opposite and a list default invites in-place edits. Declare an immutable tuple default instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The PP decode path staged each row's drafts with a per-row Python loop. Build the destination indices with a mask over the draft columns and write them in one assign, as the scheduler's draft marshal does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Both counts are rows of the rejection sampler's int32 verdict, already contiguous as spec_counts, so copy that once instead of syncing per row. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
pack() concatenated the received block and each own capture into a new tensor every step. Preallocate the payload instead: own taps' capture buffers are column views of it and received rows are copied into its leading columns, so packing allocates nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The preallocated send payload was max_num_batched_tokens rows wide and resident on every non-last stage (~470 MB per rank at pp2). Allocate it per forward instead: the first own tap's hook allocates the rows it captures, every own tap writes its columns, and pack() copies the received taps into the leading ones and hands the buffer to the send path, which clones it. A stage that owns no tap forwards the received rows as-is. PP stages run the target eagerly, so no recorded graph bakes the per-step address. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This PR adds the spec path's read-side state index, so the token count is the only reason a forking class skips checkpoints on spec decode. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Under dynamic chunking a middle chunk is acked, and that branch only queued its prefill hashes. A seq held for the chunk's checkpoint is out of the prefill scan, so nothing flushed the queue and it stalled. Both no-output branches now go through one helper that queues and flushes when the batch holds a seq. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Conflict in atom/models/kimi_k3/mono/plan.py (#2871 moved K3 mono to plain decode and ReplaySSM): take main's ReplaySSM support and keep this branch's state-fork refusal, now ahead of the spec/plain split so a forked plain decode is refused too -- this branch can put a fork on a decode, which #2871's state_slots docstring assumed never happens. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Only the stage's first own tap allocated the send buffer, so a forward where it captured nothing left the later taps writing into the previous forward's buffer (or into None). Whichever tap fires first now allocates for the forward's rows, and pack() asserts every own tap was captured. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ModelRunner reached into the drafter's private _aux_capture_spec and dspark_proposer's build_aux_capture_spec to assemble the relay. Move that into PPAuxRelay.build so the runner makes one call and keeps only the pp_aux_relay handle it uses at absorb/pack. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… lens verify_ctx_lens returns a new array, so the is_first_decode_without_local_prefill branch read batch.context_lens unshifted: a row with num_rejected > 0 in the same batch got r extra pages in kv_indptr while block_tables and positions used the shifted lengths. Read the local context_lens there, as the other two branches do. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Under pp>1 a state fork can reach decode and spec verify. Instead of the model copying the source state onto the write slot, the builder now ships `spec_state_indices_in_tensor` (the write table with the fork source in the resume column), `fused_sigmoid_gating_delta_rule_update` takes `ssm_state_indices_in` like fused_recurrent already does, and the spec conv update gets `conv_state_indices_in`. This drops `spec_honours_state_fork`, its assert, and both copies in KimiKDAAttention, and covers Qwen GDN too. ReplaySSM kernels have no read index and fold records onto the write slot's checkpoint, so its wrappers seed that checkpoint from the source (`seed_forked_slots`) and the builder parks the forked write slot's cursor at -1. Before, the cursor was left at the previous tenant's value and its records were folded into the forked state. The mono K1 refusal stays: K1 reads `st_idx` only. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Under pipeline_parallel_size > 1, only Kimi-K3 with DSpark relays the target's aux taps across stages. Every other (model, spec method) pair now fails at config time instead of mid-serving: EAGLE3's relay zip, a Qwen GDN prompt-end fork, MiniMax-M3's embed on a PPMissingLayer. The DSpark confidence scheduler and a P/D kv_consumer, which the PP loop has no remote-KV path for, are refused as well. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…idth `MLAAttention.padded_num_heads` now rounds up to a width aiter dispatches, and the plugin's MLA layer inherits it, so K3 at TP4 pads q from 24 to 32 heads. The plugin builder still planned the work descriptors for 24. Round it the same way, as aiter_mla.py's builder does. The sparse builder and layer keep max(heads, 16): they pad by repeat_interleave, which only divides evenly at that width, and they already agree with each other. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
PP refuses PAGE-backed checkpoints, so every class that applies forks and keeps a positive successor_room: a rung (or the prompt-end anchor) sits strictly inside the prompt, and `start + chunk < num_prompt_tokens` always holds. Filing checkpoints by rung hash at schedule time, which would remove the hold altogether, is left as a follow-up. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`PPAuxRelay.arm` repeated `Drafter.arm_aux_capture`'s hook mechanism and
skipped its checks. Both now go through drafter.py:
- `install_aux_taps` validates the spec's ids (each in {-1} U [0, n),
strictly increasing) and hooks a subset of them. The last stage hooks
the drafter's buffers, other stages the relay's send buffer.
- `_aux_hook` skips draft forwards and hands its writer
`ctx.ubatch_token_offset`. The relay writes at that offset and grows its
buffer for a later ubatch instead of reallocating over the first.
- `Drafter.aux_capture_spec(config)` is each drafter class's config-only
spec. A drafterless stage reads it off `factory.drafter_class(config)`,
so the relay no longer imports DSpark's builder.
`PP_AUX_KEY` is now the model-neutral "pp_aux".
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The relay allocates a fresh buffer per step and hands it to the send path, so cloning it again only doubles the copy. async_send_intermediate_tensors takes the keys the caller hands over and sends those as they are. The received aux block is now a run_model local instead of runner state, and PPAuxRelay is imported at the top of the runner. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Every PP stage runs the verify window, so the runner's is_spec_run / has_drafter pair collapses back to the single use_spec it had, now set on every stage (has_drafter's readers sit on the deferred path, which PP never takes). The scheduler's spec_decode_local was the same flag under a second name; its readers take use_spec directly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…idth The scheduler's is_deferred_out (pp_size == 1) and tokenIDProcessor's (not PP, not a recurrent-state P/D producer) agreed only because a producer never speculates. Both now come from defers_output(config): the runner passes the attention family it already resolved, the scheduler reads it off the config. GDN sizes its verify window from model_runner.num_spec_tokens, which every PP stage carries, instead of the drafter's mtp_k with a num_speculative_tokens fallback; config already refuses a speculative config without num_speculative_tokens, so the fallback's ValueError was unreachable. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… case Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…der Dynamo The stage-0 relay hook runs inside the compiled backbone. `_filled.add` mutated a set Dynamo did not create, so Dynamo replayed it with clear()+update(); with cudagraphs on, an `update` in the rewritten bytecode triggers the decorator's depyf import, which crashed PP2 warmup on images without depyf. Rebind a frozenset instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The undeferred step already brings the sampled ids to the host. Each row keeps its accepted drafts plus one token before the -1 tail, so both counts are its length; drop the verdict's own .cpu().numpy(). The ragged sampler test pins that layout against the verdict. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`torch.tensor(list, device=cuda)` is a pageable H2D, which PyTorch follows with a stream sync. Pin the list and copy it non_blocking, as the other host-staged indices do. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Conflicts: scheduler.py imports (keep both), model_support_guide.md (main's image-token line, then ours). Two of main's new tests needed follow-ups: tokenIDProcessor lets defers_output resolve attn_family itself, so it is read only for a P/D producer as before; test_aiter_flydsl_gdn's metadata stub gains has_state_fork. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The merge dropped it, so defers_output re-derived the family from config.hf_config and test_spec_deferred_status_queue's producer cases failed. The runner stub in test_h2d_runner_publication gains attn_family. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Under schedule-time advancement a middle chunk's live is_partial_prefill flips once its final chunk is scheduled, so the middle chunk's postprocess appended mtp_k + 1 placeholders to a seq still in prefill. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Under PP an earlier chunk of the same seq can still be in flight when the chunk ending on a rung is scheduled; its completion released the hold, so the next chunk could run before the rung's checkpoint was filed. The hold now records the held chunk's end, and only a batch reaching it releases. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A producer never speculates, so the Kimi-K3 + DSpark gate refused GLM-5.2's PP prefill + MTP Atomesh cells for nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The kernel resumes from column num_accepted - 1, so a forked verify with a bonus read a recycled write slot; only column 0 held the source. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
GDNStateMixin.spec_honours_state_fork() defaults to False; only Kimi-K3's KDA forward (kimi_linear) reads spec_state_indices_in_tensor and opts in, so glm5_next, Qwen3-Next/3.5 and Qwen4Exp keep the assert. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
UBatchWrapper.forward takes no intermediate_tensors, so TBO never ran under PP; refuse it in config instead of keeping a grow path nothing reaches. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
seed_forked_slots indexed with a GPU bool mask, a host sync in every linear-attention layer of a fork step. The builder already knows the (write, source) pairs: ship them once per step in the pinned copy that parks write_pos, and index with them directly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The lowest own tap opens the forward's send buffer and every own tap writes its columns; the reset and the _filled check read state that the compiled forward bakes to the warmup trace, so they guarded nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…der PP pp_speculative_methods sits beside mono_decode_installers, so the next model wired to the relay adds a row there instead of editing Config. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
_pp_checkpoint_hold -> _checkpoint_holds (keyed on advance_on_schedule), holds_checkpoint_hostage -> carries_held_checkpoint, and _defer_prefill_hashes -> _queue_or_flush_prefill_hashes, since it flushes the whole queue; mark the hold as scaffolding #2709 replaces. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Read seq.num_placeholder_tokens, which postprocess refills to one full window, instead of re-deriving the run's width from the scheduled count. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Drafter.take_relayed_aux keeps the received [N, k*H] block and aux_for serves it ahead of the own-tap buffers, so the last stage no longer splits it per tap for compute_draft_kv to concatenate back (DSpark skips the cat when it is the only block). The relay takes the drafter class's spec on every stage instead of the drafter's private _aux_capture_spec, and the runner no longer hands it _aux_buffers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
_aux_capture_spec now parses dspark_target_layer_ids once, through aux_capture_spec, and drops its unused target_model parameter. PPAuxRelay keeps only spec (no layer_ids / hidden_size copies), drops the dead _armed guard, and raises the missing-payload error from one place. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Taps owned by an earlier PP stage arrive through the relay, so the last stage gives them a 0-row buffer and no hook (~470 MB/rank for K3 at pp2). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Both now ask whether the tap's module is a PPMissingLayer, the embedding included, instead of the relay keying the -1 tap on is_first_rank. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…P stage Kimi-K3.1-DSpark ships no embed_tokens, so under PP the last stage's draft loaded nothing into the embedding it built for itself. The target now keeps its embedding on the last stage when a drafter runs there, and the draft shares it on every stage, as it already did without PP. The `-1` aux tap stays first-stage only: the last stage's copy never runs the target forward. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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.
Motivation
Follow-up to #2365. Main's
ARCH_ALIASESmapsKimiK31ForConditionalGenerationonto Kimi-K3, so the PP DSpark path in #2365 now accepts Kimi-K3.1 too. Under PP, #2365 has the draft on the last stage load the embedding from its own checkpoint.Inferact/Kimi-K3.1-DSparkhas noembed_tokens, and the loader does not fail on a missing weight, so the draft would run with an embedding that was never loaded.Technical Details
KimiLinearModelkeeps a realembed_tokenson the last stage whenspeculative_configis set, so the target checkpoint's table loads there.KimiK3DSparkalways shares the target's embedding, the same as without PP. The file goes back tomain's version.tap_is_resident(-1)is now true only on the first stage. The last stage's copy of the embedding never runs in the target forward, so the-1tap must still come through the relay.Test Plan
tests/test_pp_aux_relay.py: the last-stage test now holds a real embedding copy and checks that the-1tap still gets no buffer and no hook.test_pp.py,test_dspark.py,test_kimi_k31_config.py,models/test_kimi_k3_mono_plan.py,test_gdn_state_fork.py.Test Result
All of the files above pass in the
atom_1031container (24/99/82/3/67/7 passed). No GPU run on this branch: Kimi-K3 pp2 × tp4 needs all 8 GPUs.🤖 Generated with Claude Code