[simveit]

Before: in three corner cases, MMA-Sim returns different bits from the hardware on Hopper (H200), Blackwell (B300, B200) and RTX Blackwell. First, when a nonzero sum rounds to zero, the simulator returns −0, but the GPUs return +0. Second, with a zero f16 accumulator, the simulator anchors the fused sum at 2^-22 (Blackwell and RTX Blackwell mma.sync) or 2^-133 (wgmma, tcgen05.mma), while the GPUs anchor it at 2^-21. Hopper mma.sync already uses 2^-21 since #17. Third, tcgen05.mma kind::f8f6f4 and kind::mxf8f6f4 with e2m1 (FP4) operands crash with Unsupported dtype: torch.uint8, and the simulator does not accept FP6 (e3m2, e2m3) operands at all. Random-input tests do not hit the first two cases.

After: all three cases match the hardware bit for bit, and FP6 and FP4 work in kind::f8f6f4 and kind::mxf8f6f4. New kernels in src/mmasim_kernels let test_nv_blackwell.py test these instructions. Every other result is unchanged.

How:

• fdpa.py: ldexp_and_normalize returns +0 for every NVIDIA rounding mode. The AMD RNE-FP32 path is unchanged.
• sim.py: MMA on Blackwell and RTX Blackwell, WGMMA and TCGen05MMA use e_zero = -21 for an f16 C, like Hopper MMA after Release v2.0.3 #17. f32 keeps -133. Peichen Xie (@peichenxie) measured -21/-133 on RTX Blackwell.
• isa_*.py, helper.py and fdpa.py: the ISA classes keep the PTX type names (a_ptx_type, b_ptx_type), because e3m2, e2m3 and e2m1 are all held one per uint8. MMA_T_FDPA and MMA_ST_FDPA decode FP6 from bits 5:0 and e2m1 from bits 5:2 of that byte (PTX ISA §9.7.15.5.14, §9.7.17.10.4.3), with subnormal exponents -2 (e3m2), 0 (e2m3) and 0 (e2m1). Packed mxf4/nvf4 is unchanged.
• tcgen05mma.h, tcgen05mma_sm100a.cu and blackwell.py: 9 new tcgen05.mma kernels for e3m2, e2m3 and e2m1: kind::f8f6f4 m64n8k32 with f32 and f16 output, and kind::mxf8f6f4 m128n8k32.block32. They take the same one-value-per-byte input as the simulator. New LOAD_*_FP6 and LOAD_*_FP4 macros pack each 16 elements into 12 or 8 bytes plus zero padding (PTX ISA §9.7.17.10.4.4). On B300 the packing is little-endian: FP6 element j sits at bits 6j..6j+5, and FP4 is low nibble first.
• tests/directed/test_nv_directed.py checks these cases against the hardware bits and runs on CPU. It includes one e3m2 case that checks the -2 subnormal exponent and that bits 7:6 are ignored.

Reproducers

Both reproducers use main's own mmasim_kernels wrappers. Each prints the GPU bits next to the simulator bits for a single dot product.

import torch

def vector(dtype, k, values):
    x = torch.zeros(k)
    x[: len(values)] = torch.tensor(values)
    return x.to(dtype)

def bits(d):
    raw = d.view(torch.int16 if d.dtype == torch.float16 else torch.int32).item()
    return f"{raw & (0xFFFF if d.dtype == torch.float16 else 0xFFFFFFFF):#x}"

def run(kernels, sims, arch, cases):
    for name, a, b, c in cases:
        kind, shape_and_type = name.split(".", 1)
        kernel, sim = kernels[kind][shape_and_type], sims[kind](arch, shape_and_type)
        a, b = vector(kernel.a_type, kernel.k, a), vector(kernel.b_type, kernel.k, b)
        c = torch.tensor(c, dtype=kernel.c_type)
        print(f"{name:32} gpu {bits(kernel.dpa(a, b, c).cpu()):>10}  sim {bits(sim.dpa(a, b, c)):>10}")

1. A zero result is +0

A negative sum below the smallest output step rounds to zero.

from mmasim.nv_ptx import MMA, WGMMA, TCGen05MMA
from mmasim_kernels.nv_ptx import hopper, blackwell

f16 = [(1 + 2.0**-10) * 2.0**-5], [(1 - 2.0**-10) * 2.0**-5], -(2.0**-10)
run({"mma": hopper.mma_kernels, "wgmma": hopper.wgmma_kernels}, {"mma": MMA, "wgmma": WGMMA}, "Hopper", [
    ("wgmma.m64n8k16.f32.bf16.bf16", [-(2.0**-75)], [2.0**-75], 0.0),
    ("mma.m16n8k16.f32.bf16.bf16.f32", [-(2.0**-75)], [2.0**-75], 0.0),
    ("wgmma.m64n8k16.f16.f16.f16", *f16),
    ("mma.m16n8k16.f16.f16.f16.f16", *f16),
])
run({"mma": blackwell.mma_kernels, "tcgen05": blackwell.tcgen05mma_kernels}, {"mma": MMA, "tcgen05": TCGen05MMA}, "Blackwell", [
    ("tcgen05.m64n8k16.f16.f16.f16", *f16),
    ("mma.m16n8k16.f16.f16.f16.f16", *f16),
    ("tcgen05.m64n8k16.f32.bf16.bf16", [(1 + 2.0**-7) * 2.0**-75], [(1 + 2.0**-7) * 2.0**-75], -(2.0**-149)),
    ("mma.m16n8k16.f32.bf16.bf16.f32", [(1 + 2.0**-7) * 2.0**-75], [(1 + 2.0**-7) * 2.0**-75], -(2.0**-149)),
])

instruction	GPU	main	this PR
H200 wgmma / mma bf16 → f32	0x0	0x80000000	0x0
H200 wgmma / mma f16 → f16	0x0	0x8000	0x0
B300, B200 tcgen05 / mma f16 → f16	0x0	0x8000	0x0
B300, B200 tcgen05 / mma bf16 → f32	0x0	0x80000000	0x0

2. A zero f16 accumulator anchors at 2^-21

With C = 0, the f16 output rounds at a tie. A term at 2^-47 breaks the tie only if the anchor is 2^-22 or lower.

run({"wgmma": hopper.wgmma_kernels}, {"wgmma": WGMMA}, "Hopper", [
    ("wgmma.m64n8k16.f16.f16.f16", [2.0**-12, 2.0**-23], [2.0**-13, 2.0**-24], 0.0),
])
tie = [1.5 * 2.0**-12, 2.0**-13, -(2.0**-24)], [2.0**-11, 2.0**-12, 2.0**-23], 0.0
run({"mma": blackwell.mma_kernels, "tcgen05": blackwell.tcgen05mma_kernels}, {"mma": MMA, "tcgen05": TCGen05MMA}, "Blackwell", [
    ("tcgen05.m64n8k16.f16.f16.f16", *tie),
    ("mma.m16n8k16.f16.f16.f16.f16", *tie),
    ("mma.m16n8k8.f16.f16.f16.f16", *tie),
])

instruction	GPU	main	this PR
H200 wgmma m64n8k16	0x0	0x1	0x0
B300, B200 tcgen05 m64n8k16, mma m16n8k16 / m16n8k8	0x4	0x3	0x4

3. e2m1 operands in kind::f8f6f4

e2m1 code 1 is the subnormal 0.5. The product 0.5 · 0.5 lands on C = -(0.25 - 2^-26), so only the lowest bit of C survives. If the subnormal were normalized instead, the result would be 0x32800000. Main's blackwell.py has no e2m1 kernel, so this script builds one from main's tcgen05mma.h: it reuses the e4m3 m64n8k32 kernel with format code 5.

repro_d3.py

# B300 (sm_103a) or B200 (sm_100a): python repro_d3.py
# Main's blackwell.py has no kind::f8f6f4 e2m1 kernel, so this builds one from main's
# tcgen05mma.h (the e4m3 m64n8k32 kernel with A/B format 5 = e2m1) and wraps it in
# main's TCGen05MMA kernel class.
import ctypes
from pathlib import Path

import torch
from torch.utils.cpp_extension import load

import mmasim_kernels.nv_ptx as nv_kernels
from mmasim.nv_ptx import TCGen05MMA
from mmasim_kernels.nv_ptx import TCGen05MMA as TCGen05MMAKernel

SRC = r"""
#include "tcgen05mma.h"
extern "C"
{
    __global__ void tcgen05mma_m64n8k32_f32_e2m1_e2m1_kernel(uint32_t *d, uint8_t *a, uint8_t *b)
    {
        const uint32_t M = 64, N = 8, K = 32;
        uint32_t tid = threadIdx.x, warpid = tid / 32, laneid = tid % 32;
        uint32_t i_desc = (N >> 3 << 17) | (M >> 4 << 24), mma_barrier_phase_bit = 0;
        uint64_t a_desc, b_desc;
        __shared__ uint8_t a_smem[M * K], b_smem[N * K];
        __shared__ uint32_t d_tmem_addr;
        __shared__ uint64_t mma_barrier;

        INIT_MBARRIER();
        LOAD_A_M64K32_FP8();
        LOAD_B_N8K32_FP8();
        ALLOC_TMEM(d_tmem_addr);
        __syncthreads();
        LOAD_D_M64N8();
        __syncthreads();
        i_desc |= (1 << 4) | (5 << 7) | (5 << 10); // f32_e2m1_e2m1
        MMA("f8f6f4");
        STORE_D_M64N8();
        DEALLOC_TMEM(d_tmem_addr);
        RELINQUISH_TMEM();
    }

    void tcgen05mma_m64n8k32_f32_e2m1_e2m1(uint32_t *d, uint8_t *a, uint8_t *b)
    {
        tcgen05mma_m64n8k32_f32_e2m1_e2m1_kernel<<<1, 128>>>(d, a, b);
    }
}
"""

