You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
requirements.txt is not bumped, since no wheel contains these yet. Bump it once both land.
Kernels come from mlir-aie's factories
An operator now gives the build an aie.iron.kernelsExternalFunction through KernelObjectArtifact.from_extern, instead of its own source and flags.
mlir-aie compiles, prefixes and stamps the object.
The object's name and symbols carry a digest of its recipe. A changed recipe therefore never reuses a stale object, and fused sequences share identical recipes instead of re-prefixing them for each operator.
Extra symbols in the same object are bound with object_file.bind, e.g. softmax's mask_bf16 and mha's matmul_PV/rescale_O.
Two operators keep hand-built kernels:
flm/gemm: it picks its epilogue mode at runtime from one xclbin, and fused_mm compiles in a single mode.
swiglu_prefill_stream: its test is skipped.
Removed: IRON_AIE_KERNELS_DIR (sources now resolve through mlir-aie's MLIR_AIE_KERNEL_SOURCES), func_prefix, and iron/common/kernels.py.
Operators now pin the probed NPU as the selected device. The factories choose their architecture from the selected device and would otherwise fall back to aie2.
Llama fixes
Decode now masks softmax to the current context length instead of a running sum. Generated text used to degenerate after a few tokens. Teacher-forced against an fp32 CPU reference, top-1 agreement goes from 5/24 to 24/24.
The full-ELF callable now flushes scratch (weights, KV cache) before each dispatch. It previously flushed only its input buffer; NPU access to these buffers isn't cache-coherent, so the prefill→decode KV hand-off raced: 12 of 103 runs had different logits, now 0.
test_llama_3_2_1b_determinism requires bit-identical logits across alternating prompts.
Tests
One reference per operator. Each simple 1:1 operator has a single reference(), used both by its test and by dispatch="compare". assert_matches_reference(op, *inputs) dispatches once and judges the output under op.reference_tolerance(), the declared contract of the kernel it runs. The new references are bit-identical to the old goldens on every test parameter set, except axpy and dequant, which now round the way the kernel does.
Tolerances only tighten. Each was measured on npu2 first. relu is now exact; add/mul/axpy/dequant are held to 1 bf16 ulp. dispatch="compare" judges each step by its own kernel contract.
mlir-aie utilities are used, not re-implemented.verify_buffer uses aie.utils.verify.compare, which no longer lets NaN pass. Timing uses aie.utils.benchmark.run_iters. The trace buffer layout comes from get_trace_buffer, and trace dumps use TraceConfig. IRON tests that only re-tested these utilities are removed.
Real tests, no mocks. Tracing, conftest device gating and relative_build_dir now run for real, the first on hardware.
Metrics CSV gains a Provenance column: commit, Peano and mlir-aie versions, kernel sources and digest, device, power mode. Existing columns are unchanged.
PR Merge Checklist
The PR is rebased on the latest devel commit and pointing to devel.
Operators can now hand the compilation system an aie.iron.kernels factory's
ExternalFunction (KernelObjectArtifact.from_extern) instead of a hand-built
source + flags recipe. mlir-aie compiles, prefixes and stamps the object; its
file name and symbols carry a digest of the recipe, so a changed recipe never
reuses a stale object, and fused sequences share identical recipes instead of
re-prefixing them per operator.
- A generated MLIR module is only fresh if it links every kernel its design
was given, since a new recipe means a new object name.
- Operators pin the probed NPU as the selected device: the factories pick
their architecture from the selected device only and otherwise fall back to
aie2.
- kernels_dir resolves through mlir-aie's config (MLIR_AIE_KERNEL_SOURCES);
IRON_AIE_KERNELS_DIR is gone.
- transpose uses datamovement.transpose; gemv uses linalg.mv, with the gelu
epilogue bound from the gelu factory's object.
Co-Authored-By: Claude <noreply@anthropic.com>
rms_norm (plain and weighted), rope, softmax, and every ChanneledUnary /
BinaryElementwise operator (silu, gelu, relu, sigmoid, tanh, leaky_relu,
layer_norm, elementwise_add/mul, axpy) now take their kernels from
aie.iron.kernels. Each op builds its ExternalFunctions once in _kernel(s)
and hands the same objects to the design and to
KernelObjectArtifact.from_extern, so the design no longer re-declares the
symbol, object name, or func_prefix by hand. softmax binds mask_bf16 from
the softmax object with object_file.bind.
Co-Authored-By: Claude <noreply@anthropic.com>
The accumulator L1 type is now flat (m*n,), as the matmul, zero and
convert_copy factories all declare it. ROUND_CONV_EVEN is requested via
linalg.mm(round_conv_even=...), which needs the matching mlir-aie change.
Co-Authored-By: Claude <noreply@anthropic.com>
The QK^T product comes from linalg.mha(b_col_maj=True,
emulate_bf16_mmul_with_bfp16=True); partial_softmax, matmul_PV, rescale_O
and init_scale_buffer are bound from its object with the design's types,
and passThroughLine from a 16-bit passthrough object. Needs the matching
mlir-aie linalg.mha b_col_maj/emulate kwargs.
Co-Authored-By: Claude <noreply@anthropic.com>
mem_copy binds passThroughLine from the 16-bit passthrough object with its
bf16 line type, and no longer takes a func_prefix: the factory object is
already unique per recipe.
Co-Authored-By: Claude <noreply@anthropic.com>
flm/dequant binds q4nx_dequant_bfp from the factory's object with its
bfp16ebs8 output type. The factory also carries upstream's
--aie-pipeliner-max-stagecount=5; the output stays byte-exact.
flm/gemm stays hand-built: fused_mm compiles in a single epilogue mode,
always rounds to nearest-even and wraps mm_fused.cc in fused_mm_tile.cc,
while this operator selects among several modes at runtime from one
xclbin. It is now the only user of the aie2 lut_based_ops helper, which
moves into it from operator_bases.
swiglu_prefill_stream also stays hand-built (its test is skipped at module
level), but its silu/mul sources move to generic/, where they now live.
Co-Authored-By: Claude <noreply@anthropic.com>
compile_external_kernel() runs the compiler with cwd set to its output
directory, so a relative AIEContext build_dir (Llama uses "build_elf")
resolved twice and clang could not find the staged kernel source.
Co-Authored-By: Claude <noreply@anthropic.com>
llama_forward_pass_decode wrote a running sum of every step's context
length into softmax_vector_size, so only the first decode step was masked
correctly. From the second step on the attention softmax also covered
unwritten cache slots, and once the sum passed max_seq_len (about the 7th
token for a 293-token prompt) it covered the whole 2048-wide row. Those
slots score ~0 and soak up attention weight, which is why generated text
degenerated after a few tokens.
Write context_len, as the ELF-patching code did before #131. Teacher-forced
greedy against an fp32 CPU reference (prompt 1024, 24 tokens): top-1
agreement goes from 5/24 to 24/24, decode KL from 1.0-5.5 to <= 0.013, and
the greedy text matches the reference exactly.
Co-Authored-By: Claude <noreply@anthropic.com>
test.py only checked the exit code and scraped TTFT/TPS, so decode emitted
garbage after a few tokens for months without failing (the softmax
running-sum mask bug fixed in the previous commit).
llama_npu.py --check-accuracy feeds the NPU and llama_cpu (fp32 weights) the
reference's greedy token each step and reports KL(fp32 || NPU) of the
next-token distribution. test_llama_3_2_1b_accuracy gates on it over a
1024-char prompt and 40 steps:
prefill KL 0.074 (limit 0.1)
decode KL 0.013 (limit 0.05); 9.2 with the mask bug reintroduced
Top-1 agreement is reported but not gated: near-ties flip it (1/40 steps at
KL 0.004).
Co-Authored-By: Claude <noreply@anthropic.com>
After prefill, llama_forward_pass writes every layer's K/V cache into the
fused decode operator's scratch buffer through torch_view(), then called
scratch_buffer.to("cpu"). That syncs the other direction, and nothing
else flushes scratch (the callable only syncs its input buffer), so the
host's dirty cache lines reached DRAM whenever the CPU happened to evict
them. Decode then either read stale KV rows, or had rows it had written
overwritten one 64-byte line at a time by a late eviction.
This is the run-to-run nondeterminism in Llama's output: teacher-forced
over 40 steps, 12 of 103 runs had different logits (always starting at
decode step 1 or 2; prefill was always identical). With the flush, 0 of
103 differ, and all match the previous majority result. Latency is
unchanged (paired TTFT 0.993, decode 1.003 over 8 interleaved rounds).
Co-Authored-By: Claude <noreply@anthropic.com>
Add --check-determinism ROUNDS to llama_npu.py. It prefills two prompts
with different text, alternating, and greedily decodes a few steps from a
fresh state each round. It then counts the runs whose logits differ
bitwise from the first run of the same prompt. Alternating the prompts
matters: when a host write never reaches the device, the NPU reads the
other prompt's data instead of a leftover identical copy. That turns the
12% race fixed in the previous commit into a failure on every run: 38/38
in each of three trials with the flush reverted.
test_llama_3_2_1b_determinism runs 5 rounds of 4 tokens (about 30 s) and
requires zero differing runs.
Co-Authored-By: Claude <noreply@anthropic.com>
get_buffer() hands out writable views into the fused ELF's scratch buffer
(weights, KV caches), but _sync_inputs only flushed the input buffer. The
host runtime's own dispatch flushes every argument, and the separate-xclbin
callable goes through it; the full-ELF callable calls run_handle.start()
directly and so skipped it for scratch. NPU access to these buffers is not
cache-coherent, so an unflushed scratch write was a silent race. It was the
Llama run-to-run nondeterminism fixed at the call site in e4fed44.
_sync_inputs now flushes scratch too. With nothing dirty that transfers
nothing, and it leaves scratch marked device-resident, so reading a scratch
view after a run pulls the NPU's writes. Llama's explicit hand-off flush is
removed; the determinism test now guards the callable.
test_non_input_buffers_sync_without_explicit_flush writes a non-input
buffer, with new data each dispatch, in separate and fused modes. It also
reads a non-output buffer, with no to() from the caller. Without the flush
the fused case failed 5/5 runs.
Llama A/B, 8 interleaved rounds, prompt 1024 / 40 tokens: paired decode
1.003, TTFT 1.003, identical text.
Co-Authored-By: Claude <noreply@anthropic.com>
verify_buffer now judges through aie.utils.verify.compare under a
relative Tolerance rather than IRON's own nearly_equal. The signature
and returned mismatch indices are unchanged. The old check let a NaN
output pass; compare requires NaN to meet NaN and an infinity to meet
the same infinity, and that holds regardless of max_error_rate.
conftest adds a Provenance column to the metrics CSV (commit, Peano,
mlir-aie, kernel sources and digest, plus device and power mode when
the run used the NPU). Existing columns are unchanged, so the CI pretty
scripts still read them. When device-gated tests are collected and
there is no NPU runtime, it stops with mlir-aie's probe reason (exit 4,
e.g. "xrt-smi not on PATH ...") instead of a wall of failures.
The hand-rolled perf_counter timing in gemm and the swiglu tests now
uses aie.utils.benchmark.run_iters.
The infrastructure tests that only re-tested mlir-aie utilities
(comparison, benchmark, sequence_subviews, sequence_output_sync) are
removed; mlir-aie's own test suite covers them. Tests of IRON code stay.
Non-extensive sweep of iron/operators + iron/tests on Strix: 283
passed, 4 skipped, 1 failed. The one failure is
lazy_imports::test_lazy_catalog_does_not_import_mha, which asserts on
sys.modules and so fails whenever mha runs earlier in the same session,
also on the parent commit. No operator test flipped under the stricter
comparison.
Co-Authored-By: Claude <noreply@anthropic.com>
verify_buffer and run_test take a Tolerance, replacing rel_tol, abs_tol
and max_error_rate when given. Only per-element kinds are accepted, since
the failing elements are listed one by one.
Tolerances only tighten. Each was measured on npu2 first.
- relu, leaky_relu, elementwise_add, elementwise_mul, axpy and dequant
now use their kernel's contract. That is exact for relu and 1 bf16 ulp
for the rest, in place of rel 0.04 (0.01 for dequant).
- axpy's golden rounds s * A + B once, as the kernel does. The old bf16
expression rounded the product too.
- rope's golden uses the bf16 cos/sin tables the operator is actually
given. Measured against those it is within 1 ulp. Its tolerance goes
from rel 0.05 / abs 0.5 to rel 0.05 with no absolute floor.
- layer_norm goes from rel 0.1 / abs 0.1 to rel 0.1 / abs 0.05, the
tighter of the test's and the contract's bounds.
Co-Authored-By: Claude <noreply@anthropic.com>
Each simple 1:1 operator now has exactly one reference, the op's
reference(). Its reference.py keeps the math as reference(...) and draws
inputs with generate_inputs(...), same seeds and draws as before; the
golden output is no longer a separate code path that can drift from what
dispatch="compare" checks against.
test_utils.assert_matches_reference(op, *inputs) shapes the inputs to the
arg spec, computes op.reference(...), dispatches once and judges the output
by op.reference_tolerance() -- the declared contract of the kernel the
operator runs -- unless the test passes a tighter tolerance. The tests of
relu, tanh, sigmoid, gelu, silu, leaky_relu, layer_norm, rms_norm, softmax,
elementwise_add/mul, axpy, dequant, mem_copy, transpose, repeat and rope
become thin wrappers over it; names, parameters and metrics are unchanged.
Tests that were stricter than the contract keep their tolerance explicitly.
The new references are bit-identical to the old goldens on every test
parameter set, except axpy (now rounds once, as the kernel does: the scalar
to bf16, fp32 multiply-add, one rounding; 0 differences on the test sets)
and dequant (the fp32 golden rounded once to bf16, the same verdict).
Also drops torch_dtype_map and the dtype= parameters that only ever took
"bf16".
Co-Authored-By: Claude <noreply@anthropic.com>
dispatch="compare" used one global rule for every step: flag the step if
both max_abs and max_rel exceeded fixed thresholds. It now judges each
step with aie.utils.verify.compare under that step's
op.reference_tolerance(), the same contract its operator test holds it to.
CompareDispatch(tolerance=...) overrides it for every step. A step whose
operator declares no element-wise tolerance (none, a bound, or one with
range_frac) falls back to relative(0.025, 1e-2) per element, which is
stricter than the old max_abs-AND-max_rel rule. The logged stats are
unchanged; a mismatch now reports the verdict and the tolerance.
The infrastructure test checks both directions on real hardware: tanh
passes compare mode under its contract and is flagged under exact().
Co-Authored-By: Claude <noreply@anthropic.com>
- gemm test: drop trace_size, which every case set to 0 and nothing read.
This changes the gemm test IDs.
- lazy_imports: check in a fresh interpreter. In the session's own, any
operator collected earlier is already imported, so the check depended
on test order. iron is installed into the environment, so no path setup.
- AGENTS.md: replace the torch_to_numpy/numpy_to_torch section (neither
exists) with the bf16 view and DEFAULT_TENSOR_CLASS.from_torch/to_torch,
and describe assert_matches_reference and the one-reference layout.
Co-Authored-By: Claude <noreply@anthropic.com>
The full-ELF sequence callable now asks mlir-aie's get_trace_buffer() for
the fused trace buffer's argument index and size, instead of assuming
argument 3 and summing the slices itself. IRON's trace_buffer_size() and
its test (trace_layout.py) go; the layout is mlir-aie's and tested there.
The tests that stubbed their way around the hardware now run it:
- tracing: dispatch a traced one-step layer_norm sequence, check tracing
leaves the result bit-identical, and that dump_traces writes the
buffer's raw 32-bit words and the Perfetto JSON.
- conftest_lazy_device: run the root conftest's device gating in a real
pytest subprocess instead of calling the hook on fake items against a
stubbed runtime. The no-runtime cases hide pyxrt from the child's
PYTHONPATH, which is what an unsourced XRT amounts to.
- relative_build_dir: pass a relative build_dir instead of monkeypatching
the working directory.
Co-Authored-By: Claude <noreply@anthropic.com>
dump_traces wrote the raw words, split the buffer into slices, decoded each,
named the files and warned of truncation itself, which duplicated
TraceConfig.write_trace/trace_to_json. trace_to_json now does the per-slice
part upstream (mlir-aie branch trace-to-json-slices), so dump_traces writes the
text with write_trace and the JSON with trace_to_json, and keeps only IRON's
knobs: the output directory, the column shift, the MLIR override and the
cycles summary.
parse_trace_buffer existed to turn the parser's SystemExit into an exception;
the parser raises ValueError now, so it goes. lowered_mlir duplicated the
callable's lookup of aiecc's lowered module, which the callable now exposes as
lowered_mlir_path.
The raw text drops trailing zero words, as TraceConfig's always has, so the
test reads it back with read_trace. A buffer without slices now writes
<tag>.json rather than <tag>_trace.json.
Co-Authored-By: Claude <noreply@anthropic.com>
Resolved the devel conflict and pushed merge commit 44b849d. The resolution preserves factory-backed kernel compilation while adopting upstream binutils/symbol-prefix helpers; CodeQL reports no alerts.
@copilot can you please address my review feedback?
Addressed in a60d170: missing Llama weights now skip only outside CI, so CI fails loudly. I also audited the reference changes; the old golden generators/inline expectations were removed, and each operator method is only a thin adapter to its single module-level reference. ExternalFunction remains necessary because mlir-aie factories return it and compile_external_kernel() requires it; KernelObjectArtifact bridges that recipe into IRON’s build graph.
Ports IRON operators to mlir-aie kernel factories, standardizes reference-based verification, and fixes Llama decode correctness and cache synchronization.
Changes:
Uses factory-built kernel artifacts with recipe-aware compilation and architecture selection.
Consolidates operator references, tolerance contracts, benchmarking, tracing, and sequence synchronization.
Adds Llama accuracy and determinism checks while correcting decode masking.
AXPY: the float32 ABI does not preserve coefficient precision: the C++ kernel explicitly broadcasts bfloat16(a). Removing that rounding from the reference would be incorrect. Kept the arithmetic, clarified the docstring, and added non-integer/cancellation regressions.
Chess remains blocked upstream: most factories, including the *_sized variants, still lack use_chess in Single-core Kernel Optimizations and Tooling Xilinx/mlir-aie#3801. Some recipes also carry Peano-only flags. Upstream needs compiler-aware recipes before IRON can propagate the context choice safely; I retained the mismatch guard rather than bypassing it.
Black, whitespace/secret checks, and fallback code review passed. Pytest collection was blocked by missing dependencies; the pinned CPU PyTorch download also failed.
Compare exact results in common dtype before identifying error indices
iron/common/test_utils.py:106
Casting the reference to the output dtype can erase the difference that made compare(..., Tolerance.exact()) fail. For example, a lower-precision output may equal the rounded reference, leaving error_indices empty and causing run_test to accept a failed exact verdict. Compare in NumPy's common dtype when reconstructing the indices.
Handle bfloat16 bit patterns before converting NumPy arrays to Torch
iron/operators/dequant/reference.py:18
This conversion does not preserve an ml_dtypes.bfloat16 array through PyTorch's NumPy bridge. expand_ref returns bfloat16, so move its uint16 bit pattern into Torch instead of asking torch.from_numpy to interpret the custom dtype; otherwise the new dequant reference can fail before comparison.
mlir-aie #3801 (f119c6947) groups aie_kernels by family instead of by
architecture: generic/, aie2/ and aie2p/ are gone, and architecture-specific
code moved into *_aie2.h / *_aie2p.h headers the family's .cc selects.
- flm.GEMM compiles fused/fused_mm_tile.cc: mm_fused.cc became the header
fused/mm_fused.h, and activations.h / zero.cc moved to common/.
- Stream ops: mm.cc is linalg/mm.cc, zero.cc is zero/zero.cc, silu and mul
live in activation/ and eltwise/. StreamKernel.subdir now names the family
directory instead of defaulting to the device one, which no longer exists.
- Comments, docstrings, README kernel links and AGENTS.md follow the new
paths.
Co-Authored-By: Claude <noreply@anthropic.com>
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
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.
hunhoffe
left a comment
Requires two mlir-aie changes that are not merged yet:
linalg.mm(round_conv_even=),linalg.mha(b_col_maj=, emulate_bf16_mmul_with_bfp16=), …).trace_to_json, and a trace parser that raises instead of exiting.requirements.txtis not bumped, since no wheel contains these yet. Bump it once both land.Kernels come from mlir-aie's factories
aie.iron.kernelsExternalFunctionthroughKernelObjectArtifact.from_extern, instead of its own source and flags.object_file.bind, e.g. softmax'smask_bf16and mha'smatmul_PV/rescale_O.fused_mmcompiles in a single mode.IRON_AIE_KERNELS_DIR(sources now resolve through mlir-aie'sMLIR_AIE_KERNEL_SOURCES),func_prefix, andiron/common/kernels.py.Llama fixes
test_llama_3_2_1b_accuracygates KL(fp32 ‖ NPU): prefill < 0.1 (measured 0.074), decode < 0.05 (measured 0.013).test_llama_3_2_1b_determinismrequires bit-identical logits across alternating prompts.Tests
reference(), used both by its test and bydispatch="compare".assert_matches_reference(op, *inputs)dispatches once and judges the output underop.reference_tolerance(), the declared contract of the kernel it runs. The new references are bit-identical to the old goldens on every test parameter set, except axpy and dequant, which now round the way the kernel does.dispatch="compare"judges each step by its own kernel contract.verify_bufferusesaie.utils.verify.compare, which no longer lets NaN pass. Timing usesaie.utils.benchmark.run_iters. The trace buffer layout comes fromget_trace_buffer, and trace dumps useTraceConfig. IRON tests that only re-tested these utilities are removed.relative_build_dirnow run for real, the first on hardware.PR Merge Checklist
develcommit and pointing todevel.