#2356 pf_wmma: a second Volta prompt-expert kernel for IQ2_S gate/up and IQ4_NL down (STRATA_PF_WMMA2=1, opt-in): prompt +2.5-4.3% on a V100 + PG500 layer split

Description

## Title

pf_wmma: a second Volta prompt-expert kernel for IQ2_S gate/up and IQ4_NL down (STRATA_PF_WMMA2=1, opt-in): prompt +2.5-4.3% on a V100 + PG500 layer split

Issue: stacked on #2355 (`v100-pf-wmma-layer-split`, which carries #1771 and the layer-split fix). This PR's changes are its last two commits, `6a7425d` and `dc6cd5f` (a comment). Like #2355, it carries #1771's conflict with main at `ed4773e` in `src/prefill/prefill.cpp`.
Opt-in, so nothing changes without the variable.

## Summary

`gemm_iq_f16_grouped_kernel` (#1771's FP16 tensor-core prompt experts) dequantizes each K chunk of 128 weight rows into shared memory before the
tensor-core math: an IQ super-block (chunk 256) through `dq_dispatch`, one warp a row; IQ4_NL and Q2_0 (chunk 128) with their own loops
(IQ4_NL: one thread per row and 64-value half). The grid tables and the weight bytes are read from global memory as each chunk needs them.
We think that on a V100 this step does not fully hide behind the tensor-core math. The kernel below spreads it over the whole block; it also differs in
tile shape, static shared memory and launch bounds, and we did not isolate which part gives the gain.

The GSQ-RCO IQ3_XXS model that `setup.py --model IQ3_XXS` installs mixes expert types per layer: gate/up is IQ2_S in 10 of 48 layers and down is
IQ4_NL in 18 of 48. The new kernel takes those two types; the other layers keep the existing kernel, so the gain below comes from those layers only.

Behind the same entry point it adds a second kernel for these two types:
- A block = 128 weight rows of one expert x 64 sorted rows, K stages of 64. 8 warps = 4 x 2; each warp holds a 32 x 32 piece of the output in 2 x 2 `wmma` 16x16x16 fragments (FP32 sums).
- Each thread dequantizes exactly one 32-value sub-block per stage (row `t % 128`, sub-block `t / 128`) from a "raw stage" in shared memory
  (one IQ2_S super-block or four IQ4_NL blocks a row) that the whole block reads as contiguous 4-byte words.
  While the tensor cores work on a stage, the next raw stage and the next activations are already in registers (one barrier pair a stage).
- The IQ2_S grid and sign masks, or a byte -> two `kvalues_iq4nl` table for IQ4_NL, live in shared memory.
- Static shared memory <= 48 KiB (no per-device attribute to set). Full 16-row output tiles go straight to `Y`; an expert's last partial tile goes through a per-warp scratch.

By design the dequantized values are the ones `dq_iq2_s` / `dq_iq4_nl` give, in natural order, so the tensor cores get the same FP16 weights in the
same k order as in the existing kernel. Every other type, shape and alignment keeps the existing kernel.

The layout (a raw stage in shared memory, one sub-block per thread per stage) follows a Volta MoE kernel we wrote for our own llama.cpp build
(local patches, not published; the comment in the code says so). The code here is new, written with `nvcuda::wmma` and Strata's own group layout (no ggml dependency).

## What changed

- `src/kernels/cuda/iq_kernels.cu`: the second kernel, and with `STRATA_PF_WMMA2=1` a branch that sends IQ2_S (K a multiple of 512) and IQ4_NL (K a multiple of 128) to it.
  The variable is read like `STRATA_PF_WMMA` (on when its first character is `1`); unset or anything else: the existing kernel.
- Other backends: all of it is inside the `#if !defined(__HIPCC__)` block of `gemm_iq_f16_grouped` (both HIP builds compile the stub that returns false,
  unchanged), and the SYCL build has no `gemm_iq_f16_grouped`.
  HIP builds, without a GPU in ROCm 7.14 containers: `-DSTRATA_HIP_GFX906=ON -DCMAKE_HIP_ARCHITECTURES=gfx906` (the `mixa3607/rocm-gfx906:7.14-complete`
  image docs/AMD_HIP.md uses; `strata` and `pf_wmma_parity`) and `-DSTRATA_ENABLE_HIP=ON -DSTRATA_ENABLE_CUDA=OFF -DCMAKE_HIP_ARCHITECTURES="gfx1100;gfx1201"` (`rocm/dev-ubuntu-24.04:7.14.1-full`;
  `strata`) both compile with 0 errors. Not run (no AMD card here). SYCL not built: it does not compile the changed file.

## Checks (Tesla V100-PCIE-32GB + Tesla PG500-216 (sm_70, 32 GB), layer split K=24, CUDA 12.9, Linux, `-DSTRATA_EXPERIMENTAL_SM60=ON`)

How we checked: a build of #2355's branch + `6a7425d` (`dc6cd5f` changes a comment only), the parity test on both cards, then the engine with and without the variable in ABBA order at two power limits.

- `pf_wmma_parity` (unset and `STRATA_PF_WMMA2=1`): 0 failures on 2 devices both times. It compares the result against a double-precision reference;
  in the three cases the new kernel takes, its max |err| equals the existing kernel's (IQ2_S 1280x2560 0.000444 of 106, IQ4_NL 1280x2560 0.000393 of 97.5,
  IQ4_NL 2560x640 0.000162 of 65.4), which fits the same weights summed in the same order.
  (The first version rounded the IQ2_S sub-scale to FP16 before multiplying and gave 0.0435: within tolerance but not the same values, so we changed it.)
