[Draft · diff only] Emulator, stage tests and measurements for the static Mamba/KDA path - #115
Draft
qichao-arlo-wang wants to merge 6 commits into
Draft
[Draft · diff only] Emulator, stage tests and measurements for the static Mamba/KDA path#115qichao-arlo-wang wants to merge 6 commits into
qichao-arlo-wang wants to merge 6 commits into
Conversation
qichao-arlo-wang
force-pushed
the
feat/static-kda
branch
13 times, most recently
from
August 28, 2026 09:22
5157e19 to
c9da62e
Compare
…iling `V_FMA_VF` at 0x3B decodes and executes -- `Vector[rd] += Vector[rs1] * fp_reg<rs2>`, one accumulate over a VRAM row against an FPRAM scalar. It joins `V_SOFTPLUS_V` and `S_MAP_FP_V` from the compiler side, and it is what lets a recurrent sweep be one instruction inside one hardware loop instead of an unrolled triple per state row. `stage_profile.rs` attributes cycles to named stages, so a lowering can be read per kernel rather than as one total. Numeric formats move into config rather than being fixed in the source. **The clock stops being a hidden constant.** `runtime_config.rs` had `PERIOD` as a `const`; it is now a `LazyLock` reading `CLOCK_PERIOD_PS` from `plena_settings.toml`, and `runner` prints it at startup and checks it against the DRAM model's own clock, requiring one period to divide the other. The shipped value is 1000 ps in both modes, which is **an assumption, not a measurement** -- no RTL has been synthesised, so no critical path has set a frequency. It is in the file rather than the source so that every second, microsecond and TPOT figure the emulator reports can be traced to a value someone chose and changed without a rebuild. The HBM2 preset also being 1 ns is a coincidence of that preset, not a design relationship. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Stage-by-stage numeric tests that run the compiler's emitted machine code on the transactional emulator and compare against the float64 reference, per stage rather than per model, so a wrong answer names the emitter that produced it. The KDA tests carry `physical_rows` in their comparison parameters and take `--value-dim`, with tile geometry from `kda_prefill_tile_shapes` rather than a local guess -- the divergence between the two is how a spill came to prefetch past its allocation. **The tolerance is restated at 100% of values.** `check_mem`'s `allclose_pass = match_rate >= 90.0` let a real defect report PASSED: the prefill output scale was applied to column block 0 only, so at `value_dim` 128 half of every token was `sqrt(key_dim)` = 11.3x too large, and 94.68% of values still sat inside the case's `atol` because the data is mostly small. The shared harness and its 90% rule are left alone; the stage tests assert every value. Under the new bar four mutations of the scale sweep that survived the old one die. `isa_analysis.py` reports what an emitted program is made of, which is where the static-versus-issued distinction and the per-opcode census later in this branch started. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… the compiler `perf_model.py` gains `kda_chunk_prefill` and `kda_recurrence_decode`; `mamba2_model.py` gives Mamba-2 the same treatment; `hybrid_model.py` is a driver that parses NemotronH's `hybrid_override_pattern` and Kimi's `linear_attn_config` and costs a model layer by layer, by the operator each layer actually runs. **Calibrated, not asserted.** `test_kda_stage_calibration.py` turns the model's cycles back into instruction counts and compares them against what the compiler emits over nine shapes: 0.87-1.14 on prefill, 0.98 on decode. Getting there found a real structural error -- the spill term did not scale with column blocks -- which tightened the band from 0.75-0.96. **The roofline mismatch is real and has a mechanism.** Measured at one point it looked like noise; swept over MLEN x bandwidth it is a trend. On a wide machine KDA costs 2.9x what attention does per unit of arithmetic. The cause is KDA's channel-wise decay: it makes each decode sweep `key_dim` lanes wide -- 6.2% utilisation at VLEN 2048 -- so at MLEN 2048 about 45% of Kimi Linear's decode time is spent by an operator using 6.2% of the lanes it is charged for, while contributing 24.7% of the FLOPs. Mamba-2 has neither problem: one scalar decay per head gives an 8192-wide operand and a chunk of 256. **Prefill and decode want different machines.** The last doubling to MLEN 2048 returns +44.0% for decode and -52.1% for prefill. The mechanism is `chunk` against `mlen`: bf16 range on `1/A = exp(chunk * |gate_lower_bound|)` caps the chunk at 17, which leaves prefill's tiles 99% padding at wide `mlen`. Priced rather than asserted: row-granular prefetch is worth 1.4x at MLEN 64 and 12.5x at 2048, and moves prefill's optimum from 1024 to 2048 -- matching decode's +23.9% exactly. (An earlier draft said 13x, from 14,173 down to 1,063. That was the count with the whole spill term zeroed; row-granular still stores the live rows, so the floor is 1,137 and the factor is 12.5x.) **One caveat that limits all of it**, stated here and in the doc: MLEN and VLEN enter the model's cycles only through `ceil(work / MLEN)` instruction counts. Latencies are constants -- `V_FMA_VF` bills `VECTOR_MUL_CYCLES` whatever VLEN is -- so a wider machine means one instruction covering 32x the data at the same cost, with no area term anywhere. "MLEN 2048 makes decode 15x faster" is mostly a restatement of that assumption. Only operator-versus-operator comparisons at a fixed width survive it, and those are what is reported. `vector_sram_banking.py` models row-interleaved banking and finds 0.9-2.3% conflict residue on the recurrence and mixer sweeps -- and notes in its own docstring that word-interleaved banking, under which a row spans `min(mlen, B)` banks and no row-sequential sweep conflicts at any stride, makes the whole question moot. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… clock reads `just` targets for the Mamba-2 and KDA stage tests so they run the same way by hand and in CI, and a workflow that builds the emulator and runs them. `flake.nix` gains `clippy` and `rustfmt` in the `rust-bin` extensions. They were missing, which is why `cargo clippy` could not run in the dev shell -- not because the toolchain refused to install them. `plena_settings.toml` gains `CLOCK_PERIOD_PS = 1000` in both modes, with a comment saying in the file itself that it is an assumption and not a measurement: no RTL has been synthesised, so nothing has set a frequency. Every second, microsecond and TPOT figure downstream traces back to it. This is also the file the compiler had never managed to read -- see the compiler-side commit -- so the values here now reach it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… them Everything measured on this branch, with the plan it came from, split by how much each number can be asked to decide. **Hard facts** -- instruction counts, HBM bytes, memory footprints. Properties of the compiled artifact, reproducible, admissible as CI gates. **Uncalibrated model output** -- cycles, microseconds, TPOT, from a timing model never calibrated against silicon. Labelled on every line, and never compared against other uncalibrated model output: two numbers from two uncalibrated models decide nothing between them. There is a second split inside the hard facts, and this document got it wrong four times. A **static** count is the size of the program image; a **dynamic** count is what gets issued, with hardware loops expanded. Both are exact; only the second says anything about cost. | claim | image | issue stream | |------------------------------------------|----------:|--------------:| | projection gather, share of a KDA layer | 0.07% | **0.94%** | | convolutions vs mixer | 58% / 42% | **12% / 87%** | | what `V_FMA_VF` bought | 97x | **3.6x** | | Kimi K3 head vs 64x64, for 4x state | 1.6x | **2.3x** | The image counts are not retracted -- they are the right instrument for the CI gates, which catch a sweep falling off the hardware-loop path, an image failure no numeric test can see. What is retracted is reading any of them as a cost. A third question the counts alone could not answer: what the issued instructions *are*. Per opcode, a Kimi K3 layer's 492,681 are 124,428 arithmetic and 368,253 scaffolding -- every `V_FMA_VF` sits in a body of five of which one computes. A post-increment operand, priced only against pointer advances that have a consumer in the same loop body (215,634 foldable, 0 not), is worth 1.78x on its own and 2.83x with the FPRAM scalar auto-advancing too, taking the arithmetic share from 25.3% to 71.5%. It does not touch the 6.2% lane utilisation that KDA's channel-wise decay causes at VLEN 2048 -- that is a width mismatch, not an encoding one, and the two multiply on a wide machine. What the document records, in order: the three opcodes and what they bought; one KDA layer in both columns; where the issue slots go; memory; numeric agreement on the emulator; the banking question and what a follow-up would have to measure; prefill past one block, including the output-scale defect and the 90% bar that nearly hid it; what is **not** measured (no RTL, no synthesis, no energy model, no whole-model compilation); hybrid cost per operator with the free-width caveat stated before the numbers; and the clock, with what it is not. Also here: that packing does not cause the projection gather -- both lowerings were built and counted, identical `M_MM`, 0.12% worse separate, same gather -- and that reading at a stride removes 4,041 of its 4,650 issued instructions for nothing. And that the compiler had never read `plena_settings.toml` at all, for two independent reasons, which is why its `mram_tile_capacity` was 16x below the machine and cost the input projection 2.79x. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
qichao-arlo-wang
force-pushed
the
feat/static-kda
branch
from
August 29, 2026 00:42
5e63bf8 to
e37286b
Compare
The compiler side of this work: three opcodes, the Mamba-2 and KDA lowerings, the layer assembly, and the budgets. The emulator here decodes and executes `V_FMA_VF`, the testbenches run what that compiler emits, and `plena_settings.toml` in this repo is the file its configuration now comes from. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
qichao-arlo-wang
force-pushed
the
feat/static-kda
branch
from
August 29, 2026 08:17
e37286b to
9a1e607
Compare
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
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.
The emulator, testbench and measurement half of lowering Mamba-2 and KDA onto PLENA's
static ISA — no cache, no descriptors, no dynamic scheduling.
What is here
V_FMA_VFat0x3B—Vector[rd] += Vector[rs1] * fp_reg<rs2>, decode, execute,timing and resource classification. Two details worth reading:
mul_scalarstartsfrom the source because there the source is the base; for an accumulate, a masked-off
head must keep what the destination already had.
touches two vector rows like the VV ops. Grouping it with
V_MUL_VFwould under-reportexactly the traffic the conversion is meant to be judged on.
A cross-repo opcode guard.
decodecarried three comments saying encodings "muststay in sync with PLENA_Compiler's
doc/operation.svh", enforced by author diligencealone. It is now a test that parses that header. It immediately found two pre-existing
gaps —
V_PS_VandC_HADAMARD_TRANSFORMare declared by PLENA and implementednowhere — which are now an explicit exemption list carrying the reason.
KDA stage tests on the transactional emulator
(
transactional_emulator/testbench/kda/), nine cases: the cumulative decay, the UTtransform, chunked prefill (single and three chained chunks), the state layout
transpose, and the assembled layer (single and four layers).
Measurements —
doc/static_path_measurements.mdcollects every number in one placewith its scope stated, plus a vector-SRAM banking study and an explicitly uncalibrated
latency report.
Numbers
Kimi K3, 93 layers: 8,681,829 static instructions, 0.202 GiB of HBM, 492 of
512 FPRAM slots. One layer executes in 23,715 cycles; four chained layers in
96,114, linear.
0.202 GiB matters because of how the emulator loads memory: it preloads HBM from a flat
file starting at offset 0, so the allocation and the file both have to span every address
the program touches. A layout whose regions sit far apart needs a span-sized allocation
whatever the live data comes to, and stops being executable long before it stops being
describable. This one sits three orders of magnitude inside the configured 16 GiB.
Scope, stated plainly: 8.68M is the state-engine path only — no projections, MoE or
embeddings — so a whole-model figure is not comparable to it. No whole-model program has
been compiled here.
The worst thing found here, and it was mine
case_layer's tolerance wasatol=2e-2against a golden whose largest value is0.01215. Every compared value already sat inside the absolute bound, so writing
zeros into the output tile scored 100%. So did skipping the
vconvolution, swappingthe
qandkgather sections, and zeroing every layer but the last.layer_chainalso compared only the last layer — the code to pack every layer existedbut was gated on
if heads == 1 else None, and the only recorded invocation was--num-heads 3. Zeroing the first three layers' inputs left the verdict bit-identical,which means the claim that "the error does not chain" was never a measurement.
Both are fixed and all four mutations now fail. The honest numbers the absolute bound was
hiding: 4.3% mean relative error, which is what bf16 gives on the read-out's 128-term
contraction. The lowering itself is right to 5e-4 in float64, checked separately.
Also corrected here
V_FMA_VF1 cycle instead of 5 — it does not start with"V_MUL"and fell through to a catch-all. A 5× under-count, in the one tool Phase 4measures with, on the one opcode being measured.
analytic_models/performance/customISA_lib.jsonhad no entry, andperf_model.pyraises
KeyErroron a miss rather than defaulting.the recurrent state and the projection in VRAM while the emulator test staged them from
HBM. Two tests disagreeing about where a layer's state lives.
Update: multi-block shapes, and a bar a real defect fails (commit
8f617d7c)The prefill and transpose cases now run at Kimi K3's 128 x 128 head against
mlen64,not only at
key_dim == value_dim == mlen. All 17 cases pass with 100% of valuesinside tolerance. Two things had to change for that to mean anything.
_write_comparisonnow passesphysical_rows. A VRAM matrix wider thanmleniscolumn-block-major with block pitch
physical_rows * mlen;check_mem's stride modealready takes that as
col_block_stridebut defaults it tonum_batches, which is thelogical row count and is smaller for every padded tile. Without it the second block is
read from the wrong place.
And the cases now assert that every value is inside the tolerance they asked for.
check_mem'sallclose_passismatch_rate >= 90.0— a shared utility with manycallers, left alone — and under it the compiler's output-scale defect printed PASSED
at a 94.68% match rate, because the data is O(1e-3) and the tolerance was 5e-2. Under the
new bar, four mutations of the column-block loops die, and only one of them needed it:
k_hat/q_haton block 0 onlyA_C * stateon block 0 onlyk / Aon block 0 onlyThe layer cases move to
atol2.5e-4 alongside theirrtol0.12, for the mirror-imagereason: three of their sixty-four outputs are near 1e-3 and bf16 delivers about 1.8e-4
absolute whatever the value, so a pure relative bound is unmeetable there. They had
been passing at 95.31% on the 90% rule rather than on the tolerance. The floor keeps the
check real — an all-zero output scores 1.6% against it.
This is the second time in this branch a tolerance larger than its signal has hidden a
real error; the first is the section below.
key_dimxvalue_dimTests
Emulator: 112 passed. Compiler suite (in the submodule): 377 passed, 2 pre-existing
failures — both verified to fail on
origin/mainas well. A third that this branch hadbeen labelling pre-existing was not: it was introduced here and CI caught it. See the
compiler PR's Tests section, and the two Correction sections at the end of
docs/superpowers/progress/2026-08-25-static-mamba-kda.md. All 17 KDA stage cases pass on the transactional emulator at 100% of values intolerance.
The stage tests are scripts, not pytest, and are not in CI — run them with:
Two things a reader should know before opening the diff
The
PLENA_Toolsbump is gone, and this description used to undersell why.It said the submodule "will show as broken here". It did more than that:
08db1bb9is on no branch of
AICrossSim/PLENA_Tools, sogit submodule update --init --recursivecannot fetch it andactions/checkoutfailed outright — all threeCI jobs died before running a step, so the emulator was never built or tested in CI
on this branch. The same would happen to any reviewer cloning it.
The gitlink is back to
6b31fe00, matchingorigin/main. Checked rather thanassumed: all six KDA stage cases pass against it, the only
check_mem.pydifferencebetween the two is a NaN/Inf rule these stage tests already exceed with their own
100% bar, and the
physical_rows/col_block_stridesupport the multi-blockcomparison needs is present at
6b31fe00already.If that Tools commit is wanted, push it to
AICrossSim/PLENA_Toolsand bump again —an unreachable pin is not something a branch can carry.
The stage tests rewrite their build's TOML to Plain BF16 KV types. Under the
shipped
[TRANSACTIONAL.PRECISION], spilled activations decode as e4m3 and the readwalks into the scale stream, whose
0x7fbytes are NaN. The emitters now refuse thisat compile time rather than producing
nan.CI, after the regroup
Both repos' CI was red and every failure was introduced by this branch. Recorded in
full in the Correction sections at the end of
docs/superpowers/progress/2026-08-25-static-mamba-kda.md; in brief:MoE stage-attribution guardstage. Passes onorigin/main; it was not the pre-existing failure this description had been calling it.Generator unit testspython3 -m pyteststeps added to a job that only installstorchandtransformers.actions/checkout(2 jobs)PLENA_Toolsgitlink, above.cargo fmt --all -- --checkop.rsandvector_machine.rs. The Nix shell shipsrustfmtbut not thecargo fmtsubcommand, and I read that as "cannot check locally".rustfmt --edition 2024 --checkworks there and reproduces CI exactly.cargo clippyis still unverified locally — it is genuinely absent from the dev shell,and CI's clippy step had never run because formatting failed ahead of it.
Rebased onto
main's scoreboard timing model (#114)This PR had no checks for a while and I reported the repo's Actions as backed up.
That was wrong:
mergeable_statewasdirty. Apull_requestworkflow runsagainst
refs/pull/115/merge, and GitHub does not create a run at all when thatref cannot be computed — which looks identical to a stalled queue from outside.
mainhad moved one commit ahead: #114, "pipelined--timing-modelscoreboardwith real DMA overlap". This branch is now rebased on top of it, and the merge
was not mechanical.
#114 deleted
classify_timing_accessfromdispatch.rsand replaced it withthe
OpAccessmodel in a newaccelerator/access.rs. Of this branch's sevenhunks in
dispatch.rs(139 lines, all insertions), four still applied — twoexecution arms and two
resource_kind_for_opcodearms — and three had to bere-derived. The three opcodes are classified in
access.rsnow:V_FMA_VFgets its own arm rather than joining theV_*_VFfamily,because it reads its destination as well as its source. It lists
vector(gp(rd), vector_tile)inreads; grouping it with the othersunder-reports its vector-SRAM traffic by a row, which is the number the FMA
conversion is judged on.
V_SOFTPLUS_VjoinsV_EXP_V/V_RECI_Vunchanged.S_MAP_FP_VisUnit::Vector, notScalar: it holds the vector SRAM readport for a whole row even though its destination is FP_MEM. Its
S_MAP_V_FPmirror is the inverse and stays
Scalar.The exhaustive match is what forced the re-derivation rather than a silent
omission — adding an opcode to
op.rsfails the build untilaccess.rsclassifies it.
main's dispatch tests moved topipeline_tests.rs, so thisbranch's
mod testsindispatch.rscarries only the two guards that still meansomething after the merge.
The merge changed no numerics. All 17 KDA stage cases pass with max errors
byte-identical to the pre-merge run —
state_transpose0.00,prefill_out6.71e-04,
prefill_state9.28e-03 at 128 x 128, and the chained cases likewise.flake.nix: clippy and rustfmt were never in the toolchainCI gates on
cargo fmt --all -- --checkandcargo clippy --workspace --all-targets -- -D warnings, and the dev shell couldrun neither —
rust-bin.stable.latest.default.overridelisted onlyrust-srcand
rust-analyzer.cargo fmt's error ("not installed for the toolchain, runrustup component add rustfmt") is cargo's generic advice and misleading here:there is no rustup in this shell. I read it as "cannot be checked locally", which
is how a formatting failure reached CI and why clippy had never run on this work
at all.
Both are in the extensions now. On the rebased tree: fmt clean, clippy clean,
129 Rust tests passing (up from 112; the difference is #114's pipeline tests
plus this branch's).
Hybrid models: cost per operator, and two results that are not the ones wanted
analytic_models/performance/hybrid_model.pyreads a per-layer type list out ofa model config, dispatches each layer to the operator it actually runs, and
reports FLOPs, cycles and HBM bytes per operator. The three sibling drivers
cannot: each is
overall_exe_cycle += block_cycles * self.num_hidden_layers, oneblock shape repeated.
PerfModelwas already operator-level, so this was theonly missing piece.
FLOPs are counted in the driver, not taken from
PerfModel. They are a propertyof the model and the cost model is a property of the machine — deriving both from
the same code would make the comparison circular, and that comparison is the
point.
The KDA cost model has an oracle, and calibrating it found an error
kda_chunk_prefillandkda_recurrence_decodeare unlike every other stage inperf_model.py: the lowering exists, so the compiler can be asked how manyinstructions it emits. Setting every opcode latency to 1 turns the formula into a
count of issues, and
test_kda_stage_calibration.pycompares that across nineshapes.
The calibration did not confirm the formula, it broke it. The model scaled with
key_dimat half the compiler's rate, worst atkey_dim128 /value_dim64(0.75). The cause was specific: a VRAM matrix wider than
mleniscolumn-block-major, so a spill's zero-fill and store walk every block, and
the term billed one. Fixing that moved the range from 0.75–0.96 to 0.87–1.14.
Nothing else in this repository could have found that.
The test states what it cannot see, which matters as much. Zeroing each term:
At instruction-count granularity a matmul is one
M_TMMissue, so the terms thatdominate cycles are invisible to the oracle and stay as unvalidated as every
other stage in the file.
Result 1 — the mismatch is real, and sweeping context was the wrong axis
An earlier version of this description said the opposite: that the FLOP share and
the cycle share track each other and no operator is disproportionately expensive.
That was measured at the shipped configuration only, and context was swept
instead of the axes that matter. Context changes which operator dominates
without changing what any of them costs per unit of arithmetic:
--sweepmoves machine width and HBM bandwidth instead — neither is known, so asingle point reports the plane rather than the design. Kimi Linear, decode:
MLEN64, any bandwidthMLEN2048, 2048 B/cycleOn a narrow machine the three sit within 0.6–1.1 and nothing stands out. On a
wide one KDA costs 2.9× what attention does per unit of arithmetic. Even the
single-point reading had the signal in it — 1.27 against 0.80 — and it was read
as noise around 1.0 rather than the narrow end of a trend.
Result 2 — the answer depends on which
MLENis realANALYTICMLEN 2048TRANSACTIONALMLEN 64At short context the wide machine makes KDA look disproportionately expensive and
the narrow one makes it exactly proportionate — a qualitative disagreement
produced by a config field.
plena_settings.tomldeclaresMLEN2048 against64,
BLEN128 against 4,HLEN128 against 16. Mixing them is not survivableeither: overriding
MLENto 64 while leavingHLENat 128 makesflash_attentiondivide bymlen // hlen= 0.No times are reported anywhere.
runtime_config.rshard-codesPERIODat 1 nswith no stated basis, so cycles stay cycles.
How the configs were obtained
Fetched from Hugging Face, not written, and both layer-type encodings read out of
each model's own source rather than inferred — NemotronH's
pattern_mapping = {"M": "mamba", "E": "moe", "*": "attention"}with-forMLP, and Kimi's
is_kda_layerbeing(layer_idx + 1) in kda_layers, one-basedagainst a zero-based index. Also read rather than assumed: NemotronH's MLP is
down(act(up(x))), two matmuls and no gate, against Kimi's gated three — a gatedcount applied to NemotronH would overstate every feed-forward layer by 50%.
KdaShape.kimi_k3()matches none of them at model scale (hidden 7168 / 96 heads/ 93 layers against the real 2304 / 32 / 27). The per-head shape does match, so
the kernel measurements above stand; the model-scale label does not.
Correction: the ISA cost is three instructions, not one
Relative to
main, which stops at0x38, this work adds three:V_SOFTPLUS_V(
0x39) andS_MAP_FP_V(0x3A) for the Mamba path, andV_FMA_VF(0x3B) for bothrecurrences. Earlier text here said one, counting only
V_FMA_VF— the one theperformance argument turns on — and omitting the other two.
The argument was never about the count. All three are ordinary fixed-function ops with
their operands named in the instruction word, which is the property a statically
scheduled machine needs; what it cannot have is an instruction that fetches its own work
descriptor or a cache that decides residency at run time.
aten/tests/test_no_state_engine.pypins all three by slot and keeps
0x3Dand0x3Ffree.The mechanism, and the clock
Why KDA is the expensive one, and Mamba-2 is not
Vector-lane utilisation, and it is arithmetic rather than speculation. A vector op
covers
ceil(width / VLEN) * VLENlanes and useswidthof them:VLEN64VLEN2048value_dim= 128head_dim × state_size= 8,192KDA's decay is per key channel, so every key row carries its own scalar and
must be its own vector operation — 128 lanes wide, whatever the machine is.
Mamba-2's decay is one scalar per head, so its whole state block scales in a
single op and fills any width.
That is the same fact this branch has turned on from the start: the channel-wise
decay is what makes KDA more expressive than a scalar-gated SSM, and it is also
what makes it waste a wide vector unit. Nemotron-3 Super confirms the negative
half — its Mamba ratio moves only 0.60 to 0.84 across the whole plane while its
MoE sits flat at 1.09.
So the finding is not "recurrent operators are disproportionately expensive".
It is "an operator whose vector operand is narrower than the machine is".
The clock is a configuration value now
runtime_config.rscarriedconst PERIOD: Duration = Duration::from_nanos(1)—a bare constant with no stated basis, and every second, microsecond and TPOT
figure the emulator can produce is its consequence. It reads
[<MODE>.CONFIG.CLOCK_PERIOD_PS]now, defaulted to the same 1000 ps so nothingmoves, with the assumption written where it can be read: no RTL has been
synthesised, so no critical path has set a frequency.
Two things beyond moving the value:
It is printed at startup with its provenance. Every run begins
Clock: 1000 ps (1.000 GHz) from CLOCK_PERIOD_PS -- an assumption, not a synthesised frequency. DRAM model: 1000 ps (1x).Its relationship to the DRAM clock is asserted, not assumed.
stage_profile.rshad already recorded the hazard: a comparison there "onlyheld because the DRAM tCK happened to equal PERIOD; any preset or frequency
change made it fail". They are equal, and that is a property of the HBM2 preset
rather than of the design.
runnerrequires one period to divide the other andrefuses to start otherwise — verified by setting it to 700 ps and watching it
refuse.
Still no times are reported anywhere. Making a number configurable does not make
it known.
Verification: 129 Rust tests,
cargo fmtandcargo clippyclean, 48 Pythonguard and analytic tests, ruff clean, and all 17 KDA stage cases at 100% of values
in tolerance — max errors byte-identical to before the clock change, which is what
a 1000 ps default should give.
How wide should the machine be
The ratio grid says an operator's share of the time gets worse on a wider
machine. It does not say whether the machine is faster, which is the question a
width decision turns on.
--sweepreports that too.Width and bandwidth have to scale together, or width buys nothing. Marginal
return of doubling
MLEN, as a percentage of total decode cycles removed, KimiLinear:
The zeros are the finding. Below the diagonal a wider machine removes no
cycles at all — the memory server is already the limit and the extra lanes idle.
On Nemotron-3 Super the same cells go slightly negative (−0.1% to −0.5%): a
wider machine is fractionally slower when bandwidth does not follow it.
Where width stops paying depends on which recurrent operator the model runs.
At matched bandwidth, the last doubling tested (1024 → 2048):
MLEN64 → 2048The Mamba model is still returning nearly half per doubling at the widest point
tested. The KDA model has decayed to a quarter, and the second column is why: as
the machine widens KDA goes from a quarter of decode time to nearly half, while
contributing 24.7% of the FLOPs and running at 6.2% lane utilisation. Mamba-2's
share barely moves because its state block fills any width.
So at
MLEN2048, roughly 45% of Kimi Linear's decode time is spent by anoperator using 6.2% of the lanes it is charged for.
What this does not say. There is no area model here, so "wider is better
wherever it still helps" is trivially true and not a recommendation. What is not
trivial is the shape: the decay rate is set by the narrowest operand in the
model, so two hybrids with the same layer census and different recurrent
operators sit at different points on the same curve.
Prefill and decode want different machines — on one model, not the other
Total cycles at matched bandwidth, Kimi Linear:
MLENDecode's optimum is the widest point tested. Prefill's is 1024, and 2048 is 52%
worse. A single fixed width cannot be right for both.
Nemotron-3 Super does not have this problem. Both its modes want the widest
machine tested — prefill +59.2% and decode +44.0% at the last doubling. So the
divergence is not a property of hybrid models, of recurrent mixers, or of
prefill. It is a property of KDA.
The mechanism:
chunkagainstmlen, and a bf16 range limitDecomposing KDA prefill by machine width, one chunk and one head, instructions
issued:
MLENThe real work gets cheaper as the machine widens — 1,696 down to 1,063, which
is what wider tiles should do. The spill term grows from 754 to 13,110 and swamps
it.
Every spill zero-fills the rows past its live data, because
load_sub_matrix_*prefetches a whole
mlen × mlenMRAM block unconditionally andk_block_countselects whole blocks and cannot trim a partial one. The fill scales with
mlenwhile the live data scales with
chunk— and KDA's chunk cannot grow to meetit:
kda_chunk_check_rangecaps it at 17 because1/Areachesexp(chunk × |gate_lower_bound|)and overflows bf16 past that. Atmlen2048that is 2,048 rows zeroed to protect 16, six times per chunk.
Mamba-2's chunk is 256 and nothing bounds it numerically — the SSD scan has no
reciprocal decay to overflow — which is why its prefill keeps improving where
KDA's turns around.
Both of KDA's costs on a wide machine come from the same design choice. The
decay is channel-wise on the key axis, which makes each decode sweep 128 lanes
wide (6.2% utilisation at
VLEN2048) and forces the chunk form to divide by aper-channel cumulative decay, whose reciprocal caps the chunk and leaves prefill's
tiles 99% padding. The property that makes KDA more expressive than a
scalar-gated SSM costs it in both modes, by two different routes.
The ISA consequence, priced rather than asserted
kda_chunk_prefill(row_granular_prefetch=True)models the instruction set thatdoes not have that constraint — a prefetch taking a row count. It is off by
default and a test pins that: a hypothetical leaking into the default would
move the oracle comparison against a compiler emitting a different instruction
set, and the calibration would be measuring the wrong machine.
Instructions issued, one chunk and one head:
MLENWorth almost nothing on a narrow machine and a great deal on a wide one,
which is the shape a term scaling with
mlenhas to have. Anyone quoting 12.5×as "the value of this instruction" should say which column they mean first.
(An earlier version of this section said 13×, from 14,173 to 1,063. That was the
count with the whole spill term zeroed; row-granular prefetch still stores the
live rows, so the floor is 1,137 and the factor is 12.5×. Predicting a number and
then measuring it is how the difference showed up.)
And it makes the two modes agree on width
Total prefill cycles, Kimi Linear, matched bandwidth:
MLENThe optimum moves from 1024 to 2048 and the last doubling returns +23.9%
instead of −52.1% — which is decode's +23.9% to the digit. With the fill removed,
both modes are limited by the same dense work and the width disagreement
disappears.
So the mode divergence is not a fact about KDA's algorithm. It is a fact
about KDA's algorithm on this instruction set: a chunk capped at 17 by bf16
range, times a prefetch granularity of one whole block. The first cannot be
moved. The second can.
That is the opposite shape from a descriptor-driven state instruction — it asks
for less to be implied by one instruction, not more. Nothing here implements
it.
Correction: width is free in this model, so absolute speedups are not findings
MLEN,BLEN,HLENandVLENare the tile sizesdoc/plena_isa_spec.mddefines — the matrix machine's, the systolic array's, the partitioned array's,
and the vector machine's. "Narrow" above is
TRANSACTIONAL(64/4/16/64), "wide"is
ANALYTIC(2048/128/128/2048).They enter cycles only through instruction counts. In
perf_model.py:and the latency is a constant —
V_FMA_VFbillsVECTOR_MUL_CYCLESwhateverVLENis,M_MMbillsBLEN. So a wider machine covers 32× the data perinstruction at the same cost, with no area term anywhere in the model.
That makes the absolute numbers above near-tautological. "Widening from 64 to
2048 makes decode 15× faster" mostly restates the assumption; it should not be
quoted as a result, and the module docstring now says so where a reader meets it.
What survives is the comparison between operators. KDA and Mamba-2 are priced
by the same idealisation, so when one saturates on width at 23.9% marginal return
and the other is still at 44.0%, that difference is a property of the two
operators and not of the free throughput both were granted. Every result reported
from this driver is of that shape — a ratio between operators, or a turning
point — never an absolute speedup.
A second inconsistency, found while writing that down. The emulator's matrix
core is a hard-coded
big_4x1024— 4,096 PEs inmatrix_core.rs, independent ofMLENentirely. The two halves model width by different mechanisms: instructioncount on the analytic side, a fixed PE array on the emulator side. That is
separate from the
MLENdisagreement inplena_settings.tomland equallyunresolved.