Skip to content

finegrained-kernels: add FP8 weight-only inference - #1241

Open
IlyasMoutawwakil wants to merge 27 commits into
mainfrom
finegrained-weight-only-fp8
Open

IlyasMoutawwakil wants to merge 27 commits into
mainfrom
finegrained-weight-only-fp8

Conversation

@IlyasMoutawwakil

@IlyasMoutawwakil IlyasMoutawwakil commented Oct 16, 2026 •

Copy link
Copy Markdown
Member

Adds FP8 weight-only inference through activation_format="bf16" for dense, batched, grouped and MoE entrypoints, with per-tensor or coarse-block weight scales and BF16/FP16 activations. This covers ModelOpt's FP8_PB_WO recipe.

  • Batched decode and small 2D calls apply weight scales inline.
  • At 1,024 tokens, 2D dispatch switches to GPU weight dequantization followed by native GEMM, or the existing BF16 GEMM when fused gate/bias/output semantics require it.
  • Grouped prefill dequantizes active experts once and reuses the existing BF16 grouped GEMM and fused epilogues. Its temporary buffer reserves all expert planes; inactive planes are neither decoded nor read. No permanent BF16 weight cache.

Matched B200 CUDA-graph checks, including dequantization:

Workload Previous inline New path
Dense, 1,024 tokens, N8192/K6144 171 µs 126 µs
Dense, 8,192 tokens, N8192/K6144 1,519 µs 633 µs
Complete MoE prefill, 8,192 tokens, E16/H4096/I6400/top2 4.51 ms 3.88 ms

These use block128 scales and bounded tuning, and are microbenchmarks rather than full-model throughput. Stronger tuning made 512-token dense calls effectively tied, motivating the conservative threshold. Compiled native Torch remains competitive: 616 µs versus 633 µs for the large dense case, and can still win decode.

Extends the existing main benchmark with eager/compiled runtime-dequantization baselines, resident-BF16 dense controls and optional vLLM Marlin comparisons. Tests reuse Problem/MoEProblem; no separate weight-only test framework. Existing operator/MoE suites and focused final-source compilation, CUDA-graph, autocast and output-dtype checks pass.

Related issue

Requested by a maintainer for huggingface/transformers#49658.

Checklist

  • Linked to a discussed issue
  • Tested locally
  • New/changed functionality has test coverage
  • Mostly or completely generated with an LLM

@github-actions

Copy link
Copy Markdown

This PR changes source code for the following kernel(s): finegrained-kernels.

If you have the kernel-bot build permissions, then trigger /kernel-bot build finegrained-kernels yourself. Otherwise, a maintainer will review the PR and will take an appropriate action. Thanks for your patience.

@github-actions github-actions Bot added the chore Version bumps, releases, misc maintenance label Oct 16, 2026
@IlyasMoutawwakil IlyasMoutawwakil changed the title finegrained-kernels: add FP8 weight-only GEMMs finegrained-kernels: add FP8 weight-only inference Oct 17, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

chore Version bumps, releases, misc maintenance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant