Pull requests / #2058

#2058 Verify window: the routed experts' gate/up with the group's q8_1 rows in shared memory (STRATA_EXPDEC=1, opt-in, bit-exact)

open · @imanu86 · 0 comentarios · En GitHub

BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindows

Descripción

## Title
Issue: none (an opt-in speed change).

## Summary
`STRATA_EXPDEC=1` (opt-in, off by default) runs the verify window's VRAM-resident routed experts' **gate/up** through a
kernel that keeps the group's q8_1 activation rows in shared memory and computes several rows per lane. Same kernels
otherwise, same launch order, and the output is **bit for bit** the default kernels' (`expdec_bench` checks it with
`memcmp` on every cell). CUDA only; HIP and the prompt path are not touched.

## Why
`native_gu_multi_kernel` already decodes each weight part once per block and applies it to the group's 1-4 entries,
but the apply re-reads each entry's q8_1 block from global memory with 9 scattered 4-byte loads (36-byte blocks: a
warp's load touches 16-32 sectors). Measured with nsys on a 131k decode (RTX 2080 Ti, stage 1 of a layer split), an
extra row of the window costs about 60% of a whole expert - ~1.6 us against ~2.6 us per expert and layer - although it
reads nothing from DRAM. That is most of the gate/up kernel's slope with the window's rows (+0.86 ms per row).

