feat(sm89): GeGLU silu-fold megakernel for Qwen3-VL prefill (opt-in)

[heiheiha798]

Summary

First sm_89 (RTX 4090) megakernel for the Qwen3-VL FP8 prefill path. Fuses the gate_up GEMM + SwiGLU silu_mul + per-token block-128 FP8 quantization into a single launch, eliminating the [S, 2*inter] BF16 transient the baseline two-launch path (gate_up GEMM → silu_mul_merged_to_fp8) materializes. Mirrors #134's SwiGLU epilogue fold (sm_120) and the sm_100 flashrt_megakernel_geglu "gate stays in smem" pattern, ported to sm_89's hand-written cp.async + ldmatrix.x4 + mma.m16n8k32 idiom (no TMA / no tcgen05).

sm_89 previously had zero megakernels (sm_100/110 have encoder GeGLU; sm_120 has und_ffn_v5t/action_ffn_v6t/tinyfp8 + #134). #139 already optimized the decode path; this addresses the untouched prefill path.

Opt-in, additive-only, default path byte-identical: env FLASH_RT_SM89_SILU_FOLD_PREFILL=1 (or use_silu_fold_prefill=True), shape-gated per-inter so it only fires in the launch-bound regime where it wins.

What's in here

• fp8_bs_geglu_silu_fold_kernel (csrc/gemm/fp8_bs_gemm_device.cuh) — two-pass per CTA: pass 1 accumulates gate over K and stores silu(gate) as BF16 into a smem gate buffer; pass 2 reuses the same A/B smem staging, accumulates up, then the epilogue reads gate from smem, forms v = bf16(bf16(silu(gate))*up) (matching silu_mul_merged's two bf16 roundings), per-row amax-reduces over the 128-col quant block, and quantizes to FP8. No grid_barrier (BLOCK_N==128 = one quant block, amax is intra-CTA). Shipped tile 32x128_w4_s1.
• fp8_bs_geglu_silu_fold_apersist_kernel — A-persistent interleaved experiment (both gate+up acc in registers, B reloaded per k-iter). Kept as a bench-only tile; the two-pass s1 won the e2e sweep. Documented as an alternative explored.
• fp8_bs_gemm_kernel RESID variant (Phase 1, commit a17f971) — D = bf16(acc) + resid epilogue-fold, if constexpr (RESID) dead-strips so baseline stays byte-identical. Validated the epilogue-fold idiom on sm_89 but not wired (prefill's residual is already norm-fused → ~4µs, not worth it). Kept as bench-only + documented.
• Frontend wiring (qwen3_vl_fp8_sm89.py): use_silu_fold_prefill opt-in + per-inter S-band gate (2B [96,192], 8B S=128); outside the band the baseline two-launch path runs unchanged.

Why the shape gate (the ncu diagnosis)

ncu + a multi-agent diagnosis isolated why a naive two-pass loses on sm_89:

• s2 occupancy collapse: gate_smem (8KB) + 2 cp.async stages (40KB) → 49.7KB > 48KB opt-in tier → 1 CTA/SM, 8.33% occupancy. s1 (29.7KB, 3 CTA/SM) recovers.
• two-pass structural tax: STAGES=1 has no cp.async/compute overlap → 2× pipeline drain.
• gate_up is compute-bound at large S (8B: 71% of FP8 peak). Fusion that doesn't reduce MMA FLOPs only helps where the eliminated launch/transient is not overlapped behind compute — i.e. the launch-bound small-S regime.

So the gate only fires where it wins; above the crossover it falls back. This is the honest ceiling on a compute-bound 4090. (A follow-up survey confirmed silu-fold is the last high-value sm89 megakernel — decode is graph-captured so launch-fusion can't win there, and the remaining prefill candidates repeat this small-S-only pattern for smaller transients.)

Measurements (RTX 4090, 30-iter median, CUDA-event)

Methodology note (per CONTRIBUTING §Performance Measurement): these are eager (non-graph-captured) prefill latencies — the _layer_forward_prefill_fp8_blockscaled path timed with CUDA events around the full layer-forward loop (host issues every launch; not a cudaGraphLaunch replay). The prefill path is not graph-captured in this frontend, so eager is the honest metric. The decode path is graph-captured and is untouched by this PR.

E2E prefill (text path, single layer-forward, eager):

	off	on	Δ
2B S=128	6.10 ms	5.64 ms	+7.5%
2B S=192	6.10 ms	5.54 ms	+9.2%
8B S=128	13.97 ms	13.45 ms	+3.7% (stable over 3 runs)

Outside the gated S-band: no regression (falls back to baseline).

Correctness

• E2E prefill cosine 0.9992–1.0 vs baseline across S; argmax preserved on all S (including saturated-FP8 real activations).
• One fewer bf16(acc) rounding than baseline (fused keeps the accumulator in fp32); documented, more precise.
• Replicates silu_mul_merged_to_fp8's exact two-bf16-rounding sequence on the silu product.
• Fixed an out_scale write bug found via real-data testing (a single-thread guard left rows 1..BM-1 unwritten → garbage scale; cosine was scale-invariant so synthetic data hid it, real saturated-FP8 data exposed it).

Compliance (CONTRIBUTING.md / exec_contract §9)

• Additive-only: fp8_bs_gemm_kernel and the production dispatcher are untouched; new behavior is new kernels + an opt-in env-gated frontend branch. Default path byte-identical.
• Kernel bindings & CMake ownership: bench-only tiles under FLASHRT_ENABLE_QWEN3_VL_DEV_KERNELS; the production binding fp8_bs_geglu_silu_fold_sm89_fp8out is under ENABLE_SM89_BLOCK_FP8_GEMM (same gate as its source object). No unconditional m.def references gated symbols. Binding comment documents argument shapes.
• hasattr/getattr discipline: use_silu_fold_prefill resolves the tile via getattr(fvk, ..., None) — an optional fast path with a correct baseline fallback, not hardware routing (sm_89 is enforced by the frontend class's capability check).
• No grid_barrier, no >48KB smem on the shipped tile, no GEMM-internal double-buffer.

Test

cmake -B build -S . -DGPU_ARCH=89 -DFLASHRT_BUILD_QWEN3_VL=ON \
      -DFLASHRT_ENABLE_QWEN3_VL_DEV_KERNELS=ON
cmake --build build -j --target flash_rt_qwen3_vl_kernels
PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 pytest -q \
  tests/test_qwen3_vl_fp8_sm89.py \
  tests/test_build_inventory.py \
  tests/test_generic_kernel_bindings.py \
  tests/test_install_smoke.py tests/test_load_model_use_fp8_kwarg.py \
  tests/test_calibration_helpers.py
python scripts/test_geglu_silu_fold_correctness.py --M 128 256 512 1024
FLASH_RT_SM89_SILU_FOLD_PREFILL=1 python scripts/smoke_qwen3_vl_fp8_sm89.py \
    --text-only --checkpoint <Qwen3-VL-2B-Instruct-FP8> \
    --text-s 128 --max-prefill-seq 128

All passing locally (RTX 4090, sm_89, CUDA 12.4, torch 2.6.0+cu124, flashrt-libero env): test_qwen3_vl_fp8_sm89 5/5, test_build_inventory 5/5, test_generic_kernel_bindings 7/7, CONTRIBUTING basic smoke set 55/55.

Unsupported hardware: the GeGLU silu-fold kernel is sm_89-only (Ada FP8 mma.m16n8k32, no .kind::f8f6f4). Not built/validated on sm_80/86/87/120 — the ENABLE_SM89_BLOCK_FP8_GEMM gate keeps it out of those builds.

Tracking issue: heiheiha798#1.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.