[diff-only] Mamba + KDA support: common state engine, hybrid DSE - #113
Closed
qichao-arlo-wang wants to merge 15 commits into
Closed
[diff-only] Mamba + KDA support: common state engine, hybrid DSE#113qichao-arlo-wang wants to merge 15 commits into
qichao-arlo-wang wants to merge 15 commits into
Conversation
Two reporting defects in the layout ablation and the KDA reference. _row_major() rebuilt the baseline with replace(), which carried the skewed plan's physical_values_per_token over. The skewed layout widens each group span so field origins stay bank aligned, so the baseline was charged for padding it does not need - 50,688 against its own 49,248 values per token on KDA - and the published SRAM column flattered the layout under test. Cycle columns are unaffected: read service stays 53,176 -> 14,904 on Mamba and 536,544 -> 430,560 on KDA. The KDA reference folded the recurrent state and the conv state onto one storage parameter, but the descriptor carries state_precision and conv_state_precision independently and the shipped Kimi configuration is FP32 state with BF16 conv. That combination had no CPU reference at all. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Three changes to the layout engine, all about things no test could see. The lane widths were constants here while the Compiler swept them, so a non-default sweep would have silently kept 8/4/8 packets and diverged every bank counter. They now arrive in the descriptor trailer. mapping() and address() carried two copies of the same skew placement rule, and only mapping() feeds the bijection check and the CRC, so a divergence would have surfaced as a runtime "reads unwritten cell" far from its cause. address() is now the only placement site. The row comes from the physical linearisation and the bank from the field-local index, which agree only when the group span and every field origin are bank aligned. That precondition was unwritten and a violation reported the unrelated "layout aliases two sources". Every existing Mamba and KDA test configures the short convolution as identity on the newest tap with conv_bias absent, so the ring-buffer shift, the per-channel weight stride and the bias add were all unobservable - deleting any of them kept the suite green. Verified by mutation: removing the shift now fails. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… freeze fifo_capacity_values reaches only fifo_peak_values = min(burst, capacity) in the Rust flow; backpressure is a closed form over the spill width and the producer burst. No capacity can move the reported cycle count, so "64 and 256 produce the same result" is a tautology of that model rather than a measurement. The depth is justified by the analytic ProjectionFifoSpillModel, which steps a real queue and reports a measured high-watermark. Point both documents at it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Every point in the quick grid held hbm_bytes_per_cycle at 64 while sweeping Matrix width, state lanes, banks, FIFO depth and cache capacity. At that bandwidth the memory server is 97.75% busy on Nemotron and 96.75% busy on Kimi, so it absorbs nearly every on-chip win: the parameter that decides most of decode time was the only one never varied, and every other sensitivity in the report was quoted at a single unverified value. Adding the sweep changes what the ablations mean. The layout win rises monotonically with bandwidth - Nemotron 0.0540% at 64 B/cycle to 2.6950% at 8192 - so the headline 0.0540% cannot be quoted without its bandwidth. The conclusion survives the sweep: even at an 8 TB/s-class server the layout is worth 2.6950% on Nemotron and 0.2883% on Kimi, so it is not a system-level lever, and that is a property of the design rather than an artefact of a conservative placeholder. The two load-bearing mechanisms also invert between 512 and 1024 B/cycle. Saving state bytes matters most while bandwidth is scarce; on-chip bank conflicts only surface once it is not. A reader given only the 64 B/cycle column would put on-chip area in the wrong place. Kimi stays weight-streaming bound across the whole range: 36.757 ms TPOT and 18.66% HBM utilisation even at 8192 B/cycle. That single-device verdict does not depend on the bandwidth assumption. The pinned report hash and candidate count move accordingly, and the new section is covered by monotonicity assertions rather than by the file hash alone. 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. |
`ci.yml`'s lint job fails at `ruff format --check .`: 40 files would be reformatted. All 40 are files this branch touched and none are pre-existing -- `origin/main` is clean under the same ruff. Formatting only: `ruff check .` was already passing, so no lint rule was ever violated. Semantics-preserving, checked rather than assumed. The AST of all 173 tracked Python files is byte-identical before and after (`ast.dump`, which omits positions), so nothing in this diff can change behaviour. This only unblocks the rest of that job. Its "Run hybrid pre-RTL model guards" step runs after the format check and has therefore never executed on this branch, so whatever it reports next will be new information rather than a regression. The other failing job, `MoE stage-attribution guard`, is not this branch's fault and is not fixed here: it reads only `transactional_emulator.yml` and so cannot see the two guards this branch correctly wired into `ci.yml`. The fix belongs to that guard and is in PLENA_Simulator#112. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`test_every_guard_file_is_wired_into_ci` and `test_every_guard_is_reachable_from_ci` read only `.github/workflows/transactional_emulator.yml`. That it was the only workflow running guards from this tree was true when written and enforced by nothing. `feature/mamba-kda-support` adds a pytest step to `ci.yml` and wires two testbench guards into it. Both checks reported those guards as never run by CI: a correctly wired guard failing the check whose whole purpose is to notice unwired ones. One of the two was `test_hybrid_isa_freeze_sync.py`, the cross-repo guard for the frozen opcode allocation -- so the false positive landed on the wiring it is most costly to mistrust, and the reading it invites is "the guard is noisy", which is how a real one gets ignored later. Both now read the whole `.github/workflows` directory. Concatenating is deliberately permissive: the question is whether *some* job names the file, and which one is not these checks' business. `test_the_workflow_scan_reads_every_workflow_file` pins it. Nothing else would: narrowing the scan makes both checks *more* likely to pass, and only the guards named in the dropped file start reading as unwired -- which is the failure, not the signal. Verified against the failing tree, not just in the abstract: the fixed guard passes on `feature/mamba-kda-support` at e93b832. Injections kept honest -- dropping a guard from every workflow still fails; naming it only in `ci.yml` now passes; re-narrowing the scan fails the new test. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> (cherry picked from commit 72e6ac2) Applied partially. The original also retargets `test_every_guard_is_reachable_from_ci`, which does not exist on this branch -- it arrives with PLENA_Simulator#112, which is where this fix belongs. Cherry-picked here only so this branch's CI reflects its own state rather than a defect in a check it inherited; the duplicate resolves when #112 lands on main.
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.
13 commits · 138 files · +44,334 / −293 · tip
e93b832f(2026-08-23)What it adds
transactional_emulator/src/state_engine/— a common recurrent-state engineexecuting both Mamba-2 and KDA behind one
X_STATEcontract: descriptor decode,banked head-tile layout, residency cache, timing, precision, and per-algorithm
functional cores (
functional/mamba.rs,functional/kda.rs).analytic_models/performance/— Nemotron-3 and Kimi K3workload models, routing/MoE-event/sensitivity/formal sweeps, GPU microprofile
validators, and an HBM-bandwidth sweep.
models/nemotron3/mamba_connected_test.py,models/kimi3/kda_connected_test.py,models/kimi3/connected_blocks_test.py,bit-exact against CPU references.
test_x_state_contract_sync.py— cross-repo guard that this repo's generatedcontract matches the Compiler's spec, including a tampering test so a passing
check is evidence rather than a no-op.
Companion
Compiler side: AICrossSim/PLENA_Compiler#76. This branch's gitlink pins it.
Read alongside
doc/COMMON_STATE_ENGINE_DESIGN_ZH.md— why one engine serves both, and theGPU profiling that reshaped the design
doc/connected_hybrid_validation.md— what is validated and at what scale