[PR #337] Fix ue8m0 packing: mask mantissa bits when extracting fp32 exponents #372

New issue

Open

opened 2026-10-26 21:16:57 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:16:57 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/337
Author: @yhyang201
Created: 5/19/2026
Status: 🔄 Open

Base: main ← Head: fix/ue8m0-packing-mantissa-leak

📝 Commits (1)

16f167d Fix ue8m0 packing: mask mantissa bits when extracting fp32 exponents

📊 Changes

1 file changed (+12 additions, -9 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/smxx_layout.cuh (+12 -9)

📄 Description

transpose_and_pack_fp32_into_ue8m0 and pack_fp32_into_ue8m0 pack 4 fp32 exponent bytes into one uint32 using shift-and-OR:

packed |= (values[0] >> 23u);   // byte 0: OK
packed |= (values[1] >> 15u);   // byte 1: mantissa leaks into byte 0
packed |= (values[2] >>  7u);   // byte 2: mantissa leaks into bytes 0-1
packed |= (values[3] <<  1u);   // byte 3: mantissa leaks into bytes 0-2

When fp32 scale factors have non-zero mantissa (not exact powers of two), mantissa bits leak into adjacent exponent fields. For example, packing four copies of 0.006975 (exp=0x77) produces 0x77ffffff instead of 0x77777777.

This causes NaN in fp8_einsum when callers pass plain float32 scales.

The torch fallback get_mn_major_tma_aligned_packed_ue8m0_tensor_torch is correct: it uses >> 23 then .to(kUInt8), which naturally truncates, but the CUDA kernel path doesn't truncate.

Fix: extract each exponent with (val >> 23) & 0xFF and shift to the correct byte position, in a shared pack_4_fp32_exponents helper.

Reproduction

import torch
from deep_gemm.utils.layout import get_mn_major_tma_aligned_packed_ue8m0_tensor

# Packing bug: non-power-of-2 fp32 value gets mantissa leaked into packed result
sf = torch.full((4, 8), 0.006975, device="cuda", dtype=torch.float32)
packed = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf)
assert packed[0, 0].item() == 0x77777777, f"got {hex(packed[0, 0].item())}"  # fails: 0x77ffffff

# Power-of-2 value works because mantissa is zero
sf2 = torch.full((4, 8), 0.0078125, device="cuda", dtype=torch.float32)
packed2 = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf2)
assert packed2[0, 0].item() == 0x78787878  # passes

---

<sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/337 **Author:** [@yhyang201](https://github.com/yhyang201) **Created:** 5/19/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/ue8m0-packing-mantissa-leak` --- ### 📝 Commits (1) - [`16f167d`](https://github.com/deepseek-ai/DeepGEMM/commit/16f167d0fceb501e906c2471924f94faf15ac45a) Fix ue8m0 packing: mask mantissa bits when extracting fp32 exponents ### 📊 Changes **1 file changed** (+12 additions, -9 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/smxx_layout.cuh` (+12 -9) </details> ### 📄 Description `transpose_and_pack_fp32_into_ue8m0` and `pack_fp32_into_ue8m0` pack 4 fp32 exponent bytes into one uint32 using shift-and-OR: packed |= (values[0] >> 23u); // byte 0: OK packed |= (values[1] >> 15u); // byte 1: mantissa leaks into byte 0 packed |= (values[2] >> 7u); // byte 2: mantissa leaks into bytes 0-1 packed |= (values[3] << 1u); // byte 3: mantissa leaks into bytes 0-2 When fp32 scale factors have non-zero mantissa (not exact powers of two), mantissa bits leak into adjacent exponent fields. For example, packing four copies of 0.006975 (exp=0x77) produces 0x77ffffff instead of 0x77777777. This causes NaN in `fp8_einsum` when callers pass plain float32 scales. The torch fallback `get_mn_major_tma_aligned_packed_ue8m0_tensor_torch` is correct: it uses `>> 23` then `.to(kUInt8)`, which naturally truncates, but the CUDA kernel path doesn't truncate. Fix: extract each exponent with `(val >> 23) & 0xFF` and shift to the correct byte position, in a shared `pack_4_fp32_exponents` helper. ## Reproduction ```python import torch from deep_gemm.utils.layout import get_mn_major_tma_aligned_packed_ue8m0_tensor # Packing bug: non-power-of-2 fp32 value gets mantissa leaked into packed result sf = torch.full((4, 8), 0.006975, device="cuda", dtype=torch.float32) packed = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf) assert packed[0, 0].item() == 0x77777777, f"got {hex(packed[0, 0].item())}" # fails: 0x77ffffff # Power-of-2 value works because mantissa is zero sf2 = torch.full((4, 8), 0.0078125, device="cuda", dtype=torch.float32) packed2 = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf2) assert packed2[0, 0].item() == 0x78787878 # passes --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:16:57 +00:00

bsakaguchi referenced this issue

2026-10-26 21:17:35 +00:00

[PR #403] Fix SM120 scale-factor layout transformation #421