Add Quant Sparse Flash mla (QSMLA) operator - #39
Conversation
![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: 82983fbbb7
ℹ️ 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".
| int* metadata, | ||
| float* softmax_lse) | ||
| { | ||
| constexpr int Db = D / kTd; |
There was a problem hiding this comment.
When Td_block does not divide D (both are exposed as Makefile parameters), this floors the number of D tiles and silently skips the tail columns in both QK accumulation and output stores; for example D=500 Td_block=64 processes only the first 448 dimensions and leaves O[:,448:] unstored/wrong. Add a divisibility static_assert or a remainder tile path before accepting arbitrary D/Td_block values.
Useful? React with 👍 / 👎.
| if fail_count == 0: | ||
| print("\n=== RESULT: PASS ===") | ||
| return 0 |
There was a problem hiding this comment.
If the NPU dump is truncated or has extra elements but the shared prefix is within tolerance, this branch still prints PASS because length mismatches only emit a warning. In that scenario the verifier can approve an incomplete QSMLA result, so the length mismatch should contribute to failure before returning success.
Useful? React with 👍 / 👎.
| #define S1 64 | ||
| #define S2 128 |
There was a problem hiding this comment.
The new Makefile target exposes s1, s2, Tm, Tk, window, and scale knobs, but the CPU reference hard-codes the default shape here. Running a non-default benchmark such as make TESTCASE=quant_sparse_flash_mla s1=128 still produces a 64-row golden (and other non-default knobs produce a different attention pattern), so the verifier no longer matches the generated ELF configuration.
Useful? React with 👍 / 👎.
| const int Qb = (s1 + kTm - 1) / kTm; | ||
| const int Kb = (s2 + kTk - 1) / kTk; |
There was a problem hiding this comment.
Because s1, s2, Tm, and Tk are Makefile knobs, a non-multiple such as s2=96 Tk=64 creates a tail KV block, but the tile loads/stores still use full kTm/kTk shapes with no valid-row masking. The last block will read rows beyond the allocated Q/KV buffers and can also store past the output for a partial Q block; either require divisibility or add remainder tile handling.
Useful? React with 👍 / 👎.
| typedef float f32_t; | ||
|
|
||
| static void init_deterministic_f32(f32_t* data, int count, int seed) { | ||
| for (int i = 0; i < count; ++i) { |
There was a problem hiding this comment.
The NPU test stores the deterministic inputs after casting each value to __half, but the golden path keeps the same decimal values as full floats here. Since most 0.01-spaced values are not exactly representable in FP16, the verifier compares the kernel's result for rounded inputs against a different FP32 problem; the CPU reference should round/dequantize the inputs the same way as the NPU test or consume the dumped test inputs.
Useful? React with 👍 / 👎.
| const int Qb = (s1 + kTm - 1) / kTm; | ||
| const int Kb = (s2 + kTk - 1) / kTk; | ||
|
|
||
| const float scale = softmax_scale; |
There was a problem hiding this comment.
For callers that pass actual quantized Q/KV with non-unit q_descale or ori_kv_descale, the kernel currently scales logits only by softmax_scale and never applies either descale buffer, so both the QK scores and the final V contribution are computed in quantized units rather than real values. Either apply the descales (including the V-side scale) or reject/rename this path as supporting only pre-dequantized inputs.
Useful? React with 👍 / 👎.
| int* seqused_ori_kv, | ||
| float* sinks, | ||
| int* metadata, | ||
| float* softmax_lse) |
There was a problem hiding this comment.
When a caller supplies a non-null softmax_lse, this output pointer is accepted but never written after the pass-1 tMax/tSum reduction. Any backward or validation path that consumes LSE will read stale data even though the kernel had the information needed to store m + log(l), so either fill it or require the argument to be null.
Useful? React with 👍 / 👎.
| } | ||
| BENCHEND; | ||
|
|
||
| return 0; |
There was a problem hiding this comment.
When this testcase is built with the repository's res_check=on flow, RES_CHECK is defined but the program returns without writing out to a result file, unlike the neighboring FA tests that emit res.bin. That makes standard result collection see no QSMLA output unless the user separately uses the dump-memory path, so add a guarded writeBinaryFile for the output buffer.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b2e7f31f13
ℹ️ 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".
| int hi = diagonal + win_right; | ||
| for (int kv = 0; kv < s2; ++kv) { | ||
| bool valid = (kv >= lo) && (kv <= hi); | ||
| mask[q * s2 + kv] = valid ? 0.0f : -1e30f; |
There was a problem hiding this comment.
When the exposed shape/window knobs produce a query row whose window does not intersect [0, s2) (for example s1=128 s2=64 wleft=1 wright=1 for the early query rows), every mask entry in that row is this finite -1e30f. The online softmax then takes the row max as -1e30f, exponentiates all masked logits to 1, and returns an average of V instead of a fully masked/zero result, so add a guard/static assertion or explicit empty-row handling.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5e04af77a1
ℹ️ 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".
| } | ||
| } | ||
| // TLOAD mask tile from blockMask (stack buffer) | ||
| using gmBlockMask = global_tensor<uint32_t, RowMajor<kTm, kTk>>; |
There was a problem hiding this comment.
With the default kTk=32, blockMask only contains kTm * 1 packed words, but this RowMajor<kTm, kTk> view tells TLOAD that each row has 32 uint32_t columns. The load therefore reads subsequent rows as extra columns and then runs past the stack buffer, so the TSEL mask applied to the very first QSMLA configuration is corrupted; expose the buffer as [kTm, maskWordsPerRowBlock] or expand it to the full tile shape before loading.
Useful? React with 👍 / 👎.
| ifeq ($(TESTCASE), quant_sparse_flash_mla) | ||
| SRC_FILE += $(TEST_ROOT)/$(CASE_SRC_DIR)/quant_sparse_flash_mla.cpp |
There was a problem hiding this comment.
This direct TESTCASE branch makes QSMLA buildable with make, but I checked benchmark/one-level-arch/test/kernel/fa/compile.all (the script run by compile_all.sh for FA) and it has no TESTCASE=quant_sparse_flash_mla entry. The new operator is therefore skipped by the repository's batch build/regression path, so toolchain or kernel breakages here won't be caught unless someone remembers to invoke it manually.
Useful? React with 👍 / 👎.
| DEFINES += -DTsoftmax_scale=$(softmax_scale) | ||
| DEFINES += -DTwleft=$(wleft) | ||
| DEFINES += -DTwright=$(wright) | ||
| TARGET = $(ELF_HEAD)/$(TESTCASE)_s1$(s1)_s2$(s2)_D$(D)_Tm$(Tm)_Tk$(Tk)_Td$(Td_block).elf |
There was a problem hiding this comment.
When wleft, wright, or softmax_scale are changed, the generated code changes but the TARGET path remains identical for the same shape/tile sizes. Sweeping window/scale experiments will overwrite or mislabel the previous ELF (and any CHK_DIR derived from the target name), so include these semantic parameters in the target name.
Useful? React with 👍 / 👎.
| // TSEL 要求 dst/mask/src 同 tile_shape, 但模拟器内部 mask 按 uint32 位打包读 | ||
| // 这里用 float 类型满足编译器约束, 实际数据是 uint32 位掩码 | ||
| // tile Cols 设为 kTk (与 score tile 同形状), 模拟器只读前 maskWordsPerRowBlock 个 uint32 | ||
| using tileMask = Tile<Location::Vec, float, kTm, kTk, BLayout::RowMajor>; |
There was a problem hiding this comment.
This float mask tile is filled from packed uint32_t words rather than from an elementwise TCMP-style mask with one condition value per score lane. With the default sliding window most columns in each [kTm,kTk] score tile should be masked, but the packed word is converted into a single float value in the first lane instead of setting the corresponding per-column mask lanes, so TSEL cannot reliably replace the intended scores with tNegInf; build a full elementwise mask tile or use the TADD mask path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f068f9369f
ℹ️ 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".
| int* metadata, | ||
| float* softmax_lse) | ||
| { | ||
| constexpr int Db = D / kTd; |
There was a problem hiding this comment.
When callers use the Makefile knobs with D not divisible by Td_block, this floors the number of D tiles; the QK accumulation loops only visit dd < Db and the final store loop writes only those same blocks, so a run such as D=500 Td_block=64 silently ignores and never stores the last 52 columns. The fresh evidence is that the current test driver includes quant_sparse_flash_mla_onepass_pto.hpp, so the truncation also affects the one-pass path that is actually built; add a divisibility assertion or a remainder-tile path here too.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7f378237cb
ℹ️ 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 Qb = (s1 + kTm - 1) / kTm; | ||
| const int Kb = (s2 + kTk - 1) / kTk; |
There was a problem hiding this comment.
When the exposed Makefile knobs make s1 or s2 not divisible by kTm/kTk, these rounded-up block counts let the one-pass kernel enter a tail block while every Q/K/V/mask load and output store still uses a full tile. For example s2=96 Tk=64 makes the last KV iteration read past the allocated KV buffer, and a partial Q block can store past out; add valid-row/column masking or require divisibility. The current driver instantiates this one-pass variant, so the tail-tile issue affects the path that is actually built.
Useful? React with 👍 / 👎.
| } | ||
|
|
||
| // --- Step 2: scale + mask --- | ||
| TMULS(tW, tW, scale); |
There was a problem hiding this comment.
In the one-pass variant, logits are scaled only by softmax_scale even though the API accepts q_descale and ori_kv_descale. When callers pass actual quantized Q/KV with non-unit descales, both QK and the V contribution are accumulated in quantized units rather than real values, producing incorrect output; apply the Q/K descales and the V-side scale here or reject non-unit descales. The current driver instantiates this one-pass variant, so this omission affects the built path too.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: df775c1d59
ℹ️ 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".
| print(f"passed = {sum(passed)}/32768 ({100*sum(passed)/32768:.6f}%)") | ||
| print(f"failed = {32768-sum(passed)}") |
There was a problem hiding this comment.
When any element is outside tolerance, this compare script only prints the pass/fail counts and then exits with status 0, so any automation that invokes it as the QSMLA verifier will accept a wrong NPU dump as long as the files can be unpacked. Make the failed count or NaN count drive a non-zero exit status after reporting the metrics.
Useful? React with 👍 / 👎.
| int* seqused_ori_kv, | ||
| float* sinks, | ||
| int* metadata, | ||
| float* softmax_lse) |
There was a problem hiding this comment.
When callers pass a non-null softmax_lse to the one-pass kernel, the function computes the final online-softmax tMax/tSum state but never stores m + log(l) to this pointer, so validation or backward paths consuming LSE will read stale data. Fresh evidence is that the current test driver includes and calls quant_sparse_flash_mla_swa_onepass_pto, so the LSE omission also affects the path being built.
Useful? React with 👍 / 👎.
| import math | ||
| import struct | ||
|
|
||
| path = "SuperNPUBench/benchmark/one-level-arch/test/kernel/fa/src/" |
There was a problem hiding this comment.
When this verifier is run from the repository root or from its own src directory, this hard-coded relative prefix looks for SuperNPUBench/... under the current directory rather than next to the script, so it cannot find the golden/NPU dumps produced in the normal checkout. Use script-relative paths or CLI arguments so the comparison works outside one specific parent-directory layout.
Useful? React with 👍 / 👎.
| using gmO = global_tensor<odttype, RowMajor<s1, D>>; | ||
|
|
||
| using tileQ = TileLeft<qdtype, kTm, kTd>; | ||
| using tileKV = TileRight<kvdtype, kTk, kTd>; |
There was a problem hiding this comment.
When the two-pass TSEL kernel is selected, TMATMUL(tW_out, tQ, tK) reduces over tQ's kTd columns, so the right operand must expose kTd rows and kTk output columns. With the default kTd=64, kTk=32, this declaration loads a 32x64 tile instead, causing the QK matmul to read the wrong orientation/OOB tile rows for half of the D reduction and corrupt the logits; mirror the one-pass ColMajor<D, s2> + TileRight<kTd, kTk> view here.
Useful? React with 👍 / 👎.
| using gmMask = global_tensor<float, RowMajor<s1, s2>>; | ||
|
|
||
| using tileQ = TileLeft<qdtype, kTm, kTd>; | ||
| using tileKV = TileRight<kvdtype, kTk, kTd>; |
There was a problem hiding this comment.
When using the TADD fallback, the QK TMATMUL still needs the KV tile shaped as [kTd, kTk] because tQ has kTd columns, but this declaration makes the right operand [kTk, kTd]. For the default 64-wide D block and 32 KV columns, the matmul indexes beyond the declared right-tile rows and accumulates transposed/corrupt scores, so this fallback should use the same transposed KV view as the one-pass implementation.
Useful? React with 👍 / 👎.
| with open(path + "qsmla_onepass_npu_out.bin", "rb") as f: | ||
| actual = struct.unpack("<32768e", f.read()) | ||
|
|
||
| with open(path + "qsmla_golden.bin", "rb") as f: | ||
| golden = struct.unpack("<32768f", f.read()) |
There was a problem hiding this comment.
When the benchmark is built with non-default exposed dimensions such as s1=128 or a different D, the NPU and golden files no longer contain 32768 elements, but the verifier still unpacks exactly the default sizes. That makes otherwise valid non-default QSMLA runs fail in the compare step (or compare the wrong amount if this is later changed to tolerate sizes), so read the expected element count from arguments/metadata or compute it from the same shape knobs.
Useful? React with 👍 / 👎.
1315e2f to
a9855a6
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a9855a6b7f
ℹ️ 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".
| qdtype* q = (qdtype*)(((uint64_t)qp & ALIGN_MASK) + ALIGN); | ||
| kvdtype* kv = (kvdtype*)(((uint64_t)kvp & ALIGN_MASK) + ALIGN); | ||
|
|
||
| odttype* out = (odttype*)MAP_MEM_BASE; |
There was a problem hiding this comment.
When this testcase is run through the existing make sim path, out points at the hard-coded 0x4000802000 address instead of an allocated buffer. The common simulator invocation only provisions -s 4096M, and the one-level memory map/stack are around 0x60000000..0xb0000000, so the first TCOPYOUT writes outside the memory that the other FA tests use and can fault or make the result undumpable. Allocate and align an outp buffer like the neighboring FA drivers, or explicitly map this address before using it.
Useful? React with 👍 / 👎.
|
|
||
| path = "SuperNPUBench/benchmark/one-level-arch/test/kernel/fa/src/" | ||
|
|
||
| with open(path + "qsmla_onepass_npu_out.bin", "rb") as f: |
There was a problem hiding this comment.
The Makefile-built QSMLA driver currently includes and calls quant_sparse_flash_mla_swa_tadd_pto, but this verifier always opens qsmla_onepass_npu_out.bin. When a TADD run is being validated, the script will either fail to find the produced dump or, worse, compare a stale one-pass dump against the golden while the active TADD output is never checked; make the NPU dump path match the selected variant or pass it as an argument.
Useful? React with 👍 / 👎.
| tileSum tInvSum; | ||
| TRECIP(tInvSum, tSum); | ||
|
|
||
| for (int dd = 0; dd < Db; ++dd) { |
There was a problem hiding this comment.
With the built TADD variant and the default D=512, kTd=64, this outer dd loop makes pass 2 recompute the full QK logits once for each of the 8 output D blocks, even though the probabilities are identical for every V slice. That turns the advertised two QK sweeps into 9 sweeps and makes benchmarked QSMLA runtime scale with Db unnecessarily; compute each (i,j) probability tile once and reuse it across the V blocks, or explicitly document this as a deliberate benchmark tradeoff.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 07a4130701
ℹ️ 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".
| int lo = diagonal - win_left; | ||
| int hi = diagonal + win_right; |
There was a problem hiding this comment.
When wleft=-1 or wright=-1 is selected, the Stage-0 validation and CPU reference define that side of the window as unbounded, but this active TADD kernel instead computes lo = diagonal + 1 or hi = diagonal - 1. Such runs therefore mask valid tokens—often the entire row—and disagree with the generated golden output; handle -1 explicitly as the [0, s2) boundary.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0402abd9c7
ℹ️ 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".
| ifeq ($(IMPL), tadd) | ||
| DEFINES += -DQSMLA_USE_TADD |
There was a problem hiding this comment.
When the benchmark is built with IMPL=tadd, this branch only defines QSMLA_USE_TADD, but that macro is never consumed anywhere in the repository; quant_sparse_flash_mla.cpp unconditionally includes and invokes the one-pass implementation. Consequently TADD experiments silently benchmark the wrong kernel, so the driver should use this define to select the TADD include and entry point.
Useful? React with 👍 / 👎.
| "one-pass D-tail support is implemented in the next shape-generalization step"); | ||
| constexpr int Db = D / kTd; | ||
|
|
||
| float mask_buf[s1 * s2]; |
There was a problem hiding this comment.
When the new BSND dispatcher runs one of the listed typical_bsnd_swa_* shapes (TileM=64, S2=131072), its WorkConfig instantiates this automatic buffer with 8,388,608 floats, or 32 MiB. The one-level runtime configures only a 9 MiB per-CPU stack in test/common/src/chip_def.h, so entering the kernel overruns the stack before attention begins; generate mask blocks on demand or place the mask in appropriately sized mapped scratch memory.
Useful? React with 👍 / 👎.
| init_deterministic(kv, B*s2*N2*D, 2); | ||
|
|
||
| BENCHSTART; | ||
| if constexpr (N1 == 1 && N2 == 1) { |
There was a problem hiding this comment.
When B > 1 with N1 == N2 == 1, this condition selects the direct config kernel even though that implementation constructs only [s1,D] and [s2,D] tensor views, never reads Config::B, and therefore processes only batch 0. The remaining batches' outputs are left unwritten, so restrict this shortcut to B == 1 or use the BSND dispatcher for multi-batch configurations.
Useful? React with 👍 / 👎.
| print(f"mean_abs= {sum(errors) / count:.9f}" if count else "mean_abs= 0.000000000") | ||
| print(f"nan_npu = {sum(math.isnan(float(value)) for value in actual)}") | ||
| print(f"nan_ref = {sum(math.isnan(float(value)) for value in golden)}") | ||
| return 0 if passed_count == count else 1 |
There was a problem hiding this comment.
When both the actual and golden files are empty and no explicit shape is supplied, passed_count and count are both zero, so this vacuous equality returns success after reporting passed = 0/0. A pair of truncated or unpopulated result files can therefore pass the verifier; reject a zero element count before evaluating the comparison.
Useful? React with 👍 / 👎.
1 participant
No description provided.