Skip to content

[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
mainfrom
feat/static-kda
Draft

[Draft · diff only] Emulator, stage tests and measurements for the static Mamba/KDA path#115
qichao-arlo-wang wants to merge 6 commits into
mainfrom
feat/static-kda

Conversation

@qichao-arlo-wang

@qichao-arlo-wang qichao-arlo-wang commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator

Draft — for reading the diff only. Not intended to be merged from here.
No review is being requested and nothing here is ready to land. Paired with
AICrossSim/PLENA_Compiler#77, which carries the compiler half; the two must be read
together, and the gitlink in this branch points at that branch's head.

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_VF at 0x3BVector[rd] += Vector[rs1] * fp_reg<rs2>, decode, execute,
timing and resource classification. Two details worth reading:

  • The masked path starts from the destination, not the source. mul_scalar starts
    from the source because there the source is the base; for an accumulate, a masked-off
    head must keep what the destination already had.
  • It gets its own timing arm rather than joining the single-read VF family, because it
    touches two vector rows like the VV ops. Grouping it with V_MUL_VF would under-report
    exactly the traffic the conversion is meant to be judged on.

A cross-repo opcode guard. decode carried three comments saying encodings "must
stay in sync with PLENA_Compiler's doc/operation.svh", enforced by author diligence
alone. It is now a test that parses that header. It immediately found two pre-existing
gaps — V_PS_V and C_HADAMARD_TRANSFORM are declared by PLENA and implemented
nowhere — 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 UT
transform, chunked prefill (single and three chained chunks), the state layout
transpose, and the assembled layer (single and four layers).

Measurementsdoc/static_path_measurements.md collects every number in one place
with 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 was atol=2e-2 against a golden whose largest value is
0.01215. Every compared value already sat inside the absolute bound, so writing
zeros into the output tile scored 100%
. So did skipping the v convolution, swapping
the q and k gather sections, and zeroing every layer but the last.

layer_chain also compared only the last layer — the code to pack every layer existed
but 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

  • The profiler charged V_FMA_VF 1 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 4
    measures with, on the one opcode being measured.
  • analytic_models/performance/customISA_lib.json had no entry, and perf_model.py
    raises KeyError on a miss rather than defaulting.
  • The layer's HBM figure counted one seventh of the traffic: the measuring helper put
    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 mlen 64,
not only at key_dim == value_dim == mlen. All 17 cases pass with 100% of values
inside tolerance.
Two things had to change for that to mean anything.

_write_comparison now passes physical_rows. A VRAM matrix wider than mlen is
column-block-major with block pitch physical_rows * mlen; check_mem's stride mode
already takes that as col_block_stride but defaults it to num_batches, which is the
logical 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's allclose_pass is match_rate >= 90.0 — a shared utility with many
callers, 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:

mutation match rate caught by the 90% rule?
output scale on block 0 only (the real defect) 94.68% no
k_hat/q_hat on block 0 only 0.00% yes
A_C * state on block 0 only 55.79% yes
k / A on block 0 only 85.80% yes

The layer cases move to atol 2.5e-4 alongside their rtol 0.12, for the mirror-image
reason: 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.

case key_dim x value_dim max abs error
state layout transpose 128 x 128 0.00
prefill, out 128 x 128 6.71e-04
prefill, state 128 x 128 9.28e-03
three chunks chained, out 128 x 128 6.10e-04
three chunks chained, state 128 x 128 7.32e-03

Tests

Emulator: 112 passed. Compiler suite (in the submodule): 377 passed, 2 pre-existing
failures — both verified to fail on origin/main as well. A third that this branch had
been 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 in
tolerance.

The stage tests are scripts, not pytest, and are not in CI — run them with:

nix develop . --command bash -c 'cd transactional_emulator/testbench/kda && \
  uv run python kda_stage_test.py --case layer --num-heads 3 --key-dim 128'

Two things a reader should know before opening the diff

  1. The PLENA_Tools bump is gone, and this description used to undersell why.
    It said the submodule "will show as broken here". It did more than that: 08db1bb9
    is on no branch of AICrossSim/PLENA_Tools, so git submodule update --init --recursive cannot fetch it and actions/checkout failed outright — all three
    CI 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, matching origin/main. Checked rather than
    assumed: all six KDA stage cases pass against it, the only check_mem.py difference
    between the two is a NaN/Inf rule these stage tests already exceed with their own
    100% bar, and the physical_rows / col_block_stride support the multi-block
    comparison needs is present at 6b31fe00 already.

    If that Tools commit is wanted, push it to AICrossSim/PLENA_Tools and bump again —
    an unreachable pin is not something a branch can carry.

  2. 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 read
    walks into the scale stream, whose 0x7f bytes are NaN. The emitters now refuse this
    at 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:

repo failure cause
compiler MoE stage-attribution guard three emitters this branch adds defaulted their stage. Passes on origin/main; it was not the pre-existing failure this description had been calling it.
compiler Generator unit tests sixteen python3 -m pytest steps added to a job that only installs torch and transformers.
simulator actions/checkout (2 jobs) the unreachable PLENA_Tools gitlink, above.
simulator cargo fmt --all -- --check three hunks in op.rs and vector_machine.rs. The Nix shell ships rustfmt but not the cargo fmt subcommand, and I read that as "cannot check locally". rustfmt --edition 2024 --check works there and reproduces CI exactly.

cargo clippy is 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_state was dirty. A pull_request workflow runs
against refs/pull/115/merge, and GitHub does not create a run at all when that
ref cannot be computed — which looks identical to a stalled queue from outside.

main had moved one commit ahead: #114, "pipelined --timing-model scoreboard
with real DMA overlap". This branch is now rebased on top of it, and the merge
was not mechanical.

#114 deleted classify_timing_access from dispatch.rs and replaced it with
the OpAccess model in a new accelerator/access.rs. Of this branch's seven
hunks in dispatch.rs (139 lines, all insertions), four still applied — two
execution arms and two resource_kind_for_opcode arms — and three had to be
re-derived. The three opcodes are classified in access.rs now:

  • V_FMA_VF gets its own arm rather than joining the V_*_VF family,
    because it reads its destination as well as its source. It lists
    vector(gp(rd), vector_tile) in reads; grouping it with the others
    under-reports its vector-SRAM traffic by a row, which is the number the FMA
    conversion is judged on.
  • V_SOFTPLUS_V joins V_EXP_V / V_RECI_V unchanged.
  • S_MAP_FP_V is Unit::Vector, not Scalar: it holds the vector SRAM read
    port for a whole row even though its destination is FP_MEM. Its S_MAP_V_FP
    mirror 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.rs fails the build until access.rs
classifies it. main's dispatch tests moved to pipeline_tests.rs, so this
branch's mod tests in dispatch.rs carries only the two guards that still mean
something 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_transpose 0.00, prefill_out
6.71e-04, prefill_state 9.28e-03 at 128 x 128, and the chained cases likewise.

flake.nix: clippy and rustfmt were never in the toolchain

CI gates on cargo fmt --all -- --check and
cargo clippy --workspace --all-targets -- -D warnings, and the dev shell could
run neither — rust-bin.stable.latest.default.override listed only rust-src
and rust-analyzer. cargo fmt's error ("not installed for the toolchain, run
rustup 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.py reads a per-layer type list out of
a 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, one
block shape repeated. PerfModel was already operator-level, so this was the
only missing piece.

FLOPs are counted in the driver, not taken from PerfModel. They are a property
of 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_prefill and kda_recurrence_decode are unlike every other stage in
perf_model.py: the lowering exists, so the compiler can be asked how many
instructions it emits. Setting every opcode latency to 1 turns the formula into a
count of issues, and test_kda_stage_calibration.py compares that across nine
shapes.

modelled / measured
prefill, chunk 4–16 × key and value 64–128 0.87 – 1.14
decode, one step at 128 × 128 0.98

The calibration did not confirm the formula, it broke it. The model scaled with
key_dim at half the compiler's rate, worst at key_dim 128 / value_dim 64
(0.75). The cause was specific: a VRAM matrix wider than mlen is
column-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:

prefill, whole spill term removed        0.70   caught
prefill, UT forward substitution removed 0.91   NOT caught
prefill, all seven matmuls removed       0.91   NOT caught
decode,  scalar sweep overhead removed   0.39   caught
decode,  the FMA itself removed          0.69   caught

At instruction-count granularity a matmul is one M_TMM issue, so the terms that
dominate 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:

context KDA %FLOP %cycles ratio attention %FLOP %cycles ratio
4,096 24.7% 31.5% 1.27 8.8% 7.1% 0.80
1,048,576 2.1% 2.9% 1.37 92.1% 91.4% 0.99

--sweep moves machine width and HBM bandwidth instead — neither is known, so a
single point reports the plane rather than the design. Kimi Linear, decode:

operator MLEN 64, any bandwidth MLEN 2048, 2048 B/cycle
KDA 0.96 1.80
MoE 1.11 0.76
attention 0.64 0.63

On 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 MLEN is real

context KDA ratio, ANALYTIC MLEN 2048 KDA ratio, TRANSACTIONAL MLEN 64
4,096 1.27 0.96
1,048,576 1.37 1.46

At 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.toml declares MLEN 2048 against
64, BLEN 128 against 4, HLEN 128 against 16. Mixing them is not survivable
either: overriding MLEN to 64 while leaving HLEN at 128 makes
flash_attention divide by mlen // hlen = 0.

No times are reported anywhere. runtime_config.rs hard-codes PERIOD at 1 ns
with 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 - for
MLP, and Kimi's is_kda_layer being (layer_idx + 1) in kda_layers, one-based
against 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 gated
count 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 at 0x38, this work adds three: V_SOFTPLUS_V
(0x39) and S_MAP_FP_V (0x3A) for the Mamba path, and V_FMA_VF (0x3B) for both
recurrences. Earlier text here said one, counting only V_FMA_VF — the one the
performance 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.py
pins all three by slot and keeps 0x3D and 0x3F free.


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) * VLEN lanes and uses width of them:

