Repository navigation
riscv(rotaryembed): optional RotaryEmbed+Permute fusion via a layout flag - #7083
HougeLangley wants to merge 11 commits into
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
Part of a series: SpaceMiT K3 (RISC-V) support and optimisationThis PR is one slice of what used to be a single large change (#7037), split up
Suggested order: #7048 → #7079 → #7080, and #7083 → #7082. #7081 and #7084 中文说明本 PR 是原先单个大改动（#7037）拆分出来的一部分，便于逐个评审。依赖关系： 拆分的理由：原 PR 混合了构建选项、GEMM、注意力、四个层内核和两个纯重构， |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e75c0d01e3
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| const int seqlen = input_hc_swapped ? bottom_blob.c : bottom_blob.h; | ||
| const int num_heads = input_hc_swapped ? bottom_blob.h : bottom_blob.c; |
There was a problem hiding this comment.
With parameter 1 enabled on an x86, ARM, MIPS, or LoongArch build (or Vulkan), the registered RotaryEmbed_* override never executes this branch: those forward methods still create the output like the input and treat h as sequence and c as heads (for example, src/layer/x86/rotaryembed_x86.cpp:35-51). Thus a fused input whose h is heads and c is sequence is processed in the old orientation; during one-token decoding it even indexes the one-row cos/sin caches as though they had heads rows. Implement the swapped path in those overrides or delegate to RotaryEmbed::forward when the flag is set.
Useful? React with 👍 / 👎.
| int RotaryEmbed::load_param(const ParamDict& pd) | ||
| { | ||
| interleaved = pd.get(0, 0); | ||
| input_hc_swapped = pd.get(1, 0); |
There was a problem hiding this comment.
When a fused model carrying 1=1 is passed through ncnnoptimize, load_param stores the flag here, but the RotaryEmbed case in tools/modelwriter.h:2698-2704 writes only 0=interleaved. The saved model therefore reloads this field as 0 even though its preceding Permute has already been removed, silently reinterpreting the input axes and corrupting output shapes and values; the new field must also be emitted by ModelWriter.
Useful? React with 👍 / 👎.
…oss backends Two problems with the new input_hc_swapped parameter: * ModelWriter only emitted \" 0=interleaved\", so a fused model saved through ncnnoptimize reloaded the flag as 0 while its preceding Permute had already been removed -- the axes were then reinterpreted and the output was wrong. * Only the generic forward() implemented the swapped layout. x86, ARM, LoongArch and MIPS have their own forward() overrides that still treat h as the sequence and c as the heads, so on those builds the fused graph was processed in the old orientation and, for one-token decoding, indexed the one-row cos/sin caches as if they had num_heads rows. Those overrides now delegate to the generic implementation when the flag is set.
|
Thanks for the review -- both are real and are fixed in the branch. 1. The layout flag was not serialised by ncnnoptimize Correct. 2. Only the generic forward handled the swapped layout Correct, and worth stating plainly: on x86, ARM, LoongArch and MIPS the Vulkan is not covered by this change; the fused path is produced by the CPU Verification
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e086245b22
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (input_hc_swapped) | ||
| return RotaryEmbed::forward(bottom_blobs, top_blobs, opt); |
There was a problem hiding this comment.
When a fused model runs on ARM with fp16 or bf16 storage, this return is taken before the 16-bit dispatch below; the base method allocates a 2-byte-element output but unconditionally reads and writes rows through float* (src/layer/rotaryembed.cpp:39,59-73,85-99). This misinterprets the input/cache data and writes four bytes per element beyond the output rows. The same ordering affects bf16 on x86, MIPS, and LoongArch, so swapped mode needs layout-aware 16-bit handling rather than directly invoking the fp32 base implementation.
Useful? React with 👍 / 👎.
…puts Follow-up to the delegation added earlier: the generic implementation's fused branch is fp32/pack1 only, so delegating a 16-bit fused input to it (as the x86, ARM, LoongArch and MIPS overrides now do) would reinterpret the storage as float and write twice the row width into a 16-bit output blob. The base implementation now converts such input to fp32/pack1, runs the same fused code, and converts the result back, so every backend that delegates is correct for fp16/bf16 as well. Non-fused paths are untouched.
|
Second round -- correct, and fixed in the branch. The delegation I added answered "route swapped inputs through every backend" but The base implementation now handles it in one place: when the fused flag is set Builds clean with clang 24; greedy decoding output over qwen3-0.6b is |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 40e6991806
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (opt.use_bf16_storage && bottom_blob.elembits() == 16) | ||
| cast_bfloat16_to_float32(bottom_blob, q32, opt); |
There was a problem hiding this comment.
On ARMv8.2, Option starts with use_fp16_storage=true, so the documented opt.use_bf16_storage=true setup leaves both flags enabled; NetPrivate::convert_layout and RotaryEmbed_arm::forward select FP16 first in that case (src/net.cpp:399-405 and src/layer/arm/rotaryembed_arm.cpp:37-49). The newly added fallback instead interprets every 16-bit input/cache as BF16 here (and emits BF16 at lines 82-85), so fused models in this configuration turn ordinary FP16 bit patterns into unrelated values. Determine the selected storage format with the same precedence as the backend rather than testing use_bf16_storage alone.
Useful? React with 👍 / 👎.
| else | ||
| cast_float16_to_float32(bottom_blob, q32, opt); |
There was a problem hiding this comment.
When input_hc_swapped=1 is called with an FP32 Mat whose elempack is greater than 1, the guard is entered solely because of the packing, but this unconditional else invokes a Cast configured from FP16. It consequently reads the 32-bit float buffer as 16-bit values before convert_packing, corrupting the entire result; keep q32 as the original FP32 tensor in this case and only use this cast for actual 16-bit FP16 storage.
Useful? React with 👍 / 👎.
| in32[1] = cos32; | ||
| in32[2] = sin32; | ||
| std::vector<Mat> out32(1); | ||
| int ret = forward(in32, out32, opt); // fp32/pack1 now: normal path below |
There was a problem hiding this comment.
Under allocator pressure on a fused 16-bit or packed input, any preceding cast or convert_packing can leave q32, cos32, or sin32 empty because these void helpers discard the underlying -100 return. Calling forward anyway makes an empty q32 re-enter this conversion branch indefinitely (bits and pack are both zero), while an empty cache is later dereferenced; validate every temporary and return -100 before this recursive call.
Useful? React with 👍 / 👎.
… and failure paths Third review round on this branch, all three in the conversion fallback added last round: * fp16-vs-bf16 precedence. Option enables use_fp16_storage by default on ARMv8.2, so a build that also sets use_bf16_storage has both flags on and the backends pick fp16 first. The fallback tested bf16 first, misreading fp16 inputs as bf16. It now matches the backend order. * fp32-but-packed inputs entered the guard because of the packing, yet were still run through the 16-bit cast, reading a 32-bit buffer as fp16. Packed fp32 input now skips the cast and only goes through convert_packing. * The cast/convert_packing helpers return void and leave an empty Mat on allocation failure; recursing with it re-entered the guard forever (an empty Mat is neither fp32 nor packed). Every temporary is validated and the function returns -100 instead.
|
Third round -- all three are in the conversion fallback I added last round, 1. Preserve FP16 dispatch when both storage flags are enabled (P1) Correct. Option enables use_fp16_storage by default on ARMv8.2, so a build that 2. Bypass the FP16 cast for packed FP32 inputs (P2) Correct -- an fp32/pack>1 input entered the guard because of the packing and was 3. Reject failed temporary conversions before recursing (P2) Correct -- the cast helpers return void and leave an empty Mat on failure, and Builds clean with clang 24; greedy decoding output over qwen3-0.6b is |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6c3deb7a4b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (opt.use_fp16_storage) | ||
| cast_float16_to_float32(bottom_blob, q32, opt); | ||
| else | ||
| cast_bfloat16_to_float32(bottom_blob, q32, opt); |
There was a problem hiding this comment.
On x86 (and MIPS/LoongArch), enabling use_bf16_storage leaves use_fp16_storage=true by default, but those backends select BF16 solely from use_bf16_storage; after swapped mode delegates here, this test decodes BF16 bits as FP16 and later emits FP16, corrupting fused BF16 inference in the normal setup. Fresh evidence relative to the prior ARM precedence report is the x86 dispatch in src/layer/x86/rotaryembed_x86.cpp, which selects BF16 even when both flags are set, so the conversion must follow the format actually selected by the backend rather than the global FP16 flag.
Useful? React with 👍 / 👎.
| if (opt.use_fp16_storage) | ||
| cast_float32_to_float16(t32, top_blob, opt); | ||
| else | ||
| cast_float32_to_bfloat16(t32, top_blob, opt); |
There was a problem hiding this comment.
When a fused 16-bit call reaches this post-recursion conversion and blob_allocator cannot allocate the FP16/BF16 output, these void helpers discard Cast's -100 and leave top_blob empty. There is no empty check after this block, and the following packing wrapper also discards its return, so the layer returns success with no output instead of propagating OOM; validate the final converted and repacked blob before returning.
Useful? React with 👍 / 👎.
|
Update: the parts I had marked "compile-checked only" are now measured on real A layer-level test drives
And the ModelWriter change, end to end on the same machine: a param file with So the three rounds of findings on this branch -- backend routing, 16-bit |
…format choice Fourth review round on this branch: * The fused 16-bit conversion decoded by testing use_bf16_storage first, but the backends disagree on precedence: ARM picks fp16 when both flags are set while x86/MIPS/LoongArch select bf16 solely from use_bf16_storage. Decoding with the wrong one corrupts fused bf16 inference on those targets. Each delegating backend now adjusts the Option to the format it would have selected itself before handing off, so the generic conversion decodes the bits as they are. * The final cast-back and repack are void helpers that leave an empty Mat on allocation failure; the layer then returned success with no output. Both are now validated and return -100.
|
Fourth round -- both fixed, and both are now measured on real ARM hardware. 1. Decode BF16 according to the backend-selected format (P1) Correct, and thank you for the x86 pointer -- the backends genuinely disagree: Option opt2 = opt;
if (bottom_blobs[0].elembits() == 16)
{
opt2.use_fp16_storage = <this backend's fp16 condition>;
opt2.use_bf16_storage = opt2.use_fp16_storage ? 0 : <its bf16 condition>;
}
return RotaryEmbed::forward(bottom_blobs, top_blobs, opt2);(ARM: 2. Return OOM when final output conversion fails (P2) Correct -- the final cast-back and repack were void helpers whose failures left Measured on a Radxa Orion O6 (Cortex-A720/A520, Fedora 44, gcc 16.2.1, asimdhp +
and the ModelWriter change round-trips through ncnnoptimize with |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 716f5fc400
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| cast_bfloat16_to_float32(cos_cache, cos32, opt); | ||
| } | ||
| else | ||
| cos32 = cos_cache; |
There was a problem hiding this comment.
When input_hc_swapped=1 is invoked with a packed 2-D cos/sin cache (for example, pack4 over a sequence length divisible by four), this alias and the corresponding sin alias preserve the cache packing even though q32 is unpacked before recursion. The recursive path then iterates the unpacked sequence length while row(i) treats each packed cache group as one row, so it consumes interleaved lane data and reads beyond the physical cache rows once i >= cache.h; convert both caches to pack1 before recursing.
Useful? React with 👍 / 👎.
…urses The fused fallback unpacks q to pack1 but forwarded the cos/sin caches with their original packing. The recursive fp32/pack1 path iterates the unpacked sequence length and indexes the cache with row(i), so a packed 2-D cache (e.g. pack4 over a sequence length divisible by four) would feed interleaved lanes and be read past its physical rows once i >= cache.h. Convert both caches to pack1 alongside q; the existing empty checks right below cover allocation failures of these unpacks.
|
Fifth round -- fixed, with a reproducer. Unpack caches before the packed fallback recurses (P2) Correct. The fallback unpacked Reproducer (K3, clang-24): fused path (
End-to-end on qwen3-0.6b (this path is not exercised by that flow, so the |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d9551508c8
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| opt2.use_fp16_storage = support_fp16_storage && opt.use_fp16_storage ? 1 : 0; | ||
| opt2.use_bf16_storage = opt2.use_fp16_storage ? 0 : (opt.use_bf16_storage ? 1 : 0); | ||
| } | ||
| return RotaryEmbed::forward(bottom_blobs, top_blobs, opt2); |
There was a problem hiding this comment.
On ARMv8.2, when the input uses FP16 storage and use_fp16_arithmetic=true (the default), the normal path dispatches to forward_fp16sa, whose vector body performs the multiply/FMA operations in FP16. Delegating here instead takes the generic conversion path, upcasts q/cos/sin to FP32, computes in FP32, and rounds only when casting the output back, so enabling 1=input_hc_swapped changes the requested arithmetic mode and output bits for what is intended to be a layout-only Permute fusion. Preserve the forward_fp16sa arithmetic in a swapped-layout implementation.
Useful? React with 👍 / 👎.
Delegating the fused (input_hc_swapped) case to the generic implementation upcasts fp16 storage to fp32, computes in fp32, and rounds on the way back -- on ARMv8.2 with use_fp16_arithmetic (the default) that silently changes the requested arithmetic mode for what is a layout-only Permute fusion. Add forward_fp16sa_swapped, mirroring forward_fp16sa's fp16 multiply/FMA with the swapped indexing of the generic fused path (sequence on c, heads on h, output back to the standard layout; interleaved and full-width caches included). The dispatch requires pack1 16-bit storage for all three inputs; fp16s/bf16/packed/mixed-width inputs keep delegating to the generic conversion path, whose fp32 compute matches those modes' semantics.
|
Sixth round -- fixed with a native implementation. Preserve ARM FP16 arithmetic in swapped mode (P2) Correct. With fp16 storage + Verified on a Radxa Orion O6 (ARMv8.2, asimdhp+bf16+sve+i8mm, gcc 16.2.1), |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4a00e66692
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| /* input_hc_swapped=1：输入是 [w, c, h]（吸收掉前面的 Permute(order_type=2)）， | ||
| * 序列在 c 轴、头在 h 轴；输出仍写标准 [w, h, c]。 */ | ||
| const int seqlen = input_hc_swapped ? bottom_blob.c : bottom_blob.h; | ||
| const int num_heads = input_hc_swapped ? bottom_blob.h : bottom_blob.c; |
There was a problem hiding this comment.
When input_hc_swapped=1 and the head count exceeds the sequence length (especially one-token decoding), this makes q range over bottom_blob.h, while the loop still unconditionally evaluates bottom_blob.channel(q) before selecting the swapped pointer. For q >= bottom_blob.c, Mat::channel forms a pointer beyond the input allocation even though head is later ignored, so the common decode path has undefined behavior; construct that view only in the normal-layout branch.
Useful? React with 👍 / 👎.
In swapped mode num_heads is bottom_blob.h, which exceeds bottom_blob.c whenever the head count exceeds the sequence length (one-token decoding is the common case). The loop unconditionally evaluated bottom_blob.channel(q) before selecting the swapped pointer, forming an out-of-bounds view for q >= bottom_blob.c even though the head Mat goes unused. Construct it only for the normal layout.
|
Seventh round -- fixed. Avoid forming normal-layout channels in swapped mode (P2) Correct. In swapped mode Verified: the fused-layout micro suite (fp32/fp16/bf16/dual-flag/non-fused/ |
Summary
riscv(rotaryembed): optional RotaryEmbed+Permute fusion via a layout flag
The q/k projection is followed by a RotaryEmbed and, in the exported graphs, by a
Permute that only exists to satisfy the layout RotaryEmbed expects. This adds an
opt-in parameter (
1=input_hc_swapped) that lets the layer read the input in itsexisting layout instead, so the Permute can be dropped at export time.
中文摘要
riscv(rotaryembed): 通过布局标志实现可选的 RotaryEmbed+Permute 融合
q/k 投影之后接 RotaryEmbed，而在导出的计算图中还会多一个仅仅为了满足 RotaryEmbed
期望布局的 Permute。本 PR 新增可选参数（
1=input_hc_swapped），让该层直接按输入现有布局读取，从而在导出时省掉这个 Permute。
Changes
src/layer/rotaryembed.cpp,rotaryembed.h: the new parameter and the two layoutpaths in
forwardResults
End-to-end
Measured on a SpaceMiT K3, A100 cluster (
ai-run, VLEN=1024), 8 threads,qwen3-0.6b, llama-benchy 0.4.0. Each data point is the median of 3 repeats
(
--pp 128 512 --tg 32 128 --depth 0 --runs 3), with each variant builtstandalone and served by its own process, all other benchmarks stopped.
Note on these numbers: the run-to-run spread of the full-model prefill figure on
this board is up to ~30%, which is larger than what removing one Permute
contributes, so the table above cannot resolve the effect either way. The two
figures that would normally be quoted as opposite signs (+17.5% and -4.3% on the
same prompt length) are the spread itself, not a result, so no chart is shown for
this PR. The change is a structural one: it lets the exporter drop a Permute that
only exists to satisfy the layout, and the follow-up RVV kernel PR (#7082) builds
on the flag it introduces.
中文摘要（数据）
How the numbers were taken
this board vary by up to ~30% on the prefill figure
-march=rv64gcv_zba_zfh_zvfh_zicbopplus the SpaceMiT extensions where the variant needs them
against that same process, so the weight-packing cost is not re-measured
Correctness
The generated text is byte-identical before and after this change (greedy
decoding, same prompt and tokenizer), so the optimisation does not alter results.
No change to any model file format.
Note
The parameter is off by default, so existing models behave exactly as before and
no re-export is required. Gains depend on how much of the graph the Permute
represents; the numbers above are end-to-end.