flm.GEMM: take M, K, N and the activation as runtime parameters - #200
Merged
Merged
Conversation
The FastFlowLM harness registers one mm.xclbin per model and swaps instruction streams, against a budget of 16 xclbins for the whole model. This operator baked M, K and N into the core loop bounds, into which columns it built, and into the memtile's B buffer, so it needed one xclbin per shape: 11 (K, N, activation) combinations times up to 16 chunk lengths. Five values move into an L1 buffer the runtime sequence writes and each core reads once its barrier opens: the column's work and drain counts, M/256, K/512, and the activation. All columns are now always built, and one with no work for a shape drains its share of the A broadcast instead. The epilogue tests its mode once per chunk, outside the vector loop, so each mode keeps a branch-free inner loop; which modes it can select between stays a build-time choice, since each costs program memory. The device body is then a function of the tiling alone, verified byte for byte across every E2B shape and activation. So the xclbin is built from a module emitted at a reference shape and the per-shape build produces only the instruction stream, and the two carry different artifact stems. The core releases its barrier straight after reading the parameters. wait_for_value emits LockAction.Acquire, which does not leave the lock consumed, so without the release a core that runs twice does not wait the second time and reads the previous dispatch's parameters. Releasing before the work is safe because the sequence cannot set the barrier again until it has drained this dispatch's C. The repo's other barrier users never wait twice -- mha puts its infinite loop inside the wait, softmax writes the same parameters every dispatch -- so this does not arise there. test_one_xclbin_serves_every_shape is the regression test; the parametrised tests cannot catch it, because the aie_context fixture reconfigures the array between cases. Dropping B residency is what this costs, and it is not cheap: 12.5% at M=512, 16.4% at M=1024 and 19.3% at M=2048 on NPU2. Residency sizes the memtile buffer from K and replays it M/256 times through a buffer descriptor's repeat count, so it carries both K and M into the configuration. Restoring it needs a replay mechanism that carries neither. It does lift a cap: the repeat count expands into the memtile's BD chain at 2 blocks per replay and exceeded its 48-block limit at M=4096, so no shape with K <= 2048 would build there -- 7 of Gemma4 E2B's 10 projections. All of them build now. Verified on this base over all 12 distinct E2B prefill projections at M=256 and M=4096, dispatched back to back on one loaded xclbin with no reset in between, each checked against a CPU reference; peak 12.2 TFLOP/s at M=4096 K=12288 N=1536. Plus the operator's own suite, 17 non-extensive tests.
The merge of andrej/flm-gemm-rtp was textually clean but left the README self-contradictory: the new intro says M, K, N and the activation are runtime parameters, while the FLM-compatibility note I had added in 92080d4 still claimed they were baked in at compile time and called RTP-selectability "follow-up work" -- which 81e2006 had just done. - Drop that note entirely; the "Runtime parameters" section it would have pointed at now covers the same ground correctly. - Fix the same stale claim in the shipped-overlay section, and record what IS still build-time there: which activations the epilogue can select between, since each one compiled in costs program memory. - design.py: "None of these four" was already stale from 92080d4, which added a fifth bullet. Co-Authored-By: André Rösti <an.roesti@gmail.com>
The RTP work left clamp entirely compile-time, so _config_tag carried the bound values and every distinct pair forked a whole xclbin -- clamp=(-2,2) and clamp=(-4,4) built twice over. Split it the way epilogue_modes already splits activations: the CAPABILITY stays build-time, the SELECTION and the values go runtime. - Whether a clamped path exists at all is still -DMM_FUSED_CLAMP, because the clamped instantiation costs program memory and a build that never clamps should not carry it. That bit stays in _config_tag. - clamp_enabled and the bounds become RTP words 5-7. Bounds are floats but npu_write_rtp writes i32 only, so they travel as raw bit patterns and the kernel casts them back with __builtin_bit_cast -- memcpy leaves an unresolved external call in the compiled object rather than folding to a register move. - epilogue_body gains a CLAMP template parameter so the clamped and unclamped inner loops both stay branch-free; epilogue_dispatch picks between them once per chunk, and only compiles the clamped one when the capability is on. Deliberately NOT done: making clamp_enabled a plain runtime branch. That would double epilogue_body instantiations (mode x clamp), and program memory is the exact constraint epilogue_modes exists to manage. Verified on npu2: all four kernel variants (aie2/aie2p x clamp on/off) compile, 36/36 flm/gemm iter0 tests pass, and the new test_one_xclbin_serves_every_clamp_bound confirms three different bound pairs run back to back on one loaded xclbin while an unclamped build still resolves to a different configuration. Co-Authored-By: André Rösti <an.roesti@gmail.com>
hunhoffe
force-pushed
the
flm-gemm-rtp-v2
branch
from
September 11, 2026 22:25
d6fe009 to
92350c4
Compare
Contributor
CI Test Resultsbb0d09b (2026_09_15_22_10_53) IRON - CI SummaryExamplesiron/applications/llama_3.2_1b
Smalliron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/flm/gemm
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/mha
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Krackan - SmallIRONTested on iron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/flm/gemm
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/mha
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Krackan - ExamplesIRONTested on iron/applications/llama_3.2_1b
Phoenix - SmallIRONTested on iron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/flm/gemm
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Phoenix - ExamplesIRONTested on Trend tables omitted, the comment hit GitHub's size limit. Full report in the workflow run. |
A core can HOLD a B chunk across m_chunk accumulators instead of releasing
it after one, so DDR reads B m_row_blocks/m_chunk times. That recovers the
traffic the removed B residency used to save, without residency's cost:
residency sized the memtile buffer from k_iters and replayed it
m_row_blocks times, putting both K and M in the device configuration,
whereas m_chunk is a configuration constant.
Defaults to 1 (off). It is a tested knob, not a default, because
interleaved A/B at K=1024 N=4096 does not show a consistent win:
M m_chunk=1 min/med m_chunk=2 min/med min med
512 525.4 / 622.7 502.1 / 566.3 -4.4% -9.1%
1024 959.0 / 1081.6 1092.4 / 1162.9 +13.9% +7.5%
2048 1902.8 / 1986.0 1837.1 / 1940.0 -3.5% -2.3%
B's traffic does fall, but the per-unit A descriptors m_chunk forces
(a_split) appear to eat it. The M=1024 regression is not monotonic in
n_units and is unexplained -- that is the open question, and the reason
this is off by default.
How the interleave works, since it is the non-obvious part: the core
consumes A as (k, b_iter, mc, band), with mc INSIDE b_iter. A second A
fifo would express that directly but needs a third core input DMA channel
against a hardware limit of two. So the memtile A object holds m_chunk
stacked tiles and the forward's dims_to_stream emits them interleaved,
which fits the memtile BD's four dimensions only because mc's stride
(M_TILE*K_TILE) exactly equals the row-group dimension's size*stride and
the two merge.
A partial group is inexpressible -- stride 0 inner is rejected, stride 0
outermost is the BD repeat count (one object per repetition, not one in
total), and sub-object fills do not coalesce, all three confirmed on
hardware -- so op.py resolves m_chunk to 1 when it would not divide
m_row_blocks, or when the group's ROWS*M_TILE*K stride would overflow the
shim BD's 20-bit step.
Also stops forcing OVERLAP=1 on every split block: only a block spanning
more than one window awaits inside itself. Worth -3.5% at M=2048 alone.
37/37 iter0 tests pass on npu2.
Co-Authored-By: André Rösti <an.roesti@gmail.com>
Moves off the 1.4.3.dev85 snapshot onto the tagged release. llvm-aie is unchanged at 22.0.0.2026090701+3e93bf7b, which is the pin mlir-aie v1.4.3 itself names in utils/peano-requirements.txt, so the two stay in step. Tagged wheels live under their own tag's asset page rather than latest-wheels-4 (which carries only the .dev builds), hence the extra find-links. Co-Authored-By: Claude <noreply@anthropic.com>
CT_MAX_K_FOR_N is carried to the kernel as -DMM_FUSED_CT_K but appeared in neither artifact name. This repo's build cache keys on filename and mtime rather than on source or flags, so an object or xclbin built at one ct_max_k silently satisfied a request for another. Naming it in _kernel_object alone is not enough, and the reasoning that let that hole stand is worth recording: tile_ma usually moves with ct_max_k, but tile_ma is caller-overridable, so tn128/ma16 is reachable at two different ct_max_k values. An xclbin built by an experiment at ck=128 was then served to test_gemm_tile_options[tn128-ma16], which wants ck=32, and it returned NaN rather than an error. An artifact name must cover every input to THAT artifact; do not argue one field is implied by another unless the implication holds for every reachable configuration, overrides included. Co-Authored-By: Claude <noreply@anthropic.com>
A parameter word is not free. Each costs ~66 ns per core and the sequence
writes ROWS*COLS = 32 of them, so every word is ~2.06 us of dispatch latency
-- measured by padding the buffer at a fixed core count (12 words 108.6 us,
24 words 135.2, 48 words 182.7). Against a ~107 us floor that is most of a
short-prefill dispatch.
Five of the ten were dead or duplicated:
* the clamp trio is now sent only by a clamp-capable build. Where no
clamped path is compiled in -- the default, and every real projection --
those three words were written on every dispatch and never read.
* n_chunks / n_units are sent only when m_chunk > 1. They are
m_row_blocks // M_CHUNK and each other, so at the shipped M_CHUNK = 1 the
core reads m_row_blocks instead. No on-core arithmetic.
and n_work / n_drain are now derived rather than sent, from a raw N word plus
the tile's own column. Branch-free, and both divisors are powers of two, so
it lowers to sdiv-by-constant and leaves no __divsi3 (verified in the .elf,
not just the .ll). Note ScalarValue overloads add/sub/mul/floordiv/mod but
NOT the shift operators, so this uses // rather than >>.
The column index is per-tile static data, deliberately not a constant folded
into the program: the 32 core programs differ today only in symbol names, and
baking it into code would make them differ in instructions, foreclosing a
future one-program xclbin.
rtp_layout() sizes the buffer from (clamp_capable, m_chunk), both of which are
already in _config_tag, so the word count cannot vary within a configuration
and ship-once-use-many is preserved -- verified by loading one xclbin and
dispatching shapes that disagree on M, K, N, on which columns sit a block
out, and on the activation.
Also guards CT_MAX_K_FOR_N: it reads like a tuning table but is load-bearing
for correctness, and a wrong value fails silently (err/mass 3.45e-02 at
tile_n=64/ct_k=64, NaN at tile_n=128). Root cause not found; pack_b is ruled
out by test, its permutation round-trips at ct_k 128, 64 and 32. Unverified
pairs now raise rather than miscompute.
30-shape suite: -3.5% median at M=256 (best -11.3%, E2B/kv), within noise at
M >= 1024 -- the saving is a constant ~12 us. Accuracy bit-identical on all
30 shapes; frozen-reference control drifted +0.11%.
Co-Authored-By: Claude <noreply@anthropic.com>
--no-short is not a real pytest option. What the harness actually wants is --iterations 1: each test already averages ITERS dispatches over ROUNDS interleaved rounds, so conftest's default of 5 repeats the whole matrix five times for nothing. Also records that -s must not be passed when the CSV is wanted -- the reporter parses captured stdout, so disabling capture yields a CSV with no metric columns. Co-Authored-By: Claude <noreply@anthropic.com>
Where K or N is 10240 the row-block stride overflows the shim BD's 20-bit
iteration step, so that leg is issued as one transfer per row-block. Two shim
resources bound how many may be outstanding and neither is modelled by the
toolchain: BD ids (16/tile, freed without a completion check) and the channel
task queue (4 deep, pushed unconditionally).
That bound was enforced by windowing -- issue four, await the whole window,
issue the next four, then the same again per column-block. It is correct, but
TaskGroup.finish() emits dma_await_task, so each of those is a real barrier and
the channel drains to EMPTY at every window and column-block boundary. On a
DDR-rate-bound design those bubbles are the entire cost of the split path. The
OVERLAP value computed for this path was never read by it.
Retire the oldest transfer as the next is issued instead: the same number stay
in flight, the queue bound is enforced directly rather than by draining, and
the channel stays full across both kinds of boundary. Interleaved 8 rounds x 30
iters, bit-exact against the previous sequence on every shape:
E4B/gateup M1024 4119.4 -> 3617.7 -12.2%
E4B/gateup M2048 7664.8 -> 7182.4 -6.3%
E4B/down M1024 3649.2 -> 3553.4 -2.6%
E4B/down M2048 7196.6 -> 6978.7 -3.0%
The 24 shapes that do not split are untouched, measured at -0.1% median over
the full suite. One xclbin still serves every shape.
A unit is weighted by the C descriptors it drains (M_CHUNK under c_split), not
counted as one: counting units would overrun the 4-deep queue by exactly
M_CHUNK, and overrunning it hangs rather than diagnoses.
Also drops three knobs that no longer have a consumer -- the overlap parameter
(never passed), col_work/col_drain (superseded by the core deriving its own
trip counts) and the mt_*_bytes sizes (orphaned when B residency went) -- and
the prose describing mechanisms this design no longer has.
Co-Authored-By: Claude <noreply@anthropic.com>
hunhoffe
changed the title
hunhoffe
commented
Sep 14, 2026
hunhoffe
left a comment
hunhoffe
left a comment
Collaborator
Author
There was a problem hiding this comment.
More fine-tuning needed.
This operator had accumulated about one line of prose per line of code, and
the RTP work made it worse rather than better -- design.py went from 429
comment and docstring lines to 633 while the code grew by 120.
Most of that was not explaining the code. It was measurement tables that
belong in README.md, rationale for alternatives that were tried and rejected,
and restatements of the line underneath. Several facts were stated two or
three times over: that the build cache keys on filename, that a second A fifo
would want a third input DMA channel, that C drains in row-block order even
under M_CHUNK.
What is kept is the class of comment that is expensive to lose, because the
failure it describes is silent: the BD-id recycle hazard that corrupts rather
than faulting, the shim task-queue depth that hangs rather than diagnosing,
_VERIFIED_CT_K where a wrong entry computes the wrong answer, and the memtile
placement pin. Those are stated once, at the code they constrain.
design.py 42% -> 33% prose op.py 45% -> 35%
test.py 41% -> 33% benchmark.py 42% -> 35%
mm_fused.cc 46% -> 36%
Also corrects documentation that had gone stale. test.py still described the
split legs as retired in windows, which they have not been since the rolling
retire landed, so its docstring, an inline comment and two entries in the
shape table were all describing a mechanism that no longer exists.
test_gemm_split_leg_windowing is renamed to test_gemm_split_leg_bounds for the
same reason. Drops an unused `import os`.
No functional change: the AST with docstrings stripped is identical except for
that import and the two renames. 185/185 pass.
Co-Authored-By: Claude <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Instruction-cache collisions and DMA queue overrun can produce stale clamp bounds, incorrect activations, or device hangs.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Moves FLM GEMM shape, activation, and clamp bounds into runtime parameters so configurations can share an xclbin.
Changes:
- Adds runtime-parameterized GEMM execution and epilogue dispatch.
- Introduces rolling split-transfer retirement.
- Pins MLIR-AIE v1.4.3 and expands regression coverage.
File summaries
| File | Description |
|---|---|
requirements.txt |
Pins the required MLIR-AIE release. |
iron/operators/flm/gemm/test.py |
Tests shared xclbins and runtime parameters. |
iron/operators/flm/gemm/README.md |
Documents runtime parameters and performance tradeoffs. |
iron/operators/flm/gemm/op.py |
Separates configuration and instruction artifacts. |
iron/operators/flm/gemm/design.py |
Implements RTP handling and rolling transfers. |
iron/operators/flm/gemm/benchmark.py |
Updates benchmark guidance and commentary. |
aie_kernels/generic/mm_fused.cc |
Adds runtime epilogue and clamp dispatch. |
Review details
- Files reviewed: 6/7 changed files
- Comments generated: 3
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
The clamp-capable NPU1 build fails aiecc's measured_stack_sizes edge:
error: stack_size is absent, so this core uses the device default of
1024 bytes, but it needs 1088 bytes
which is every Phoenix failure in CI. The activation LUT path plus the
epilogue's clamp vectors put it 64 bytes over the device default, and
nothing had set stack_size, so the default applied.
2048 with the exact requirement named in the comment, since aiecc reports
it if a change ever outgrows this. It also comes off the L1 budget rather
than being left for aiecc to catch: _default_l1 was free to hand a buffer
the bytes the stack needs, and relying on a build failure to notice is
only tolerable while nothing has to fit in the gap. Verified not to move
(tile_ma, b_depth) for any tile_n on either device, so the geometry and
the measured performance are unchanged.
Co-Authored-By: Claude <noreply@anthropic.com>
Review feedback: the find-links line above stands on its own. Co-Authored-By: Claude <noreply@anthropic.com>
emit_split appended a unit's TaskGroup and only then retired down to SHIM_TASK_QUEUE, but fill/drain pushes the DMA task immediately while TaskGroup.finish() is what emits the await. So each unit's transfers went onto the channel with four already outstanding -- a transient fifth against a queue documented as four deep, and one more than test_gemm_split_leg_bounds' 4 + 2 + 2 descriptor arithmetic assumes. Retire to SHIM_TASK_QUEUE - unit_cost first instead, so the push happens with room for it. The channel still never drains to empty, which is the property worth -12.4% on these shapes: three transfers stay in flight while the oldest is awaited. Reported by Copilot on #200. Build-verified on both devices at M=512 K=10240 N=10240, where both legs split -- a miscounted bound fails to close its task groups at build time. Not verified on hardware: this box cannot dispatch, so CI is the gate. Co-Authored-By: Claude <noreply@anthropic.com>
Four review comments on #200, all about what belongs in the xclbin and what belongs in the instruction stream. MM_FUSED_CLAMP forked the build between clamped and unclamped callers, which reads as "a build per clamp" and is the odd one out: the activation is runtime-selected with every mode compiled in, so the clamp should be too. It now is. There is no unclamped instantiation to compile out -- an unclamped dispatch sends (-inf, +inf), and min(x, +inf) / max(x, -inf) leave every finite value bit-identical, so this costs no accuracy. Gone with it: the CLAMP template parameter, epilogue_dispatch, _clamp_capable, the cl axis in the config tag and the kernel object, and the clamp_enabled RTP word and kernel argument. The bounds are two always-sent words rather than a conditional trio, so an unclamped caller pays ~4 us of dispatch it did not before and a clamping one saves ~2. The README's -3.5%-at-M=256 figure was measured at four words and is annotated rather than restated, since it has not been re-measured. _config_tag folded into config_name: it had one caller besides config_name itself, and name now composes on config_name instead. name also carries the clamp bounds, as raw bit patterns. They are immediates in the runtime sequence and the build cache keys on filename and mtime, so without them a second operator at the same shape with different bounds is served the first one's instruction stream and silently clamps to the first one's values. _epilogue_mask ORs rather than sums, so two copies of a mode cannot carry into the neighbouring mode's bit, and __post_init__ now rejects an epilogue that epilogue_modes omits instead of letting it reach the kernel's default arm and apply no activation at all. Also trims the epilogue comment block, which review found excessive, and drops test.py's late import of get_target_model, already at module top. Build-verified on npu1 and npu2 across every tile_n, with and without a clamp, and with gelu for the LUT archive path. Not verified on hardware: this box ships a cp310 pyxrt that no mlir_aie 1.4.3 wheel matches, so CI is the gate on the numerics. Co-Authored-By: Claude <noreply@anthropic.com>
hunhoffe
marked this pull request as ready for review
September 16, 2026 19:32
andrej
approved these changes
Sep 16, 2026
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.
Moves flm.GEMM's shape out of the device configuration and into the instruction stream, so one xclbin serves every projection of a model instead of one per shape. All 30 E2B + E4B projection shapes now resolve to a single configuration, asserted by
test_one_xclbin_serves_every_shape. M, K, N, the activation and the clamp bounds are runtime parameters.The parameter buffer is 4 words rather than 10. A word costs ~2.06 µs of dispatch latency across the 32 cores, which is most of a short-prefill dispatch against a ~107 µs floor, so the clamp trio is sent only by a clamp-capable build and n_work/n_drain are derived on-core from a raw N word plus the tile's own column. Separately, split legs (K or N = 10240, where the row-block stride overflows the shim BD's 20-bit iteration step) now retire the oldest transfer as the next is issued instead of draining a whole window and then a whole column-block — worth up to -12.4% on those shapes.
The cost is memtile B residency, which sized the memtile buffer from k_iters and replayed it m_row_blocks times, putting both K and M into the configuration. Against the pre-RTP build, short prefill gets faster (M=256 median -5.5%) and long prefill pays for B's re-reads (M=1024 +13.1%, M=2048 +22.0%). This is a deliberate trade rather than a regression to fix before merge. Recovering it needs a replay count driven from the instruction stream; that is now reachable via
aiex.dma_channel_reset_forin the pinned v1.4.3, but the buffer would still be sized from k_iters. Follow-up, not a blocker.Accuracy is identical to pre-RTP.
PR Merge Checklist
develcommit and pointing todevel.