deep-gemm: fused per-token FP8 cast in Triton (fused_per_token_cast_to_fp8), version 3 by IlyasMoutawwakil · Pull Request #1326 · huggingface/kernels-community

IlyasMoutawwakil · 2026-11-12T14:14:29Z

Adds deep_gemm.fused_per_token_cast_to_fp8, which does what utils.per_token_cast_to_fp8 does in one Triton kernel, with bit-identical output.

utils.per_token_cast_to_fp8 is upstream's PyTorch reference, used in its tests. transformers calls it on every DeepGEMM forward: the FP8/FP4 linear, both activations of the grouped experts, and Mega MoE. Each call costs:

about a dozen kernels;
a zero-filled padded copy of the input;
fp32 temporaries the size of the input;
with packed UE8M0 scales, a host sync in its packing check.

	`utils.per_token_cast_to_fp8`	`fused_per_token_cast_to_fp8`
Kernels	~12	1 (Triton)
Host sync	with packed UE8M0	none
Outputs	allocated	allocated, or written in place into given `out` / `out_sf` (e.g. views of the Mega MoE symmetric buffer, so no `.copy_`)
Compiled extension / CUDA toolkit	no	no

Same arguments as the reference: use_ue8m0, gran_k, use_packed_ue8m0. Input is bfloat16 or float32. Optional out / out_sf are added.
Local to the Hub build (torch-ext/deep_gemm/fused_cast.py). The vendored upstream sources are unchanged.
To match bit for bit, the reciprocal uses tl.math.div_rn. tl.fdiv lowers to the approximate div.full.f32, and that is enough to move some values across an E4M3 rounding boundary.
This adds a new function, so version goes to 3, and the tests load version 3.

Results

Bit-identical to utils.per_token_cast_to_fp8 in all 60 cases checked on a B200:

bfloat16 and float32 input;
group 32 and 128 with packed UE8M0, UE8M0 and fp32 scales;
shapes 1×512, 7×4096, 129×7168, 64×1000 (a partial last group) and 3000×4096;
an all-zero group and very large values.

Timings are on a B200 with hidden 4096, bfloat16, group 32 and packed UE8M0 (the Mega MoE recipe):

Tokens	`utils.per_token_cast_to_fp8` eager	Fused eager	Fused in a CUDA graph
1	128 µs	6.4 µs	1.2 µs
64	129 µs	6.3 µs	1.5 µs
512	128 µs	8.2 µs	2.0 µs
4096	335 µs	15.5 µs	6.3 µs
16384	1145 µs	38.7 µs	31.4 µs

Tests

tests/test_kernels_ci.py (kernels_ci) checks two things:

bit-exactness against the reference over the recipes and shapes above;
writing into views of a larger buffer, leaving the rest of the buffer untouched.

🤖 Generated with Claude Code

…version 3 `deep_gemm.utils.per_token_cast_to_fp8` is upstream's PyTorch reference: about a dozen kernels, a zero-filled padded copy of the input, fp32 temporaries its size, and a host sync in the packed-UE8M0 check. This adds the same cast as one CUDA kernel, bit for bit (bfloat16 or float32 input, group 32 or 128, fp32 / UE8M0 / packed-UE8M0 scales), writing into caller-given outputs such as views of the Mega MoE symmetric buffer. It is compiled ahead of time in its own `[kernel.per_token_cast]` section, also for pre-Hopper archs, so it needs neither the JIT runtime nor a CUDA toolkit. A new function, hence version 3. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

github-actions · 2026-11-12T14:14:46Z

This PR changes source code for the following kernel(s): deep-gemm.

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

Same API and bit-exact output (60/60 cases against utils.per_token_cast_to_fp8 on B200), without a compiled extension. Writes straight into the given out/out_sf views. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

vasqu · 2026-11-12T14:40:26Z

 import kernels

-deep_gemm = kernels.get_kernel("kernels-community/deep-gemm", version=2)
+deep_gemm = kernels.get_kernel("kernels-community/deep-gemm", version=3)


quite a lot of this lol, wondering whether this can be simplified at some point using some helpers but not here

agreed, a shared load_kernel() in the tests' conftest (or a VERSION constant) would make bumps one-liners. leaving it for a separate pr

did it here after all: tests/generators.py loads the kernel once (the only get_kernel call), and every test file does from generators import deep_gemm, so the next bump is one line

vasqu · 2026-11-12T14:42:17Z

do we use licences here?

no per-file headers for local additions in this repo, only vendored upstream files keep theirs. it's covered by deep-gemm's LICENSE (DeepGEMM's MIT), and it reimplements their per_token_cast_to_fp8

vasqu · 2026-11-12T14:52:27Z

+        tl.store(sf_ptr + row * stride_sf + groups, scale, mask=group_mask)
+
+
+def fused_per_token_cast_to_fp8(x, use_ue8m0, gran_k=128, use_packed_ue8m0=False, out=None, out_sf=None):


I assume this is essentially 1:1 to the deepgemm one but with slight adjustments

Lets shorten the docstrings to point that out + the main diff, no need to overcomplicate things

yes, it's utils.per_token_cast_to_fp8 1:1 (bit-identical), the differences being one kernel and writing into out / out_sf. shortened the docstrings to say just that

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The version is pinned in one place, so a bump touches one line. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

github-actions Bot added the chore Version bumps, releases, misc maintenance label Nov 12, 2026

IlyasMoutawwakil changed the title ~~deep-gemm: fused per-token FP8 cast (fused_per_token_cast_to_fp8), version 3~~ deep-gemm: fused per-token FP8 cast in Triton (fused_per_token_cast_to_fp8), version 3 Nov 12, 2026

IlyasMoutawwakil requested a review from vasqu November 12, 2026 14:30

vasqu previously approved these changes Nov 12, 2026

View reviewed changes

deep-gemm: shorter fused cast docstrings

5b103c4

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

IlyasMoutawwakil dismissed vasqu’s stale review via 5b103c4 November 12, 2026 16:43

vasqu previously approved these changes Nov 12, 2026

View reviewed changes

deep-gemm: load the kernel once in the tests (generators.py)

e832658

The version is pinned in one place, so a bump touches one line. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

IlyasMoutawwakil dismissed vasqu’s stale review via e832658 November 12, 2026 16:58

deep-gemm: mark test_lazy_init's kernel import as intentional

ac4b438

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

vasqu approved these changes Nov 12, 2026

View reviewed changes

Navigation Menu

Repository navigation

deep-gemm: fused per-token FP8 cast in Triton (fused_per_token_cast_to_fp8), version 3 - #1326

deep-gemm: fused per-token FP8 cast in Triton (fused_per_token_cast_to_fp8), version 3#1326
IlyasMoutawwakil wants to merge 5 commits into
mainfrom
deep-gemm-fused-cast

IlyasMoutawwakil commented Nov 12, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Nov 12, 2026

Uh oh!

vasqu Nov 12, 2026

Uh oh!

IlyasMoutawwakil Nov 12, 2026

Uh oh!

IlyasMoutawwakil Nov 12, 2026

Uh oh!

vasqu Nov 12, 2026

Uh oh!

IlyasMoutawwakil Nov 12, 2026

Uh oh!

vasqu Nov 12, 2026

Uh oh!

IlyasMoutawwakil Nov 12, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		tl.store(sf_ptr + row * stride_sf + groups, scale, mask=group_mask)


		def fused_per_token_cast_to_fp8(x, use_ue8m0, gran_k=128, use_packed_ue8m0=False, out=None, out_sf=None):

Conversation

IlyasMoutawwakil commented Nov 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Results

Tests

Uh oh!

github-actions Bot commented Nov 12, 2026

Uh oh!

vasqu Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

IlyasMoutawwakil Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

IlyasMoutawwakil Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

IlyasMoutawwakil Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

IlyasMoutawwakil Nov 12, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

IlyasMoutawwakil commented Nov 12, 2026 •

edited

Loading