[diff-only] Mamba + KDA support: descriptor-driven state ISA, hybrid lowering - #76
Draft
qichao-arlo-wang wants to merge 17 commits into
Draft
[diff-only] Mamba + KDA support: descriptor-driven state ISA, hybrid lowering#76qichao-arlo-wang wants to merge 17 commits into
qichao-arlo-wang wants to merge 17 commits into
Conversation
The three packet widths the skew is tuned against - head, head-dim and state-dim lanes - lived only as constants inside the emulator while the Compiler carried them on the projection plan as swept DSE parameters. Nothing tied the two together, and the mapping CRC could not: it hashes the source-to-cell mapping, which does not depend on packet shape. Any sweep off 8/4/8 would have kept producing 8/4/8 packets on the emulator side, so every bank service and stall counter would have diverged from the DSE that motivated the layout, with no test failing. Field records already stop at byte 248 of the 256-byte descriptor, so the lanes go in a trailer there and the "unused bytes must be zero" rule now covers 200..248 plus the five bytes after the trailer. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The 0x38 sub-target that lets the policy register and the correction-bias address share one opcode is written independently in the routed-MoE emitter and in the emulator, and so is the packed expert shift. Neither the opcode table nor the descriptor contract covers them, so a one-sided edit would have written the wrong control register at runtime rather than failing a check. state_contract.py already walks the emulator source with --simulator-root; teach it these two constants as well. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The tail loop chose a layer type with a conditional whose two arms were both MAMBA, which reads as an unfinished dispatch: if a non-Mamba layer ever emitted an X_STATE event it would have been silently relabelled. Only Mamba layers emit state events, so state the invariant and fail on violation. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Collaborator
Author
|
These two PRs exist to make the diff easy to read and comment on. They are not going to be merged as-is. Opened as a pair so the cross-repo change can be reviewed together:
Please treat them as read-only review surfaces:
If any of this does get proposed for merge later, it should be as separate, scoped PRs with their own CI, not by un-drafting these. |
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
2 participants
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.
Michael's
feature/mamba-kda-support, opened as a draft so the change is browsableand commentable per-line. Ownership stays with @michael; nobody should merge or push
here on the basis of this PR existing.
17 commits · 104 files · +21,161 / −567 · tip
8bc274b(2026-08-23)What it adds
A descriptor-driven recurrent-state ISA shared by Nemotron-3 Mamba-2 and Kimi K3 KDA:
X_STATE(opcode0x3D) — one command per recurrent-state operation, 7 subops(
PRELOAD / RESET / PREFILL / STEP / COMMIT / EVICT / FENCE), 256-byte little-endiandescriptor with a common header and a per-algorithm payload.
L_SCATTER_M(opcode0x3F) — materialize a Matrix result into a programmablebanked producer→consumer layout (
ROW_MAJOR / TRANSPOSE / MAMBA_SKEW / KDA_SKEW / CUSTOM)before the matching
X_STATE.C_SET_TOPK_REGextension — atargetoperand plus sigmoid-normalized routeweights and correction-bias ranking, for Kimi K3's no-aux-loss MoE. Backward
compatible: omitting
targetreproduces the original machine word byte for byte.aten/state/,aten/mamba/,aten/kda/,aten/kimi3/,aten/nemotron3/— contract,lowering, projection scatter, and capacity-aware schedulers.
tools/state_contract.py— generates and checks the wire format against the Simulator.Companion
Simulator side is the paired draft in
PLENA_Simulator; its gitlink pins this branch.Read alongside
doc/nemotron3_mamba_isa.md— normative descriptor and lifecycledoc/plena_isa_spec.md— the three ISA changesdoc/kimi3_connected_scaling.md— the measured static-emitter boundary