sm89 prefill FFN megakernel(epilogue-fold;Phase 2 silu-fold 已 wire,opt-in,2B +7.5%/8B +3.7%)

[heiheiha798]

背景

sm_89(RTX 4090)是本机主力开发卡。repo 里 megakernel 覆盖:sm_100/110 有 encoder GeGLU,sm_120 有 und_ffn_v5t/action_ffn_v6t/tinyfp8 + flashrt-project#134 的 Qwen3 prefill NVFP4 epilogue 融合。sm_89 零个 megakernel。 flashrt-project#139 已把 decode 优化透,prefill 路径没动。

方向

走 epilogue-fold(仿 flashrt-project#134 + action_ffn_v6t),排除 v5t grid_barrier 移植。 epilogue-fold 省 HBM 来回,无 grid_barrier 不碰 occupancy 死穴。

Phase 0 baseline(2026-07-17, 2B 实测)

• 2B text prefill:S=128→5.8ms, S=256→5.9ms, S=512→8.6ms, S=1024→16.4ms。
• launch overhead 实测(graph vs loop):S=128 1.8ms(30.9%),S≥256 被 GPU exec overlap 掉,prefill 是 compute/HBM-bound。
• 结论:launch-fusion 收益只在 S≤128;大 S 得靠消灭中间 HBM。

Phase 1 — residual-fold(✅ idiom 验证完成,不 wire)

把 residual_add fold 进 down GEMM epilogue。commit a17f971。正确性 cos 1.0。不 wire:prefill 的 residual 已被 residual_add_rms_norm_to_fp8 三合一融合,real 收益太小(~4µs)。

Phase 2 — GeGLU silu-fold(✅ 已 wire,opt-in,commit 9a27090)

把 gate_up GEMM + silu_mul + per-token block-128 FP8 quant 压进单 launch,消灭 [S,2*inter] BF16 transient。仿 sm100 flashrt_megakernel_geglu 的 gate-in-smem + flashrt-project#134 SwiGLU fold,移植到 sm89 手写 cp.async+ldmatrix+mma.m16n8k32(无 TMA/tcgen05)。

两个 kernel 变体(均 additive,baseline fp8_bs_gemm_kernel 不动)

• fp8_bs_geglu_silu_fold_kernel(two-pass:gate→smem→up):shipped tile 32x128_w4_s1。单 B smem region 跨 pass 复用,一次只活一个 acc(~70 regs)。
• fp8_bs_geglu_silu_fold_apersist_kernel(A-persistent interleaved 实验):双 acc 在寄存器,B 每 k-iter 重载。bench-only 保留;e2e sweep 仍是 two-pass s1 赢。

ncu 根因诊断(为什么朴素 two-pass 输)

• s2 occupancy 崩:gate_smem(8KB)+ 双 cp.async stage(40KB)→ 49.7KB > 48KB opt-in → 1 CTA/SM,8.33% occupancy。s1(29.7KB,3 CTA/SM)恢复。
• two-pass 结构税:STAGES=1 无 overlap,2× pipeline drain,2× A re-load。
• tile 计算密度:32x128 = 160 B/mma vs baseline 128x128 = 64 B/mma(2.5× 差),compute-bound 区间吃瘪。
• B 流量翻倍是 RED HERRING:总 B = 2N·K 两边相同。
• 结论:8B compute-bound(71% peak),HBM 省下来没地方花;2B 56% peak + K 短 + B<L2,fused 反赢。

shape-gate(opt-in,additive,default byte-identical)

env FLASH_RT_SM89_SILU_FOLD_PREFILL=1。gate 按 inter:

• 2B(inter=6144):S ∈ [96, 192] 走 fused
• 8B(inter=12288):S = 128 走 fused
• 其余 S 回退 baseline 两-launch 路径

正确性

e2e prefill cosine 0.9992-1.0 vs baseline,argmax 全 S 保持(含 saturated-FP8 真实激活)。比 baseline 少一次 bf16(acc) rounding(文档化,更精确)。修了一个 out_scale 写 bug(单线程 guard 导致 row 1..BM-1 未写 → scale 垃圾,fp8 正确)。

E2E 收益(4090, 30-iter median, fold on vs off)

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

band 之外无回归(回退 baseline)。tests/test_qwen3_vl_fp8_sm89.py 5/5 通过。

排除项

• v5t grid_barrier 多 phase:Phase 0 后排除。
• GEMM 内部双缓冲/warp-spec:warp-spec 实验已排除。
• residual-fold wire:Phase 1 验证后排除(已被 norm 预融合)。
• interleaved 双 B smem:73.6KB s2 → 1 CTA/SM,occupancy 死,排除(留 apersist 实验变体)。

环境

• conda env flashrt-libero(torch 2.6.0+cu124, python 3.10)。
• build:GPU_ARCH=89 + FLASHRT_BUILD_QWEN3_VL=ON + FLASHRT_ENABLE_MOTUS=ON + FLASHRT_ENABLE_QWEN3_VL_DEV_KERNELS=ON。
• 权重:/data/pretrained_models/Qwen3-VL-{2B,8B}-Instruct-FP8(8B 从 modelscope 下)。

分支

feat/sm89-prefill-ffn-megakernel(从 main @ a4e01b0)。Phase 1 a17f971,Phase 2 9a27090。

[heiheiha798]

Shipped as upstream PR flashrt-project#151 (draft): flashrt-project#151

Final state:

• Phase 1 residual-fold: idiom validation, not wired (commit a17f971)
• Phase 2 GeGLU silu-fold: wired, opt-in, 2B +7.5% / 8B +3.7% in the gated small-S regime, zero regression outside (commit 9a27090)

Remaining-space survey (multi-agent): diminishing returns. silu-fold was the last high-value sm89 megakernel — it hit the biggest transient ([S,2*inter]) in the only regime where fusion wins on a compute-bound 4090 (launch-bound small-S). The three remaining candidates (attn O→FP8 fold, qk_norm_rope fold, ViT RESID_BIAS) all repeat silu-fold's shape-gated pattern for smaller transients with higher complexity/risk (FA2 fork violates add-only; qk_norm_rope epilogue is 3-4× silu-fold's complexity; ViT RESID_BIAS is real but ~1-2% of ViT only). Closing this issue; further sm89 work would pivot to GEMM-side (persistent kernel / tile-tuning), not fusion.