diff --git a/configs/vllm_kimi_k3_dspark_stage1.yaml b/configs/vllm_kimi_k3_dspark_stage1.yaml index 07768a2d..ed3e1ecb 100644 --- a/configs/vllm_kimi_k3_dspark_stage1.yaml +++ b/configs/vllm_kimi_k3_dspark_stage1.yaml @@ -86,7 +86,9 @@ inference: inference_fetch_batch: 8 max_sample_pool_size: 32 store_last_hidden_states: true - # Must match target_layer_ids in the draft config above. + # Must match target_layer_ids in the draft config above. These are the layers + # whose *output* is captured; the engine shifts them by +1 into the capture ids + # the model uses, so 7 means "the residual stream after layer 7 ran". aux_hidden_states_layers: [7, 31, 47, 63, 87] last_hidden_states_prenorm: true diff --git a/docs/kimi_k3_attn_res_capture.md b/docs/kimi_k3_attn_res_capture.md new file mode 100644 index 00000000..673f4bab --- /dev/null +++ b/docs/kimi_k3_attn_res_capture.md @@ -0,0 +1,251 @@ +# Capturing Kimi-K3 Hidden States under AttnRes + +This is the background you need before touching Kimi-K3 auxiliary hidden-state +capture, and the record of two bugs we fixed in `vllm_k3.patch`. It is written +in the order that makes the problem make sense rather than the order we found +things in. + +## The thing that makes K3 different + +In an ordinary transformer, "the hidden state after layer K-1" is a tensor. It +is threaded from one layer to the next, and every reader sees the same value. + +Under AttnRes, K3 threads three things instead: + +- `prefix_sum` — the running sum within the current block +- `block_residual`, the *bank* — the blocks committed so far, `(T, num_blocks, H)` +- a pending `delta`, the previous layer's MLP output, not yet folded in + +and no reader consumes them directly. Every consumer — the next layer's +attention front-end, and the model's own output-side aggregation — first runs a +softmax mixture over the bank and the prefix. From the kernel: + +``` +sources = bank[0..num_blocks-1] + [prefix] # num_blocks + 1 of them +logit(v) = (v · (norm_w ⊙ proj_w)) * rsqrt(mean(v²) + eps) +mixed = Σ softmax(logits)_i · v_i # a convex combination +``` + +The query vector `norm_w ⊙ proj_w` comes from the *consumer's* own +`self_attention_res_norm.weight` and `self_attention_res_proj.weight`. + +**So "the hidden state after layer K-1" is not a property of the residual +stream. It is defined by whichever layer reads it.** Every consequence below +follows from that sentence. + +## Three quantities people confuse + +| Quantity | What it is | Who reads it | +|---|---|---| +| `prefix_sum + pending_mlp_out` | the running prefix on its own | **nobody** | +| `mixed` | the softmax mixture, using layer K's score weights | the tap, and the drafter | +| `mixed` normed by layer K's `input_layernorm` | layer K's attention input | only layer K | + +The drafter is trained against the second row. The third row is layer K's +private input preprocessing — the drafter has its own `fc_norm` +(`"fc_norm": true` in `configs/draft_models/kimi_k3_dspark_mla.json`), so +feeding it a value already normed by the target model stacks two unrelated +normalisations. + +The first row is what the pre-tap code recorded, and what a PP boundary used to +fall back to. It is not an approximation of the second row; it is one of +`num_blocks + 1` sources. + +## What a "tap" is + +A tap is a probe on the residual stream, in the plumbing sense: it reads the +signal without interrupting it. Concretely it is the same `attn_res` call the +consumer would make, with the output norm turned off: + +```python +attn_res(prefix, None, bank, score_norm, score_proj, None, num_blocks=..., + block_write_idx=-1, ...) +# ^^^^ output_norm_weight = None -> returns `mixed` +``` + +Three things make it a pure read, and all three matter: + +- `delta=None` — the kernel's `if HAS_DELTA: tl.store(prefix_ptr, ...)` does not fire +- `block_write_idx=-1` — `if WRITE_BLOCK: tl.store(blocks_ptr, ...)` does not fire +- the `prefix` passed in is `prefix_sum + pending_mlp_out`, freshly computed, not + the live `prefix_sum` + +Verified directly (`.tmp/attnres-pp/verify_tap_is_pure_read.py`): after the call, +`prefix` and the bank are bitwise unchanged, the output aliases neither, and +running it twice gives the same answer. + +This is why the tap cannot reuse the consumer's call. The consumer's call has +side effects — it folds the delta back into the prefix and commits a block into +the bank — and it returns the *normed* value. `tests/models/kimi_k3/test_eagle3.py` +pins both conditions (`delta is None`, `block_write_idx == -1`), because an +end-to-end PP comparison cannot catch a regression here: both sides would do the +same corrupting tap and the damage would cancel. + +## Why the tap is a second computation at all + +`attn_res` has exactly one output pointer: + +``` +149: output = mixed +154: if APPLY_OUTPUT_NORM: +162: output = mixed * rstd * output_norm_weight +164: tl.store(output_ptr + ..., output, ...) +``` + +It writes either the mixture or the normed value, never both — even though at +line 162 both are live in registers. That single-exit interface, not the maths, +is why a tapped layer computes the mixture twice. + +Note the shape of the fix that is *not* available: unfusing the norm into a +separate kernel would force `mixed` out to memory and back. These kernels are +entirely memory-bound, so that costs two extra passes over `(T, H)` on all 93 +layers × 2 calls: + +| | passes per forward | traffic at T=32768, H=7168, bf16 | +|---|---|---| +| fused (today) | 1536 | 672 GiB | +| norm as a separate kernel | 1908 | 835 GiB (+24%) | + +The cheap version is a second optional output pointer, used only on tapped +layers: `+5` passes to write `mixed`, `-40` passes because the separate tap call +disappears, net **-15.3 GiB**. That needs both the Triton kernel and +`csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu` changed — changing only Triton +would force tapped layers off the native path and make the model's arithmetic +depend on the capture config, which is the bug in the next section. So it means +recompiling vLLM, which our patch stack is built to avoid. Not done. + +## Bug 1: the tap at a pipeline boundary + +If capture id `K` is exactly a non-last stage's `end_layer`, layer `K` lives on +the next rank, so its score weights are not here. + +The upstream behaviour ([vLLM #50487](https://github.com/vllm-project/vllm/pull/50487), +still open) is to fall back to the running prefix. In the language above, that +forces the softmax one-hot onto a single source and discards the rest. For a +stage ending at layer 48 with `attn_res_block_size=8` that is 1 of 7 sources. In +a harness driving the real kernel with synthetic weights, cosine similarity +between the fallback and the real mixture ran **0.05 to 0.60** — a different +feature, not a nearby one. + +The damage is that the exported feature becomes a function of the pipeline +partition. The same capture id yields the mixture at pp=1 and the raw prefix at a +pp that happens to cut there, while the drafter at serving time always consumes +the mixture. It is silent: right shape, right dtype, no NaNs. + +**Fix: keep a copy of the two vectors.** Everything else the tap needs is already +local, including one detail that makes it exact — this stage's +`num_attn_res_blocks` is `cdiv(end_layer, block_size)`, which is precisely layer +`end_layer`'s `prev_valid_blocks`, so the bank is already the right length rather +than merely present. What is missing is two `hidden_size` vectors, unsharded and +unquantised, about **28KB** together. + +They are redirected out of the checkpoint by `_maybe_boundary_attn_res_name` +before the PP filter drops them as belonging to a layer this rank does not own, +and initialised to NaN so a checkpoint that never supplies them fails loudly +instead of exporting a mixture over uninitialised memory. + +Rejecting boundary ids was our first attempt. It does not survive mostly-PP +sharding: for 93 layers with capture ids at {8, 32, 48, 64, 88} under the default +`get_pp_indices`, the first collision is at pp=10 (id 64), pp=12 has none, pp=16 +has two, and from pp=40 on all five are stage ends. + +## Bug 2: the kernel dispatch at a pipeline boundary + +Found while measuring bug 1, and unrelated to capture — it affects plain serving. + +The native fused op is only eligible when, among other conditions, +`delta is not None` and `block_write_idx < 0`. The handoff folds the delta into +the prefix and the receiving stage started its layers with no delta, so **the +first layer of every non-first stage ran the Triton fallback** while the same +layer at pp=1 ran the native op. They agree to within bf16 rounding and then +compound, so the arithmetic — every tap after that point, not just the boundary +one — depended on where the pipeline was cut. + +**Fix: hand that layer an explicit zero delta.** Block-write layers never reach +the native path from either side, so the filler is skipped for them and a +block-aligned partition allocates nothing. + +## What was verified + +On GB300 against the real `attn_res` kernel, comparing a two-stage run against a +single-stage one at every possible split point: + +| stack | splits | boundary tap | all taps | final hidden states | +|---|---|---|---|---| +| 24 layers / block 4 | 23 | 23 | 23 | 23 | +| 32 layers / block 8 | 31 | 31 | 31 | 31 | + +Bit-identical, not close. Before the bug 2 fix, end-to-end identity held at 4 of +31 splits, and those four were exactly the block-write-aligned cuts — which is +also the independent confirmation of the dispatch mechanism. + +Caveat on provenance: the harness reproduces K3's AttnRes bookkeeping and drives +the real kernel, but with synthetic weights. Bit-identity is a structural claim +and holds regardless; the cosine range above is illustrative. + +Test-suite effect, against the pinned image with `--noconftest` and +`test_latent_moe_tail.py` excluded (it fails to collect on pristine too): +`tests/models/kimi_k3` goes from 59 failed / 26 passed to 59 failed / 29 passed, +with the failing set identical line for line. Those 59 are GPU tests that cannot +run on a CPU-only box — `current_platform.device_type` is `''` there, so +`torch.device('')` raises. + +## What capture costs + +Per tapped layer, one extra `attn_res` reading `num_blocks + 1` sources. For the +93-layer stage-1 recipe: + +| capture id | 8 | 32 | 48 | 64 | 88 | total | +|---|---|---|---|---|---|---| +| source reads | 2 | 5 | 7 | 9 | 12 | 35 | + +With the write of each result that is 40 passes over `(T, H)`, against the 1536 +the model's own AttnRes calls already do — about **2.6%** of AttnRes traffic, +17.5 GiB, and a smaller fraction of the whole forward. It grows with the depth +of the tapped layer, so where you tap matters more than how many times. + +Worth keeping in proportion: K3's PP handoff carries the bank, so it is +`1 + cdiv(start_layer, block_size)` tensors of `(T, H)` rather than one. At +T=32768 and H=7168 one such tensor is 0.44 GiB: + +| cut at | tensors on the wire | per handoff | +|---|---|---| +| 24 | 4 | 1.75 GiB | +| 47 / 48 | 7 | 3.06 GiB | +| 88 | 12 | 5.25 GiB | + +pp=2 moves 3.1 GiB per step, pp=4 moves 9.2 GiB, pp=8 moves 22.3 GiB — growing +superlinearly, since deeper cuts also carry more blocks. That is over NVLink/IB +rather than HBM, so any optimisation effort belongs there long before it belongs +in capture's 17.5 GiB of local traffic. + +## Choosing where to cut + +Two independent reasons to prefer PP boundaries at multiples of +`attn_res_block_size`: + +- the first layer of the stage is then a block-write layer, which takes the + Triton path from either side — this was the zero-code workaround for bug 2 and + is still the cheapest configuration, since the filler delta is skipped +- shallower cuts move less bank + +Boundaries coinciding with capture ids are fine now. That is what bug 1's fix +buys, and it matters because the MLA layer pattern puts the natural capture ids +on the same multiples of 8 that make good block-aligned cuts. + +`aux_hidden_states_layers` in the YAML is **pre-shift**: the engine adds one, so +`7` means "the residual stream after layer 7 ran" and becomes capture id 8. + +## Where things live + +| | | +|---|---| +| model changes | `patches/vllm/nightly-7794b1e08.../vllm_k3.patch` | +| tests | `patches/vllm/nightly-7794b1e08.../tests/vllm_k3_tests.patch` (not in `series`) | +| equivalence harness | `.tmp/attnres-pp/verify_boundary_tap.py` | +| purity check | `.tmp/attnres-pp/verify_tap_is_pure_read.py` | +| dispatch probe | `.tmp/attnres-pp/verify_native_vs_triton.py` | + +Neither fix went upstream. Bug 2 would have made a clean standalone vLLM PR, and +bug 1 belongs in #50487, which currently ships the fallback this replaces along +with a test asserting it. See `UPSTREAM_EXPORT_PLAN.md` (M13) for that decision. diff --git a/patches/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/tests/vllm_k3_tests.patch b/patches/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/tests/vllm_k3_tests.patch new file mode 100644 index 00000000..3f2cee05 --- /dev/null +++ b/patches/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/tests/vllm_k3_tests.patch @@ -0,0 +1,150 @@ +Tests for vllm_k3.patch. Not listed in ../series -- the image does not ship +vLLM's test tree, so apply this by hand in a source checkout that already has +../vllm_k3.patch: + + cd $VLLM_ROOT && git apply /path/to/vllm_k3.patch + && git apply /path/to/tests/vllm_k3_tests.patch + pytest tests/models/kimi_k3/test_eagle3.py tests/models/kimi_k3/test_pp_handoff.py + +test_pp_handoff.py (new) pins the invariant behind the zero-delta fix: the first +layer of a stage must see the same kind of delta regardless of where the pipeline +was cut, because that is what decides whether attn_res dispatches to the fused +native op or the Triton fallback. It fails without the fix. The two negative +cases guard the other direction -- no filler where it would not change dispatch. + +test_eagle3.py needs two fixture updates, both consequences of the AttnRes tap +itself rather than of anything about PP: +- test_kimi_k3_uses_shared_eagle3_layer_configuration's stub predates the tap and + does not set use_attn_res, which _set_aux_hidden_state_layers now reads. +- test_kimi_linear_forward_extracts_attn_res_aux_hidden_states asserted the old + post-mixture sum. It now checks the mixture comes back from the kernel, and + additionally that the pending MLP output is folded into the prefix rather than + passed as a delta -- the kernel writes an applied delta back in place, which + would double-add it into the live residual stream. + +Measured against the pinned image, tests/models/kimi_k3 with --noconftest and +test_latent_moe_tail.py excluded (it fails to collect on pristine too): 59 +failed / 26 passed before, 59 failed / 29 passed after, with the failing set +identical line for line. The 59 are environmental -- no GPU and no vllm._C. + +diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py +index 61a24a83e..4ca193756 100644 +--- a/tests/models/kimi_k3/test_eagle3.py ++++ b/tests/models/kimi_k3/test_eagle3.py +@@ -30,6 +30,7 @@ def test_kimi_k3_uses_shared_eagle3_layer_configuration(): + torch.nn.Module.__init__(target) + model = _make_kimi_linear_model() + object.__setattr__(model, "layers", [None] * 93) ++ object.__setattr__(model, "use_attn_res", False) + language_model = SimpleNamespace( + embed_input_ids=lambda _: None, + model=model, +@@ -126,5 +127,13 @@ def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch): + + torch.testing.assert_close(output, final_hidden_states) + torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) +- torch.testing.assert_close(aux_hidden_states[1], prefix_sum + layer_hidden_states) +- assert final_attn_res.call_args.args[2] is block_residual ++ torch.testing.assert_close(aux_hidden_states[1], final_hidden_states) ++ ++ aux_call, output_call = final_attn_res.call_args_list ++ # No delta and no block write is what keeps the tap a pure read: either one ++ # would have the kernel store back into state the next PP stage consumes. ++ torch.testing.assert_close(aux_call.args[0], prefix_sum + layer_hidden_states) ++ assert aux_call.args[2] is block_residual ++ assert aux_call.args[1] is None ++ assert aux_call.kwargs["block_write_idx"] == -1 ++ assert output_call.args[2] is block_residual +diff --git a/tests/models/kimi_k3/test_pp_handoff.py b/tests/models/kimi_k3/test_pp_handoff.py +new file mode 100644 +index 000000000..b39641301 +--- /dev/null ++++ b/tests/models/kimi_k3/test_pp_handoff.py +@@ -0,0 +1,87 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++"""The delta a pipeline stage hands its first layer under AttnRes. ++ ++``attn_res`` only takes its fused native path when a delta is supplied, so ++whether the first layer of a stage sees one decides which kernel runs -- and ++that must not depend on where the pipeline was cut. ++""" ++ ++from types import SimpleNamespace ++from unittest.mock import Mock ++ ++import pytest ++import torch ++ ++from vllm.models.kimi_k3.nvidia import model as kimi_model ++from vllm.models.kimi_k3.nvidia.model import KimiLinearModel ++ ++ ++def _run_first_layer(monkeypatch, *, is_first_rank, is_block_write_layer): ++ """Run one AttnRes stage and return the delta its first layer received.""" ++ start_layer = 0 if is_first_rank else 4 ++ prefix = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) ++ layer = Mock( ++ return_value=( ++ torch.zeros_like(prefix), ++ prefix, ++ torch.zeros(prefix.size(0), 1, prefix.size(1)), ++ ) ++ ) ++ layer.is_block_write_layer = is_block_write_layer ++ ++ model = object.__new__(KimiLinearModel) ++ object.__setattr__(model, "use_sequence_parallel", False) ++ object.__setattr__(model, "use_attn_res", True) ++ object.__setattr__(model, "num_attn_res_blocks", 1) ++ object.__setattr__(model, "aux_hidden_state_layers", ()) ++ object.__setattr__(model, "start_layer", start_layer) ++ object.__setattr__(model, "end_layer", start_layer + 1) ++ object.__setattr__(model, "layers", [None] * start_layer + [layer]) ++ monkeypatch.setattr( ++ kimi_model, ++ "get_pp_group", ++ lambda: SimpleNamespace(is_first_rank=is_first_rank, is_last_rank=False), ++ ) ++ ++ model.forward( ++ input_ids=None, ++ positions=torch.tensor([0, 1]), ++ intermediate_tensors={"hidden_states": prefix, "residual": None}, ++ inputs_embeds=prefix if is_first_rank else None, ++ ) ++ return layer.call_args.kwargs["hidden_states"], prefix ++ ++ ++def test_mid_block_stage_start_receives_a_zero_delta(monkeypatch): ++ delta, prefix = _run_first_layer( ++ monkeypatch, is_first_rank=False, is_block_write_layer=False ++ ) ++ ++ assert delta is not None ++ assert delta.shape == prefix.shape ++ assert delta.dtype == prefix.dtype ++ assert not delta.any() ++ ++ ++@pytest.mark.parametrize( ++ "is_first_rank,is_block_write_layer", ++ [ ++ # Block-write layers take the Triton path from either side, so a filler ++ # delta would buy nothing. ++ (False, True), ++ # The first stage has no handoff and no pending delta to recover. ++ (True, True), ++ ], ++) ++def test_no_filler_delta_where_it_would_not_change_dispatch( ++ monkeypatch, is_first_rank, is_block_write_layer ++): ++ delta, _ = _run_first_layer( ++ monkeypatch, ++ is_first_rank=is_first_rank, ++ is_block_write_layer=is_block_write_layer, ++ ) ++ ++ assert delta is None diff --git a/patches/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/vllm_k3.patch b/patches/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/vllm_k3.patch index a64e2a59..6f339e30 100644 --- a/patches/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/vllm_k3.patch +++ b/patches/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/vllm_k3.patch @@ -28,6 +28,31 @@ logs the active mode once at configuration time: Once the base image carries PR #50487, drop this patch but keep exporting VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=1, or the default flips back to prefix_only. +Two fixes on top of #50487 make the capture invariant to the pipeline +partition. Background and derivations: docs/kimi_k3_attn_res_capture.md. + +PP boundary: the mixture for capture id ``k`` uses layer ``k``'s +``self_attention_res_norm`` / ``_proj``, which for a non-last stage's +``end_layer`` sit on the next rank. #50487 falls back to the running prefix +there, a different feature entirely (cosine 0.05-0.60 against the real mixture). +Everything else the tap needs is local, so a non-last AttnRes stage now keeps a +copy of those two hidden-size vectors (~28KB, unsharded), redirected out of the +checkpoint by ``_maybe_boundary_attn_res_name`` before the PP filter drops them. +Rejecting the boundary ids instead was the previous behaviour here; it does not +survive mostly-PP sharding, since for the 93-layer stage-1 recipe every +intermediate id is a stage end from pp=40 on. + +Zero delta after a handoff, found while measuring the above: ``attn_res`` only +takes the fused native path when a ``delta`` is supplied, and the handoff folds +the delta into the prefix, so the first layer of every non-first stage ran the +Triton fallback -- ~0.008% of elements differ (max relative 7.6e-4) and the +error compounds through the rest of the stack. An explicit zero delta restores +the dispatch, and is skipped for block-write layers, which never reach it. + +Verified on GB300 against the real kernel: the boundary tap is bit-identical to +the PP=1 capture at every split point, and with the zero delta the whole +two-stage run is bit-identical end to end (31 of 31 split points, from 4). + Note: the v0.22.1 patch is not needed on this nightly: - flash_attn rotary import fix: upstream, the import is now guarded with ``suppress(ModuleNotFoundError)`` in @@ -64,13 +89,58 @@ Apply: "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py -@@ -1095,6 +1095,85 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): +@@ -1066,6 +1066,23 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): + if self.use_attn_res: + self.output_attn_res_norm = PPMissingLayer() + self.output_attn_res_proj = PPMissingLayer() ++ # An aux tap at ``end_layer`` needs that layer's mixture weights, ++ # which live on the next rank. Two unsharded hidden-size vectors ++ # are cheaper to copy than the tap is to move. ++ self.boundary_attn_res_norm = RMSNorm( ++ config.hidden_size, eps=config.rms_norm_eps ++ ) ++ self.boundary_attn_res_proj = ReplicatedLinear( ++ config.hidden_size, ++ 1, ++ bias=False, ++ quant_config=None, ++ prefix=f"{prefix}.boundary_attn_res_proj", ++ ) ++ # Sentinel: a checkpoint that never supplies these must fail ++ # rather than export a mixture over uninitialised memory. ++ self.boundary_attn_res_norm.weight.data.fill_(float("nan")) ++ self.boundary_attn_res_proj.weight.data.fill_(float("nan")) + + world_size = get_tensor_model_parallel_world_size() + assert config.num_attention_heads % world_size == 0, ( +@@ -1095,6 +1112,106 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): } ) + def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + super()._set_aux_hidden_state_layers(layers) + if self.use_attn_res: ++ if ( ++ self._aux_attn_res_stream ++ and not get_pp_group().is_last_rank ++ and self.end_layer in layers ++ ): ++ missing = [ ++ f"layers.{self.end_layer}.self_attention_res_{kind}.weight" ++ for kind, param in ( ++ ("norm", self.boundary_attn_res_norm.weight), ++ ("proj", self.boundary_attn_res_proj.weight), ++ ) ++ if not torch.isfinite(param).all() ++ ] ++ if missing: ++ raise ValueError( ++ f"Kimi-K3 aux layer id {self.end_layer} is the end of PP " ++ f"stage [{self.start_layer}, {self.end_layer}), so it " ++ f"needs a copy of layer {self.end_layer}'s AttnRes " ++ "mixture weights, but the checkpoint supplied no " ++ f"{' and no '.join(missing)}." ++ ) + # Emitted once, at configuration time. Which layers are tapped and + # which convention is in force are the two things you need to + # confirm from a running process, and neither is recoverable from @@ -127,12 +197,12 @@ Apply: + score_proj = self.output_attn_res_proj + num_blocks = self.num_attn_res_blocks + else: -+ # Last layer of a non-final pipeline stage: the consumer lives on -+ # the next rank and the output-side aggregation only exists on the -+ # 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 ++ # Consumer is layer ``end_layer`` on the next rank. This is exact ++ # rather than approximate because ``num_attn_res_blocks`` is ++ # ``cdiv(end_layer, block_size)`` -- that layer's ``prev_valid_blocks``. ++ score_norm = self.boundary_attn_res_norm ++ score_proj = self.boundary_attn_res_proj ++ num_blocks = self.num_attn_res_blocks + + return attn_res( + prefix, @@ -150,7 +220,22 @@ Apply: def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) -@@ -1161,7 +1240,10 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): +@@ -1146,6 +1263,14 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): + block_residual[:, : residual.size(1), :].copy_(residual) + prefix_sum = hidden_states + hidden_states = None ++ if ( ++ not get_pp_group().is_first_rank ++ and not self.layers[self.start_layer].is_block_write_layer ++ ): ++ # ``attn_res`` only takes its fused path when a delta is supplied, ++ # and the handoff folded this stage's away. Without a stand-in ++ # the arithmetic depends on where the pipeline was cut. ++ hidden_states = torch.zeros_like(prefix_sum) + residual = block_residual + + for layer_idx, layer in enumerate( +@@ -1161,7 +1286,10 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): if (layer_idx + 1) in self.aux_hidden_state_layers: if self.use_attn_res: assert prefix_sum is not None @@ -162,3 +247,39 @@ Apply: else: assert residual is not None aux_hidden_state = hidden_states + residual +@@ -1213,6 +1341,21 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): + return hidden_states, aux_hidden_states + return hidden_states + ++ def _maybe_boundary_attn_res_name(self, name: str) -> str | None: ++ """Map layer ``end_layer``'s two mixture weights onto the boundary copies. ++ ++ Everything else about that layer still belongs to the next rank. ++ """ ++ if not self.use_attn_res or get_pp_group().is_last_rank: ++ return None ++ for suffix, target in ( ++ ("self_attention_res_norm.weight", "boundary_attn_res_norm.weight"), ++ ("self_attention_res_proj.weight", "boundary_attn_res_proj.weight"), ++ ): ++ if name.endswith(f"layers.{self.end_layer}.{suffix}"): ++ return target ++ return None ++ + def load_weights( + self, + weights: Iterable[ +@@ -1290,6 +1433,13 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue ++ boundary_name = self._maybe_boundary_attn_res_name(name) ++ if boundary_name is not None: ++ # Must precede the PP filter below, which would drop these. ++ param = params_dict[boundary_name] ++ default_weight_loader(param, loaded_weight) ++ loaded_params.add(boundary_name) ++ continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue