feat(normalization): add tree-reduction static RMS Norm - #144
wangyuascend-spec wants to merge 3 commits into
Conversation
|
The two blockers for this PR are now resolved on the default branches:
Rebuild against current |
Replace the three mxquant variants (mxquant [128,64], mxquant_1024_assembly and mxquant_1024_streaming [128,1024]) with a single [512,256] operator that matches the dynamic_mx_quant solution's tail OCP shape: four PEs split M (128 rows each), each 32-row block is quantized one 32-wide MX group at a time (8 groups/block, 4 blocks/PE), and the eight per-group E8M0 columns are packed four-at-a-time into two U32 words per row (two full-CELL 128 B scale stores per row block instead of eight 1/4-full ones). The byte order stays plain row-major [512,8] (scale(row, group) = scales[row*8 + group]) and the test/reference contract is unchanged; the tests are reworked to the same single variant. Why max(x,-x) instead of abs: on the current gfrun a unary TEPL op writes its CUBE_M32 destination densely instead of through the CELL indexer, which permutes the [32,32] tile (LinxISA/SuperScalarModel issue #678). TSUB/TMAX are CUBE-cell aware and max(x,-x) is bit-exact for every input, so the scale/quant result is identical to an abs-based one. The removed variants used TABS and were only correct on other model builds. Performance: this kernel is TLSU bound. Each 32-wide group is TLOADed/TSTOREd with its own strided view into the [32,256] row block, i.e. 64 B payload bursts at a 512 B stride (32 B at 256 B on the store side), both well under the 256 B cacheline, so per-access efficiency is burst/256; the non-contiguous stores additionally pin the SCB. The fix is one contiguous [32,256] TLOAD (512 B rows) plus TPARTVIEW subviews for the per-group reductions and TASSEMBLY/B.ASSEMBLE on the store side, which is blocked today: TPARTVIEW's SubTileView is not yet a general tile operand (LinxISA/Linx-TileOP-API issue #144) and gfsim does not model B.ASSEMBLE. The disabled #if 0 variant in mxquant.hpp records the blocked contiguous form. Bandwidth model (from kernels/multi_thread/tlsu_bandwidth_usage.md): * one TLSU op costs ~ fixed overhead + bytes/256; load ~15 cyc, store ~31 cyc; * only burst >= CACHELINE(256 B), or stride == burst, fills a cacheline; a smaller burst costs burst/256 efficiency and still pays a tag/beat per cacheline; * measured on the same [512,256] MX quant: the contiguous [1,16384] reshape is ~1367 cyc total, while the per-group [512,256] strided form is ~12676 cyc (TLSU Active Wall ~99%, STQ peak 15, SCB 30-46 of 64, store-side back-pressure into the frontend); * store is dearer than load (31 vs 15) and is pinned by the SCB global-visibility window, so cutting store count / making stores contiguous has the largest payoff. Verification (res_check=on; gfrun -s softcore.multiThreadNum=4 -f <elf>): payload 131072/131072, scale 4096/4096 bytes, byte-exact; input_readback.bin == input.bin; both goldens non-zero; perturbing one golden byte is detected at the expected offset; the default kernel-only build compiles.
4240b5e to
6618d6a
Compare
6618d6a to
d32897c
Compare
2 participants
Summary
Base
Validation
Compilation currently reaches TileOP TROWSUM/TROWSUM_ASS lowering but is blocked by LinxISA/llvm-project#94:
The earlier TROWSUM physical-shape assertion is resolved by TileOP PR #136. No ELF or gfrun result is available for this new testcase until llvm-project#94 is fixed.
Dependency