- Engine: the GSQ-RCO IQ3_XXS model and pack that `setup.py --model IQ3_XXS` installs (`ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF`), with setup's
  launch settings (`--mmap-experts`, image input on, `--spec 4 --spec-min-p 0.5`), 256K context, int8 KV, layer split K=24, `STRATA_SM70_TABLE=1`,
  `STRATA_PLE_PROBE=0` (skips the start-up PLE read check and keeps direct reads). Every expert resident on both cards.
  A fresh engine start per run, on, off, off, on (ABBA), at 170 W and at 250 W.
  Prompt tok/s (prompt tokens / time to the first content token, measured on the client). Fresh 26.7K, 52K and 104K prompts, then another 26.7K and 52K (nothing reused):

  | prompt | 170 W off (2 runs) | 170 W on (2 runs) | ratio of means | 250 W off (2 runs) | 250 W on (2 runs) | ratio of means |
  |---|---|---|---:|---|---|---:|
  | 52,019 | 2,191.7 / 2,195.8 | 2,276.6 / 2,277.3 | +3.8% | 2,360.8 / 2,361.9 | 2,439.7 / 2,449.2 | +3.5% |
  | 103,869 | 2,267.2 / 2,267.5 | 2,332.5 / 2,325.4 | +2.7% | 2,457.1 / 2,457.0 | 2,511.4 / 2,525.4 | +2.5% |
  | 26,692 (second) | 1,905.5 / 1,907.0 | 1,986.4 / 1,985.0 | +4.2% | 2,038.4 / 2,040.7 | 2,127.5 / 2,126.0 | +4.3% |
  | 52,020 (second) | 2,191.4 / 2,193.6 | 2,261.0 / 2,273.9 | +3.4% | 2,362.7 / 2,356.0 | 2,453.0 / 2,446.4 | +3.8% |

  Every run with the variable is above every run without it; the two runs of a setting are within 0.6%. The first 26.7K after start (26,693 tokens)
  is left out of the table because one run (170 W off, 1,517.1) was far below its pair; all eight values: 170 W off 1,517.1 / 1,764.4, on 1,827.4 / 1,826.1;
  250 W off 1,864.3 / 1,869.8, on 1,939.3 / 1,938.6.
  The kernel runs on the prompt path only; decode is not part of this comparison.
- The long-prompt recall check (answer a passphrase placed in the middle) passed in all 40 requests (8 starts x 5).

## Extra Notes

- Only some layers of this model take the kernel (see Summary). Extending the same layout to other types (IQ3_S, IQ2_XS, IQ2_XXS, IQ3_XXS, Q2_0) might gain more; not done here.
- Disclosure: the change and this text were prepared with an AI assistant (Claude); the numbers are from runs on the machine above.

🤖 Generated with [Claude Code](https://claude.com/claude-code)