major, minor = torch.cuda.get_device_capability()
assert major == 10, "needs a datacenter Blackwell GPU (B200 or B300)"
src = Path(__file__).with_name("repro_d3_kernel.cu")
src.write_text(SRC)
lib = ctypes.CDLL(
    load(
        "repro_d3_kernel",
        [str(src)],
        extra_cuda_cflags=[f"-arch=compute_{major}{minor}a", f"-code=sm_{major}{minor}a"],
        extra_include_paths=[str(Path(nv_kernels.__file__).with_name("kernels"))],
        is_python_module=False,
    )
)
lib.tcgen05mma_m64n8k32_f32_e2m1_e2m1.argtypes = [ctypes.c_void_p] * 3
kernels = {"m64n8k32.f32.e2m1.e2m1": TCGen05MMAKernel(
    "Blackwell", "m64n8k32.f32.e2m1.e2m1", lib.tcgen05mma_m64n8k32_f32_e2m1_e2m1)}


def smem_e2m1(codes, k):  # SMEM .b4x16_p64: 16 codes in 8 bytes (low nibble first), 8 zero bytes
    x = torch.zeros(k, dtype=torch.uint8)
    for i, code in enumerate(codes):
        x[i // 16 * 16 + i % 16 // 2] |= code << (4 * (i % 2))
    return x


def container_e2m1(codes, k):  # simulator: one code per uint8, in bits 5:2
    x = torch.zeros(k, dtype=torch.uint8)
    x[: len(codes)] = torch.tensor(codes, dtype=torch.uint8) << 2
    return x


def bits(d):
    return f"{d.view(torch.int32).item() & 0xFFFFFFFF:#x}"


def run(kernels, sims, arch, cases):
    for name, a, b, c in cases:
        kind, shape_and_type = name.split(".", 1)
        kernel, sim = kernels[shape_and_type], sims[kind](arch, shape_and_type)
        c = torch.tensor(c, dtype=kernel.c_type)
        gpu = bits(kernel.dpa(smem_e2m1(a, kernel.k), smem_e2m1(b, kernel.k), c).cpu())
        try:
            s = bits(sim.dpa(container_e2m1(a, sim.k), container_e2m1(b, sim.k), c))
        except AssertionError as e:
            s = f"error: {e}"
        print(f"{name:30} gpu {gpu:>10}  sim {s}")


# e2m1 code 1 = 0.5 (subnormal); a * b = 0.25 = 0.01b * 2^0, so the last bit of C is lost
run(kernels, {"tcgen05": TCGen05MMA}, "Blackwell", [
    ("tcgen05.m64n8k32.f32.e2m1.e2m1", [1], [1], -(0.25 - 2.0**-26)),
])

GPU	GPU result	main	this PR
B300 (sm_103a)	0x33000000	Unsupported dtype: torch.uint8	0x33000000
B200 (sm_100a)	0x33000000	Unsupported dtype: torch.uint8	0x33000000

Hardware

• H200 (sm_90a): driver 580.105.08, run with main's hopper.py unchanged.
• B300 (sm_103a): driver 610.43.02, CUDA 13.0, torch 2.14.1+cu130. Main's blackwell.py was compiled with -code=sm_103a.
• B200 (sm_100a): same driver and toolchain as B300, run with main's blackwell.py unchanged.

Each GPU ran 3 fresh processes per simulator. The outputs were bit-identical across runs. tests/equivalence/test_nv_hopper.py passes with this PR on H200, and test_nv_blackwell.py passes on B300 and B200, both before and after the e2m1 commit.

The +0 rule also covers Volta through Ada and RTX Blackwell, which were not measured. A fixed-point sum carries no sign, so I expect the same behaviour there.

The final commit, which adds FP6, the new kernels and the RTX Blackwell change, was tested on B300 only. test_nv_blackwell.py now covers 31 instructions, including the 9 new ones, and all of them match bit for bit.

This PR was prepared with AI assistance (Claude Code). I reviewed the changes and the hardware results.

🤖 Generated with Claude Code

[simveit]

@microsoft-github-policy-service agree

[Peichen Xie (peichenxie)]

Thank you for your PR! I can confirm Bullets 1 and 2. I also tested e_zero on RTX Blackwell, and the results were -21/-133 as well.

For Bullet 3, could you add support for FP6 and include the corresponding instructions in src/mmasim_kernels? That way, the existing test scripts can test these instructions.

[simveit]

For Bullet 3, could you add support for FP6 and include the corresponding instructions in src/mmasim_kernels? That way, the existing test scripts can test these instructions.

included the new instructions :) tried to align with existing conventions, please lmk if there is something more i can change