Add single-dispatch layer-by-layer multi-head attention#91
Draft
andrej wants to merge 1 commit into
Draft
Conversation
andrej
commented
Apr 6, 2026
Collaborator
Author
There was a problem hiding this comment.
Can we reuse the reference from the existing mha? (Note: does not include RoPE and Q, K, V projections, but some code reuse should be possible.)
Contributor
📊 Test Results for Test Example Applications1d87fe8 (2026_04_07_21_05_39) IRONCLADTested on
📈 Trends (vs main branch) for Test Example Applications1d87fe8 (2026_04_07_21_05_39) IRONCLAD Trendsllama_3.2_1b
llama_3.2_1b_prompt_1024_tokens_1
llama_3.2_1b_prompt_1024_tokens_40
llama_3.2_1b_prompt_13_tokens_1
llama_3.2_1b_prompt_13_tokens_40
llama_3.2_1b_prompt_2048_tokens_1
llama_3.2_1b_prompt_2048_tokens_40
|
Contributor
CI Test Results55c43fb (2026_07_09_22_19_28) IRON - CI SummaryExamplesiron/applications/llama_3.2_1b
Krackan - ExamplesIRONTested on iron/applications/llama_3.2_1b
Trends: IRON Trendsiron/applications/llama_3.2_1btest_llama_3_2_1b[llama_3.2_1b_prompt_1024_tokens_1]
test_llama_3_2_1b[llama_3.2_1b_prompt_1024_tokens_40]
test_llama_3_2_1b[llama_3.2_1b_prompt_13_tokens_1]
test_llama_3_2_1b[llama_3.2_1b_prompt_13_tokens_40]
Phoenix - ExamplesIRONTested on Trends: IRON Trends |
andrej
commented
Jul 9, 2026
Comment on lines
+68
to
+74
| // stats[0] = running max | ||
| // stats[1] = running sum (of exp(x - max)) | ||
| // --------------------------------------------------------------------------- | ||
|
|
||
| void softmax_partial_stats_impl(bfloat16 *restrict input, | ||
| bfloat16 *stats, | ||
| const int32_t vector_size) |
Collaborator
Author
There was a problem hiding this comment.
For readability, can we make stats a struct rather than an array with hard-coded indices? Also pointer should never alias with the input so we can annotate with *restrict.
Collaborator
Author
There was a problem hiding this comment.
We'll have struct softmax_stats *stats as the function argument but if that's hard to do from the design.py it's fine if it doesn't match from there and is just a void * there in the function signature
| return; | ||
| } | ||
|
|
||
| // --------------------------------------------------------------------------- |
Collaborator
Author
There was a problem hiding this comment.
Make sure code comment style aligns with rest of the file, not sure these // --- are used elsewhere
Comment on lines
+163
to
+172
| // --------------------------------------------------------------------------- | ||
| // Online (partial / tiled) softmax helpers | ||
| // | ||
| // These three kernels implement a two-pass online softmax that processes a row | ||
| // in sub-tile chunks, keeping running max and sum statistics in a small local | ||
| // buffer (`stats`). Layout of the stats buffer (bfloat16[16], only [0..1] | ||
| // used): | ||
| // stats[0] = running max (scaled by log2e) | ||
| // stats[1] = running sum (of exp2(x*log2e - max)) | ||
| // --------------------------------------------------------------------------- |
Collaborator
Author
There was a problem hiding this comment.
Apply same comments as in aie_kernels/aie2/softmax.cc: comment style, struct instead of array for stats
Comment on lines
+189
to
+226
| // --- Phase 1: find local max (scaled by log2e) ------------------------- | ||
| float local_max = -INFINITY; | ||
| auto it_in1 = aie::cbegin_restrict_vector<SM_VEC_LEN>((bfloat16 *)input); | ||
| for (int i = 0; i < elem_iters; i++) { | ||
| input_bf16 = *it_in1++; | ||
| scaled_accum = aie::mul(input_bf16, log2e_vec); | ||
| float chunk_max = aie::reduce_max(scaled_accum.to_vector<bfloat16>()); | ||
| if (chunk_max > local_max) { | ||
| local_max = chunk_max; | ||
| } | ||
| } | ||
|
|
||
| // --- Phase 2: update running max, rescale running sum ------------------ | ||
| float old_max = (float)stats[0]; | ||
| float old_sum = (float)stats[1]; | ||
|
|
||
| if (local_max > old_max) { | ||
| // New max is larger — rescale the old sum by exp2(old_max - new_max) | ||
| aie::vector<float, SM_VEC_LEN> diff_vec = | ||
| aie::broadcast<float, SM_VEC_LEN>(old_max - local_max); | ||
| aie::vector<bfloat16, SM_VEC_LEN> corr = aie::exp2<bfloat16>(diff_vec); | ||
| old_sum = old_sum * (float)corr[0]; | ||
| old_max = local_max; | ||
| } | ||
|
|
||
| // --- Phase 3: accumulate exp2(input * log2e - max) for this chunk ------ | ||
| aie::vector<bfloat16, SM_VEC_LEN> max_val_vec = | ||
| aie::broadcast<bfloat16, SM_VEC_LEN>((bfloat16)old_max); | ||
|
|
||
| auto it_in2 = aie::cbegin_restrict_vector<SM_VEC_LEN>((bfloat16 *)input); | ||
| for (int i = 0; i < elem_iters; i++) { | ||
| input_bf16 = *it_in2++; | ||
| scaled_accum = aie::mul(input_bf16, log2e_vec); | ||
| exp_in_accum = aie::sub(scaled_accum, max_val_vec); | ||
| aie::vector<bfloat16, SM_VEC_LEN> exp_val = | ||
| aie::exp2<bfloat16>(exp_in_accum.to_vector<float>()); | ||
| exp_val_accum = add(exp_val_accum, exp_val); | ||
| } |
Collaborator
Author
There was a problem hiding this comment.
Make this use the same on-line approach as in aie_kernels/aie2/softmax.cc
Comment on lines
84
to
90
| void partial_softmax_alias_bf16(bfloat16 *restrict input_vector, | ||
| bfloat16 *restrict output_vector, | ||
| bfloat16 *restrict scale_buffer, | ||
| const int32_t vector_size, | ||
| const int32_t row_idx, | ||
| const int32_t num_rows, | ||
| const bfloat16 scale) |
Collaborator
Author
There was a problem hiding this comment.
There appears to be a lot of code duplicated in the new partial softmax functions below. Factor out shared functionality into a separate function used across both, forcing inlining so it does not affect performance -- or, if sharing code is not easily possible, explain what the differences between the two implementations are.
Comment on lines
+82
to
+87
| // The tile is interpreted as a `vector_size`-wide horizontal strip of the | ||
| // per-head (S, S) attention-score block; idx[0] is the strip's starting | ||
| // column within that block, idx[1] is the strip's row within the block. | ||
| // The kernel implements the causal mask in-place by writing `a` to elements | ||
| // strictly above the diagonal and copying y -> z everywhere else. This | ||
| // avoids materialising an H*S*S mask buffer entirely. |
Collaborator
Author
There was a problem hiding this comment.
This function is a general 'triangular fill', not specific to attention. Make the comments reflect this. Consider factoring this out into a separate operator since it has less to do with AXPY, or making it more AXPY-like by allowing a scale or constant addition be applied within the triangular masked region (again using C++ templates).
Comment on lines
+26
to
+34
| """Operator that fuses multiple MLIROperators into one. | ||
|
|
||
| Args: | ||
| dispatch: Dispatch strategy for the fused operator. | ||
| ``"auto"`` (default) selects ``"fused"`` on NPU2 and | ||
| ``"separate"`` on NPU1. ``"fused"`` uses a single-ELF | ||
| dispatch (requires NPU2). ``"separate"`` compiles each | ||
| sub-operator to its own xclbin and invokes them sequentially. | ||
| """ |
Collaborator
Author
There was a problem hiding this comment.
Fusion's a bit of a big term for what this does. It still reconfigures the array between each operator. Let's say it fuses multiple operators into a "single dispatch." (except for NPU1 where full ELF is not yet supported where this still degenerates into multiple dispatches). The sequential feature would be used on NPU2 for debugging and benchmarking the impact of single-dispatch.
Comment on lines
-55
to
-57
| for obj in objs: | ||
| obj.filename = f"op{idx}_{obj.filename}" | ||
| obj.prefix_symbols = f"op{idx}_" |
Collaborator
Author
There was a problem hiding this comment.
Why was this removed? Disambiguating function names between the sub-operators seems desirable. An explanation / reasoning for removing it is a valid response to this if there are good reasons to remove this.
Comment on lines
-85
to
-86
| if len(op.get_kernel_artifacts()) > 0: | ||
| mlir_artifact.generator.kwargs["func_prefix"] = f"op{idx}_" |
Collaborator
Author
There was a problem hiding this comment.
Why was this removed? Disambiguating function names between the sub-operators seems desirable. An explanation / reasoning for removing it is a valid response to this if there are good reasons to remove this.
Comment on lines
+275
to
+277
| Mirrors the pattern from ``chain_swiglu_artifacts`` in | ||
| ``iron/operators/swiglu_base.py``: each unique operator gets its own | ||
| xclbin + insts compiled separately, linked via ``--xclbin-input``. |
…lxl) Add a layer-by-layer (LxL) prefill MHA operator built as an OperatorSequence, and rename the existing data-flow (DF) MHA operator for clarity: - Rename iron/operators/mha -> iron/operators/mha_prefill_df (DF = data-flow: single fused kernel). - Add iron/operators/mha_prefill_lxl (LxL = layer-by-layer: chained sub-operators dispatched back-to-back). - Extend axpy with scale / scalar_add / causal-mask modes and softmax with partial/online softmax, used by the LxL design. - Register the benchmark marker in an operator-local conftest.py. Update README and operator registrations to the new names.
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
1 participant
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.
"Naive" alternative implementation for multi-head attention from the currently checked-in data-flow design. This is a simple layer-by-layer implementation, but it uses the single-dispatch mechanism to fuse it all into one MLIR file and save on CPU roundtrips and XRT overheads.
Includes two variants:
Q,K,V. This matches the functionality of the checked-in dataflow MHA.