[GH-ISSUE #365] [BUG] Non-deterministic numerical problem of k_grouped_fp8_gemm_nt_contiguous #174

New issue

Open

opened 2026-10-26 21:14:10 +00:00 by bsakaguchi · 3 comments

bsakaguchi commented

2026-10-26 21:14:10 +00:00

Owner

Originally created by @FFGGSSJJ on GitHub (Jun 23, 2026).
Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/365

I encountered a numerical problem when using k_grouped_fp8_gemm_nt_contiguous for wgrad computation in FP8 training. Specifically:

when one of the expert gets exactly 128 tokens
when one channel (column) of the tensor a that belongs to this expert are all zeros

the computed result for that expert will be wrong non-deterministically. And what interesting is that if I use different scale factor value for the all-zero channel during per-channel cast, the error might disappear.

I can reproduce the error using the following minimal unit test on my GH200 machine and deep_gemm version '2.6.1':

import torch
import deep_gemm

_KS = [256, 128, 128, 256]
_HIDDEN = 2048   # n: grad_y channels / wgrad rows
_FFN = 2048      # H: activation channels / wgrad cols
_AMAX_FLOOR = 0

def _per_channel_cast_to_fp8(x, gran_k=128, sf_val=1e-30):
    m, n = x.shape
    x_view = x.view(-1, gran_k, n)
    amax = x_view.abs().float().amax(dim=1).view(-1, n)
    if _AMAX_FLOOR > 0:
        amax = amax.clamp(_AMAX_FLOOR)
    sf = amax / 448.0

    # NOTE: handle cases when the channel contains all zeros
    sf = torch.where(sf==0, torch.full_like(sf, sf_val), sf).clamp_min(1e-30)
    fp8 = (x_view * (1.0 / sf.unsqueeze(1))).to(torch.float8_e4m3fn).view(m, n)
    return fp8, sf

def _pack_kmajor(fp8, ks, mcols):
    """Per-expert k-major pack, mirroring generators.py::generate_k_grouped_contiguous."""
    out = torch.empty(sum(ks) * mcols, dtype=fp8.dtype, device=fp8.device)
    prefix = 0
    for k in ks:
        out[prefix * mcols:(prefix + k) * mcols] = fp8[prefix:prefix + k].T.flatten()
        prefix += k
    return out

def _make_operand(x, ks, sf_val=1e-30):
    fp8, sf = _per_channel_cast_to_fp8(x, 128, sf_val)
    return _pack_kmajor(fp8, ks, x.shape[1]), sf.T


def _build_operands(device, seed=0):
    """Synthetic grad_y / activation with the minimal trigger structure."""
    g = torch.Generator(device=device).manual_seed(seed)
    m = sum(_KS)
    grad_y = torch.randn(m, _HIDDEN, device=device, dtype=torch.bfloat16, generator=g) * 1e-7
    s = torch.randn(m, _FFN, device=device, dtype=torch.bfloat16, generator=g) * 0.4

    offsets, off = [], 0
    for k in _KS:
        offsets.append(off)
        off += k
    
    # NOTE: setup one all-zero channel for 128-token expert
    for e, k in enumerate(_KS):
        if k == 128:
            grad_y[offsets[e]:offsets[e] + k, 0] = 0
    return grad_y, s, offsets

def test_k_grouped_fp8_gemm_is_deterministic(sf_val):
    """The grouped FP8 wgrad GEMM must return the same correct result every call."""
    device = torch.cuda.current_device()
    iters = 10000
    num_experts = len(_KS)

    grad_y, s, offsets = _build_operands(device)
    grouped_layout = torch.tensor(_KS, dtype=torch.int32, device=device)

    # per-channel cast
    fp8_grad_y = _make_operand(grad_y, _KS, sf_val)
    fp8_s = _make_operand(s, _KS, sf_val)
    torch.cuda.synchronize()

    # reference wgrad
    ref = torch.zeros(num_experts, _HIDDEN, _FFN, device=device, dtype=torch.float32)
    for e, k in enumerate(_KS):
        if k == 0:
            continue
        gy_e = grad_y[offsets[e]:offsets[e] + k].float()
        s_e = s[offsets[e]:offsets[e] + k].float()
        ref[e] = gy_e.t() @ s_e
    ref_norms = torch.linalg.vector_norm(ref.reshape(num_experts, -1), dim=1)
    thresholds = (100.0 * ref_norms).clamp_min(1.0)  # blow-up = >=100x the true grad

    failures = []
    for it in range(iters):
        out = torch.zeros(num_experts, _HIDDEN, _FFN, device=device, dtype=torch.float32)
        torch.cuda.synchronize()
        deep_gemm.k_grouped_fp8_gemm_nt_contiguous(
            fp8_grad_y, fp8_s, out, _KS, grouped_layout, out, recipe=(1, 1, 128),
            use_psum_layout=False,
        )
        torch.cuda.synchronize()
        out_norms = torch.linalg.vector_norm(out.reshape(num_experts, -1), dim=1)
        bad = torch.nonzero(out_norms > thresholds).flatten().tolist()
        if bad:
            failures.append((it, bad, out_norms.max().item()))

    if failures:
        first_it, experts, worst = failures[0]
        print(
            f"\nk_grouped_fp8_gemm_nt_contiguous is nondeterministic: "
            f"{len(failures)}/{iters} calls blew up (first at iter {first_it}, "
            f"experts {experts}); worst output norm {worst:.3e} vs reference max "
            f"{ref_norms.max().item():.3e}. Operands were cast once and are identical "
            f"across calls (single stream, fully synced); amax_floor={_AMAX_FLOOR} sf_val={sf_val}."
        )

if __name__ == "__main__":
    for sf_val in [1, 1e-4, 1e-8, 1e-30]:
        test_k_grouped_fp8_gemm_is_deterministic(sf_val)

with outputs:

k_grouped_fp8_gemm_nt_contiguous is nondeterministic: 1749/10000 calls blew up (first at iter 8, experts [2]); worst output norm 7.451e+03 vs reference max 1.311e-03. Operands were cast once and are identical across calls (single stream, fully synced); amax_floor=0 sf_val=1.

k_grouped_fp8_gemm_nt_contiguous is nondeterministic: 181/10000 calls blew up (first at iter 127, experts [2]); worst output norm 1.093e+00 vs reference max 1.311e-03. Operands were cast once and are identical across calls (single stream, fully synced); amax_floor=0 sf_val=0.0001.

Originally created by @FFGGSSJJ on GitHub (Jun 23, 2026). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/365 I encountered a numerical problem when using `k_grouped_fp8_gemm_nt_contiguous` for wgrad computation in FP8 training. Specifically: - **when one of the expert gets exactly 128 tokens** - **when one channel (column) of the tensor $a$ that belongs to this expert are all zeros** the computed result for that expert will be wrong non-deterministically. And what interesting is that if I use different scale factor value for the all-zero channel during per-channel cast, the error might disappear. I can reproduce the error using the following minimal unit test on my GH200 machine and deep_gemm version '2.6.1': ```py import torch import deep_gemm _KS = [256, 128, 128, 256] _HIDDEN = 2048 # n: grad_y channels / wgrad rows _FFN = 2048 # H: activation channels / wgrad cols _AMAX_FLOOR = 0 def _per_channel_cast_to_fp8(x, gran_k=128, sf_val=1e-30): m, n = x.shape x_view = x.view(-1, gran_k, n) amax = x_view.abs().float().amax(dim=1).view(-1, n) if _AMAX_FLOOR > 0: amax = amax.clamp(_AMAX_FLOOR) sf = amax / 448.0 # NOTE: handle cases when the channel contains all zeros sf = torch.where(sf==0, torch.full_like(sf, sf_val), sf).clamp_min(1e-30) fp8 = (x_view * (1.0 / sf.unsqueeze(1))).to(torch.float8_e4m3fn).view(m, n) return fp8, sf def _pack_kmajor(fp8, ks, mcols): """Per-expert k-major pack, mirroring generators.py::generate_k_grouped_contiguous.""" out = torch.empty(sum(ks) * mcols, dtype=fp8.dtype, device=fp8.device) prefix = 0 for k in ks: out[prefix * mcols:(prefix + k) * mcols] = fp8[prefix:prefix + k].T.flatten() prefix += k return out def _make_operand(x, ks, sf_val=1e-30): fp8, sf = _per_channel_cast_to_fp8(x, 128, sf_val) return _pack_kmajor(fp8, ks, x.shape[1]), sf.T def _build_operands(device, seed=0): """Synthetic grad_y / activation with the minimal trigger structure.""" g = torch.Generator(device=device).manual_seed(seed) m = sum(_KS) grad_y = torch.randn(m, _HIDDEN, device=device, dtype=torch.bfloat16, generator=g) * 1e-7 s = torch.randn(m, _FFN, device=device, dtype=torch.bfloat16, generator=g) * 0.4 offsets, off = [], 0 for k in _KS: offsets.append(off) off += k # NOTE: setup one all-zero channel for 128-token expert for e, k in enumerate(_KS): if k == 128: grad_y[offsets[e]:offsets[e] + k, 0] = 0 return grad_y, s, offsets def test_k_grouped_fp8_gemm_is_deterministic(sf_val): """The grouped FP8 wgrad GEMM must return the same correct result every call.""" device = torch.cuda.current_device() iters = 10000 num_experts = len(_KS) grad_y, s, offsets = _build_operands(device) grouped_layout = torch.tensor(_KS, dtype=torch.int32, device=device) # per-channel cast fp8_grad_y = _make_operand(grad_y, _KS, sf_val) fp8_s = _make_operand(s, _KS, sf_val) torch.cuda.synchronize() # reference wgrad ref = torch.zeros(num_experts, _HIDDEN, _FFN, device=device, dtype=torch.float32) for e, k in enumerate(_KS): if k == 0: continue gy_e = grad_y[offsets[e]:offsets[e] + k].float() s_e = s[offsets[e]:offsets[e] + k].float() ref[e] = gy_e.t() @ s_e ref_norms = torch.linalg.vector_norm(ref.reshape(num_experts, -1), dim=1) thresholds = (100.0 * ref_norms).clamp_min(1.0) # blow-up = >=100x the true grad failures = [] for it in range(iters): out = torch.zeros(num_experts, _HIDDEN, _FFN, device=device, dtype=torch.float32) torch.cuda.synchronize() deep_gemm.k_grouped_fp8_gemm_nt_contiguous( fp8_grad_y, fp8_s, out, _KS, grouped_layout, out, recipe=(1, 1, 128), use_psum_layout=False, ) torch.cuda.synchronize() out_norms = torch.linalg.vector_norm(out.reshape(num_experts, -1), dim=1) bad = torch.nonzero(out_norms > thresholds).flatten().tolist() if bad: failures.append((it, bad, out_norms.max().item())) if failures: first_it, experts, worst = failures[0] print( f"\nk_grouped_fp8_gemm_nt_contiguous is nondeterministic: " f"{len(failures)}/{iters} calls blew up (first at iter {first_it}, " f"experts {experts}); worst output norm {worst:.3e} vs reference max " f"{ref_norms.max().item():.3e}. Operands were cast once and are identical " f"across calls (single stream, fully synced); amax_floor={_AMAX_FLOOR} sf_val={sf_val}." ) if __name__ == "__main__": for sf_val in [1, 1e-4, 1e-8, 1e-30]: test_k_grouped_fp8_gemm_is_deterministic(sf_val) ``` with outputs: ``` k_grouped_fp8_gemm_nt_contiguous is nondeterministic: 1749/10000 calls blew up (first at iter 8, experts [2]); worst output norm 7.451e+03 vs reference max 1.311e-03. Operands were cast once and are identical across calls (single stream, fully synced); amax_floor=0 sf_val=1. k_grouped_fp8_gemm_nt_contiguous is nondeterministic: 181/10000 calls blew up (first at iter 127, experts [2]); worst output norm 1.093e+00 vs reference max 1.311e-03. Operands were cast once and are identical across calls (single stream, fully synced); amax_floor=0 sf_val=0.0001. ```

bsakaguchi commented

2026-10-26 21:14:12 +00:00

Author

Owner

@hiSandog commented on GitHub (Jun 25, 2026):

The repro is nicely narrowed down. For triage, the interesting boundary seems to be the combination of exactly 128 tokens for one expert and an all-zero column in a; that would make a very good regression case because it exercises the grouped kernel path without relying on a large model. It may also be worth noting whether changing only the scale for the zero column alters the failing lane/column consistently, since that could help distinguish stale shared-memory data from a scale/dequantization edge case.

@hiSandog commented on GitHub (Jun 25, 2026): The repro is nicely narrowed down. For triage, the interesting boundary seems to be the combination of exactly 128 tokens for one expert and an all-zero column in `a`; that would make a very good regression case because it exercises the grouped kernel path without relying on a large model. It may also be worth noting whether changing only the scale for the zero column alters the failing lane/column consistently, since that could help distinguish stale shared-memory data from a scale/dequantization edge case.

bsakaguchi commented

2026-10-26 21:14:12 +00:00

Author

Owner

@Functionhx commented on GitHub (Jul 9, 2026):

I traced the root cause to a PTX-level race condition in sm90_fp8_gemm_1d1d.cuh. Have a PR ready if you'd like to take a look.

@Functionhx commented on GitHub (Jul 9, 2026): I traced the root cause to a PTX-level race condition in sm90_fp8_gemm_1d1d.cuh. Have a PR ready if you'd like to take a look.

bsakaguchi commented

2026-10-26 21:14:12 +00:00

Author

Owner

@LiRunGuo commented on GitHub (Sep 22, 2026):

I checked this on an H200 (SM90, 132 SMs, driver 580.159.03, CUDA 13.0, PyTorch 2.11.0+cu130). It reproduces on 2.6.1, and it no longer reproduces after #343 (1f6f3f3), including on current main.

I used the script from this issue with two changes: 2000 iterations per sf_val instead of 10000, and an extra check that each call's output is bitwise equal to the first call's output. A call counts as bad if it fails either check.

Build	Bad calls out of 2000, for `sf_val` = 1, 1e-4, 1e-8, 1e-30
`54e2261` (2.6.1, before #343)	2000, 1999, 1999, 1999
`54e2261` + only the `sm90_fp8_gemm_1d1d.cuh` change from #343	0, 0, 0, 0
current `main` (`78b6900`)	0, 0, 0, 0

On 2.6.1 the output differs between calls in almost every call. With sf_val=1, the first call already blows up (output norm 8.1e3 vs 1.3e-3 for the reference). The C++ extension is identical between 54e2261 and 1f6f3f3, so the only difference between the first two rows is the tensormap drain added in #343.

So this looks fixed on main by #343. #375 adds a further drain on top of #343; I did not test it, but for this repro #343 alone is enough. I only tested on H200, not GH200, so if it still reproduces for you on a recent main, the commit and GPU would help.

@LiRunGuo commented on GitHub (Sep 22, 2026): I checked this on an H200 (SM90, 132 SMs, driver 580.159.03, CUDA 13.0, PyTorch 2.11.0+cu130). It reproduces on 2.6.1, and it no longer reproduces after #343 (1f6f3f3), including on current `main`. I used the script from this issue with two changes: 2000 iterations per `sf_val` instead of 10000, and an extra check that each call's output is bitwise equal to the first call's output. A call counts as bad if it fails either check. | Build | Bad calls out of 2000, for `sf_val` = 1, 1e-4, 1e-8, 1e-30 | |---|---| | 54e2261 (2.6.1, before #343) | 2000, 1999, 1999, 1999 | | 54e2261 + only the `sm90_fp8_gemm_1d1d.cuh` change from #343 | 0, 0, 0, 0 | | current `main` (78b6900) | 0, 0, 0, 0 | On 2.6.1 the output differs between calls in almost every call. With `sf_val=1`, the first call already blows up (output norm 8.1e3 vs 1.3e-3 for the reference). The C++ extension is identical between 54e2261 and 1f6f3f3, so the only difference between the first two rows is the tensormap drain added in #343. So this looks fixed on `main` by #343. #375 adds a further drain on top of #343; I did not test it, but for this repro #343 alone is enough. I only tested on H200, not GH200, so if it still reproduces for you on a recent `main`, the commit and GPU would help.

bsakaguchi referenced this issue

2026-10-26 21:15:42 +00:00

[PR #174] [MERGED] support compile with cuda 13.0 #278