qsmla: TransB #257 adaptation + dynamic-shape v2 + single-PE CSA kernel (FP16/HIF8) - #151
Conversation
TileOP-API PTO-ISA#104 (01d62c4) redefined the cooperative Shared-B TransB control as a physical-storage-shape declaration per pto-spec #257 (ASL BundleMatrixSharedBPrimarySchemaLegal): TransB=0 -> physical [N, K] (K contiguous) TransB=1 -> physical [K, N] (N contiguous) QSMLA stages both K and V as SharedMatrixRight<kvdtype, kTk, kTd> (one KV token per row, D contiguous), so: - QK (Q @ K^T): the physical [kTk, kTd] layout IS the [N, K] form for the K^T operand -> drop transpose_b() (TransB=0) - PV (P @ V): the same physical layout is the [K, N] form for the V operand (N contiguous) -> add transpose_b() (TransB=1) Without this adaptation the kernel fails the compile-time effective-K and cooperative-D shape static_asserts against TileOP-API >= 748ac69. Verified with SuperScalarModel fix/gfrun-shared-tmatmul-layout-257 @ 6a60fb5 (the model-side TransB fix, de3ea07f), TileOP-API 64ef554, LLVM 10abd6e+: gfrun 4-PE on swa_small / csa_small / ori_sparse_small all 32768/32768 (100%) with max_abs bitwise identical to the pre-PTO-ISA#104 baselines. Refs: PTO-ISA/pto-spec#255, pto-spec#257, LinxISA/Linx-TileOP-API#103, LinxISA/Linx-TileOP-API#104, LinxISA/Linx-TileOP-API#122
Same convention as mega_moe_sim.cpp: write the pass value to the test-finisher address after BENCHEND so gfsim --dump-memory can complete (requires SuperScalarModel#665 finisher observation fix). Numerical verification remains gfrun R2=0. Not part of PR#141 (which only contains the TransB adaptation); this is a separate local commit for gfsim performance collection.
…ariant v2
Add quant_sparse_flash_mla_v2.hpp, a variant of the unified five-mode
four-PE TADD kernel where all Vec-engine tiles use DYNAMIC
ValidRow/ValidCol instead of compile-time constants:
- Valid regions are set at runtime via the Tile(VR, VC) constructor per
KV block, so B.DIM encodes the register form ("B.DIM %[reg], 0")
instead of the immediate form ("B.DIM zero, %c[imm]").
- tileW's ValidCol tracks the actual number of valid KV tokens in each
block (instead of the full kTk), so TROWMAX/TROWSUM and broadcasts
operate on the true valid region. The software mask is retained for
pass-2 PV correctness, but the hardware now also sees the correct
geometry.
- Shared tiles (cooperative TMATMUL operands) and the Cube accumulator
stay static: PTO v0.58 requires compile-time valid shapes for
cooperative matrix operations (TMATMUL static_assert).
- gmGatherKV and iterators are unchanged (they consume the physical
shape).
Enabled with QSMLA_DYNAMIC_SHAPE=on (unified four-PE IMPLs only).
Default build is unchanged.
Verified on ops-20260915 (main 4b6ae0a) + SuperScalarModel 6a60fb5 +
TileOP-API 64ef554 + llvm 7c1388d: csa_small / swa_small /
ori_sparse_small gfrun (HIF8/BF16, 4 threads) all pass 32768/32768
with max_abs bit-identical to the static kernel (0.004676545 /
0.004920054 / 0.009070974, atol=rtol=2e-2). ELF encoding check:
95 register-form B.DIM + 211 immediate-form C.B.DIMI.
Add quant_sparse_flash_mla_csa_1pe.hpp, a single-PE Local CUBE implementation of the CSA mode with the same dual-source semantics as the unified four-PE kernel (ORI window + CMP TopK logical sources, two-pass online softmax), plus the tadd.hpp Tk16 pre-transposed rebuild (square-B local-tile contract, TTRANS retired) and build wiring (IMPL=csa_tadd_1pe, embedded csa_small input). dtype dual state via if constexpr: - FP16: descale=1, no P quantization (FP16 output) - HIF8: aligned with the official aclnnQuantSparseFlashMla contract (per-tensor descale on scores, P*16 -> HIF8, PV*kvDescale, output/16, BF16 output); HIF8 CMP gather uses uint32 carriers because LinxV5 cannot legalize divergent scalar i8/HIF8 copies. Verification (csa_small: B=1 S1=1 oriS2=128 cmpS2=64 N1=64 D=512 cmpTopK=40 ratio=4, gfrun R2=0): - FP16: 32768/32768 pass @ atol=rtol=1e-3, max_abs=9.07e-06 - HIF8: 32768/32768 pass @ atol=rtol=2e-2, max_abs=4.68e-03
![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: af884527fd
ℹ️ 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".
| static_assert(kTk == kTd, | ||
| "unpatched-headers single-PE tadd requires square B " | ||
| "tiles (Tk == Td_block)"); |
There was a problem hiding this comment.
The existing single-PE invocation still defaults to Tk=32 and Td_block=64 in the test Makefile (and is documented that way at the top of this header), so every IMPL=tadd QSMLA_DTYPE=FP16 build that relies on those established defaults now fails this new assertion. Please make the tadd defaults square (for example, conditionally default Td_block to 32) or preserve a compatible non-square path so the existing benchmark remains buildable.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Valid finding, fixed in c534397. Two notes on the fix: (1) Td_block ?= 32 alone would still leave the default build broken in a worse way — the [32,32] profile compiles but violates the pto-spec 0.58.6 2048-byte row-reduction source contract (TileOP has no compile-time check, so it faults the model at runtime; see the CSA kernel note on the same constraint). So IMPL=tadd now defaults to the validated Tk16 square profile (Tk ?= 16 / Td_block ?= 16, other IMPLs keep the legacy defaults), and tadd.hpp gains the missing kTm*kTk*4 <= 2048 static_assert so an explicit Tk=32 build fails with a clear message. Verified: default IMPL=tadd QSMLA_DTYPE=FP16 now builds and passes 32768/32768 @ 1e-3 against the tadd golden, max_abs identical to the 09-15 archive.
| alignas(64) kvdtype kv_t[D * s2]; | ||
| for (int t = 0; t < s2; ++t) | ||
| for (int d = 0; d < D; ++d) | ||
| kv_t[d * s2 + t] = ori_kv_ptr[t * D + d]; |
There was a problem hiding this comment.
When N1 > 1 selects quant_sparse_flash_mla_swa_tadd_bsnd_pto, its run_full_rows/run_tail_rows helpers call this function separately for every work item and head chunk, so this whole-sequence transpose is repeated even though those calls share the same batch KV. For the supported B=8, S1=4, N1=128, N2=1, Tm=32 shape, the identical D × S2 input is transposed 128 times, making scalar preprocessing dominate the benchmark as S2 grows; hoist or cache the transpose at the BSND dispatcher/batch level.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Valid finding, fixed in c534397. The BSND dispatcher now owns one batch-shared whole-sequence transpose (change-detected per batch — work items decode batch-major, and with N2==1 the kv offset depends only on batch) and passes it into config_pto through a new optional ori_kv_t_ptr parameter; the N1==1 entry point keeps the in-function transpose. Verified on B=2/S1=2/S2=64/N1=128: output bit-identical before/after, Total Inst 54,013,996 -> 41,140,992 (-23.8%), total blocks -45%; csa_tadd_1pe regression unchanged (100% @ 1e-3, instruction count identical).
… K^T) Two review findings on the tadd path, both fixed: 1. The legacy Tk=32/Td_block=64 defaults now fail the new square-B static_assert, so every default IMPL=tadd build broke. The defaults are not just non-square: the [32,32] profile also violates the pto-spec 0.58.6 2048-byte row-reduction source contract (TileOP has no compile-time check, so it compiled but faulted the model's runtime contract). Fix: IMPL=tadd now defaults to the validated Tk16 square profile (Tk=16/Td_block=16, 2048-byte tW), other IMPLs keep the legacy defaults, and tadd.hpp gains the missing row-reduction static_assert so an explicit Tk=32 build fails with a clear message instead of a runtime fault. 2. The BSND dispatcher re-transposed the whole [D, S2] KV sequence for every head chunk (e.g. 128 identical transposes for the B=8/S1=4/N1=128 shape) because config_pto owned the transpose. Fix: the BSND dispatcher now owns one batch-shared transpose (change-detected per batch; work items decode batch-major) and passes it to config_pto through a new optional ori_kv_t_ptr parameter. The N1==1 entry point keeps the in-function transpose. Verification (gfrun R2=0): - IMPL=tadd default build now compiles (was: static_assert failure) and passes 32768/32768 @ 1e-3 against the tadd golden, max_abs identical to the 09-15 archive (1.22041e-04) - BSND B=2/S1=2/S2=64/N1=128: output bit-identical before/after; Total Inst 54,013,996 -> 41,140,992 (-23.8%), blocks -45% - csa_tadd_1pe FP16 regression: 100% @ 1e-3, instruction count unchanged (7940132); the review fixes only touch the tadd path
|
Both Codex review findings are addressed in c534397 (see inline replies):
All gfrun runs report R2=0. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c534397fbb
ℹ️ 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".
| Tm ?= 32 | ||
| Tk ?= 32 | ||
| Td_block ?= 64 | ||
| ifeq ($(IMPL), tadd) |
There was a problem hiding this comment.
A build selecting the newly added IMPL=csa_tadd_1pe without explicit tile overrides takes the else defaults Tk=32 and Td_block=64, but quant_sparse_flash_mla_csa_1pe_pto requires kTk == kTd and also limits kTm * kTk * 4 to 2048 bytes. Consequently, the documented csa_tadd_1pe configuration fails both new static assertions with the Makefile defaults; this implementation needs the same validated 16×16 defaults as the other single-PE path.
Useful? React with 👍 / 👎.
| tileKRight tK; | ||
| auto gK = itOriT(dd, ori_blk_begin + j); | ||
| TLOAD_CUBE(tK, gK); |
There was a problem hiding this comment.
When Config::S2 is not divisible by kTk and the SWA window reaches the end of the ORI sequence, the final iterator load still reads a full kTk columns from ori_t, although that scratch buffer is allocated with an exact row stride of s2. The load therefore crosses into the next transposed row (and past the allocation for the final D row); the later mask cannot prevent the out-of-bounds access, and the corresponding full-width V load can also let an invalid NaN contaminate P·V despite a zero probability. Pad ORI storage/rows to a tile-aligned capacity as is already done for CMP, or explicitly reject non-aligned S2 values.
Useful? React with 👍 / 👎.
| ifeq ($(QSMLA_DYNAMIC_SHAPE), on) | ||
| ifeq ($(filter $(IMPL),$(QSMLA_UNIFIED_IMPLS)),) | ||
| $(error QSMLA_DYNAMIC_SHAPE requires a unified four-PE IMPL) | ||
| endif | ||
| DEFINES += -DQSMLA_USE_DYNAMIC_SHAPE |
There was a problem hiding this comment.
When the same unified configuration is built once normally and once with QSMLA_DYNAMIC_SHAPE=on, this flag changes the selected kernel but is not represented anywhere in TARGET. Both builds therefore produce the identical ELF path, so the second build overwrites the first and benchmark/result collectors cannot distinguish the static and dynamic implementations. Add a dynamic/static variant suffix to the target name when enabling this define.
Useful? React with 👍 / 👎.
2 participants
Summary
Consolidates the three QSMLA work streams into one PR with a chronological commit history. Supersedes #141 and #149 (identical content for the shared commits; rebase onto current main is conflict-free).
qsmla: adapt Shared-B TransB usage to pto-spec #257qsmla: add gfsim test-finisher writegfsim --dump-memorycan completeqsmla: add dynamic-shape (DYNAMIC ValidRow/ValidCol) four-PE kernel variant v2QSMLA_DYNAMIC_SHAPE=on(default build unchanged)qsmla: add single-PE CSA kernel (FP16/HIF8 dual dtype, Local CUBE)1. TransB adaptation (from #141)
TileOP-API #104 redefined the cooperative Shared-B
TransBcontrol as a physical-storage-shape declaration per pto-spec #257. The QSMLA kernel still used the legacy math-transpose-flag reading and failed the compile-time effective-K / cooperative-D shapestatic_asserts since TileOP-API748ac69.2. gfsim test-finisher (from #141 branch)
Same convention as
mega_moe_sim.cpp: write the pass value to the test-finisher address after BENCHEND so gfsim--dump-memorycan complete. Numerical verification remains gfrun R2=0.3. Dynamic-shape v2 (from #149)
quant_sparse_flash_mla_v2.hpp: dynamic-shape variant of the unified five-mode four-PE TADD kernel. Vec tiles use DYNAMIC ValidRow/ValidCol so the hardware only sees the live tile geometry. All new code sits behindifeq/#ifdefguards; the default build is unchanged.4. Single-PE CSA kernel (new)
quant_sparse_flash_mla_csa_1pe.hpp: single-PE Local CUBE implementation of the CSA mode with the same dual-source semantics as the unified four-PE kernel (ORI window + CMP TopK logical sources, two-pass online softmax), square-B tiles (Tk==Td, unpatched Local-B contract), row-reduction source <= 2048B, and tiled four-segment MM1 (no generic lambdas: cross-source tile capture triggers raw TSTORE stack spill).aclnnQuantSparseFlashMlacontract — per-tensor descale on scores, P16 -> HIF8, PVkvDescale, output/16, BF16 output. HIF8 CMP gather uses uint32 carriers (LinxV5 cannot legalize divergent scalar i8/HIF8 copies).quant_sparse_flash_mla_tadd.hpp: Tk16 pre-transposed rebuild (TTRANS retired, square-B static_assert)Verification (csa_small: B=1 S1=1 oriS2=128 cmpS2=64 N1=64 D=512 cmpTopK=40 ratio=4, gfrun R2=0)
IMPL=csa_tadd_1pe QSMLA_DTYPE=FP16IMPL=csa_tadd_1pe QSMLA_DTYPE=HIF8Both ELFs rebuilt from the consolidated branch tip and re-verified against the CPU golden (
qsmla_compare.py).