perf(gdn-prefill): ldmatrix A+B GDN-projection GEMM — 2.1x kernel, warm-TTFT median -5.8% / p90 -10.3%

[tbraun96]

Summary

The GDN-projection prefill GEMM (fp8_gemm_t) was the single largest kernel in warm prefill (34% of warm GPU time) and was badly untuned: ncu measured 10.7% SM throughput with a 56% MIO-queue stall — it read every A+B MMA fragment with ~32 scalar shared-memory loads per K-step.

This ports the proven int8_gemm_8w_ldmab ldmatrix.x4 A+B fragment loads to FP8. E4M3 shares the m16n8k32 8-bit fragment layout with s8, so the loads transfer verbatim; only the MMA (s8→e4m3) and the epilogue (direct f32 accumulate — fp8_gemm_t is unscaled, the scale is folded into the weights upstream) change. 8 ldmatrix.x4 replace 32 scalar loads.

Result: 2.1× on the kernel (14.6 → 30.5 TFLOP/s, cosine 1.000000) and the first prefill change to move end-to-end latency above the noise floor.

Kernel

	before (fp8_gemm_t)	after (fp8_fp8_gemm_ldmab)
TFLOP/s @ GDN in_proj (N=10240, K=5120, M=512)	14.6	30.5 (+109%)
ncu MIO stall	74 cyc	44 cyc
kernel duration	3.30 ms	1.84 ms
cosine vs fp8_fp8_gemm_t	—	1.000000 (bit-exact)

End-to-end (agentic subset, seed 42, all runs 174/174, IoU 0.6264 identical)

Same-box A/B (same image, only the env toggled) and cross-box replication:

warm TTFT	baseline (scalar)	ldmab	delta
median	1481	1397	−5.7%
avg	1655	1537	−7.1%
p90	2307	2074	−10.1%
max	5145	4270	−17.0%

Cross-box confirms: ldmab dgx1 1397/1537/2074 · ldmab dgx3 1385/1522/2054 · baseline dgx1 1481/1655/2307 · baseline dgx2 1470/1644/2296. Averaged: median −5.8%, avg −7.3%, p90 −10.3%.

Accuracy gates

GDN is drift-sensitive, so this was gated hard:

• KL logit drift: 100% token agreement (519/519), mean_KL = max_KL = 0.000000 — bit-identical to baseline (10 diverse prompts, top-5 logprobs).
• Microtest cosine 1.000000 vs the validated fp8_fp8_gemm_t.
• Agentic IoU 0.6264, identical across all four A/B runs.
• CI: radix_tree 41 pass / 0 fail. Coherence (factual / haiku / math / parallel tool-calls) pass.

Stability: full golden e2e

Full agentic run on the golden config (unsloth Qwen3.6-27B-NVFP4, qwen3_coder parser, --disable-tool-grammar true, 128 slots / interval 32, MTP --num-drafts 1 bf16, util 0.85, bf16 KV, 32k), ldmab default-on:

1007 samples issued / 1007 completed / 0 failed, duration 5714s, TTFT median 1318 ms / avg 1591 ms, agentic score 0.6202.

For reference, llama.cpp full-run on the same suite is 1431 median / 1666 avg — so Atlas now wins warm-TTFT median (−7.9%) and avg (−4.5%) on top of the existing decode and wall-clock leads.

Default-on

ATLAS_FP8_LDMAB is an opt-out (=0 falls back to the scalar path), not an opt-in — the win is unconditional and accuracy-neutral, so it should not sit behind a flag. Routing lives in ops::fp8_gemm_n128 (K must be a multiple of 32, the ldmab K-tile); it quantizes the BF16 activation to E4M3 once into a persistent grow-only scratch, then launches the ldmatrix GEMM.

Commit notes for the reviewer

• ad92cd66 / ec1159a8 / cc2ec62f — the win (kernel, routing, default-on). These are the commits that matter.
• 026d7aa9 — the DIAG-dump gating fix this PR originally carried (finalize_last fired a full-vocab 248k logits D2H + sort on every deep prefill turn; now env-gated only). Production hygiene, unchanged and still included.
• 006ea421 (half-batched B-frag register pre-stage) and 8c7c5b94 (4-stage cp.async pipeline + the TFLOP/s microtest harness) touch fp8_gemm_t, which is now only the fallback path. Both are cosine-verified but measured perf-neutral end-to-end; they were in the tree for the gated build. 8c7c5b94 is still WIP-labeled. Recommend squash-merge, or I can rebase those two out if you'd rather ship only the ldmab commits.

What this does NOT fix

p90 at full-suite scale (2436 ms) still trails llama's 2276 ms. That residual is not this kernel — it's fixed per-turn overhead and the attention-at-depth path. Separately verified as dead ends this cycle: fp8-KV / KV-compression (neutral-to-worse → attention is not bandwidth-bound), and a deeper NST=4 pipeline on the ldmab kernel itself (correct at cosine 1.0 once a loop-counter/threadIdx shadow was fixed, but a regression at every shape — the 32KB smem halves occupancy and the 2-stage double-buffer already hides the latency).

[SeedSource]

[BETA] Automated PR review — SeedSource AI pipeline

Advisory only, not a merge gate. AI-generated (internal risk-triage, Atlas serve gate, GLM-5.2 review), attributed per Atlas convention.

Triage

• Risk: low - serve-test-required (triage advisory): False - path-heuristic (authority): True
• Rationale: Removed diagnostic syncs based on context length, making them env-gated to improve TTFT performance.
• Serving-path files: crates/spark-model/src/model/trait_impl/prefill_b/finalize_last.rs

Serve gate

• Build: 31s - Qwen3.6-35B-A3B-FP8 on GB10 node1
• Top-k KL vs main: 0.0009 (no shift vs main) - decode delta -0.2%
• Coherence probe: coherence=pass ping=10/10 (note: our test_coherence.py is flaky here; KL is the reliable serve signal)

GLM-5.2 code review

Correctness: Sound. Both gates correctly preserve the env-var path; the auto-trigger removal is semantically clean. No behavioral change when ATLAS_DIAG_GEMMA4 is unset (the new default — no sync/readback).

Risk: Low-to-moderate. This is a default-behavior change: previously, any prefill with chunk_start + chunk_len > 16384 silently emitted diagnostics (sync + D2H + sort on 248k logits). If anyone in CI/staging relied on those dumps appearing unconditionally for long-context runs, they'll silently disappear. Verify no dashboards/alerts scrape these tracing::warn lines. The env-var escape hatch remains, so on-call debugging is unaffected.

Quality: Comment is accurate and useful. Minor: is_ok_and(|v| v == "1" || v == "true") is duplicated at both sites — not worth a helper for two call sites, but note if a third appears. Consider documenting the env var in a central diagnostics README so operators know it exists.

Verdict: Approve. Confirm no CI/log consumers depend on the auto-trigger output.

_{Pipeline: build+serve on GB10, GLM-5.2 via NVIDIA NIM, internal risk-triage classifier. External diff treated as untrusted input. Public repo.}