Learned low-rank activation codec for split boundaries

[danielwinterw]

Draft: learned low-rank activation codec for split boundaries

A per-boundary PCA codec targeting ~8x vs f16 on the inter-stage wire
(4096-wide f16 = 8 KB/token -> rank-512 int8 ~ 0.5 KB, rank-1024 ~ 1 KB):

• WireActivationDType::Lowrank: the rank rides in the high bits of the
  state header's reserved field, so frame sizes stay derivable at read time
  and existing dtypes are byte-identical.
• BoundaryCodec: orthogonal-iteration PCA (deterministic seed), per-token
  int8 quantized coefficients, .skbc file format with dimension sanity
  bounds. Fit offline with the new fit-boundary-codec subcommand from
  activations captured via SKIPPY_CAPTURE_BOUNDARY_ACTIVATIONS.
• serve-binary --boundary-codec + --activation-wire-dtype lowrank
  enable it on the standalone stage path; embedded stages can load a codec
  via MESH_LLM_BOUNDARY_CODEC.

Draft because: protocol + tooling are complete and unit-tested
(roundtrip error bounds on synthetic low-rank data), but there are no
real-model quality gates yet (perplexity delta, suffix accept-rate delta),
no package-manifest shipping of codec tensors, and auto dtype resolution
deliberately does not select lowrank until the embedded forward path is
wired. Codec identity is keyed informally for now — package-hash keying is
the planned follow-up.

[github-actions]

This pull request is currently a draft. Reviews will not take place until the PR is marked as ready for review.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[i386]

Reviewed end-to-end: protocol tagging (rank packed into the reserved field's high bits with the dtype tag in the low byte — safe since existing dtypes had zero high bits), codec encode/decode math, size arithmetic with checked multipliers, CLI/env plumbing, and the allowlist regen.

Clean aspects: strictly opt-in (Lowrank bails at startup without a codec), all existing dtype paths unchanged, decode validates rank and payload size, fit is deterministic, and there are round-trip, persistence, and degenerate-shape tests.

Two non-blocking notes for the follow-up plan:

1. Codec identity is not verified across a boundary. The decode path checks that the frame's rank matches the loaded codec's rank, but two different codecs of the same rank will silently corrupt outputs (per the doc comment, both stages must hold byte-identical tensors — nothing enforces it). Suggest hashing the codec tensors at load and advertising/verifying the digest at connection setup, or at minimum in the state header handshake.
2. Env-only config (MESH_LLM_BOUNDARY_CODEC) applies one codec to all boundaries on a stage; with multi-stage splits this will collide once per-boundary codecs exist. Fine for now given the tracked plan, just flagging the seam.

Also minor: the {}x compression print in fit_boundary_codec_cli divides width*2 by rank + 4 — the per-token f32 scale makes it rank + 4/d per row, so the printed factor is pessimistic for small d.

[i386]

This is the codec-identity gap from my review, flagged inline: the check validates only the rank, so two different codecs of the same rank on the two sides of a boundary will decode garbage silently — the doc comment's "both stages must hold byte-identical codec tensors" invariant is unenforced. Suggest hashing the codec tensors at load (e.g. SHA-256 over the .skbc payload) and advertising/verifying the digest at connection setup so mismatched codecs fail loudly instead of corrupting outputs.

[danielwinterw]

Thanks for the approve. On the codec-identity gap — agreed, and it is the one I would least want to ship silently: two same-rank codecs decoding garbage with no error is exactly the failure mode that looks like a model quality problem. Hashing the .skbc payload at load and verifying the digest at connection setup is the right shape, and it fits the package-hash keying already listed as a follow-up in the PR description.

Not doing it in this PR: the digest wants to ride the stage handshake alongside the other capability negotiation, and this PR is deliberately inert until the embedded forward path is wired (auto never resolves to lowrank). I would rather land the negotiation and the forward path together than add a half-enforced check now. Tracking it as the blocking item before this comes off draft.

[i386]

Reviewed at origin/feat/lowrank-boundary-codec against its base feat/runahead-verify-windows. Draft, so I'm focusing on the wire contract rather than the missing quality gates — those are correctly listed in the description.

The framing decision is good: putting the rank in the header so frame sizes stay derivable at read time is the right call, and it's what keeps read_stage_message a single pass. I checked the allocation path and it's sound — the Lowrank branch sits above the MAX_STAGE_ACTIVATION_BYTES check, so a peer-supplied k can't drive an unbounded vec![0; activation_bytes]. Worth calling out because k is the first wire-controlled input to the frame-size computation; n_embd has always been local config.

Three things, in order of how much I care.

1. A codec mismatch is silently wrong output, not an error

decode catches a rank mismatch, because payload size is a strict function of k and the length check will fail. It cannot catch a content mismatch: two stages holding same-rank .skbc files fitted from different calibration runs will encode and decode happily and produce garbage activations, with no error anywhere on the path. The model just gets worse, at a boundary, in a way that looks like a quality regression rather than a config bug.

You flag this ("codec identity is keyed informally for now — package-hash keying is the planned follow-up"), and package-hash keying is the right end state. But the gap between "informally keyed" and "silently wrong" is the part I'd close before this leaves draft, because it's the failure that will cost someone a day. A cheap interim: hash the .skbc bytes to 32 bits and exchange the fingerprint once at connection setup, refusing the connection on mismatch. That's independent of package-manifest work and doesn't need a frame-format change.

2. dtype() lost validation it used to have

pub fn dtype(self) -> io::Result<WireActivationDType> {
    WireActivationDType::try_from(self.reserved & 0xFF)
}

reserved used to be exactly dtype as i32, so anything outside 0..=2 was rejected as "unknown activation wire dtype". With the mask, the high bits are now silently dropped for every dtype, not just Lowrank. reserved = 0x0100 was a hard error before and now parses as F32. That's a real loss of frame validation on a field that was previously fully constrained, and it applies to every existing dtype rather than just the new one.

Suggest requiring the high bits to be zero unless the low byte is Lowrank — it preserves the old strictness for F32/F16/Q8 and costs one comparison.

3. This needs its own generation, if #1409 lands first

#1409 defines generation 5 as the compatibility contract, and this PR changes the meaning of reserved and adds a dtype tag on top of it. As a stacked branch that's fine — but if #1409 merges before this does, a shipped gen-5 peer predates the reserved change, and this needs generation 6 rather than riding on 5.

The failure mode is safe either way (an older peer reads reserved = 3 | (k << 8) and errors on TryFrom), but "the connection fails at runtime" and "the peer was excluded during split planning" are different operator experiences, and the whole point of the generation token is to get the second one. Worth deciding now which way the stack lands.

Nits

• write_stage_message calls state.lowrank_k().map_err(io::Error::other)? purely for its error, discarding the value. It reads as dead code — a named validate_lowrank() (or at minimum a comment saying it's a validation call) would stop someone deleting it as unused.
• decode reads the per-token scale off the wire without checking it's finite. A NaN or inf scale produces NaN activations that propagate into the model rather than failing the frame. Not a memory-safety issue, and arguably the codec-identity check above subsumes it, but a is_finite() guard is one line.
• encode's per-row scale is max_abs / 127.0. I walked the degenerate inputs — denormal max_abs, inf, NaN dot products — and they all land safely via clamp plus Rust's saturating float-to-int cast, so no panic. Just noting I checked; no change wanted.

[danielwinterw]

Thanks for going through the allocation path — good to have that confirmed rather than assumed.

2. dtype() lost validation. Fixed in 19b189c and you're right that this was the worst of the three: it silently weakened a field that had been fully constrained, for every dtype rather than just the new one. dtype() now requires zero high bits unless the low byte is Lowrank, with a test that reserved = 0x0100 fails as it did before.

3. Generation. Bumped this branch to generation 6. Reasoning matches yours: #1409 is the base and merges first, a release can be cut between the two, and at that point a shipped gen-5 peer predates the reserved change. Better to exclude it at split planning than to have it fail on a frame. README says why generation 6 exists.

Nits: validate_lowrank() names the write-path validation call so it no longer reads as a discarded value; decode rejects a non-finite per-token scale rather than letting NaN into the model. Noted on the degenerate-input walk of encode — I'd reached the same conclusion via the saturating cast, and it's useful to have it independently checked.

1. Codec identity — the one I'm not doing here, and I want to be straight about why. I agree it's the failure that costs someone a day, and your interim (32-bit fingerprint over the .skbc bytes, exchanged once at connection setup, refuse on mismatch) is the right shape and genuinely independent of package-manifest work. I've held off because it needs a handshake exchange on the standalone path, which today has no negotiation at all — that's the same gap I just documented on #1409 — and I'd rather add that once, deliberately, than bolt a codec-specific exchange onto a path that will need a general one. It is the blocking item before this leaves draft and I've said so in the PR description.

If you'd rather see the fingerprint land now on the grounds that a silent-garbage failure outweighs the tidiness of doing negotiation once, say so and I'll do it — that's a reasonable call and I don't feel strongly enough to argue it.

[ndizazzo]

Closing as superseded by #1482 (commit 7369cfd). This draft encodes Lowrank and its rank through WireActivationDType and StageStateHeader.reserved. #1482 removed the dtype selector and reserved field, made raw F32 the only stage activation representation, and bumped the state and protocol contracts. The branch also records codec identity negotiation and real-model quality gates as blockers before it could leave draft. Reviving low-rank boundary compression now needs a fresh design against the current raw-F32 wire rather than merging this implementation.