Add one-level RMSNorm kernel with tiling harness and precision check. - #26
Add one-level RMSNorm kernel with tiling harness and precision check.#26wangyuascend-spec wants to merge 5 commits into
Conversation
Introduce fp16 RMSNorm (fp32 compute) under one-level-arch, with tiling_info specified in the harness and scripts to generate/compare golden outputs. Co-authored-by: Cursor <cursoragent@cursor.com>
![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: ac0c5f6919
ℹ️ 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".
| template <typename TileVec> | ||
| inline void rsqrt_newton(TileVec &out, TileVec &a) { | ||
| TileVec x, t1, t2; | ||
| TRECIP(x, a); |
There was a problem hiding this comment.
For rows whose mean square is far from 1, this initializes the reciprocal-square-root Newton iteration with 1/a instead of a value near 1/sqrt(a). The update x *= 1.5 - 0.5*a*x*x only converges from a suitable rsqrt seed; with small activations, such as a row of 0.1s where a≈0.01, the iteration overflows or changes sign, so RMSNorm emits infinities or wildly scaled outputs even though the golden path computes a finite 1/sqrt(mean_sq + eps). Use a true rsqrt/sqrt-based seed or another stable approximation.
Useful? React with 👍 / 👎.
| // const int64_t tile_n = gN; | ||
| /*以下写法能好*/ | ||
| const int64_t tile_m = 1; | ||
| const int64_t tile_n = 512; |
There was a problem hiding this comment.
When tiling[1]/tiling[3] describes anything other than the default 512 columns, this hard-coded tile_n makes the kernel silently process only 512 columns while still dividing the row sum by gN. For example, gN=1024 leaves the second half of each row unstored and normalizes the first half with a denominator based on only half the data; gN<512 makes TLOAD read past the logical row. Either enforce gN == tile_n == 512 up front or loop over all N tiles and accumulate the full row sum before storing.
Useful? React with 👍 / 👎.
|
|
||
| : "${COMPILER_DIR:?Set COMPILER_DIR to the linx_blockisa_llvm_musl bin directory}" | ||
|
|
||
| make TESTCASE=rms_norm COMPILER_DIR="$COMPILER_DIR" DType=__half diss |
There was a problem hiding this comment.
This per-case compile script is not reachable from the advertised benchmark/one-level-arch/compile_all.sh path: I checked that script's explicit operator list and it still omits test/kernel/normalization/rms_norm. In CI or benchmark runs that rely on the full compile script, the new kernel will never produce an ELF or get disassembled/validated, so add a compile_operator entry for this directory.
Useful? React with 👍 / 👎.
Static RowMajor shape only accepts a data pointer; keep compile-time Valid and avoid unsigned size_t wrap on i*gN. Co-authored-by: Cursor <cursoragent@cursor.com>
Switch tiles/GM to dynamic Valid (-1) and construct with active dims so tiling-driven shapes work on the register B.DIM path. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 21ba32b3e9
ℹ️ 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".
| using tile_h = Tile<Location::Vec, dtype, tM, tN, BLayout::RowMajor, -1, -1>; | ||
| using tile_f = Tile<Location::Vec, float, tM, tN, BLayout::RowMajor, -1, -1>; | ||
| // ValidCol=1; Cols=32 → 128B PE-local (TSize min). | ||
| using tile_v = Tile<Location::Vec, float, tM, 32, BLayout::RowMajor, -1, 1>; |
There was a problem hiding this comment.
With the one-level Makefile the kernel is compiled with -D__linx, and the __linx implementations of TRECIP and TROWSUM still static_assert that tile ValidRow/ValidCol are not DYNAMIC. These aliases set the RMSNorm data and vector tiles to -1 valids, so the advertised compile.all path fails as soon as TROWSUM(sqrsum, squared) or TRECIP(x, a) is instantiated instead of producing an ELF; keep the one-level tile valids static or avoid those unsupported ops for dynamic-valid tiles.
Useful? React with 👍 / 👎.
| // tiling_info: {g_m, g_n, tile_m, tile_n} | ||
| int64_t tiling_info[4] = {16, 512, 1, -1}; |
There was a problem hiding this comment.
The data generator writes tiling_info.bin and exposes --g-m, --g-n, and tile arguments, but this driver always reads and computes the hard-coded 16 x 512 shape. If someone generates a non-default RMSNorm case and then runs the res_check path, the binary input size no longer matches what this code reads, and the kernel either fails the read and runs on uninitialized stack data or processes the wrong prefix with the wrong stride; load the generated tiling metadata or remove the shape options from the check path.
Useful? React with 👍 / 👎.
Handle full tiles in the main loop and the trailing ValidRow block separately; remove early-return shape checks from the kernel. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0ae99dbba7
ℹ️ 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 int64_t gM = tiling[0]; | ||
| const int64_t gN = tiling[1]; | ||
| const int64_t tile_m = tiling[2] > 0 ? tiling[2] : tM; |
There was a problem hiding this comment.
When the tiling metadata requests tile_m > 1 (the generator exposes --tile-m and the public API documents tiling[2]), this value is passed as active_row even though all RMSNorm tile aliases are compiled with tM = 1. The dynamic-valid tile constructors do not resize the backing tile, so TLOAD/TSTORE and the elementwise ops iterate multiple rows against one row of tile storage, corrupting memory or skipping rows; either keep tile_m fixed to tM or instantiate tile types with the requested row capacity.
Useful? React with 👍 / 👎.
Co-authored-by: Cursor <cursoragent@cursor.com>
1 participant
Introduce fp16 RMSNorm (fp32 compute) under one-level-arch, with tiling_info specified in the harness and scripts to generate/compare golden outputs.