Skip to content

Add opt-in FP8 activation kernels for SM120 E2M1 prefill MoE - #13

Closed
qu0b wants to merge 1 commit into
ormandj:mainfrom
qu0b:feat/fp8-prefill-moe
Closed

qu0b wants to merge 1 commit into
ormandj:mainfrom
qu0b:feat/fp8-prefill-moe

Conversation

@qu0b

@qu0b qu0b commented Oct 12, 2026

Copy link
Copy Markdown

Large eager batches in the SM120 E2M1-K32 MoE path currently use BF16 activation GEMMs. This adds an opt-in FP8 activation path using native E2M1 tensor-core operands, per-row activation scales and FP32 application of the existing K32 weight scales. It preserves the prepared weight representation; activation quantization remains lossy.

The kernel uses a six-stage asynchronous copy pipeline, pairs CTA barriers, skips padded M16 blocks and vectorizes activation loads and the final route reduction. SGLANG_B12X_FP8_PREFILL_MIN_ROWS enables it above a token-expert row threshold; zero remains the default. CUDA graph capture continues using W4A16. This applies to the existing E2M1-K32 path, not the block-FP8 weights used by the default GLM checkpoint.

benchmarks/sm120_fp8_prefill.py provides synthetic numerical checks and interleaved single-GPU timings against W4A16. It covers partial blocks, uneven routing, zero weight scales and masked routes. It reports activation error against a PyTorch FP8 emulation and the BF16 kernel. These checks do not establish model quality or end-to-end serving performance. No public image was built or qualified by this change.

Measured on one NVIDIA RTX PRO 6000 Blackwell Max-Q at a 300 W power limit, using a separate validation environment with PyTorch 2.14.1+cu130 and CUDA 13.0. Synthetic shape: 288 experts, hidden size 4096, intermediate size 1024, top-k 8. Each entry is the median of 90 CUDA-event samples in three alternating-order rounds after compilation and warm-up.

Tokens W4A16 (ms) FP8 activations (ms) Kernel time reduction
1,024 2.028 1.705 15.9%
2,048 2.887 2.073 28.2%
4,096 5.197 3.640 30.0%
8,192 9.832 6.799 30.8%

The synthetic FP8-reference relative L2 error was 0.86–0.95%; error against W4A16 was about 4.6%. The benchmark permits small split-K repeat differences and rare adjacent-code FP8 rounding ties. These are numerical observations on synthetic inputs, not model-quality guarantees.

Validation: nine selection and kernel-contract tests pass; synthetic GPU checks pass including int32/int64 masked routes; release-contract and source-patch verification pass. A GPU trace confirms execution of the FP8 fused kernel, row quantizer and vectorized route sum. An independent comparison with the unchanged W4A16 package found identical outputs at two tested sizes and relative L2 error of 0.00052% at the third.

Developed with AI assistance.

@qu0b

qu0b commented Oct 12, 2026

Copy link
Copy Markdown
Author

Superseded by direct native-source contributions: FlashInfer format support flashinfer-ai/flashinfer#6355 and FP8 prefill flashinfer-ai/flashinfer#6356. The independent scheduler admission improvement is sgl-project/sglang#43844. The FlashInfer port was adapted to upstream main and remeasured on SM120, including public-API reference tests and a cooperative-launch fix for small batches. Closing this integration-patch PR to keep upstream review focused on the source implementations.

@qu0b qu0b closed this Oct 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant