[DO NOT MERGE] Yanwen/release14 quant shaders - #13
Open
xuyanwen2012 wants to merge 28 commits into
Open
xuyanwen2012 wants to merge 28 commits into
xuyanwen2012 wants to merge 28 commits into
Conversation
Large prefills (e.g. 8B @ 2048 tokens) pack >2.56s of GPU work into a single 128-node command buffer submission, tripping the sgpu job watchdog. Setting the env var submits every N nodes instead, without adding a blocking stall (submits are non-blocking; execute() fences once at the end) — a measurement aid until the real fix lands driver-side.
Differential Revision: D113962415 Pull Request resolved: pytorch#21444
…o release/1.4
Aggregates the best-known-validated quantized-linear coopmat work, scattered
across several worktree branches, onto a release/1.4 base:
- Ports dev's e2e-validated shipped tile geometry for 4w (specs/036,
128x128x16) and 8da4w (specs/027, 64x32x32 dbuf2) verbatim into
linear_qw_coopmat.glsl/yaml and linear_dq8ca_qw_coopmat.glsl/yaml.
- Fixes two independent bugs found via ETDump (neither shader ever actually
dispatched before these fixes, on any branch tested this session):
- can_use_q4gsw_coopmat's dim>2 rank check rejected the real model's
rank-3 [1, M, K] activations; replaced with dev's leading-dims-numel
check (a [1, M, N] buffer is bit-identical to [M, N] when the leading
dim is 1).
- Q4gswLinear.cpp (upstream) silently hijacks the et_vk.linear_q4gsw.default
op registration with its own tiled-only q4gsw_linear_gemm__* shaders,
making QuantizedLinear.cpp's linear_q4gsw_coopmat path dead code for 4w.
Same bug as memory quant-perf-rebase-orphaned-4w-coopmat; applied the
same known fix (restore QuantizedLinear.cpp's registration, disable the
colliding one).
- Ports texture-dbuf4's texture-storage IO capability (specs/040/041) for
both 4w and 8da4w, so WMMA can dispatch on the canonical texture3d _embq_
PTE instead of requiring the retired buffer-only export path. Gated behind
ET_VK_TEXTURE_COOPMAT=1 + ET_VK_Q4GSW_COOPMAT_VARIANT=tsweep_dbuf4_t128x128k16g22s32
/ ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_t64x32k32g12s64 -- off by default,
so behavior without the env vars is unchanged.
ETDump-confirmed on M51 (secondary board, SUMD change-66942): 100% of 112 real
prefill GEMM dispatches hit the coopmat kernel for both quant modes. Prefill
tok/s (median of 5, Llama 3.2 1B, 2048-token prompt): 4w 1064.5 (1.43x vs
release-1.4 tiled baseline), 8da4w 969.2 (1.27x). Correctness not yet run
through the small-shape coopmat bench -- only "coherent, non-garbage output"
observed so far.
Authored with Claude.
![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: f60080d04a
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+101
to
+103
| for (const char* prefix : kTsweepPrefixes) { | ||
| if (v.rfind(prefix, 0) == 0) { | ||
| return true; |
There was a problem hiding this comment.
When either variant environment variable contains a value such as the documented tsweep_t.../tsweep_dbuf1_t... namespaces, this prefix-only check accepts it even though the repository generates only tsweep_dbuf4 shaders; an eligible dispatch then fails in VK_KERNEL_FROM_STR. A typo such as tsweep_t is also accepted and reaches unchecked std::stoul/grid indexing in parse_tsweep_tile, terminating graph setup instead of falling back as the comment promises. Accept only complete, generated variant tokens before parsing or selecting them.
Useful? React with 👍 / 👎.
Comment on lines
420
to
422
| add_storage_type_suffix(kernel_name, graph->storage_type_of(output)); | ||
| add_storage_type_suffix( | ||
| kernel_name, graph->storage_type_of(packed_int_weight)); |
There was a problem hiding this comment.
When ET_VK_TEXTURE_COOPMAT is set without a sweep variant and an eligible linear uses Texture3D input/output, this suffix construction requests linear_q4gsw_coopmat_texture3d_* (or its dq8ca equivalent). Neither base coopmat YAML defines a Texture3D variant—the only such variants are in the new dbuf4 sweep YAML—so shader lookup fails instead of running the experiment or falling back. Texture I/O needs to require a variant that actually has the requested storage combination, or the base variants must be generated.
Useful? React with 👍 / 👎.
Comment on lines
+346
to
+350
| const int64_t csh_bytes = | ||
| int64_t(sg_grid_y) * kMmaM * tile_n * int64_t(sizeof(uint16_t)); | ||
| const int64_t limit = | ||
| graph->context()->adapter_ptr()->max_compute_shared_memory_size(); | ||
| if (csh_bytes >= limit) { |
There was a problem hiding this comment.
For Texture3D sweep variants, comparing only Csh against the device limit does not enforce the stated requirement that it fit on top of Ash and Bsh. For example, tsweep_dbuf4_t256x256k16g14s32 allocates 24,576 bytes for Ash, 16,896 for Bsh, and 32,768 for Csh: on a 64-KiB device this check passes because 32,768 is below the limit, although the shader requires 74,240 bytes and can trigger the pipeline/hang behavior this guard is intended to prevent. Compare the total shared-memory allocation with the limit.
Useful? React with 👍 / 👎.
Full precompiled-variant sweep (125 texture3d tiles, e2e prefill tok/s on real llama_main, not microbench) on M51 (secondary board, SUMD change-66942) found tsweep_dbuf4_t128x16k64g12s64 beats the prior default (t64x32k32g12s64) by 11.8-17.3% across 1B/3B/8B, with no rank flip. ETDump- confirmed genuine coopmat dispatch on the real prefill path both before and after this change. 4w's existing default (t128x128k16g22s32) was re-confirmed as still optimal in the same sweep -- no change there. Only takes effect when ET_VK_TEXTURE_COOPMAT=1 is set; that master switch stays opt-in. ET_VK_DQ8CA_COOPMAT_VARIANT still overrides this default when explicitly set. Authored with Claude.
… too ET_VK_TEXTURE_COOPMAT=1 alone (no explicit ET_VK_Q4GSW_COOPMAT_VARIANT) crashed on every model size: the eligibility gate accepts texture3d storage unconditionally for q4gsw, but the empty-variant kernel name resolves to the bare "linear_q4gsw_coopmat" shader, which was never compiled with a texture3d IO_STORAGE variant (only the tsweep_dbuf4-suffixed one was) -- "Could not find ShaderInfo with name linear_q4gsw_coopmat_texture3d_texture2d_half". Same fix as dq8ca_coopmat_variant() (prior commit): default the empty case to tsweep_dbuf4_t128x128k16g22s32 -- same geometry as the shipped buffer default, already re-confirmed #1 in the M51 tile sweep, just resolving to a kernel name that actually has a texture3d build. Authored with Claude.
…ically wrong The M51 tile-sweep winner shipped as the 8da4w default two commits ago (tsweep_dbuf4_t128x16k64g12s64) failed test_llama_microbench --correctness-only across every texture3d case, including the rank-3 (real model shape) ones -- roughly 70% of output elements mismatched in a structured per-row pattern. e2e prefill tok/s and ETDump dispatch confirmation are not a substitute for a real numeric correctness check; the sweep never ran one, and should have before this shipped. Checked and ruled out the M>K/2 int_input_sums undersizing described in upstream issue pytorch#21423: that buffer is marked unused in both the old and new tsweep_dbuf4 shaders identically, so it isn't the discriminator here. Root cause of the new tile's specific failure not yet isolated -- most likely a genuine indexing bug in the dbuf4 template's spec-resolved code at that exact tile shape (M=128, N=16, K=64, 1x2 grid, sub=64). Reverts to t64x32k32g12s64, now verified via test_llama_microbench --scheme=8da4w --storage=texture3d --correctness-only: 0 failures, real coopmat dispatch confirmed via the harness's own dispatch log (not just "produces coherent text"). Also ports test_llama_microbench.cpp onto this branch (from texture-dbuf4, where it originated) plus its CMakeLists registration, so this check can be re-run here going forward instead of depending on a cross-branch binary. Authored with Claude.
tsweep_dbuf4_t64x64k32g24s32 passed a correctness-first tile sweep and looked like a rep-confirmed win on 1B/3B/8B across two boards, but repeating test_llama_microbench --correctness-only back-to-back on the same board/driver with no code change showed intermittent failures (1/10 runs, most shapes wrong, no crash/error) vs. 0/6 for the current default over the same window. A single correctness-bench pass is not sufficient evidence for a coopmat tile. Net effect: no default change, comment updated to record why and to flag that the sweep's other 11 "passing" candidates are equally unverified.
Brings over the SDPA QK^T/AV coopmat shaders and SDPA.cpp gating from dev (originally bundled with the linear coopmat port in 8ed2a70, plus the later default-on flip from 573d44d) -- only the linear half of that work had made it into release14-quant-shaders so far (f60080d), leaving SDPA on the tiled path even with ET_VK_TEXTURE_COOPMAT=1 set. Same gating as dev: enabled by default on capability-eligible devices (cooperative_matrix + subgroup_size==64), ET_VK_DISABLE_COOPMAT=1 remains the shared kill switch. No texture3d IO work was needed -- SDPA's q/k/v/attn_weights tensors are already buffer+half in the real e2e graph, which is what the existing buffer-storage-only coopmat shaders require. Validated via ETDump on primary M51: sdpa_compute_attn_weights_coopmat and sdpa_compute_out_coopmat now dispatch (previously silently fell back to tiled). 1B/4w prefill 1075 -> 1508 tok/s stacking this on top of the existing linear WMMA. Authored with Claude Code.
run_sdpa_suite only measures timing for sdpa_compute_attn_weights_coopmat / sdpa_compute_out_coopmat -- it never reads output data back or checks it against a reference. Add a correctness gate: builds the same graph shape directly via ComputeGraph (SymInt support the TestCase framework lacks), at small coopmat-tile-aligned shapes (S=128, D=64, matching the 1B model's real Q_H/KV_H config as one case), computes a causal GQA-aware fp32 CPU reference matching the shader's exact head-mapping (kv_h = q_h / (Q_H/KV_H)), and confirms both coopmat shaders actually dispatched (not a silent tiled fallback, which would pass numerically too). Verified 10 back-to-back --sdpa-correctness-only runs, 0 failures across both cases (20/20) -- the repeat-run discipline the linear tile sweep's flaky-tile incident earlier today established as mandatory for any coopmat correctness claim, not just a single pass.
…DATED) Adds linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr, a fork of the dbuf4 tile-sweep kernel that stages A through cooperative-matrix registers -- coopMatLoad from global, coopMatStore into LDS -- instead of a hand-rolled per-thread ivec4 copy. Ported from shaders/shmem_double_buf4-tr.comp on vk_cooperative_matrix_perf's gemm-ubm branch. Only the A half of the reference's idea is portable here. B is already column-major in LDS, so the reference's "tr" property is not a delta for it, and B cannot be coopmat-staged at all: the int4 weights need a nibble-extract that coopMatLoad cannot do, and a coopmat's per-lane layout is opaque so one cannot be built from unpacked registers. B staging is left byte-identical. A needed a new activation layout. The stock 4h4w packing is not row-major -- element [m4*K4 + k4] is an ivec4 whose component selects one of 4 rows, so the flat index is m4*(4*K4) + k4*4 + r, which is not affine in the row index and cannot be addressed by any RowMajor/ColumnMajor coopMatLoad (ColumnMajor also fails on contiguity, since a uint packs 4 K-values rather than 4 M-values). quantize_and_pack_4w_with_group_sums emits the same quantization and group sums into the existing kPackedInt8_4W layout instead, which is plain row-major. write_block in linear_int8_input_block.glslh assumes an ivec4-typed output, so it gets a SKIP_BLOCK_WRITE_HELPERS opt-out; no other shader is affected. Layout and kernel are chosen together. Both quantized_linear_impl (graph build time) and pick_linear_dqa_qw_shader (dispatch time) now go through dq8ca_coopmat_dispatch_eligible(), so a shape that falls back to the tiled path still gets the 4h4w layout that path expects. The predicate is guarded by dq8ca_variant_wants_rowmajor_a() first, so on the shipped default it short-circuits and nothing about the stock path changes. NOT the default and NOT validated on hardware: reachable only via ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4tr_t<M>x<N>k<K>g<SGX><SGY>s<sub>. Verified so far only that it builds, that all 9 kernel variants and 16 packer variants land in spv.cpp, and that the SPIR-V matches the intent -- against dbuf4 at the same tile, OpCooperativeMatrixMulAddKHR is unchanged at 8 while LoadKHR goes 16 -> 24 and StoreKHR 4 -> 12, exactly A_TILES_PER_SG=4 times the prologue and main-loop staging sites.
…orrection Reduces kernel time for the 8B prefill GEMM shapes by 33.6% and raises efficiency from 28.9% to 43.5% of the architectural peak, measured on M51 (s5e9975) at pinned 980/5333/934 against main SUMD 0a88330954. That closes 58% of the gap to the reference dense-int8 vk_cooperative_matrix_perf dbuf4 shader, which reaches 54.2% on the same board, shape and clocks while doing neither int4 unpacking nor dynamic activation quantization. The starting point was a pipeline-dump ISA analysis of the shipped kernel with its group loop isolated from prologue and epilogue. That loop is 98.6% of the dynamic instruction stream, and only 4.2% of it is the matrix multiply -- a 23:1 ratio of overhead to WMMA. Two interventions moved that, one did not. The first, in linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl, applies the activation zero-point correction and the per-row activation scale once after the group loop instead of once per quantization group. The per-group term factors exactly: out = ifs * [ SUM_g wsc_g*acc_g - izp * SUM_g wsc_g*wsum_g ] The right-hand sum depends on no activation data, so it is accumulated once in the prologue from t_weight_scales and t_weight_sums -- both already bound, so this needs no new binding, no dispatch-signature change, no prepack work and no export-format change. It removes the per-group multiply and subtract, drops the wsum_sh shared array and its ping-pong, and lets izp/ifs be loaded after the loop rather than held live across it: vgpr_count 136 -> 108, v_sub_nc_u32 16 -> 0, WMMA unchanged. 18.5% faster. The second retiles from 64x32x32 to 128x64x32 with a 4x2 subgroup grid at wave32, worth a further 16.2%. Note the reference shader's own winning geometry, 128x128x64, is our *worst* valid candidate here (+31%), so that result does not transfer. Also note 128x128x64 at a 4x2 grid is not merely slow but numerically wrong -- A_ACTIVE_THREADS is 512 against a workgroup of 256, so half of A is never staged; it needs a 4x4 grid. The earlier 12/12 correctness failure at that tile was this, not register pressure. The third, in ...dbuf4zpn.glsl, widens int4 to int8 byte-parallel. The four nibbles are already one per byte, and v^8 is exactly the four-bit two's complement of v-8 because -8 == +8 (mod 16), so only a per-byte sign extension remains, done with sgn*0x1E (0x08*0x1E == 0xF0, and sgn <= 0x08080808 so it cannot carry across bytes). A naive nib-0x08080808 would borrow across byte lanes whenever a nibble is below 8. Nibble-category ops fall 92 -> 36, bit-identical output, 2.7% faster. What did not work is recorded so it is not retried: making the B shared-memory stride a power of two, to remove the integer multiplies it forces into every B address, is 4.3% slower at stride 4 and 21.8% slower at stride 8. The bank-conflict skew is worth far more than the twelve multiplies it costs. The ...dbuf4zpb.glsl variant keeps that skew templated so the measurement is reproducible. Every variant is additive and reachable only through ET_VK_DQ8CA_COOPMAT_VARIANT; the shipped default is untouched. Each passed 12 consecutive correctness runs with zero failures before any timing was taken, which this kernel family needs -- two tile defaults were shipped and reverted on 2026-08-18, one deterministically wrong and one wrong in 1 of 10 identical runs. Timings are 5 reps with clocks read back before and after, against a baseline whose spread is 0.03-0.14%. Also fixes linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl, which was committed in 7de604d failing 12/12: it bound the packed activations as a 32-bit int array while loading a coopmat<int8_t> from them. That mismatch is silently wrong from a StorageBuffer on this driver, though it is fine through Workgroup storage, which upstream already relies on. Binding it as int8_t gives 14/14. The dbuf4trm and dbuf4trd variants are the bisect probes that isolated it. One caveat on the whole exercise: instruction counts located the candidates but predicted none of the outcomes. The hoist beat its instruction count by 2x, the nibble change fell short of its, and the stride change went the wrong way entirely. Every claim here rests on an end-to-end measurement, not a count.
…esce B store Three real, on-device-validated changes to the dq8ca (8da4w) coopmat linear kernel, found via the dq8ca-dequant-unpack-ablation investigation into why this kernel trails a dense-int8 reference by a wide margin: - Promote the shipped tile from t64x32k32g12s64 to t128x128k64g81s64 (already ahead pre-fix; +11.8-13.2% relative on top of the other two changes). - Drop the classic "+1" anti-bank-conflict skew on B's LDS stride (B_STRIDE_U32): measured slower than no padding on this hardware/driver (+2.98pp on its own). - Make the B-operand LDS write coalesced: invert the write-address formula so consecutive threads write consecutive LDS words, instead of 4 words apart. The read side and every dequantized value are unchanged; only which thread writes which element (+0.7-0.8% on top of the other two). Combined: 8B prefill efficiency of int8 theoretical peak goes from 28.93% (original) to 36.97%+ (now shipped) on the primary validation board. Real e2e prefill: 1B 1532.9, 3B 616.9, 8B 299.2 tok/s (ETDump-confirmed coopmat dispatch, ET_VK_TEXTURE_COOPMAT=1). Validated to this kernel family's own stated bar: 10 consecutive --correctness-only passes, zero failures, for every (model x storage) combination -- 1B/3B/8B x buffer/texture3d, 60+ runs total. Decode-regime validation and a second board/driver are still open. Full investigation, ablation chain, and the (still open) question of why the B-operand LDS read itself is so much more expensive than a dense-int8 reference's is in openspec/changes/dq8ca-dequant-unpack-ablation/results/README.md. Authored with Claude.
…efault tsweep_dbuf4zpg_t128x64k32g42s32 combines zp-hoist, byte-parallel int4 unpack, static branch elision, and loop-invariant B-staging-index hoisting with the B_STRIDE_U32 skew removal and coalesced B-store the prior default already shipped. Validated on the sibling dq8ca-uvec4-redesign branch (cut from this branch at the same base commit) at 46.49-46.50% efficiency of int8 peak on 8B, vs. 36.97% for tsweep_dbuf4_t128x128k64g81s64 -- +20.5/20.9/21.2% relative on 8B/3B/1B. Re-verified on this branch's own build: 48/48 correctness cases pass (3 consecutive runs), and a real e2e run confirms the scored ETDump block dispatches this shader for 67.4% of leaf GPU time, 344.7 tok/s prefill on 8B (xgpusw-debug08, canonical main-fafb46ae9c0d driver, maxpin 980/5333/934). Removed the shaders and dispatch-table prefixes for the superseded single-intervention isolation variants (-tr/-trm/-trd, -zp/-zpn/-zpb/ -zpx, -zpi/-zpk) this promotion folds together or supersedes -- this branch ships only the validated default; the experimental siblings live on dq8ca-uvec4-redesign and dq8ca-arch-redesign.
…d patches Adds env-var-driven tile-sweep support to the SDPA QK^T/attn*V coopmat shaders (ET_VK_SDPA_ATTN_COOPMAT_VARIANT / ET_VK_SDPA_OUT_COOPMAT_VARIANT in SDPA.cpp), mirroring the mechanism QuantizedLinear.cpp already uses for dq8ca/q4gsw. Default behavior is unchanged -- kSdpaAttnDefaultDims/ kSdpaOutDefaultDims reproduce the prior hardcoded tile exactly. Only the shipped default tile is compiled per shader (t128x64k32g22s64 for QK^T, t64x64k32g22s64 for attn*V). A wider sweep (~40 variants per shader) was tried but found unvalidated by its own admission -- that doesn't belong in this branch (final shipped shaders only), so it's trimmed here; a real sweep belongs on an experimental branch. Also includes shader-lab's two sanctioned additive patches: --json microbench output (test_llama_microbench.cpp) and ET_VK_EXTRA_INSTANCE_LAYERS (vk_api/Runtime.cpp), needed to enable Vulkan debug layers on an Android native binary. Verified on xgpusw-debug08/00000b750f413c33 (canonical main-fafb46ae9c0d driver): 48/48 correctness cases pass (3 consecutive runs, unchanged from before -- this harness's correctness gate doesn't cover SDPA), and a real e2e run confirms coherent output at the default SDPA tiles, 343.9 tok/s prefill on 8B (matching the already-promoted dq8ca default).
Replaces tsweep_dbuf4zpg_t128x64k32g42s32 with tsweep_dbuf4zpgtr_t128x64k32g42s32 as the shipped dq8ca coopmat linear default. Same tile, B-staging, and zp-hoist as dbuf4zpg; only the A-operand LDS staging changes, from a per-thread scalar scatter to a coopMatLoad(global)->coopMatStore(LDS) sequence. Validated on the sibling dq8ca-tr-staged-a-on-zpg branch (cut from this branch @ 1f3322c): -6.90%/-6.82%/-6.76% kern_us on 8B/3B/1B (46.50% -> 49.94% efficiency of int8 peak on 8B), 349.7 -> 364.5 tok/s real e2e prefill on 8B (+4.2%), vgpr_count 133 -> 128. Correctness: 10/10 consecutive clean (buffer) + 6/6 consecutive clean (texture3d) across all three model sizes, backed by an exhaustive host-side address-equivalence proof that the old and new A-staging schemes write the identical LDS address set. A follow-up tile re-sweep against this shader's own lower register pressure found no better tile. See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg and openspec/changes/coopmat-tr-tilesweep-4w-port for the full record.
… defaults barrier() alone does NOT order shared-memory stores against a subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver; it needs an explicit memoryBarrierShared(). The symptom is recognisable and quiet: one stale MMA_M-row band of the A operand (16 consecutive rows, all WG_TILE_N columns), silently wrong output, no crash and no DEVICE_LOST, on roughly 2.5% of runs. Found 2026-09-02 in sdpa_compute_out_coopmat.glsl. Both shipped coopmat linear defaults had this defect on their double-buffer staging path, as did the 8da4w env-var fallback. The memoryBarrierShared() pairs already present in these files all guarded wcorr_sh / bias_sh / Csh_out -- none guarded Ash_int8/Bsh_int8, i.e. none guarded the store -> coopMatLoad ordering point that runs every chunk. linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl (8da4w default) +2 linear_q4gsw_coopmat_tsweep_dbuf4.glsl (4w default) +3 linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl (8da4w fallback) +2 All three now have 0 bare barrier() calls. The diff is exactly 7 memoryBarrierShared() lines plus comments; QuantizedLinear.cpp is comment-only. COST: NONE, measured, not assumed. Interleaved same-binary A/B (unfenced clones of both defaults compiled alongside the fixed ones, alternating round-by-round so drift hits both arms equally), 8B prefill, buffer, 4 rounds each: 8da4w fenced 70625.1 us vs unfenced 70624.6 us -> +0.001% 4w fenced 95740.4 us vs unfenced 95741.9 us -> -0.002% Observed spread across all 8 runs was 0.037% (8da4w) and 0.017% (4w), and the paired per-round deltas straddle zero in both cases, so the fence cost is below the resolution of the measurement. The fence is NOT optimised away -- verified in the compiled SPIR-V rather than inferred from the timing: OpMemoryBarrier goes 1 -> 2 (8da4w) and 0 -> 2 (4w) between the unfenced and fenced builds, with OpControlBarrier unchanged at 2. Without that check a "no cost" result would be consistent with the compiler having stripped the fence, which would also mean it fixes nothing. Correctness re-gated after the driver reflash noted below: 14 PASSED / 0 FAILED for both defaults, correct shader confirmed dispatched. Measured on xgpusw-debug08 / 00000bf70c579c33, driver md5 1eea300aa3974ff8974d04808c9ff394 (SUMD main@fafb46ae9c0d, re-verified after a reflash -- a crash-to-bootloader recovery on this board silently reverted /vendor to the BSP factory driver), clocks pinned 980/5333/934.
sdpa_compute_out_coopmat produced silently wrong output on ~2.5% of runs at S=256, with no crash and no DEVICE_LOST. The previous correctness gate could not see it: both its shapes were S=128, where the shader runs only 2 M-tiles and 4 K-chunks, and the failure needs more. Diagnosis. The microbench seeds with srand(0), so every repeat feeds byte identical inputs to an identical shader on an identical driver -- and the outcome still varied run to run. That rules out a compile-time miscompile (those fail every time, as the QK^T UBO-stride bug did at 20/20) and means a race. Localizing the wrong elements pinned it: each failure corrupted exactly one MMA_M-row band -- 16 consecutive rows, ALL WG_TILE_N columns, a single head. Rows confined to one 16-row block while every column is wrong is the signature of a bad A operand, since A is indexed by row and shared across all column blocks. Those rows are Ash rows written by one subgroup and read back by coopMatLoad in several, and the band was stale even for the subgroup that wrote it -- so the store was late, not misindexed. barrier() alone was not ordering the uvec4 Ash/Bsh stores against the cooperative-matrix load path on this driver. Make it explicit with memoryBarrierShared(), as the linear coopmat kernel already does for its shared staging. Evidence: 0 failures / 300 case-runs with attn*V coopmat isolated (150 reps), and 0 / 160 on the full extended gate with both coopmat shaders (40 reps). Against the measured pre-fix rate of 5/200, zero failures in 300 has probability ~0.05%. Per-kernel cost is nil: attn*V prefill is 13002.16 us with the fence against 13002.20 us without (n=5, CoV 0.008%). Residual risk, stated rather than assumed away: only the write->read barrier got the fence. The loop's second barrier orders this chunk's coopMatLoad reads against the next chunk's stores, and if the same driver quirk applies in that direction it would reproduce an identical symptom. 460 clean case-runs do not exclude a rarer WAR variant at a ~2.5% base rate. (cherry picked from commit 8dc05ca)
The SDPA coopmat gate had two cases, both S=128. At the shipped 128x64 QK^T
tile that gives num_tiles_m=1, num_tiles_n=2, so BOTH tiles are diagonal: the
shader's whole-tile all-masked early-out never fired and an all-visible fast
path would have gotten zero coverage. A gate cannot be treated as covering a
path its shapes cannot reach, so extend it before touching any shader.
--sdpa-regions-only enumerates the QK^T tile grid and prints each tile's
region classification, then reports which of the three shader paths the gate
actually reaches. It also checks the classification at its boundaries: for a
fixed M-tile row the classes must walk visible -> diagonal -> masked with no
class recurring and no column skipped, which catches both a tile claimed by
two classes and a tile claimed by none.
Two S=256 cases are added as a separate "regions" tier. S=256 is the smallest
size reaching all three regions at the shipped tile (num_tiles_m=2,
num_tiles_n=4), and it populates both boundary transitions; S=192 is not
usable because it is not a multiple of WG_TILE_M=128 and the dispatch gate
would refuse the shader. It also gives attn*V 8 K-chunks instead of 4.
--sdpa-tier=<fast|regions|all> keeps the cheap S=128 pre-check available; all
three tiers measure 222/355/406 ms per pass, so the 10+ repeat discipline
stays practical even for the full gate.
--sdpa-force-fallback runs the same shapes with coopmat disabled. This proves
the gate's dispatch assertion is load-bearing (the tiled path is numerically
correct, so it reports mismatches=0 and the gate must still fail it), and it
is the control that separates a wrong shader from a wrong harness.
Extending the gate immediately exposed a pre-existing intermittent failure in
sdpa_compute_out_coopmat at S=256 (4/100 case-runs) with QK^T clean at 0/100,
and the tiled path clean at 0/40 on identical shapes. The two
ET_VK_SDPA_DISABLE_{QK,OUT}_COOPMAT switches here are the diagnostic that
attributed it and are needed to run a QK^T-isolated gate while attn*V remains
flaky. They are temporary scaffolding and must be removed before promotion.
(cherry picked from commit 9d31d99)
At the 2048-prefill workload with the shipped 128x64 tile there are 512 QK^T workgroups: 240 are entirely above the diagonal and already early-out, 32 genuinely straddle it, and 240 are entirely BELOW it -- fully visible, nothing to mask. Those 240 were still paying the full masked epilogue: a coopMatStore into the Csh scratch, a barrier, and a per-element scalar read-mask-write over the whole 128x64 tile. That is 240 of the 272 workgroups that execute (88%) doing an LDS round trip that cannot mask anything. Classify the tile and send the fully-visible case straight to global, as the linear kernel's buffer epilogue does. The fully-visible condition is the complement of the existing all-masked one; both holding at once would require WG_TILE_M + WG_TILE_N < 2, and a tile matching neither falls through to the per-element path, so the three classes are exhaustive and disjoint. The address arithmetic reproduces the Csh path's global address exactly. The store's stride cannot be the UBO-derived aw_row_width. coopMatStore miscompiles on the Xclipse/AMD-PAL compiler when its stride derives from a UBO value -- the same bug sdpa_compute_out_coopmat.glsl works around for its output stride -- and measured deterministically here: 20/20 case-runs wrong with the UBO value, 0/20 with a compile-time one. Unlike the linear kernel's out_N_arg this width is not static; resize_sdpa_attn_weights_node recomputes it from the input_pos symint on every resize. So it is passed as a spec constant resolved at node construction and the fast path is entered ONLY when that baked value still equals the live width, falling through to the compile-time-stride Csh path otherwise. Correctness therefore never depends on the guess: a chunked prefill simply does not get the fast path. The entry condition also requires the tile to lie wholly inside both extents, since the store writes whole MMA tiles unchecked. The dispatch gate already guarantees that (and this shader's staging reads depend on it unguarded), so under that gate no workgroup loses the fast path -- but it makes an out-of-extent write impossible by construction rather than by appeal to the gate. Per-kernel QK^T prefill time, 3 suite reps each, dispatch confirmed every rep: 8B 9505.6 -> 8178.6 us (-14.0%), 3B 7199.2 -> 6194.0 us (-14.0%), 1B 6806.5 -> 6167.4 us (-9.4%). Between-rep spread is 0.05-0.30% for 8B/3B, two orders of magnitude below the effect; 1B's fast-path reps spread 6.4% (5922.8-6319.8) so its magnitude is not resolved by 3 reps, only its direction. Correctness: 0/60 case-runs over 15 reps of the extended gate with attn*V held on the tiled path (its coopmat shader has a pre-existing intermittent failure at S=256, unrelated to this change). The fast path is confirmed live rather than assumed: perturbing only its store fails S=256 and leaves S=128 untouched -- matching the enumerated claim that no S=128 tile is fully visible -- with both first mismatches at s=128, the first row of the fully-visible M-tile row. (cherry picked from commit 013ef5a)
The row-wise softmax between the two SDPA GEMMs read and reduced over the whole context_len row, three times, for every row. But the QK^T shader has already written -inf to every element with c > s + input_pos, and a masked element cannot contribute: exp(-inf - max) is 0 in the sum, it cannot be the row max unless the entire row is masked, and it normalizes to exactly 0. So row s only needs the first (s + input_pos + 1) columns. Over a full prefill the rows are a triangle, not a rectangle, so this halves the bytes read. The store still covers the whole row, because attn*V stages every chunk with chunkK < context_len and multiplies it by V -- a stale tail would be multiplied as if it were attention weight. But the tail is written as zero WITHOUT loading the input, so pass 3 loses its tail read too, and the zero store is spread across all 64 workers rather than serialized on worker 0 the way the existing straddling-texel path is. Truncation is gated on HAS_INPUT_POS. Fused SDPA has no input_pos and takes its mask from an attn_mask bias instead, so no truncation is valid there. 8B prefill, per-kernel, n=5: 13971.1 -> 9451.9 us, -32.3%, with the two ranges disjoint ([13895.7,14048.7] vs [9405.0,9499.1]). Correctness 0 failures / 80 case-runs on the extended gate. Also adds a softmax bucket to test_llama_microbench's --sdpa suite, which only timed qk and av. That measurement is why this change is trustworthy: an ETDump put softmax at 443.3 -> 444.5 ms, i.e. no effect, but a single ETDump capture turned out to carry +-20% capture-to-capture noise (the same pair of captures moved attn*V by +20.5% across a change that the querypool measures at 0.00% with CoV 0.008%). The querypool path with repeats resolves a 32% effect the ETDump could not see at all. (cherry picked from commit 67facb0)
attn*V walked all num_k_chunks_arg chunks of the context for every output tile. Two thirds of that work could not contribute anything. num_k_chunks_arg is max_context_len/WG_TILE_K, sized from the KV cache, so on a ctx3072 PTE it is 96 chunks even when the live context_len is 2048. The 32 beyond-context chunks staged zeros -- and then ran the MMA on them anyway. Worse, P is the softmax output of a causally masked score matrix, so P[s, c] == 0 for every c > s + input_pos exactly (QK^T writes -inf, softmax normalizes it to zero). A chunk whose lowest context index already exceeds the highest row in the M-tile is therefore all zeros, and a zero A tile contributes 0*V == 0. Skipping it is value-preserving, not an approximation. Bounding the loop by both facts takes the chunk-loop iterations at the 2048 prefill from 32 M-tiles x 96 chunks = 3072 down to 1056 (34.4%): the work is a triangle over M-tiles, not a rectangle, and the beyond-context rectangle is gone entirely. 8B prefill, per-kernel, n=5: 13002.2 -> 4560.4 us, -64.9% (2.85x), ranges disjoint. The 2.9x predicted by the iteration count matches the 2.85x measured, so the model of where the time went is right. Correctness 0 failures / 120 case-runs on the extended gate. REJECTED on the way here: dbuf4 double-buffered staging, the intervention this shader was originally slated for. It measured 13002.2 -> 13124.7 us, +0.94%, ranges disjoint -- a small but real regression. attn*V moves ~1.04 GB per dispatch in ~13.0 ms, i.e. ~80 GB/s, which is bandwidth-bound, and double buffering hides latency rather than bandwidth. It also doubled LDS 9728 -> 19456 B, halving co-residency from 6 to 3 workgroups per CU. Reducing traffic was the answer; overlapping it was not. (cherry picked from commit d132d44)
ET_VK_SDPA_DISABLE_QK_COOPMAT and ET_VK_SDPA_DISABLE_OUT_COOPMAT existed for one purpose: attributing an intermittent wrong-output failure to one of the two SDPA GEMMs by forcing the other onto its tiled path. That worked -- it put the fault in attn*V (4/100 case-runs) with QK^T clean (0/100) -- and the underlying missing shared-memory barrier is now fixed, so there is nothing left to isolate and no reason to keep a debug env var in the dispatch gate. Gate after removal: 0 failures / 100 case-runs, with both coopmat shaders confirmed dispatching on every rep (which is now the only mode, so a silent fallback would have shown up as a dispatch warning rather than passing). (cherry picked from commit 846d217)
Completes the memoryBarrierShared() work for SDPA. 8dc05ca fixed only the write->read barrier in attn*V and said so; this covers the rest. sdpa_compute_attn_weights_coopmat.glsl (QK^T) +3 -> 3 fenced, 0 bare sdpa_compute_out_coopmat.glsl (attn*V) +1 -> 2 fenced, 0 bare Two of the QK^T sites are the WRITE->READ direction, i.e. the same direction that produced the observed ~2.5%-of-runs stale-A-band bug in attn*V, not the untested WAR direction: QK^T:266 after the Bsh[...] staging stores, before the MMA coopMatLoad QK^T:349 after coopMatStore into Csh, before the epilogue reads it QK^T:297 WAR -- this chunk's reads vs the next chunk's stores attn*V:260 WAR -- same So QK^T was exposed in the dangerous direction. The prior note that it was "measured clean 0/100, not proven safe" understated it. VALIDATION STATUS: build verified, device validation PENDING. Two builds (fenced/unfenced) were produced and every .spv compared: of 1620 shaders only the 2 SDPA ones differ, so build variance is excluded as a confounder. The fences are confirmed present in the compiled SPIR-V rather than assumed -- OpMemoryBarrier 0 -> 3 (QK^T) and 1 -> 2 (attn*V), OpControlBarrier unchanged. The on-device correctness gate and the interleaved A/B have NOT run: the target board (xgpusw-debug08 / 00000bf70c579c33) crashed to bootloader twice and came back with a third, unrecognised driver hash. Do not treat this commit as perf-validated until that A/B is recorded. The equivalent fences on the linear defaults measured +0.001% / -0.002% (7892151), so a cost is unlikely.
…v var
ET_VK_TEXTURE_COOPMAT began as an experiment hook (specs/040/041), off by
default, documented as "Off by default, so buffer dispatch is byte-identical".
That reasoning only holds for buffer-storage PTEs. The embq PTEs this branch
targets place the linear weights in TEXTURE storage, so with the hook off the
*_texture3d_* coopmat variants were rejected and the entire quantized linear
path -- 72.4% of prefill GPU time -- silently fell back to the tiled shader.
No linear WMMA ran at all unless the caller happened to know about the flag.
Flip the default to enabled. ET_VK_TEXTURE_COOPMAT=0 (or "false"/"off") still
restores the old buffer-dispatch-only behavior for baseline measurement; unset,
or any other value, enables it.
SDPA needed no change -- it was already on by default, gated only on
supports_cooperative_matrix() && subgroup_size() == 64, with
ET_VK_DISABLE_COOPMAT as the kill switch.
8B/8da4w, 2048 prefill, embq ctx3072, xgpusw-debug08/00000b750f413c33, driver
main@fafb46ae9c0d (md5 1eea300aa39..), BSP CP2A.260605.016 20260831.130326,
clocks maxpin 980/5333/934, 3 reps each:
no env vars, BEFORE this change 174.5 / 174.8 / 174.9 -> 174.7 tok/s
ET_VK_TEXTURE_COOPMAT=1 + NODE_THRESHOLD=32
367.5 / 366.8 / 368.0 -> 367.4 tok/s
no env vars, AFTER this change 366.9 / 367.9 / 368.0 -> 367.6 tok/s
2.10x, entirely from the default. The post-change no-env run matches the
env-forced run to 0.05%, and needs no ET_VK_EXECUTE_NODE_THRESHOLD -- that
watchdog workaround is not required on this driver.
ETDump verification (no env vars set), leaf GPU 5484.0 us over 1363 events:
72.4% linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr_t128x64k32g42s32_texture3d_texture2d_half
5.7% sdpa_attn_weights_softmax_buffer_half
4.8% sdpa_compute_attn_weights_coopmat_t128x64k32g22s64_buffer_buffer_half
2.7% sdpa_compute_out_coopmat_t64x64k32g22s64_buffer_buffer_half
Both SDPA GEMMs dispatch on the _coopmat path as buffer_buffer_half, so the
buffer+fp16 gate in SDPA.cpp is satisfied even though the linear path is on
texture IO. No tiled sdpa_compute_attn_weights / sdpa_compute_out appears --
there is no silent fallback left. Dumps archived at
.artifacts/rel14qs-{coopmat,defaultwmma-noenv}-2026-09-08.etdp.
Not validated: decode (prefill-only runs, --max_new_tokens=1) and models other
than 8B/8da4w.
…l rejected
Optimization attempts against gemm-ubm's shmem_double_buf4-tr3, which is 1.39x
faster than our shipped dq8ca kernel on M51 (2231.8 vs 1601.4 us at
2048x1024x4096, 47.94% vs 66.82% of int8 WMMA peak, profiler off, same
board/driver/clocks). Shipped default is UNCHANGED; every variant is opt-in via
ET_VK_DQ8CA_COOPMAT_VARIANT and none is a default.
A-layout (dbuf4zpgtr3): port tr3's shared-A layout -- scalar int8_t storage,
row-major over the full chunk, A_ROW_PAD_I8 bytes of per-row pad (yaml
parameter). Correctness 12/12 clean; MEASURED SLOWER at every pad value:
p16 +4.73%, p8 +4.71%, p0 +48.78%. Axis decomposition:
- layout+element-type alone (p0, LDS identical to baseline): +48.78%
- adding 8B of pad recovers 44.1 pp
- 16B (tr3's own value) adds nothing over 8B (+0.019%) -- tr3 over-pads for
M51 and the extra 2 KiB of LDS is a pure loss
- vs the earlier zpgtrp attempt (same layout+pad but `shared uint`, +11.50%),
the element-type change alone is worth 6.8 pp
Mechanism: slab-major gives each 16x16 A fragment a row stride of exactly 16
bytes, so the fragment is 256 CONTIGUOUS bytes; row-major full-chunk makes it
strided (32/40/48 B). Padding fixes the bank pattern but cannot restore
contiguity. tr3's row-major A is a constraint of its own staging, not a virtue.
VGPR 128 unchanged, scratch 0, LDS 14080 -> 16128 (still 4 WG/CU).
B-staging: five variants moving ownership from (block,col) slots spread over 4
threads to one thread per contiguous run, to cut the measured 3.25x static ds_*
gap.
- bv4 (uvec4 array, 4-wide), bw (no retype, 4-wide), bwr (bw, reversed store
order): all KNOWN-INCORRECT, deterministic 3/3 failure on the
num_groups==2 shapes only (K=256); K=128/2048/4096 pass. Cause isolated to
4-wide-per-thread ownership with half the workgroup idle. Store order is
not the trigger (bwr) and neither is the uvec4 retype (bw, no retype,
fails identically). Deliberately NOT registered so they cannot be selected.
- bw2 (2-wide, all 256 threads active): CORRECT 14/14 x3, ~0% (-0.02%/-0.115%)
- bw3 (uvec2 retype + bw2 ownership): CORRECT 14/14 x3, +0.34%. Also proves
an 8x packing ratio is fine for coopMatLoad, which retroactively clears the
retype as bv4's cause.
The transformation demonstrably landed and bought nothing: scalar ds_store_b32
11 -> 3, paired ds_store_2addr_b32 10 -> 13, total ds ops 104 -> 100, static
instructions 1647 -> 1607. bw2 and bw3 emit byte-identical LDS op mixes, so the
compiler was already merging bw2's adjacent scalar writes.
Why it could never have worked, now measured: our LDS instruction count is
dominated by LOADS, not stores -- ours 80 loads / 24 stores against TR3's
24 / 8. ds_load_2addr_b32 alone is 72 and is the largest single ISA category.
The 80 loads include the six auxiliary shared arrays (izp_sh, ifs_sh, wsc_sh,
wcorr_sh stride-0 broadcasts, bias_sh, Csh_out) that TR3 has no analogue for.
That is the next target; B-store width is closed.
Full measurement record, counter data and provenance:
openspec/changes/dq8ca-vs-tr3-iso-tile-counters/results.md
…riants Housekeeping: prunes the tsweep GLSL/YAML files that no longer serve any purpose now that this branch ships a single validated default per family (q4gsw: tsweep_dbuf4_t128x128k16g22s32; dq8ca: tsweep_dbuf4zpgtr_t128x64k32g42s32) -- - dbuf1/dbuf2/dbuf3 and the bare "tsweep_t" prefixes: dead in both the q4gsw and dq8ca prefix lists (no GLSL ever existed for them in this branch), so selecting one could only ever crash confusingly at shader lookup instead of failing validation cleanly. - dq8ca dbuf4 / dbuf4zpg: the 2026-08-26 and 2026-08-28 defaults, each twice superseded. - dq8ca dbuf4zpgtr3: correctness-clean but measured +4.71% slower. - dq8ca dbuf4zpgbw2/bw3: correctness-clean, ~0% delta, no reason to keep as a build target. - sdpa_compute_attn_weights_coop / sdpa_compute_out_coop: unreferenced by any remaining code path. QuantizedLinear.cpp's prefix lists and dq8ca_coopmat_variant()'s invalid-token fallback are updated in lockstep so an env var can never name a shader that no longer exists. Full history (promotion dates, measured deltas, why each variant was rejected) is preserved in the comments below and in the deleted files themselves, on branch yanwen/release14-quant-shaders-archived-2026-09-15. Authored with Claude Code (Sonnet 5). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…sw tile-sweep yaml
Second pass, same goal as the previous commit: this branch ships only the
production shaders, minimal enough to serve as a future release branch.
- linear_qw_coopmat.{glsl,yaml}, linear_dq8ca_qw_coopmat.{glsl,yaml}: removed.
Zero real dispatch references anywhere in the codebase -- the only mentions
were two comment lines in QuantizedLinear.cpp citing their tile geometry as
historical provenance for two constants. Comments updated to explain what
those constants are now that the files they cited are gone, and
kDq8caQ4gswCoopmatDims's value corrected to the real current shipped tile
(it held a stale, pre-08-26 geometry -- harmless in practice since it's an
unreachable defensive fallback, but confusing to leave wrong).
- linear_q4gsw_coopmat_tsweep_dbuf4.yaml: trimmed from 97 shader_variants
(specs/041's full tile-sweep search space, ~48 geometries x storage combos)
down to 3 -- the shipped tile's (t128x128k16g22s32) buffer/texture3d/
weight-buffer storage combos. Only this one tile is ever dispatched; the
other 94 were pure sweep residue never reachable outside an explicit
ET_VK_Q4GSW_COOPMAT_VARIANT override.
Deliberately left alone: coopmat_mm.glsl / GemmCoopmat.h (Linear.cpp /
Matmul.cpp's general fp16 coopmat path) -- real shared backend
infrastructure, not one of our experimental variants, just permanently
gated off on M51 via !is_integrated_gpu(). Out of scope for this cleanup.
Rebuilt clean (install + custom_ops/test_llama_microbench) after this change;
device re-verification not run for this commit.
Full pre-cleanup history for everything removed today is preserved on
branch yanwen/release14-quant-shaders-archived-2026-09-15 (snapshotted
before the first housekeeping commit, so it also covers this one).
…pped Companion to 1731d59, which should have included these two but a fatal git-add error on an already-staged path silently excluded them from that commit: - QuantizedLinear.cpp: the comment/constant fixes for the now-deleted linear_qw_coopmat.yaml / linear_dq8ca_qw_coopmat.yaml citations, per 1731d59's own message. - linear_q4gsw_coopmat_tsweep_dbuf4.yaml: the 97->3 shader_variants trim, same commit. No content difference from what was already described and rebuild-verified in 1731d59 -- this just lands the two files that commit's message says are there but isn't.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
[PLEASE REMOVE] See CONTRIBUTING.md's Pull Requests for ExecuTorch PR guidelines.
[PLEASE REMOVE] If this PR closes an issue, please add a
Fixes #<issue-id>line.[PLEASE REMOVE] If this PR introduces a fix or feature that should be the upcoming release notes, please add a "Release notes: " label. For a list of available release notes labels, check out CONTRIBUTING.md's Pull Requests.
Test plan
[PLEASE REMOVE] How did you test this PR? Please write down any manual commands you used and note down tests that you have written if applicable.