cuda: revive 2 dropped kernels with FMA-contraction fixes (stacked on #2)

[TrevorS]

PR3 draft: cuda: revive 2 previously-dropped kernels with FMA-contraction fixes (stacked on #2)

Summary

Two new commits stacked on mtp-beats-plain-kernels (PR #2). Revives kernel optimizations that had been dropped from PR #2 due to numerical drift, after root-causing each to a specific FMA-contraction issue and applying a targeted fix.

The new commits:

1. 584de5e — cuda: shared-weight Q8 matmul kernel for small-N batched (MTP verifier), originally dropped from PR cuda: small-N batched kernel polish (stacked on #1) #2 for MTP-strict drift, now correctly gated so it only replaces the bit-identical warp8 fallback (not the cuBLAS path).
2. af16691 — cuda: fuse head_rms_norm + rope_tail on Q (decode + batched paths), originally dropped for long-context long_memory_archive greedy drift, now uses __fmul_rn as an FMA-contraction barrier to reproduce the sequential reference's bit pattern.

What's still NOT in this PR (and why)

• routed_moe small-N decode-LUT path (the K5 candidate from PR cuda: small-N batched kernel polish (stacked on #1) #2 investigation) — drift root-caused but fix has unacceptable blast radius. The K5 down path uses dev_dot_q2_K_q8_K_block (single-pair variant: pre-computes dall = y_d * x_d then dall * isum - dmin * summs), while the upstream reference at n_tokens=2 uses dev_dot_q2_K_q8_K_block8 (8-pair batched variant: inlines yd * xd * isum[p] - yd * xmin * summs[p]). Under --use_fast_math, nvcc emits different FMA contractions for these two formulations, producing different fp32 outputs from the same INT32 partial sums. Closing the gap would require either (a) a new K5-only down kernel inlining the 8-pair math but looping over the 6 expert slots, or (b) modifying the shared single-pair kernel — but dev_dot_q2_K_q8_K_block has 8 other callers across the MoE down stack (lines 9289, 9322, 9348, 9374, 9456, 9491, 9515, 9553), so a change there risks regressing every batched MoE path. Both options exceed the reasonable scope for this PR. Worth a separate, focused PR if/when someone wants to chase it.

Root-cause details

K4 — shared-weight Q8 matmul

Originally fired unconditionally at n_tok=2..4, replacing both the cuBLAS Gemm/Sgemm path and the matmul_q8_0_preq_batch_warp8_kernel fallback. The cuBLAS path uses tensor-core tiling that is numerically different from any custom Q8 warp kernel; the warp8 fallback uses the same per-block wscale * xs * dot operand order as the share-warp kernel, so K4 IS bit-identical to warp8 but NOT to cuBLAS.

Fix (7-line gating predicate at ds4_cuda.cu:5996-6053):

if (blocks <= 32u &&
    (!g_cublas_ready ||
     (cuda_q8_f32_ptr(...) == NULL && cuda_q8_f16_ptr(...) == NULL))) {
    // share-warp dispatch
}

The cache-pointer checks short-circuit on label/dim mismatch and respect the same allowed-cache predicates as the cuBLAS dispatch, so the gate fires precisely on the weights that cuBLAS would have skipped (cache miss → fallback to warp8). On those weights, share-warp matches warp8 bit-for-bit. Weights that hit cuBLAS stay on cuBLAS.

K1 — head_rms_norm + rope_tail fuse

The naive fused kernel applied the rms scale to tail[i] / tail[i+1] inline in the same register-level expression as the RoPE cos/sin multiply. Under --use_fast_math nvcc fused scale * tail[i] * c and scale * tail[i+1] * s into single FMAs, which differs from the sequential reference path where head_rms_norm_kernel writes xr[i] *= scale to fp32 memory and rope_tail_kernel reads it back. The ULP-scale per-pair drift compounds across 30 layers × high pos0 to flip argmaxes on long-context tests.

Fix (2 inline operand changes at ds4_cuda.cu:2416-2425):

const float x0 = __fmul_rn(tail[i], scale);
const float x1 = __fmul_rn(tail[i+1], scale);

__fmul_rn is a hard FMA-contraction barrier that forces a single-rounded fp32 multiply, reproducing the bit pattern of the sequential reference's memory store-then-load.

Headline numbers (DGX Spark / GB10, ds4flash IQ2XXS-w2Q2K)

Single-prompt smoke benches (./ds4 -n 256 -p "knight" --temp 0), 3-run mean:

Mode	PR #2 base	this PR
Plain decode	16.52	16.55
MTP batched-verifier strict	13.74	13.80

Marginal absolute gains because the kernels target paths that don't dominate single-token decode. The engineering value is the root-causing and fixing, not the t/s number — the K1 and K4 fixes are general-purpose techniques (FMA barriers, dispatch gating on cache-pointer availability) that future kernel work can reuse.

Tested against

• make clean && make cuda-spark — clean, no new warnings
• ./ds4_test --all — only pre-existing logprob-vectors short_code_completion failure (also on upstream/main); tensor-equivalence summary: capture_fail=0 logits_fail=0 greedy_fail=0 top1_mismatch=0
• make cuda-regression — pre-existing build error (also on upstream/main), unchanged
• make cpu — clean build, no new warnings
• Plain 3-run stability: byte-identical coherent output, 16.5 ± 0.03 t/s
• MTP STRICT 3-run: 13.8 t/s on first 2 runs, 10.86 t/s on a third run impacted by external GPU workload (not commit-induced)
• Byte-equality diff vs upstream/main for plain decode (n=64, "knight") — 0 bytes
• Byte-equality diff vs upstream/main for MTP BATCH_VERIFY=1 STRICT=1 (n=64, "knight") — 0 bytes
• Hardware: NVIDIA DGX Spark (GB10 / sm_121), driver 580.142, CUDA 13.0
• Model: DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf (80.76 GiB)

Pre-existing non-determinism observation

Independent of this PR: the mtp-beats-plain-kernels (#2) baseline and upstream/main both show intermittent long_memory_archive greedy_fail in roughly 1 of 4 ds4_test --all runs. Source: MoE down path uses atomicAdd at n_tokens >= 128 (prefill chunk size), which is not float-deterministic. The 5 dropped-kernel runs and 3 PR3-final runs show this same distribution — neither K1 nor K4 introduces it. Flagged here because it surfaced during the careful investigation that produced this PR; it's a separate concern worth its own ticket.

Out of scope / follow-ups

• K5 (routed_moe small-N decode-LUT path) revival — would need a focused PR on the dev_dot_q2_K_q8_K_block/block8 family
• The atomicAdd non-determinism above
• Combined-forward MTP path (downstream fork) — separate follow-up
• Captured-graph spec decode support — separate subsystem PR

[TrevorS]

Superseded by the reframed 2-PR stack (#11 + #12), which tells the same Spark/GB10 + MTP combined-forward story more concisely, rebased on current upstream/main, with the exploratory paths dropped.