Skip to content

ET backend: N-split infrastructure for Q4_0, Q4_K, Q8_0 MUL_MAT kernels#20

Open
CodeDoes wants to merge 12 commits into
aifoundry-org:etfrom
CodeDoes:clean-et
Open

ET backend: N-split infrastructure for Q4_0, Q4_K, Q8_0 MUL_MAT kernels#20
CodeDoes wants to merge 12 commits into
aifoundry-org:etfrom
CodeDoes:clean-et

Conversation

@CodeDoes

@CodeDoes CodeDoes commented Jul 21, 2026

Copy link
Copy Markdown

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 by n_lo .. n_hi instead of 0 .. N.

int64_t n_lo = 0;
int64_t n_hi = N;
if (N >= NSPLIT_MIN_N) {
    const int64_t shire_id = hart_id >> 6;        // 0..31
    const int64_t N_per_shire = (N + 31) / 32;
    n_lo = shire_id * N_per_shire;
    n_hi = min(n_lo + N_per_shire, N);
}

Files changed

  • mul_mat_Q4_0.c — add NSPLIT_MIN_N gate + n_lo/n_hi before every for (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 simple n_lo/n_hi slice added on top of the existing M-split; the separate q8_dot_tile / q8_dot_compute_x2_aligned gather 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_f32 destination 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

  • The body of the 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.
  • Note the Q8_0 threshold is 100000 here vs 10000 in 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 to M <= 4 && N >= 32. Expect to reconcile the two thresholds when merging.

vidas and others added 12 commits July 13, 2026 07:35
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants