feat: pipeline-parallel hidden-state export, a patched vLLM image, and the K3 DSpark MLA draft - #165
Merged
Merged
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 79d7318ad0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…es per PP stage Hidden-state extraction under pipeline parallelism needs every stage to hand back the aux layers it owns, and no released vLLM does that: upstream builds the capture list on the last rank only, and `ExtractHiddenStatesModel` does not implement `SupportsPP`, so `ModelConfig.verify_with_parallel_config` rejects `pipeline_parallel_size > 1` while `SpeculativeConfig` is still being constructed. The following commits are therefore unrunnable without a patched vLLM, and `pyproject.toml`'s plain `vllm>=0.18.0` cannot express that. Pin the substrate instead of naming a branch: base image `vllm/vllm-openai:nightly-7794b1e08bf505ff28664515ffaaeeec955ab796` (`0.26.1rc1.dev353+g7794b1e08`, vLLM main @ 2026-08-05) plus an ordered patch stack, so every claim about PP behaviour is reproducible from this tree alone. `vllm_pp_hidden_states.patch` carries the per-stage capture, the prefill-only extract path and the PP completion barrier, and adds `SupportsPP` plus a `make_empty_intermediate_tensors` factory to `ExtractHiddenStatesModel`, which is what gets past the config-time rejection above. It is an upstream candidate: a shrinking diff after a rebase onto a newer nightly is the signal that upstream has absorbed part of it, which is why it stays separate from `vllm_k3.patch` (the Kimi-K3 model and its env flag, ours indefinitely). The patch is scoped to the export path: it touches the capture model, the schedulers, the input processor, `ModelRunnerOutput`, the extract proposer and the GPU model runner, and no draft-model file. It was originally cut together with unrelated draft-model changes belonging to a different feature on a different release path; those have been removed. Nothing in the remaining hunks references them, so the two are genuinely independent rather than merely separable. The Dockerfile reads a `series` file rather than globbing `*.patch` so the stack stays explicit as patches are added or renamed, and so each patch's upstream intent is recorded next to it; the ordering itself is not a correctness requirement. Both patches edit `vllm/models/kimi_k3/nvidia/model.py` — the K3 one inserting `_capture_aux_hidden_stream` around line 1095, the PP one generated against the tree the K3 one produces — but the edited regions are disjoint and `patch` resolves the +79-line shift as an offset, so either order applies with no fuzz and yields the same tree. Verified by applying them in reverse. Patches were forward-ported by committing the original working tree on its base (vLLM PR #50000 head `0498dc7ea9204f18`) and cherry-picking onto the nightly, so git performs a real three-way merge instead of context-fuzz matching. No conflicts, and the resulting diffstat was identical to the original per file even though upstream had drifted on 7 of the 12 files it then covered. Verified by building this Dockerfile from a pristine base: both patches apply with no fuzz and no rejected hunks, in the resulting image `supports_pp(ExtractHiddenStatesModel)` is `True` and `ModelRunnerOutput` carries the `prefill_only` field, and the equivalence check added in the next commit passes bitwise inside that image rather than in a hand-assembled container. Dropping the `fa` extra from the install is part of making the image runnable rather than a cleanup. `flash-attn-4` pins `nvidia-cutlass-dsl==4.6.0.dev0` where this base ships `4.6.0`, and pulls `apache-tvm-ffi` past the version vLLM's flashinfer was built against; every worker then aborts during memory profiling with `tvm::ffi::Error: TypeAttr __ffi_repr__ is already registered for type index 132`, so the image builds and imports cleanly and only fails once it reaches a GPU. Constraining the install to the base versions instead does not work: flash-attn-4's pin is exact, and the image is not pip-reproducible anyway, since its torch requires `nvidia-nccl-cu13==2.29.7` while the image ships `2.30.7`. The extra only enables the optional flash_attention draft backend, whose import is already guarded and which `training.attention_backend` does not select by default; without it the install leaves the flashinfer and CUTLASS packages untouched. `tests/vllm_pp_hidden_states_tests.patch` is deliberately absent from `series`: the Dockerfile's COPY glob is flat and the installed dist-packages tree has no `tests/` directory. It is kept alongside for anyone patching a vLLM source checkout. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
…ssemble them
Under pipeline parallelism no single rank holds the whole residual stream, so the
connector cannot publish one concatenated object per request. Each stage now
writes only the aux layers it owns as separate `<key>_layer<id>` objects, and the
scheduler side advertises a `pp_layer_manifest` naming every fragment and its
role so the trainer can put the sample back together.
Ownership comes from vLLM's own `get_pp_indices`, which returns a half-open
range, while our capture ids use post-layer semantics — capture id `k` is the
residual stream after layer `k-1` ran, so it belongs to the owner of layer `k-1`,
hence `start_layer < layer_id <= end_layer` rather than `<=` on both sides. That
boundary was enumerated exhaustively rather than reasoned about: every
`(pp_rank, position)` claim for layer counts {32, 35, 36, 37, 61} across pp_size
{1, 2, 3, 4, 6, 8}, with default-shaped and boundary aux ids, against the real
`get_pp_indices` including its uneven partitions (61 layers over 8 stages splits
[7,7,8,8,8,8,8,7]). Every position is claimed exactly once, none twice. The
`layer_id == 0 and pp_rank == 0` branch is load-bearing only if an aux id of 0 is
ever configured — without it, id 0 would be claimed by nobody — and is dead today
only because the engine shifts every id to >= 1.
`MooncakeDataset._load_vllm_pp_layers` is the reassembly counterpart, and stays
the single place that knows the fragment layout. It concatenates in manifest
order, cross-checks each fragment's `input_ids` against the first, and removes
fragment keys rather than the base key, which PP never writes.
PP=1 keeps its existing single-object layout. Making the fragment layout
universal is the better end state and is left for follow-up, because it changes
what every existing producer writes and deserves its own review.
Three fail-closed changes come with this, all of the same shape — two code paths
independently derive one fact and disagreement is silent:
Placement matters more than coverage, so store reachability moved from
`save_kv_layer` to the constructor. Measured on Qwen3-0.6B at PP=2, raising from
`save_kv_layer` left the other stage waiting on a rank that had stopped
answering, and the caller saw `TimeoutError: RPC call to sample_tokens timed out`
wrapped in `EngineDeadError` 3m33s later, with the real message only in worker
stderr. From the constructor the same misconfiguration is a normal engine-init
failure in about 30s. The store object itself stays lazy: building it in
`__init__` registers buffers and perturbs vLLM's KV-cache memory profiling.
A writer that cannot reach the store, or whose put fails, now raises instead of
warning. Previously a whole run could complete having published nothing, and the
`raise` for PP>1 sat directly after a log line saying "skipping". Dropping
training data is wrong at every pipeline size, and no consumer can detect the
hole.
`_check_layer_layout` rejects an empty or non-ascending
`eagle_aux_hidden_state_layer_ids`, and a final id that disagrees with vLLM's
`get_total_num_hidden_layers()`. PP=1 concatenates in vLLM's capture order, which
is always ascending, while PP concatenates in configured manifest order, so a
descending list trains on a layout that depends on the pipeline size with no
error at either end. "Which layer is last" is likewise encoded twice — the engine
appends it from the HF config, the manifest labels it from vLLM's count — and
disagreement leaves PP with no `last_hidden_states` fragment while PP=1 silently
mislabels a real layer as the final one.
`OfflineSavingActor` rejects samples carrying a manifest instead of doing a plain
`get` on the base key. That key is never written under PP, so it would have
blocked waiting for it and then "cleaned up" a key that does not exist while
every real fragment leaked.
Reader-side manifest coverage checks were written, tested green, and then
removed. The manifest is built in one place by one list comprehension over
`self._layer_ids`, which `_check_layer_layout` pins at startup, so duplicate ids,
a second `last_hidden_states` entry, a reordered manifest and a short reassembly
are all unreachable. They cost a per-sample pass on the training hot path and
could never fire.
`tests/test_vllm_pp_equivalence.py` is the acceptance test, because the one
property that protects the training signal is that reassembled tensors match what
a single-stage run produces. Both arms go through production code — the connector
publishes, `MooncakeDataset._load_from_mooncake` reads — and each runs in its own
process, since vLLM does not release GPU workers within a process and two `LLM()`
objects fight over memory. On Qwen3-8B at PP=1 versus PP=2 all three prompt
lengths are bitwise identical for both `hidden_states` and `last_hidden_states`.
Bitwise is the pass criterion, which required removing batch composition as a
variable: batching the three prompts together changes GEMM shapes and moves bf16
results by 1-2 ULP (absolute 128 on a ~1e4 massive-activation element). Holding
pp fixed and varying only composition reproduces the identical deviation, so it
is not a pipeline effect — hence one request per `generate()` call in both arms,
with `--tolerant` available but not the default. Comparison metrics are computed
in float64; an fp32 dot product over these tensors, which carry ~1e4 activations
next to ~1e-2 values, loses enough precision to report a cosine above 1.
Known gaps, recorded so they are not rediscovered: `vllm_pp_complete` is a
hardcoded `True`, so the guard in the fetcher cannot fire and the real
fail-closed behaviour is the barrier in the vLLM patch; the manifest is not
self-describing, so the reader reverse-engineers fragment shapes and hardcodes
int64/bf16 defaults instead of using the shapes `put()` already returns; each
fragment stores a redundant `input_ids` copy; and a store failure that first
appears mid-run still propagates through the slow RPC-timeout path above, because
that failure is genuinely dynamic and there is nowhere better to raise.
Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
… per-stage capture gate Three engine-side changes that the preceding fragment layout needs in order to be reachable and safe to reach. `vllm_pp_size` has existed in `VllmEngineConfig` since the config was written but never reached vLLM. `_resolve_parallel_sizes` now derives `tp_size` as `world_size // pp_size` instead of assuming the whole world is one TP group, and `pipeline_parallel_size` joins `_PROTECTION_ENGINE_KEYS` so `extra_args` cannot quietly contradict it. At the default `pp_size=1` this reduces to the previous `nnodes * num_gpus_per_engine`, so nothing changes for existing runs; the two existing tests already covered the arithmetic. The request-side counterpart is forwarding the connector's `pp_layer_manifest` into the sample metadata, which is what lets the trainer find the fragments. Followers on a multi-node engine now run vLLM's headless executor rather than constructing an `LLM`. vLLM multi-node multiprocessing has exactly one scheduler and EngineCore, on `node_rank=0`; building `LLM` on every node creates a second EngineCore per follower which then issues control-plane RPCs with no leader message queue. The headless path has to rebuild `KVTransferConfig` and `CompilationConfig` from their dicts by hand because it bypasses `LLM.__init__`, which is where that conversion normally happens, and leader and follower must produce field-identical `vllm_config`s or they disagree about model and KV layout. The monitor runs on a daemon thread because `start_worker_monitor(inline =True)` blocks and a Ray actor's `init()` has to return. Validated by simulating two nodes on one host — vLLM's MP rendezvous is only `master_addr`/`master_port`, so `nnodes=2` with `node_rank` 0/1 and 2 GPUs each reproduces it: the follower's `init()` returns with `_engine=None` and a live `MultiprocExecutor`, and the leader's capture path matches a HF reference forward at cos >= 0.99998. That simulation also exposed a total teardown gap, so `shutdown()` no longer looks only at `self._engine`. On a follower that attribute is `None`, so `shutdown()` was a no-op there and left the `MultiprocExecutor`, its monitor thread and all child worker processes alive holding GPU memory after the actor was done. `health_check` had the same root cause and reported a healthy follower as dead; nothing calls it on worker engines today, but it made a follower indistinguishable from a dead engine to any future monitoring. `_check_per_stage_capture_support` rejects `pp_size > 1` for target architectures not known to capture aux states on every stage. This is a correctness requirement rather than a message-quality nicety, and specifically because of the patch stack this series pins. On stock vLLM the configuration fails closed early and accurately: `ExtractHiddenStatesModel` lacks `SupportsPP`, so `verify_with_parallel_config` raises inside `create_engine_config` before weights load. Our patch is precisely what removes that guard — it widens the drafter construction condition to `is_last_rank or extract_hidden_states`, so `use_aux_hidden_state_outputs` becomes true on every rank, which is the point. But only patched models return `(IntermediateTensors, aux_hidden_states)` from a non-last rank; a generic model still returns a bare `IntermediateTensors` and discards the aux list it just built. Measured on Qwen3-0.6B at PP=2 before the gate: the run reached memory profiling and died at `_dummy_run`'s `hidden_states, _ = outputs` with `ValueError: too many values to unpack (expected 2)`, naming nothing relevant. `IntermediateTensors` defines `__getitem__`/`__len__` but no `__iter__`, and integer `__getitem__` falls into its slice branch returning another `IntermediateTensors`, so unpacking iterates unbounded and only fails on the arity check. The gate is knowingly a stopgap: a hardcoded allowlist of four architecture names will rot, and it stands in for the real question, which is whether the model emits aux states on non-last stages. It exists because the alternative today is that obscure crash. Making per-stage capture generic means changing the non-last-rank return contract for every Eagle3-capable model in vLLM, or threading aux tensors through `IntermediateTensors`, plus fixing `_dummy_run` — a much larger upstream conversation than this allowlist implies. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
…ding the patched image `configs/vllm_qwen3_8b_pp2.yaml` is the worked example for the feature and the config the end-to-end validation runs: Qwen3-8B split across two pipeline stages on two GPUs, feeding a two-GPU FSDP draft training loop. It differs from `vllm_qwen3_8b.yaml` only in setting `inference.vllm.pp_size`, which is the point — the engine derives `tp_size` as `inference_num_gpus_per_engine / pp_size`, so nothing else has to change to move an existing config onto pipeline stages. The justfile header now says which directory pairs with which backend version and how patches are picked up, because selecting the patched nightly is the one thing a reader cannot infer from the recipes: `VLLM_VERSION` looks like an ordinary version string, but only the nightly directory carries a patch `series`, and a released vLLM silently captures aux states on the last stage only. It also records `ARG_TAG_POSTFIX`, since every backend otherwise builds to the same `IMAGE_REPO:version` tag. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
torchspec-bot
force-pushed
the
export/vllm-pp-hidden-states
branch
from
August 9, 2026 20:45
79d7318 to
ee061ba
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ee061ba723
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| + # last one, so there is nothing here to mix against. Falling back | ||
| + # to the running prefix keeps the tap defined rather than reaching | ||
| + # for weights this rank does not construct. | ||
| + return prefix |
There was a problem hiding this comment.
When a Kimi-K3 run captures an aux layer that lands exactly at the end of a non-last pipeline stage after VllmEngine's +1 shift, this branch stores the raw running prefix instead of the AttnRes mixture used for the same layer in pp=1 and for non-boundary layers. That makes the exported hidden state depend on the PP partition for those layer selections, so training can silently consume pipeline-layout-specific features; reject such boundary layer ids or compute the tap on the next stage instead of falling back to prefix.
Useful? React with 👍 / 👎.
torchspec-bot
force-pushed
the
export/vllm-pp-hidden-states
branch
from
August 9, 2026 21:16
ee061ba to
fb4eb7f
Compare
5 tasks
`Inferact/Kimi-K3-DSpark` is published with `architectures: ["K3DSparkModel"]` and `model_type: k3_dspark`, and no loader here resolves either, so `AutoDraftModelConfig` raises on the checkpoint and it cannot be trained, continued, or evaluated in this repo at all. `DSparkDraftModel` cannot stand in for it: that checkpoint's attention is DeepSeek MLA — `q_a_proj` / `q_b_proj` LoRA pair, `kv_a_proj_with_mqa`, `kv_b_proj`, and a `v_head_dim` distinct from the qk dims — while `DFlashAttention` is a Qwen3-style GQA block with per-head norms, so neither the weight names nor the shapes line up. Add `K3DSparkConfig`, `K3DSparkMLAAttention`, `K3DSparkDecoderLayer` and `K3DSparkModel`. The attention block keeps DFlash's dual-source KV — Q is projected from the draft tokens only, K and V from `cat(context, draft)` with shared weights — and rotates only the `qk_rope_head_dim` side dims. Each side gets its own position ids, because Q spans draft-only positions while K spans context plus draft; V is never rotated, and the block mask keeps attention bidirectional. `_init_rope`, `_compute_softmax_scale` and the interleaved rotate helper are borrowed from the existing `DeepSeekMLAAttention` rather than duplicated, so the YaRN cache and its mscale are built one way in this repo. The published checkpoint rotates and no NoPE recipe exists for this draft, so the block carries no `mla_use_nope` knob and always builds a rotary; a knob would only add an untrained code path. `mla_use_output_gate` raises `NotImplementedError` rather than building a gate the published checkpoint has no weights for. Config normalization exists because transformers 5.x aliases `rope_scaling` onto `rope_parameters` as a property: both spellings collapse to one dict, an explicit legacy `rope_scaling` wins, a nested `rope_theta` is lifted to the attribute `_init_rope` reads while staying in the dict for serving-config round trips, and a partial yarn block gets the DeepSeek defaults filled in. The lift is load bearing rather than cosmetic — under transformers 5.x the top-level attribute is gone, and `_init_rope` reading it is what keeps the rotary base at the published `rope_theta`. `dflash.py` gains `attention_class` and `decoder_layer_class` so the K3 draft swaps its attention implementation without copying `DFlashDraftModel`'s forward, plus optional `fc_norm`: a per-target-layer RMSNorm applied before the context projection. It defaults off, so existing configs keep identical module names and numerics; the K3 recipe turns it on. Registry entries are exact-type keys because `from_config` looks up `type(config)`, and `DSparkTrainer` builds `K3DSparkModel` when handed a `K3DSparkConfig` — which still subclasses `DSparkConfig`, so the existing isinstance dispatch keeps selecting `DSparkTrainer`. Tests pin the state-dict layout to the MLA names and shapes with no `q_norm` / `k_norm` present, the YaRN mscale math, the output-gate rejection, and a forward/backward through the DSpark wrapper that reaches `q_a_proj`, `kv_b_proj`, `o_proj`, the markov head and the confidence head while the embedding stays frozen. Dispatch is covered from JSON in both directions: a K3 config resolves to `K3DSparkConfig` and remains a `DSparkConfig`. `configs/draft_models/kimi_k3_dspark_mla.json` matches the published checkpoint's config except that it enables `fc_norm` and selects its own capture layers, so it is a training recipe rather than a mirror of the release. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
The K3 DSpark draft added in the previous commit is only reachable through a config that wires it to a pipeline-parallel target, and nothing in `configs/` shows that combination: the vLLM examples are single-stage, and the multi-node examples are all SGLang. This is the recipe our own stage-1 runs use, minus the site-specific parts. The layout is 40 GPUs over 8 nodes: 32 inference GPUs as two engines of `tp_size=8` across two pipeline stages, and 8 training GPUs under FSDP `FULL_SHARD`. It therefore needs a vLLM built from `patches/vllm/<image-tag>/`, since released vLLM captures aux hidden states on the last pipeline stage only. Two couplings are easy to get wrong and are called out in the file. `aux_hidden_states_layers` must match `target_layer_ids` in the draft config, because the connector writes exactly the layers the draft's context projection expects to read. And `min_loss_tokens: 14` is not arbitrary: DSpark supervises whole seven-token diffusion blocks, so a sample shorter than two full blocks contributes no usable signal. The Mooncake block departs from the dataclass defaults in three ways that only matter at this scale. `enable_hard_pin` keeps published fragments from being evicted while the trainer is still reading them; it needs a Mooncake exposing `ReplicateConfig.with_hard_pin` and degrades to a warning otherwise. `kv_lease_ttl_s: 180` and `get_retry_max_wait_seconds: 180` raise the master lease and the consumer's fetch-retry ceiling above their 5s and 60s defaults, because a two-stage engine spread over eight nodes can leave a fragment unclaimed for far longer than a single-node run does. And `global_segment_size: 512GB` sizes the store for 7168-wide hidden states across five aux layers. Checkpoint, dataset and output paths are left as `???` following the convention in `configs/default.yaml`, and the NIC names and RDMA device are commented templates rather than values, since both are properties of the cluster rather than of the recipe. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
`DeepSeekMLAAttention._init_rope` passes `base=rope_theta` in every branch except the YaRN one, where the argument is simply absent and `LlamaYarnRotaryEmbedding`'s own default of 10000 takes over. The sibling implementation in `llama3_eagle.py` does pass it, and the surrounding branches here pass it, so this is an omission rather than a decision. It went unnoticed because the configs that reach this branch resolve `rope_theta` to 10000 anyway: transformers 5.x moved the field into `rope_parameters`, so the top-level attribute the method reads no longer exists and the `getattr` default wins. `K3DSparkConfig` is the one config class that lifts the nested value back to the attribute, which is what turns a latent mismatch into a real one — the K3 DSpark draft declares `rope_theta: 50000.0` and was silently trained against rotary frequencies for 10000, at a factor of five in every position-dependent frequency. The published `Inferact/Kimi-K3-DSpark` config carries `rope_theta` nested in its yarn block, so any run started from that checkpoint was affected; drafts trained before this commit learned against the wrong frequencies and are not comparable to ones trained after it. The existing test asserted the rotary class and its `dim` but never its `base`, which is precisely the gap that hid this. The new test pins the base to the configured value and compares `inv_freq` against a reference embedding built at 50000, asserting it differs from one built at 10000 so the assertion cannot pass vacuously. Verified by reverting the one-line change: the test fails with `10000 != 50000.0`. This does not change behaviour for any other config in the tree. The Eagle3 configs still resolve `rope_theta` to the 10000 default, so their rotary base is untouched here; making them honour their declared `rope_theta` means teaching `_init_rope` to read the nested value, which is a separate change with its own consequences for already-trained drafts. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
torchspec-bot
force-pushed
the
export/vllm-pp-hidden-states
branch
from
August 9, 2026 21:55
fb4eb7f to
c8314a5
Compare
Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
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
1 participant
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.
Summary
Hidden-state export currently only works when the target model fits in one pipeline stage. This adds pipeline-parallel export end to end, pins a patched vLLM image so the behaviour is reproducible from this tree, and makes the published
Inferact/Kimi-K3-DSparkdraft loadable so a PP target has something to train.Released vLLM cannot do this: it builds the aux-capture list on the last rank only, and
ExtractHiddenStatesModeldoes not implementSupportsPP, soModelConfig.verify_with_parallel_configrejectspipeline_parallel_size > 1whileSpeculativeConfigis still being built. The first commit therefore pins a nightly base image plus an ordered patch stack; the patch is an upstream candidate kept separate from the K3 model patch so a shrinking diff after a rebase is the signal that vLLM has absorbed it.Each stage now writes only the layers it owns as separate Mooncake fragments, records them in a
vllm_pp_layer_manifest, and the consumer reassembles them in layer order, cross-checkinginput_idsacross fragments before concatenating. Paths that cannot honour fragments fail closed rather than training on partial tensors: offline saving rejects manifest-bearing samples, andpp_size > 1is refused for architectures without per-stage capture instead of silently producing tensors no stage wrote.What's in each commit
series-ordered patch stack. The Dockerfile readsseriesrather than globbing*.patchso the stack stays explicit and each patch's upstream intent is recorded next to it; ordering itself is not a correctness requirement, since the two patches edit disjoint regions andpatchabsorbs the ~80-line shift as an offset (verified by applying them in reverse).pp_sizeplumbing, derivedtp_size, headless follower executors for multi-node engines, and the per-stage capture gate.justfilenotes for building the patched image.K3DSparkModel), which no loader here previously resolved, plusfc_normand two subclass hooks in the DFlash draft.Test plan
PP=1vsPP=2hidden states are bitwise identical across three prompt lengths (rel_l2=0,cos=1.0,max_abs=0), read back through the productionMooncakeDatasetpath —tests/test_vllm_pp_equivalence.py, run inside the image this PR's Dockerfile builds.hidden_states=(1, 128, 12288)confirming three aux layers were reassembled from both stages.supports_pp(ExtractHiddenStatesModel)isTrueandModelRunnerOutputcarriesprefill_only.ruff checkandruff format --checkclean. Full suite shows no regressions againstmainin the same container (identical pre-existing failures, +41 passing tests).Notes for review
patches/vllm/<image-tag>/;vllm_pp_hidden_states.patchis the upstream candidate andvllm_k3.patchis ours indefinitely.faextra is intentionally not installed in this image:flash-attn-4pinsnvidia-cutlass-dsl==4.6.0.dev0against the base image's4.6.0, which breaks vLLM's flashinfer stack at worker startup. It only gates the optionalflash_attentiondraft backend.configs/draft_models/kimi_k3_dspark_mla.jsonmatches the published checkpoint except that it enablesfc_normand picks its own capture layers.