operand width per vector op VLEN 64 VLEN 2048
KDA, one key row value_dim = 128 100% 6.2%
Mamba-2, one state block head_dim × state_size = 8,192 100% 100%
Kimi hidden activation 2,304 100% 56.2%
MoE intermediate 1,024 100% 50.0%

KDA'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.rs carried const 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 nothing
moves, 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.rs had already recorded the hazard: a comparison there "only
held 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. runner requires one period to divide the other and
refuses 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 fmt and cargo clippy clean, 48 Python
guard 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. --sweep reports 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, Kimi
Linear:

doubling 64 B/cyc 256 1024 4096
64 → 256 14.4% 74.6% 75.5% 75.5%
256 → 512 0.0% 6.2% 43.0% 43.0%
512 → 1024 0.0% 0.0% 36.8% 36.8%
1024 → 2048 0.0% 0.0% 13.4% 23.9%

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):

marginal return recurrent operator's share of decode cycles, MLEN 64 → 2048
Kimi Linear (KDA) 23.9% 23.9% → 44.5%
Nemotron-3 Super (Mamba-2) 44.0% 10.7% → 15.1%

The 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 MLEN 2048, roughly 45% of Kimi Linear's decode time is spent by an
operator 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:

MLEN decode (Mcycles) prefill (Mcycles)
512 8.10 1,856
1024 5.12 1,742
2048 3.89 2,650

