feat(kimi-k3): serve DSpark drafts (fc_norm + AttnRes tap) - #1031
feat(kimi-k3): serve DSpark drafts (fc_norm + AttnRes tap)#1031torchspec-bot wants to merge 3 commits into
Conversation
A TorchSpec-trained K3 DSpark draft needs two things the Inferact reference checkpoint does not, and both fail silently -- they cost acceptance rate and raise nothing. fc_norm: the draft normalizes each target tap on its own before the taps are concatenated and projected. Its five weights ride in the checkpoint; without the modules they have no destination and context_proj sees inputs on a scale it was never trained for. The incremental-projection fast path re-projects from pre-split fc columns and would skip it, so it stands down for such a draft rather than quietly disagreeing with project_target_hidden. aux_hidden_stream: K3's residual backbone mixes across blocks, so the value a consumer layer actually reads is the pre-norm AttnRes mixture, not the running prefix sum. Drafts trained against that stream must be fed it. The checkpoint declares which one it wants, the target refuses a request it cannot honour, and TOKENSPEED_DFLASH_AUX_STREAM exists to A/B the claim. Default is unchanged. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
![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: 5f1e48b5d1
ℹ️ 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".
|
|
||
| import torch | ||
| from torch import nn | ||
| from tokenspeed_kernel.ops.layernorm.triton import segmented_rmsnorm |
There was a problem hiding this comment.
In deployments where the runtime wheel is upgraded while an older tokenspeed-kernel wheel satisfying the existing tokenspeed-kernel>=0.1.3.dev0 dependency remains installed, this unconditional import fails before K3 DSpark can even serve prefix-only checkpoints because segmented_rmsnorm only exists in the kernel half of this commit. Please raise the runtime dependency floor or add a guarded fallback so the new runtime cannot be paired with a kernel that lacks the symbol.
Useful? React with 👍 / 👎.
Signed-off-by: Torchspec bot <262938024+torchspec-bot@users.noreply.github.com>
1 participant
Summary
A K3 DSpark draft needs two things the Inferact reference checkpoint does not. Both currently fail silently — they cost acceptance rate and raise nothing, which is the reason this PR makes each one declared and enforced rather than inferred.
fc_norm— the draft normalizes each target tap on its own before the taps are concatenated and projected. Its five weights ride in the checkpoint; without the modules they have no destination andcontext_projsees inputs on a scale it was never trained for. A draft that declares"fc_norm": truegets the modules and ships onefc_norm.N.weightper tap; declaring it without the weights, or shipping weights without declaring it, now fails the load rather than serving an identity-weight norm. The incremental-projection fast path re-projects from pre-splitfccolumns and would skip the norms entirely, so it stands down for such a draft instead of quietly disagreeing withproject_target_hidden.aux_hidden_stream— K3's residual backbone mixes across blocks, so the value a consumer layer actually reads is the pre-norm AttnRes mixture, not the running prefix sum. A draft trained against that stream must be fed it. The checkpoint declares which stream it wants, the target refuses a request it cannot honour, and the resolved choice is logged next to the tap ids at startup:The default is unchanged — a checkpoint that declares neither field behaves exactly as before.
TOKENSPEED_DFLASH_AUX_STREAMoverrides the checkpoint for A/B runs.Test Plan
test/runtime/test_k3_dspark_torchspec_draft.py— 24 new tests covering config defaults and validation, stream resolution and env override,_wire_aux_hidden_stream,set_dflash_aux_hidden_stream,_dspark_capture_stream(prefix / attn_res / consumer weights / last layer / no blocks), fc_norm projection ordering and itsstrict=Truelength mismatch, and the incremental-projection stand-down. All use synthetic configs — no checkpoint required, so they run in ordinary CI.End-to-end, Kimi-K3 + a TorchSpec DSpark draft, TP16 across 4× GB200 (requires #1015 to start at all):
attn_res(what the draft trained against)Reproducibility caveat, stated plainly: the TorchSpec checkpoint these numbers come from is not public, so the end-to-end table above is not something a reviewer can independently re-run today. The unit tests are self-contained. Happy to hold this PR until a checkpoint is available if you'd rather not take serving-side evidence on trust.
A full-length GSM8K run is blocked on an unrelated engine hang under sustained load in the target's KDA prefill path — it reproduces with
stream=prefix, so it predates this change. Filing separately.