Two perf changes from a fresh review of the hot paths, now fully validated on both fleet architectures (RTX 3090 sm_86, Quadro P5200 sm_61). The deployment-shape validation materially changed the story for change 1 — see below. Change 2 (server logprobs) stands as originally measured.

1. CUDA: FA vec kernel instances for D == 512 with matched quantized KV — now gated to gqa_ratio <= 4

Gemma 4 global-attention layers (head size 512) always dispatch to the tile (pre-Volta) or MMA (Turing+) FA kernels — both require K/V in F16, so with a quantized KV cache the entire per-layer global KV is dequantized into an F16 staging buffer on every decode step. The vec kernel reads q4_0/q8_0 directly but had no D=512 instances (only 64/128/256). This PR adds matched q4_0/q8_0 D=512 vec instances.

What deployment-shape validation found: the original perf cases used gqa_ratio = 4. The real Gemma 4 global layers are MQA (head_count_kv = 1, gqa_ratio = 16). At that shape the vec kernel re-reads the single KV head once per Q head (16×), while TILE/MMA pay the dequant staging once and amortize K/V reads across Q heads via the GQA optimization (8 Q-heads per pass). The vec kernel loses there — on both architectures:

Per-op, deployment shape (MQA, nh_kv=1, ratio 16, nb=1, q8_0):

(q4_0 is worse still: up to −159% on sm_61.) End-to-end this was a measurable ~4% tg regression at d32768 on the 3090, reproduced in both run orders.

Per-op, gqa_ratio = 4 (the shape where vec wins), q8_0:

(q4_0 at ratio 4: sm_86 keeps a 1.65–1.79x win; sm_61 is a wash, 1.00–1.04x — the in-kernel q4_0 dequant costs roughly what the staging pass saved on Pascal.)

Resolution (3c8a368): the D=512 vec dispatch requires gqa_ratio <= 4. The deployment shape falls back to the baseline kernels — verified per-op parity (≤2%) and e2e parity on both arches. Net effect for the current fleet: neutral on FA (no regression, no win), with the kernels + dispatch + test coverage in place for low-ratio D=512 shapes and as the substrate for the real fix.

Follow-up that would capture the actual prize: a quantized-KV-direct kernel with GQA packing (vec + ncols2, or TILE reading q8_0/q4_0 natively) would cut Gemma 4 global-layer FA traffic to ~1/4 of either current path — the per-op headroom at MQA-16 is ~4x. That is a kernel-engineering task, not a dispatch tweak.

2. Server: avoid full-vocab sort in get_token_probabilities

When a client requests logprobs (n_probs > 0), the server did a full std::sort of the entire vocabulary (262k entries for Gemma) per emitted token, then the caller linearly re-scanned the sorted vector to find the sampled token. Softmax normalization only needs max+sum (O(V), no ordering); top-k comes from std::partial_sort. O(V log V + V) → O(V + k log k) per token, identical output.

Validation

• test-backend-ops test -o FLASH_ATTN_EXT vs CPU reference: 2903/2903 pass on sm_86 (RTX 3090) and full sweep pass on sm_61 (Quadro P5200) — including new hs=512 quantized-KV cases at both gqa_ratio=4 and the deployment MQA-16 shape (q8_0/q4_0, nb={1,2,3,32}, sinks variant). Quantized-KV FA previously had zero hs=512 coverage (only 64/72).

• sm_61 run exercises the cols_per_block=2 vec instantiation (pre-Volta dispatch, ne1<=2); sm_86 exercises cols_per_block=1.

• E2e llama-bench A/B (gemma-4-12B QAT Q4_K_XL, -fa 1 -ctk q8_0 -ctv q8_0): gated build ≡ baseline within noise on sm_86; ungated build reproduced the predicted regression (3090: 58.6→56.3 tg @d32k). sm_61 gated per-op parity confirmed (228/397/752 us vs 231/404/766 baseline).

• sm_61 3-way e2e (amethyst, idle P5200, ngl=99, r=2 — stddev ≤0.31 throughout):

  tg32 @ depth	baseline	ungated vec	gated (this PR)
  d0	24.70	24.35 (−1.4%)	24.54
  d4096	22.39	21.77 (−2.8%)	22.24
  d16384	20.36	18.88 (−7.3%)	20.30
  d32768	18.11	16.02 (−11.5%)	18.12 (parity)

  The ungated dispatch would have cost the Pascal fleet 11.5% decode at 32k depth; the gate restores exact parity.

• Deployment-shape perf cases added to make_test_cases_perf so this dispatch decision stays measurable.

• generate_cu_files.py updated so regeneration reproduces the instance files byte-for-byte.

Numbers provenance

• sm_86: local RTX 3090 (desktop-loaded; per-op numbers are thousands-of-runs averages, e2e cross-checked in both run orders).
• sm_61: amethyst (idle Quadro P5200), binaries cross-built on crystal with the scripts/build-pascal-p5200.md recipe, static-linked so the deployed v3 libs can't shadow.

🤖 Generated with Claude Code