## What changed
- `src/kernels/cuda/iq_kernels.cu`: `expdec_gu_kernel<TG, R, MINB>` (gate/up IQ2_XXS, IQ2_XS, IQ3_XXS, IQ3_S, IQ2_S at
  n_embd 2560) and `expdec_down_kernel<TD, R>` (down Q2_0 at n_ff 640; IQ4_NL with `STRATA_EXPDEC_D20=1`). The group's
  entries are staged 4 at a time in a layout where a call's 8 words are two conflict-free 16-byte loads (IQ4_NL: two
  8-byte loads) and the fp16 block scale is already a float; R rows per lane share the activation words; a block takes
  `rb` rows (auto from the launch's groups) instead of 16; the IQ3_XXS / IQ3_S codebooks are staged in shared memory like
  the IQ2 ones. `Split<TG>::apply_u` is `apply` on the block's words; `Split<18>/<21>::load<true>` read the staged grid.
  `native_expert_grouped` takes the new kernels when the lever is on and the layout matches, else the default ones.
- `include/strata/kernels/iq_kernels.hpp`: `native_expert_set_expdec`, `native_expert_set_expdec_params` /
  `native_expert_expdec_params` (`NativeExpertExpdecParams`), `native_expert_expdec_takes`.
- `src/kernels/expdec_bench.cpp` (`EXCLUDE_FROM_ALL`): the default kernels against EXPDEC at T 1-4 on 10 format pairs,
  synthetic blobs read cold from a >= 96 MB pool, `memcmp` of the output, median of 5 x `reps` calls, per phase.
- `docs/EXPERTDEC.md`.

Knobs (environment, read at startup): `STRATA_EXPDEC_R` (gate/up rows per lane, 1 | 2, default 1), `STRATA_EXPDEC_DR`
(down, 1 | 2 | 4, default 4), `STRATA_EXPDEC_MINB` (gate/up `__launch_bounds__` min blocks per SM, 4 | 2, default 4),
`STRATA_EXPDEC_RB` / `_DRB` (rows per block, 0 = auto), `STRATA_EXPDEC_MINBLK` (the auto rule's minimum blocks, 800),
`STRATA_EXPDEC_DOWN` (the down too, default 0), `STRATA_EXPDEC_D20` (IQ4_NL down too, default 0). The first launch of
each (role, format) logs one line (`strata expdec: CUDA1 gate/up type 21: expdec_gu_kernel R 1 MINB 4, ... captured
into a graph`).

**The down stays separate on purpose.** In the micro-bench the Q2_0 down with EXPDEC is 0.65-0.85x of the default, but
in the engine (nsys, two captures with two block geometries) it was about 10% *slower* than the default kernel, which
already keeps the group's h rows in shared memory; IQ4_NL showed no gain in either. So `STRATA_EXPDEC=1` takes the
gate/up only and `STRATA_EXPDEC_DOWN=1` is a separate opt-in for whoever wants to measure it on another card.

## Bit for bit
For every (row, entry) the sum is the default kernel's tree (`row_dot_80_sub16`: lane t of a 16-lane sub-warp does
calls t, t+32, t+64 into `s` and t+16, t+48 into `s16`, `s = __fadd_rn(s, s16)`, butterfly xor 8-4-2-1; `row_dot_40_sub8`
/ `row_dot_20_sub16` for the downs). The same lane does the same calls in the same order with the same integer and
float operations on the same words; the entries of a group are independent, so the staging and the rows-per-lane /
rows-per-block choices only change which lane computes what. `expdec_bench`: `identical` on all 1,280 cells (10 pairs x
T 1-4 x 2 window shapes x 8 choices x 2 cards). In the engine the greedy text is unchanged.

## Measurements
Hardware: RTX 2080 Ti 22 GB (sm_75) + RTX 3060 12 GB (sm_86), Windows 11, CUDA 12.6, layer split (3060: layers 0-15,
2080 Ti: 16-47 + head), `--pipeline-windows 2`, Qwen3.8-Flash-Next IQ3_XXS pack, 131k context.

Micro-bench, gate/up phase, us per call with 20 distinct experts (the shape of a real window), `def` -> ratio with EXPDEC:

| card | format | T=1 | T=2 | T=3 | T=4 |
|---|---|---|---|---|---|
| RTX 2080 Ti | IQ3_S | 72 / 0.73x | 112 / 0.79x | 124 / 0.75x | 138 / 0.74x |
| RTX 2080 Ti | IQ3_XXS | 61 / 0.79x | 95 / 0.84x | 107 / 0.79x | 119 / 0.77x |
| RTX 2080 Ti | IQ2_XXS / IQ2_XS / IQ2_S | 42-62 / 0.75-0.98x | 74-83 / 0.92x | 86-96 / 0.86-0.87x | 97-109 / 0.84x |
| RTX 3060 | IQ3_S | 74 / 1.09x | 144 / 0.99x | 154 / 0.89x | 174 / 0.87x |
| RTX 3060 | IQ3_XXS | 78 / 0.89x | 132 / 0.92x | 143 / 0.90x | 162 / 0.82x |
| RTX 3060 | IQ2_XXS / IQ2_XS / IQ2_S | 60-68 / 0.90-1.04x | 112-118 / 0.89-1.00x | 119-136 / 0.89-0.95x | 132-149 / 0.87-0.91x |

Per-expert / per-entry cost (us per layer call, RTX 2080 Ti, IQ3_S): 5.20 / 1.26 default -> 3.73 / 0.64.

nsys in the engine (96 verify windows, median kernel time per window at T = 1 / 2 / 3 / 4):
RTX 2080 Ti gate/up 1.60 / 2.58 / 3.51 / 4.16 ms -> **1.42 / 2.07 / 2.75 / 3.20 ms (-11 / -20 / -22 / -23%)**, slope per
row +0.86 -> +0.60 ms; RTX 3060 gate/up 0.77 / 1.43 / 1.97 / 2.23 -> 0.80 / 1.21 / 1.83 / 1.84 ms.

End to end, deterministic 131k bench (forced tokens and window sizes, 1024 decode tokens, three interleaved runs
lever on / off / on): **+1.1%** decode tokens/s (+0.0 and +2.1 in the two B arms, no B arm below A), text identical.
With the companion QSA scores lever (`STRATA_QSA_ROWS=1`, separate PR) on top: +2.9% together. The expert kernels are
~20% of the critical card's window and EXPDEC takes ~20% of that, so the end-to-end number is small by construction.

Ported and validated on v0.1.43 (micro-bench bit-exact on both cards; our full stack on 0.1.43 matches our 0.1.42
daily on the forced 131k bench).

## Risks / limits
- Off by default; a format or geometry outside the list keeps the default kernel for that role; `STRATA_OLD_IQ_MMVQ=1`
  turns it off as it turns the multi kernels off. Captured graphs keep the kernels they captured.
- Shared memory per block: 11.5 KB of activations + 1-8 KB of codebook (IQ2_S: 19.7 KB); registers 64 with the default
  `MINB` 4 (116-128 with `2`). Checked with `cuobjdump -res-usage` on sm_75 and sm_86.
- Not measured on sm_80/89/120 or on HIP (CUDA only). The weight side (the 8-9 two-byte loads per call and row) is
  untouched and is the next step; T=1 gains less (the flat cost is the weight side).



**Scope.** The gain comes from VRAM-resident routed experts / decode QSA scores on Turing/Ampere setups where these kernels are a large share of the verify window. On rigs where most routed experts are off-GPU (e.g. 4x RTX 5080 16 GB, see the test in #2058) the window is dominated by stage hand-offs and host experts, and the levers measure neutral.

En el sitio

Enlaces a install, modelos, releases.