Decode'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: chunk against mlen, and a bf16 range limit

Decomposing KDA prefill by machine width, one chunk and one head, instructions
issued:

MLEN total without the spill term spill share
64 2,450 1,696 30.8%
512 4,390 1,072 75.6%
2048 14,173 1,063 92.5%

The 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 × mlen MRAM block unconditionally and k_block_count
selects whole blocks and cannot trim a partial one. The fill scales with mlen
while the live data scales with chunk — and KDA's chunk cannot grow to meet
it
: kda_chunk_check_range caps it at 17 because 1/A reaches
exp(chunk × |gate_lower_bound|) and overflows bf16 past that. At mlen 2048
that 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 VLEN 2048) and forces the chunk form to divide by a
per-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 that
does 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:

MLEN whole-block prefetch row-granular factor
64 2,450 1,801 1.4×
512 4,390 1,146 3.8×
2048 14,173 1,137 12.5×

Worth almost nothing on a narrow machine and a great deal on a wide one,
which is the shape a term scaling with mlen has 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:

MLEN whole-block prefetch row-granular
512 1,856 1,325
1024 1,742 (optimum) 676
2048 2,650 (52% worse) 515 (optimum)

The 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, HLEN and VLEN are the tile sizes doc/plena_isa_spec.md
defines — 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:

cycles = ceil(work / MLEN) × instruction latency

and the latency is a constant — V_FMA_VF bills VECTOR_MUL_CYCLES whatever
VLEN is, M_MM bills BLEN. So a wider machine covers 32× the data per
instruction 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 in matrix_core.rs, independent of
MLEN entirely. The two halves model width by different mechanisms: instruction
count on the analytic side, a fixed PE array on the emulator side. That is
separate from the MLEN disagreement in plena_settings.toml and equally
unresolved.

@qichao-arlo-wang
qichao-arlo-wang force-pushed the feat/static-kda branch 13 times, most recently from 5157e19 to c9da62e Compare August 28, 2026 09:22
qichao-arlo-wang and others added 5 commits August 29, 2026 01:42
…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>
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant