ET backend: N-split infrastructure for Q4_0, Q4_K, Q8_0 MUL_MAT kernels#20
Open
CodeDoes wants to merge 12 commits into
Open
ET backend: N-split infrastructure for Q4_0, Q4_K, Q8_0 MUL_MAT kernels#20CodeDoes wants to merge 12 commits into
CodeDoes wants to merge 12 commits into
Conversation
Reorder from contiguous write to read with atomic stores.
…ecovery (aifoundry-org#15) This PR adds Q4_K MUL_MAT to the ET (ETSOC-1) backend — a scalar kernel plus a matrix-engine kernel for prefill — improves the existing Q4_0/F16/F32 matrix-engine kernels, and fixes a Q8_0 generation regression in the uberkernel path. Verified on Llama-3.2-1B-Instruct on ETSOC-1 with flash attention enabled.
…sor-engine & vectorized dots, plus uberkernel support (aifoundry-org#16) This PR adds full ET-backend support for the Q2_K, Q3_K, Q5_K, and Q6_K super-block quantizations (previously only Q4_0/Q8_0/Q4_K existed), and substantially improves generation throughput across all K-quants — including the pre-existing Q4_K — by (a) vectorizing the generation dot product and (b) enabling the K-quant matmuls to run inside the uberkernel. For prefill, every K-quant gets a tensor-unit (matrix-engine) kernel delivering very large token/sec. Everything is validated on test-backend-ops (MUL_MAT + GET_ROWS) and by coherent end-to-end generation on Llama-3.2-1B models.
This was referenced Jul 21, 2026
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
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.
Adds an N-split sharding path to the ET Q4_0, Q4_K, and Q8_0 MUL_MAT kernels so that all 32 shires are utilized during generation when the output-column dimension (N) is large.
This is a squashed cleanup of PR #17's infrastructure (previously 4 commits: bffd006 + 3 syntax fixups). The guard threshold is intentionally set high so that tests with small N keep the old “every shire computes full output” behavior, which is what the test infra validates.
How N-split works
Old
M-split:for (m = hart_id; m < M; m += STRIDE_M). During autoregressive generation M is typically 1, so 31 of 32 harts are idle and every hart still loads the full weight matrix from DRAM.New
N-split: each shire (64 harts = 1 shire) computes a contiguous[n_lo, n_hi)column slice of the output. The loop body for every path in every kernel is now bounded byn_lo .. n_hiinstead of0 .. N.Files changed
mul_mat_Q4_0.c— addNSPLIT_MIN_Ngate +n_lo/n_hibefore everyfor (n = 0; n < N; n++)loop (K-split path, K-split-grouped path, tile-outer scalar path, simple path). Net +35 / −15.mul_mat_Q4_K.c— identical scaffold, identical gate. Net +35 / −15.mul_mat_Q8_0.c— identical scaffold, identical gate. Note the Q8_0 file here is still the simplen_lo/n_hislice added on top of the existing M-split; the separateq8_dot_tile/q8_dot_compute_x2_alignedgather rewrite (which eliminates the alignment crash) lives in PR fix: stable fgb.ps dot product and N-split for 2048 harts #21. Net +33 / −21.Why NSPLIT_MIN_N = 10000 (Q4_0/Q4_K) / 100000 (Q8_0)
At very small N, every shire hits the same
atomic_store_f32destination and the shire distribution adds 32× more cache-line invalidations than a single hart doing full N. Setting the threshold high keeps the fast path disabled until N is large enough that the bandwidth win outweighs the atomics cost. The threshold is deliberately conservative; tune it after real hardware benchmarks.Review flags
if (N >= NSPLIT_MIN_N)branches in Q8_0 does not yet duplicate the single-hart path — it still relies on the existing M-split distribution, so correctness for large-M workloads is unchanged.100000here vs10000in the Q4_0/Q4_K files; PR fix: stable fgb.ps dot product and N-split for 2048 harts #21's production path later changes the Q8_0 activation condition toM <= 4 && N >= 32. Expect to reconcile the two thresholds when merging.