From 4e19e2297268af61747f8cb373785a96ecb0d701 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Sun, 11 Oct 2026 06:49:03 +0000 Subject: [PATCH 1/4] Return +0 when an NVIDIA MMA result rounds to zero A nonzero negative sum below the smallest output step came back as -0 (0x8000 for f16, 0x80000000 for f32). H200 and B300 return +0. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Simon Veitner --- src/mmasim/mmasim/arithmetic/fdpa.py | 2 ++ tests/directed/test_nv_directed.py | 48 ++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+) create mode 100644 tests/directed/test_nv_directed.py diff --git a/src/mmasim/mmasim/arithmetic/fdpa.py b/src/mmasim/mmasim/arithmetic/fdpa.py index d6eda15..7e34572 100644 --- a/src/mmasim/mmasim/arithmetic/fdpa.py +++ b/src/mmasim/mmasim/arithmetic/fdpa.py @@ -58,6 +58,8 @@ def ldexp_and_normalize(s: DoubleTensor, e: IntTensor, rho: str) -> torch.Tensor else: # "RZ-FP32" s = torch.trunc(s * 2.0**23) * 2.0**-23 # RZ res = (s * pow2(e)).to(torch.float32) + if rho != "RNE-FP32": # NVIDIA MMA units return +0 for a zero result + res[res == 0.0] = 0.0 nans = res.isnan() if res.dtype == torch.float16: res = res.view(torch.int16) diff --git a/tests/directed/test_nv_directed.py b/tests/directed/test_nv_directed.py new file mode 100644 index 0000000..fe8e7de --- /dev/null +++ b/tests/directed/test_nv_directed.py @@ -0,0 +1,48 @@ +import torch + +from mmasim.nv_ptx import MMA, WGMMA, TCGen05MMA + +# Expected D bits are measured on H200 (Hopper) and B300 (Blackwell) + + +def vector(dtype: torch.dtype, k: int, entries: dict[int, float]) -> torch.Tensor: + x = torch.zeros(k, dtype=torch.float32) + for i, v in entries.items(): + x[i] = v + return x.to(dtype) + + +def raw(x: torch.Tensor) -> int: + if x.dtype == torch.float16: + return int(x.view(torch.int16).item()) & 0xFFFF + return int(x.view(torch.int32).item()) & 0xFFFFFFFF + + +def directed_test(sim, a: dict, b: dict, c: float, expected: int): + a_vec = vector(sim.a_type, sim.k, a) + b_vec = vector(sim.b_type, sim.k, b) + d = sim.dpa(a_vec, b_vec, torch.tensor(c, dtype=sim.c_type)) + assert raw(d) == expected, f"D_sim = {raw(d):#x}, D_gpu = {expected:#x}" + + +def test_zero_result_is_positive(): + # a nonzero negative sum that rounds to zero gives +0, not -0 + tiny = {0: -(2.0**-75)}, {0: 2.0**-75}, 0.0 + f16_tiny = {0: (1 + 2.0**-10) * 2.0**-5}, {0: (1 - 2.0**-10) * 2.0**-5}, -(2.0**-10) + bf16_tiny = ( + {0: (1 + 2.0**-7) * 2.0**-75}, + {0: (1 + 2.0**-7) * 2.0**-75}, + -(2.0**-149), + ) + for sim, (a, b, c) in [ + (WGMMA("Hopper", "m64n8k16.f32.bf16.bf16"), tiny), + (MMA("Hopper", "m16n8k16.f32.bf16.bf16.f32"), tiny), + (WGMMA("Hopper", "m64n8k16.f16.f16.f16"), f16_tiny), + (MMA("Hopper", "m16n8k16.f16.f16.f16.f16"), f16_tiny), + (TCGen05MMA("Blackwell", "m64n8k16.f16.f16.f16"), f16_tiny), + (MMA("Blackwell", "m16n8k16.f16.f16.f16.f16"), f16_tiny), + (TCGen05MMA("Blackwell", "m64n8k16.f32.bf16.bf16"), bf16_tiny), + (MMA("Blackwell", "m16n8k16.f32.bf16.bf16.f32"), bf16_tiny), + ]: + print(f"Testing {sim.arch} {sim.shape_and_type}") + directed_test(sim, a, b, c, 0) From 4be0da7897207f6c2a85e2622a3c26fff7ec4e72 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Sun, 11 Oct 2026 06:49:03 +0000 Subject: [PATCH 2/4] Use e_zero = -21 for f16 accumulators on Blackwell and in wgmma A zero f16 C anchors the fused sum at 2^-21 on B300 and B200 for mma.sync and tcgen05.mma, and on H200 for wgmma. The simulator used -22 for Blackwell mma.sync and -133 for wgmma and tcgen05.mma. Hopper mma.sync already uses -21. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Simon Veitner --- src/mmasim/mmasim/nv_ptx/sim.py | 10 ++++++---- tests/directed/test_nv_directed.py | 32 +++++++++++++++++++++++++++++- 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/src/mmasim/mmasim/nv_ptx/sim.py b/src/mmasim/mmasim/nv_ptx/sim.py index 1902e49..3ca9dbc 100644 --- a/src/mmasim/mmasim/nv_ptx/sim.py +++ b/src/mmasim/mmasim/nv_ptx/sim.py @@ -42,9 +42,9 @@ def __init__(self, arch: str, shape_and_type: str): e_zero = -20 if self.c_type == torch.float16 else -131 elif self.arch in ["Turing", "Ampere", "Ada Lovelace"]: e_zero = -21 if self.c_type == torch.float16 else -132 - elif self.arch == "Hopper": + elif self.arch in ["Hopper", "Blackwell"]: e_zero = -21 if self.c_type == torch.float16 else -133 - else: # Blackwell, RTX Blackwell + else: # RTX Blackwell e_zero = -22 if self.c_type == torch.float16 else -133 self.arithmetic_op = fdpa.MMA_T_FDPA(F, rho, L_max, e_zero) @@ -106,7 +106,8 @@ def __init__(self, arch: str, shape_and_type: str): F = 25 rho = "RNE-FP16" if self.d_type == torch.float16 else "RZ-FP32" L_max = 32 // self.a_type.itemsize - self.arithmetic_op = fdpa.MMA_T_FDPA(F, rho, L_max, e_zero=-133) + e_zero = -21 if self.c_type == torch.float16 else -133 + self.arithmetic_op = fdpa.MMA_T_FDPA(F, rho, L_max, e_zero) def dpa(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: return self.arithmetic_op.dpa(a, b, c) @@ -126,7 +127,8 @@ def __init__(self, arch: str, shape_and_type: str): F = 25 rho = "RNE-FP16" if self.d_type == torch.float16 else "RZ-FP32" L_max = 32 // self.a_type.itemsize - self.arithmetic_op = fdpa.MMA_T_FDPA(F, rho, L_max, e_zero=-133) + e_zero = -21 if self.c_type == torch.float16 else -133 + self.arithmetic_op = fdpa.MMA_T_FDPA(F, rho, L_max, e_zero) def dpa(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: return self.arithmetic_op.dpa(a, b, c) diff --git a/tests/directed/test_nv_directed.py b/tests/directed/test_nv_directed.py index fe8e7de..86d3015 100644 --- a/tests/directed/test_nv_directed.py +++ b/tests/directed/test_nv_directed.py @@ -2,7 +2,7 @@ from mmasim.nv_ptx import MMA, WGMMA, TCGen05MMA -# Expected D bits are measured on H200 (Hopper) and B300 (Blackwell) +# Expected D bits are measured on H200 (Hopper), B300 and B200 (Blackwell) def vector(dtype: torch.dtype, k: int, entries: dict[int, float]) -> torch.Tensor: @@ -46,3 +46,33 @@ def test_zero_result_is_positive(): ]: print(f"Testing {sim.arch} {sim.shape_and_type}") directed_test(sim, a, b, c, 0) + print("Tests passed!") + + +def test_e_zero_f16_accumulator(): + # a zero f16 C anchors the fused sum at 2^-21 + hopper = [ + WGMMA("Hopper", "m64n8k16.f16.f16.f16"), + ] + for sim in hopper: + print(f"Testing {sim.arch} {sim.shape_and_type}") + # 2^-25 is a tie at the f16 subnormal step, 2^-q breaks it iff q <= 46 + for q, expected in [(46, 0x0001), (47, 0x0000)]: + a, b = {0: 2.0**-12, 1: 2.0**-23}, {0: 2.0**-13, 1: 2.0 ** (23 - q)} + directed_test(sim, a, b, 0.0, expected) + blackwell = [ + TCGen05MMA("Blackwell", "m64n8k16.f16.f16.f16"), + MMA("Blackwell", "m16n8k16.f16.f16.f16.f16"), + ] + for sim in blackwell: + print(f"Testing {sim.arch} {sim.shape_and_type}") + # 1.5 * 2^-23 + 2^-25 is a tie, -2^-47 is below the 2^-46 grid + a = {0: 1.5 * 2.0**-12, 1: 2.0**-13, 2: -(2.0**-24)} + b = {0: 2.0**-11, 1: 2.0**-12, 2: 2.0**-23} + directed_test(sim, a, b, 0.0, 0x0004) + print("Tests passed!") + + +if __name__ == "__main__": + test_zero_result_is_positive() + test_e_zero_f16_accumulator() From 866bdda23edf3f81e9950906d6a0017faddd95bf Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Sun, 11 Oct 2026 07:39:02 +0000 Subject: [PATCH 3/4] Decode unpacked e2m1 operands in kind::f8f6f4 and mxf8f6f4 tcgen05.mma kind::f8f6f4 and kind::mxf8f6f4 accept e2m1, but the simulator raised "Unsupported dtype: torch.uint8". Decode one e2m1 per uint8 from bits 5:2 of the 8-bit container and keep its subnormal at exponent 0. B300 and B200 return 0x33000000 for 0.5 * 0.5 with C = -(0.25 - 2^-26), which this now matches. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Simon Veitner --- src/mmasim/mmasim/arithmetic/fdpa.py | 34 ++++++++++++++++++++++---- src/mmasim/mmasim/arithmetic/helper.py | 6 +++++ tests/directed/test_nv_directed.py | 12 +++++++++ 3 files changed, 47 insertions(+), 5 deletions(-) diff --git a/src/mmasim/mmasim/arithmetic/fdpa.py b/src/mmasim/mmasim/arithmetic/fdpa.py index 7e34572..16830e9 100644 --- a/src/mmasim/mmasim/arithmetic/fdpa.py +++ b/src/mmasim/mmasim/arithmetic/fdpa.py @@ -6,6 +6,7 @@ truncate_fp32_to_tf32, truncate_e4m3_to_ue4m3, unpack_uint8_to_fp4, + decode_uint8_to_fp4, ) DoubleTensor = Annotated[torch.Tensor, "float64"] @@ -89,9 +90,11 @@ def t_fdpa( F: int, rho: str, e_zero: int, + e_subnormal_a: int | None = None, + e_subnormal_b: int | None = None, ) -> torch.Tensor: # [...] - s_a, e_a = frexp_and_normalize(a) - s_b, e_b = frexp_and_normalize(b) + s_a, e_a = frexp_and_normalize(a, e_subnormal_a) + s_b, e_b = frexp_and_normalize(b, e_subnormal_b) s_c, e_c = frexp_and_normalize(c) s = torch.cat([s_a * s_b, s_c.unsqueeze(-1)], dim=-1) e = torch.cat([e_a + e_b, e_c.unsqueeze(-1)], dim=-1) @@ -111,11 +114,23 @@ def dpa(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor if a.dtype == torch.float32: # tf32 a = truncate_fp32_to_tf32(a) b = truncate_fp32_to_tf32(b) + e_subnormal_a = e_subnormal_b = None + if a.dtype == torch.uint8: # e2m1, whose subnormal exponent is 0 + a, e_subnormal_a = decode_uint8_to_fp4(a), 0 + if b.dtype == torch.uint8: + b, e_subnormal_b = decode_uint8_to_fp4(b), 0 K = a.shape[-1] L = min(K, self.L_max) for i in range(0, K, L): c = t_fdpa( - a[..., i : i + L], b[..., i : i + L], c, self.F, self.rho, self.e_zero + a[..., i : i + L], + b[..., i : i + L], + c, + self.F, + self.rho, + self.e_zero, + e_subnormal_a, + e_subnormal_b, ) return c @@ -137,9 +152,11 @@ def st_fdpa( F: int, rho: str, e_zero: int, + e_subnormal_a: int | None = None, + e_subnormal_b: int | None = None, ) -> torch.Tensor: # [...] - s_a, e_a = frexp_and_normalize(a) - s_b, e_b = frexp_and_normalize(b) + s_a, e_a = frexp_and_normalize(a, e_subnormal_a) + s_b, e_b = frexp_and_normalize(b, e_subnormal_b) s_c, e_c = frexp_and_normalize(c) s_alpha, e_alpha = frexp_and_normalize(alpha) s_beta, e_beta = frexp_and_normalize(beta) @@ -166,6 +183,11 @@ def dpa( alpha: torch.Tensor, beta: torch.Tensor, ) -> torch.Tensor: + e_subnormal_a = e_subnormal_b = None + if a.dtype == torch.uint8: # e2m1, whose subnormal exponent is 0 + a, e_subnormal_a = decode_uint8_to_fp4(a), 0 + if b.dtype == torch.uint8: + b, e_subnormal_b = decode_uint8_to_fp4(b), 0 K = a.shape[-1] L = min(K, self.L_max) for i in range(0, K, L): @@ -178,6 +200,8 @@ def dpa( self.F, self.rho, self.e_zero, + e_subnormal_a, + e_subnormal_b, ) return c diff --git a/src/mmasim/mmasim/arithmetic/helper.py b/src/mmasim/mmasim/arithmetic/helper.py index ec3ee37..fe9d6c7 100644 --- a/src/mmasim/mmasim/arithmetic/helper.py +++ b/src/mmasim/mmasim/arithmetic/helper.py @@ -55,3 +55,9 @@ def unpack_uint8_to_fp4(x: torch.Tensor) -> torch.Tensor: low = x & 0x0F unpacked = torch.stack([high, low], dim=-1).view(x.shape[:-1] + (-1,)) return fp4_value_table.to(x.device)[unpacked.int()] + + +def decode_uint8_to_fp4(x: torch.Tensor) -> torch.Tensor: + assert x.dtype == torch.uint8 + x = (x >> 2) & 0x0F # e2m1 sits in bits 5:2 of the 8-bit container + return fp4_value_table.to(x.device)[x.int()] diff --git a/tests/directed/test_nv_directed.py b/tests/directed/test_nv_directed.py index 86d3015..d9a4008 100644 --- a/tests/directed/test_nv_directed.py +++ b/tests/directed/test_nv_directed.py @@ -73,6 +73,18 @@ def test_e_zero_f16_accumulator(): print("Tests passed!") +def test_e2m1_f8f6f4(): + # a * b = 0.25 is 0.01b * 2^0, so C = -(0.25 - 2^-26) loses its last bit + sim = TCGen05MMA("Blackwell", "m128n8k32.f32.e2m1.e2m1") + print(f"Testing {sim.arch} {sim.shape_and_type}") + a = torch.zeros(sim.k, dtype=torch.uint8) + a[0] = 0x1 << 2 # e2m1 sits in bits 5:2 of the 8-bit container + d = sim.dpa(a, a, torch.tensor(-(0.25 - 2.0**-26), dtype=torch.float32)) + assert raw(d) == 0x33000000, f"D_sim = {raw(d):#x}, D_gpu = 0x33000000" + print("Tests passed!") + + if __name__ == "__main__": test_zero_result_is_positive() test_e_zero_f16_accumulator() + test_e2m1_f8f6f4() From 1f881509edb2da22ffb5dd103e9cb6d59a320293 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Wed, 14 Oct 2026 16:17:33 +0000 Subject: [PATCH 4/4] Support FP6 in tcgen05.mma f8f6f4 and use e_zero = -21 on RTX Blackwell RTX Blackwell anchors a zero f16 C at 2^-21 like Hopper and Blackwell, so mma.sync now uses e_zero = -21 (f16) / -133 (f32) on all three. tcgen05.mma kind::f8f6f4 and kind::mxf8f6f4 now accept e3m2 and e2m3, held one per uint8 in bits 5:0 (e2m1 stays in bits 5:2). The ISA classes keep the PTX type names, and the FDPA ops decode uint8 operands with the matching subnormal exponent (e3m2 -2, e2m3 0, e2m1 0). New tcgen05 kernels cover m64n8k32 f32/f16 and m128n8k32.block32 for e3m2, e2m3 and e2m1, so test_nv_blackwell.py tests them. The load macros pack 16 elements into 12 (FP6) or 8 (FP4) bytes plus zero padding, as PTX ISA 9.7.17.10.4.4 requires. All 31 Blackwell instructions match B300 bit for bit. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Simon Veitner --- README.md | 2 +- src/mmasim/mmasim/arithmetic/fdpa.py | 61 ++-- src/mmasim/mmasim/arithmetic/helper.py | 13 +- src/mmasim/mmasim/nv_ptx/isa_common.py | 2 + src/mmasim/mmasim/nv_ptx/isa_mma.py | 4 + src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py | 12 +- src/mmasim/mmasim/nv_ptx/sim.py | 26 +- .../mmasim_kernels/nv_ptx/blackwell.py | 24 ++ .../nv_ptx/kernels/tcgen05mma.h | 153 +++++++++ .../nv_ptx/kernels/tcgen05mma_sm100a.cu | 306 ++++++++++++++++++ tests/directed/test_nv_directed.py | 44 +-- 11 files changed, 594 insertions(+), 53 deletions(-) diff --git a/README.md b/README.md index 063d72b..522b7d9 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ Supported GPU architectures: `Volta`, `Turing`, `Ampere`, `Ada Lovelace`, `Hoppe Supported MMA instructions: `mma.sync`, `wgmma.mma_async`, `tcgen05.mma`, and `v_mfma`. -Supported data types: FP64, FP32, TF32, FP16, BF16, FP8, FP4, MXFP8, MXFP4, and NVFP4. +Supported data types: FP64, FP32, TF32, FP16, BF16, FP8, FP6, FP4, MXFP8, MXFP6, MXFP4, and NVFP4. ## How to verify the equivalence between MMA-Sim and GPU diff --git a/src/mmasim/mmasim/arithmetic/fdpa.py b/src/mmasim/mmasim/arithmetic/fdpa.py index 16830e9..d4d2262 100644 --- a/src/mmasim/mmasim/arithmetic/fdpa.py +++ b/src/mmasim/mmasim/arithmetic/fdpa.py @@ -3,10 +3,11 @@ from .. import MMAOperation, MMABlockScaleOperation from .helper import ( dtype_subnormal_exponent, + uint8_subnormal_exponent, truncate_fp32_to_tf32, truncate_e4m3_to_ue4m3, unpack_uint8_to_fp4, - decode_uint8_to_fp4, + decode_uint8_to_fp6_fp4, ) DoubleTensor = Annotated[torch.Tensor, "float64"] @@ -49,7 +50,7 @@ def ldexp_and_normalize(s: DoubleTensor, e: IntTensor, rho: str) -> torch.Tensor # as PyTorch-CPU computes FP64 -> FP32 -> FP16 internally s, e = frexp_and_normalize(s * pow2(e), e_subnormal=-14) s = torch.round(s * 2.0**10) * 2.0**-10 # RNE - res = (s * pow2(e)).to(torch.float16) + res = (s * pow2(e) + 0.0).to(torch.float16) # -0.0 + 0.0 = +0.0 elif rho == "RNE-FP32": res = (s * pow2(e)).to(torch.float32) else: # RZ @@ -58,9 +59,7 @@ def ldexp_and_normalize(s: DoubleTensor, e: IntTensor, rho: str) -> torch.Tensor s = torch.trunc(s * 2.0**13) * 2.0**-13 # RZ else: # "RZ-FP32" s = torch.trunc(s * 2.0**23) * 2.0**-23 # RZ - res = (s * pow2(e)).to(torch.float32) - if rho != "RNE-FP32": # NVIDIA MMA units return +0 for a zero result - res[res == 0.0] = 0.0 + res = (s * pow2(e) + 0.0).to(torch.float32) # -0.0 + 0.0 = +0.0 nans = res.isnan() if res.dtype == torch.float16: res = res.view(torch.int16) @@ -104,21 +103,33 @@ def t_fdpa( class MMA_T_FDPA(MMAOperation): - def __init__(self, F: int, rho: str, L_max: int, e_zero: int): + def __init__( + self, + F: int, + rho: str, + L_max: int, + e_zero: int, + a_ptx_type: str | None = None, + b_ptx_type: str | None = None, + ): self.F = F self.rho = rho self.L_max = L_max self.e_zero = e_zero + # fp6 and fp4 operands are held one per uint8, so keep their ptx types + self.a_ptx_type = a_ptx_type + self.b_ptx_type = b_ptx_type + self.e_subnormal_a = uint8_subnormal_exponent.get(a_ptx_type) + self.e_subnormal_b = uint8_subnormal_exponent.get(b_ptx_type) def dpa(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: if a.dtype == torch.float32: # tf32 a = truncate_fp32_to_tf32(a) b = truncate_fp32_to_tf32(b) - e_subnormal_a = e_subnormal_b = None - if a.dtype == torch.uint8: # e2m1, whose subnormal exponent is 0 - a, e_subnormal_a = decode_uint8_to_fp4(a), 0 + if a.dtype == torch.uint8: # fp6 or fp4 + a = decode_uint8_to_fp6_fp4(a, self.a_ptx_type) if b.dtype == torch.uint8: - b, e_subnormal_b = decode_uint8_to_fp4(b), 0 + b = decode_uint8_to_fp6_fp4(b, self.b_ptx_type) K = a.shape[-1] L = min(K, self.L_max) for i in range(0, K, L): @@ -129,8 +140,8 @@ def dpa(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor self.F, self.rho, self.e_zero, - e_subnormal_a, - e_subnormal_b, + self.e_subnormal_a, + self.e_subnormal_b, ) return c @@ -169,11 +180,24 @@ def st_fdpa( class MMA_ST_FDPA(MMABlockScaleOperation): - def __init__(self, F: int, rho: str, L_max: int, e_zero: int): + def __init__( + self, + F: int, + rho: str, + L_max: int, + e_zero: int, + a_ptx_type: str | None = None, + b_ptx_type: str | None = None, + ): self.F = F self.rho = rho self.L_max = L_max self.e_zero = e_zero + # fp6 and fp4 operands are held one per uint8, so keep their ptx types + self.a_ptx_type = a_ptx_type + self.b_ptx_type = b_ptx_type + self.e_subnormal_a = uint8_subnormal_exponent.get(a_ptx_type) + self.e_subnormal_b = uint8_subnormal_exponent.get(b_ptx_type) def dpa( self, @@ -183,11 +207,10 @@ def dpa( alpha: torch.Tensor, beta: torch.Tensor, ) -> torch.Tensor: - e_subnormal_a = e_subnormal_b = None - if a.dtype == torch.uint8: # e2m1, whose subnormal exponent is 0 - a, e_subnormal_a = decode_uint8_to_fp4(a), 0 + if a.dtype == torch.uint8: # fp6 or fp4 + a = decode_uint8_to_fp6_fp4(a, self.a_ptx_type) if b.dtype == torch.uint8: - b, e_subnormal_b = decode_uint8_to_fp4(b), 0 + b = decode_uint8_to_fp6_fp4(b, self.b_ptx_type) K = a.shape[-1] L = min(K, self.L_max) for i in range(0, K, L): @@ -200,8 +223,8 @@ def dpa( self.F, self.rho, self.e_zero, - e_subnormal_a, - e_subnormal_b, + self.e_subnormal_a, + self.e_subnormal_b, ) return c diff --git a/src/mmasim/mmasim/arithmetic/helper.py b/src/mmasim/mmasim/arithmetic/helper.py index fe9d6c7..2dbc5ec 100644 --- a/src/mmasim/mmasim/arithmetic/helper.py +++ b/src/mmasim/mmasim/arithmetic/helper.py @@ -11,6 +11,8 @@ torch.float8_e5m2fnuz: -15, torch.float8_e4m3fnuz: -7, } +uint8_subnormal_exponent = {"e3m2": -2, "e2m3": 0, "e2m1": 0} +uint8_type_bits = {"e3m2": (3, 2), "e2m3": (2, 3), "e2m1": (2, 1)} # (E, M) fp4_value_table = torch.tensor( @@ -57,7 +59,12 @@ def unpack_uint8_to_fp4(x: torch.Tensor) -> torch.Tensor: return fp4_value_table.to(x.device)[unpacked.int()] -def decode_uint8_to_fp4(x: torch.Tensor) -> torch.Tensor: +def decode_uint8_to_fp6_fp4(x: torch.Tensor, x_type: str) -> torch.Tensor: assert x.dtype == torch.uint8 - x = (x >> 2) & 0x0F # e2m1 sits in bits 5:2 of the 8-bit container - return fp4_value_table.to(x.device)[x.int()] + E, M = uint8_type_bits[x_type] + x = (x.int() & 0x3F) >> (5 - E - M) # fp6 sits in bits 5:0, fp4 in bits 5:2 + bias = 2 ** (E - 1) - 1 + e = (x >> M) & (2**E - 1) + s = (x & (2**M - 1)) + (e > 0) * 2**M + res = (s * torch.exp2(e.clamp(min=1) - bias - M)).to(torch.float16) + return torch.where(x >> (E + M) == 1, -res, res) diff --git a/src/mmasim/mmasim/nv_ptx/isa_common.py b/src/mmasim/mmasim/nv_ptx/isa_common.py index c3b2f6a..80d2864 100644 --- a/src/mmasim/mmasim/nv_ptx/isa_common.py +++ b/src/mmasim/mmasim/nv_ptx/isa_common.py @@ -10,6 +10,8 @@ "e5m2": torch.float8_e5m2, "ue8m0": torch.float8_e8m0fnu, "ue4m3": torch.float8_e4m3fn, + "e3m2": torch.uint8, # torch has no fp6 dtype + "e2m3": torch.uint8, # torch has no fp6 dtype "e2m1": torch.uint8, # torch.float4_e2m1fn_x2 is not well-implemented } diff --git a/src/mmasim/mmasim/nv_ptx/isa_mma.py b/src/mmasim/mmasim/nv_ptx/isa_mma.py index fb17468..c57a8b1 100644 --- a/src/mmasim/mmasim/nv_ptx/isa_mma.py +++ b/src/mmasim/mmasim/nv_ptx/isa_mma.py @@ -136,6 +136,8 @@ def __init__(self, arch: str, shape_and_type: str): self.m, self.n, self.k = nv_shape_to_mnk(shape) self.a_type = nv_torch_dtype[a_type] self.b_type = nv_torch_dtype[b_type] + self.a_ptx_type = a_type + self.b_ptx_type = b_type self.c_type = nv_torch_dtype[c_type] self.d_type = nv_torch_dtype[d_type] @@ -161,6 +163,8 @@ def __init__(self, arch: str, shape_and_type: str): self.packing = 2 if self.k == 64 else 1 self.a_type = nv_torch_dtype[a_type] self.b_type = nv_torch_dtype[b_type] + self.a_ptx_type = a_type + self.b_ptx_type = b_type self.c_type = nv_torch_dtype[c_type] self.d_type = nv_torch_dtype[d_type] self.s_type = nv_torch_dtype[s_type] diff --git a/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py b/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py index bbe97f8..5200375 100644 --- a/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py +++ b/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py @@ -28,14 +28,16 @@ def __init__(self, arch: str, shape_and_type: str): else: # d_type == "f16" assert a_type == b_type == "f16" else: # k == 32, kind::f8f6f4 - assert a_type in ["e5m2", "e4m3", "e2m1"] # TODO: support e3m2 and e2m3 - assert b_type in ["e5m2", "e4m3", "e2m1"] + assert a_type in ["e5m2", "e4m3", "e3m2", "e2m3", "e2m1"] + assert b_type in ["e5m2", "e4m3", "e3m2", "e2m3", "e2m1"] assert d_type in ["f32", "f16"] self.arch = arch self.shape_and_type = shape_and_type self.m, self.n, self.k = m, n, k self.a_type = nv_torch_dtype[a_type] self.b_type = nv_torch_dtype[b_type] + self.a_ptx_type = a_type + self.b_ptx_type = b_type self.c_type = nv_torch_dtype[d_type] self.d_type = nv_torch_dtype[d_type] @@ -52,8 +54,8 @@ def __init__(self, arch: str, shape_and_type: str): assert k in [32, 64] assert d_type == "f32" if k == 32: # kind::mxf8f6f4 - assert a_type in ["e5m2", "e4m3", "e2m1"] # TODO: support e3m2 and e2m3 - assert b_type in ["e5m2", "e4m3", "e2m1"] + assert a_type in ["e5m2", "e4m3", "e3m2", "e2m3", "e2m1"] + assert b_type in ["e5m2", "e4m3", "e3m2", "e2m3", "e2m1"] assert s_type == "ue8m0" assert block_size == "block32" else: # k == 64, kind::mxf4nvf4 @@ -71,6 +73,8 @@ def __init__(self, arch: str, shape_and_type: str): self.packing = 2 if k == 64 else 1 self.a_type = nv_torch_dtype[a_type] self.b_type = nv_torch_dtype[b_type] + self.a_ptx_type = a_type + self.b_ptx_type = b_type self.c_type = nv_torch_dtype[d_type] self.d_type = nv_torch_dtype[d_type] self.s_type = nv_torch_dtype[s_type] diff --git a/src/mmasim/mmasim/nv_ptx/sim.py b/src/mmasim/mmasim/nv_ptx/sim.py index 3ca9dbc..44d6b68 100644 --- a/src/mmasim/mmasim/nv_ptx/sim.py +++ b/src/mmasim/mmasim/nv_ptx/sim.py @@ -42,11 +42,11 @@ def __init__(self, arch: str, shape_and_type: str): e_zero = -20 if self.c_type == torch.float16 else -131 elif self.arch in ["Turing", "Ampere", "Ada Lovelace"]: e_zero = -21 if self.c_type == torch.float16 else -132 - elif self.arch in ["Hopper", "Blackwell"]: + else: # Hopper, Blackwell, RTX Blackwell e_zero = -21 if self.c_type == torch.float16 else -133 - else: # RTX Blackwell - e_zero = -22 if self.c_type == torch.float16 else -133 - self.arithmetic_op = fdpa.MMA_T_FDPA(F, rho, L_max, e_zero) + self.arithmetic_op = fdpa.MMA_T_FDPA( + F, rho, L_max, e_zero, self.a_ptx_type, self.b_ptx_type + ) def dpa(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: return self.arithmetic_op.dpa(a, b, c) @@ -69,7 +69,12 @@ def __init__(self, arch: str, shape_and_type: str): ) else: self.arithmetic_op = fdpa.MMA_ST_FDPA( - F=25, rho="RZ-FP32", L_max=32, e_zero=-133 + F=25, + rho="RZ-FP32", + L_max=32, + e_zero=-133, + a_ptx_type=self.a_ptx_type, + b_ptx_type=self.b_ptx_type, ) def dpa( @@ -128,7 +133,9 @@ def __init__(self, arch: str, shape_and_type: str): rho = "RNE-FP16" if self.d_type == torch.float16 else "RZ-FP32" L_max = 32 // self.a_type.itemsize e_zero = -21 if self.c_type == torch.float16 else -133 - self.arithmetic_op = fdpa.MMA_T_FDPA(F, rho, L_max, e_zero) + self.arithmetic_op = fdpa.MMA_T_FDPA( + F, rho, L_max, e_zero, self.a_ptx_type, self.b_ptx_type + ) def dpa(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: return self.arithmetic_op.dpa(a, b, c) @@ -151,7 +158,12 @@ def __init__(self, arch: str, shape_and_type: str): ) else: self.arithmetic_op = fdpa.MMA_ST_FDPA( - F=25, rho="RZ-FP32", L_max=32, e_zero=-133 + F=25, + rho="RZ-FP32", + L_max=32, + e_zero=-133, + a_ptx_type=self.a_ptx_type, + b_ptx_type=self.b_ptx_type, ) def dpa( diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/blackwell.py b/src/mmasim_kernels/mmasim_kernels/nv_ptx/blackwell.py index 3aef12f..0bb22b6 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/blackwell.py +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/blackwell.py @@ -38,13 +38,25 @@ # sm_100a fp8 f32_output lib.tcgen05mma_m64n8k32_f32_e5m2_e5m2.argtypes = [ctypes.c_void_p] * 3 lib.tcgen05mma_m64n8k32_f32_e4m3_e4m3.argtypes = [ctypes.c_void_p] * 3 +# sm_100a fp6 and fp4 f32_output +lib.tcgen05mma_m64n8k32_f32_e3m2_e3m2.argtypes = [ctypes.c_void_p] * 3 +lib.tcgen05mma_m64n8k32_f32_e2m3_e2m3.argtypes = [ctypes.c_void_p] * 3 +lib.tcgen05mma_m64n8k32_f32_e2m1_e2m1.argtypes = [ctypes.c_void_p] * 3 # sm_100a fp8 f16_output lib.tcgen05mma_m64n8k32_f16_e5m2_e5m2.argtypes = [ctypes.c_void_p] * 3 lib.tcgen05mma_m64n8k32_f16_e4m3_e4m3.argtypes = [ctypes.c_void_p] * 3 +# sm_100a fp6 and fp4 f16_output +lib.tcgen05mma_m64n8k32_f16_e3m2_e3m2.argtypes = [ctypes.c_void_p] * 3 +lib.tcgen05mma_m64n8k32_f16_e2m3_e2m3.argtypes = [ctypes.c_void_p] * 3 +lib.tcgen05mma_m64n8k32_f16_e2m1_e2m1.argtypes = [ctypes.c_void_p] * 3 # sm_100a mxf8f6f4 lib.tcgen05mma_m128n8k32_block32_f32_e5m2_e5m2_ue8m0.argtypes = [ctypes.c_void_p] * 5 lib.tcgen05mma_m128n8k32_block32_f32_e4m3_e4m3_ue8m0.argtypes = [ctypes.c_void_p] * 5 +# sm_100a mxf8f6f4 fp6 and fp4 +lib.tcgen05mma_m128n8k32_block32_f32_e3m2_e3m2_ue8m0.argtypes = [ctypes.c_void_p] * 5 +lib.tcgen05mma_m128n8k32_block32_f32_e2m3_e2m3_ue8m0.argtypes = [ctypes.c_void_p] * 5 +lib.tcgen05mma_m128n8k32_block32_f32_e2m1_e2m1_ue8m0.argtypes = [ctypes.c_void_p] * 5 # sm_100a mxf4 and nvf4 lib.tcgen05mma_m128n8k64_block32_f32_e2m1_e2m1_ue8m0.argtypes = [ctypes.c_void_p] * 5 lib.tcgen05mma_m128n8k64_block16_f32_e2m1_e2m1_ue8m0.argtypes = [ctypes.c_void_p] * 5 @@ -76,14 +88,26 @@ # sm_100a fp8 f32_output "m64n8k32.f32.e5m2.e5m2": lib.tcgen05mma_m64n8k32_f32_e5m2_e5m2, "m64n8k32.f32.e4m3.e4m3": lib.tcgen05mma_m64n8k32_f32_e4m3_e4m3, + # sm_100a fp6 and fp4 f32_output + "m64n8k32.f32.e3m2.e3m2": lib.tcgen05mma_m64n8k32_f32_e3m2_e3m2, + "m64n8k32.f32.e2m3.e2m3": lib.tcgen05mma_m64n8k32_f32_e2m3_e2m3, + "m64n8k32.f32.e2m1.e2m1": lib.tcgen05mma_m64n8k32_f32_e2m1_e2m1, # sm_100a fp8 f16_output "m64n8k32.f16.e5m2.e5m2": lib.tcgen05mma_m64n8k32_f16_e5m2_e5m2, "m64n8k32.f16.e4m3.e4m3": lib.tcgen05mma_m64n8k32_f16_e4m3_e4m3, + # sm_100a fp6 and fp4 f16_output + "m64n8k32.f16.e3m2.e3m2": lib.tcgen05mma_m64n8k32_f16_e3m2_e3m2, + "m64n8k32.f16.e2m3.e2m3": lib.tcgen05mma_m64n8k32_f16_e2m3_e2m3, + "m64n8k32.f16.e2m1.e2m1": lib.tcgen05mma_m64n8k32_f16_e2m1_e2m1, } tcgen05mma_block_scale_cuda_kernels = { # sm_100a mxf8f6f4 "m128n8k32.block32.f32.e5m2.e5m2.ue8m0": lib.tcgen05mma_m128n8k32_block32_f32_e5m2_e5m2_ue8m0, "m128n8k32.block32.f32.e4m3.e4m3.ue8m0": lib.tcgen05mma_m128n8k32_block32_f32_e4m3_e4m3_ue8m0, + # sm_100a mxf8f6f4 fp6 and fp4 + "m128n8k32.block32.f32.e3m2.e3m2.ue8m0": lib.tcgen05mma_m128n8k32_block32_f32_e3m2_e3m2_ue8m0, + "m128n8k32.block32.f32.e2m3.e2m3.ue8m0": lib.tcgen05mma_m128n8k32_block32_f32_e2m3_e2m3_ue8m0, + "m128n8k32.block32.f32.e2m1.e2m1.ue8m0": lib.tcgen05mma_m128n8k32_block32_f32_e2m1_e2m1_ue8m0, # sm_100a mxf4 and nvf4 "m128n8k64.block32.f32.e2m1.e2m1.ue8m0": lib.tcgen05mma_m128n8k64_block32_f32_e2m1_e2m1_ue8m0, "m128n8k64.block16.f32.e2m1.e2m1.ue8m0": lib.tcgen05mma_m128n8k64_block16_f32_e2m1_e2m1_ue8m0, diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/tcgen05mma.h b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/tcgen05mma.h index a307bc6..f0075cf 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/tcgen05mma.h +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/tcgen05mma.h @@ -134,6 +134,112 @@ a_desc |= 1ll << 46; \ } while (0) +// (16x2) tiles, each with (8x16) elements, packed in 12 bytes +#define LOAD_A_M128K32_FP6() \ + do \ + { \ + for (uint32_t i = 0; i < 16; i++) \ + { \ + for (uint32_t k = 0; k < 2; k++) \ + { \ + uint32_t row = i * 8 + tid / 16; \ + uint32_t col = k * 16 + tid % 16 * 8 / 6; \ + uint32_t shift = tid % 16 * 8 % 6; \ + uint8_t x = 0; \ + if (tid % 16 < 12) \ + { \ + uint8_t lo = a[row * K + col] & 0x3f; \ + uint8_t hi = a[row * K + col + 1] & 0x3f; \ + x = lo >> shift | hi << (6 - shift); \ + } \ + a_smem[i * 256 + k * 128 + tid] = x; \ + } \ + } \ + a_desc = desc_encode(__cvta_generic_to_shared(a_smem)); \ + a_desc |= desc_encode(128ll) << 16; \ + a_desc |= desc_encode(256ll) << 32; \ + a_desc |= 1ll << 46; \ + } while (0) + +// (8x2) tiles, each with (8x16) elements, packed in 12 bytes +#define LOAD_A_M64K32_FP6() \ + do \ + { \ + for (uint32_t i = 0; i < 8; i++) \ + { \ + for (uint32_t k = 0; k < 2; k++) \ + { \ + uint32_t row = i * 8 + tid / 16; \ + uint32_t col = k * 16 + tid % 16 * 8 / 6; \ + uint32_t shift = tid % 16 * 8 % 6; \ + uint8_t x = 0; \ + if (tid % 16 < 12) \ + { \ + uint8_t lo = a[row * K + col] & 0x3f; \ + uint8_t hi = a[row * K + col + 1] & 0x3f; \ + x = lo >> shift | hi << (6 - shift); \ + } \ + a_smem[i * 256 + k * 128 + tid] = x; \ + } \ + } \ + a_desc = desc_encode(__cvta_generic_to_shared(a_smem)); \ + a_desc |= desc_encode(128ll) << 16; \ + a_desc |= desc_encode(256ll) << 32; \ + a_desc |= 1ll << 46; \ + } while (0) + +// (16x2) tiles, each with (8x16) elements, packed in 8 bytes +#define LOAD_A_M128K32_FP4() \ + do \ + { \ + for (uint32_t i = 0; i < 16; i++) \ + { \ + for (uint32_t k = 0; k < 2; k++) \ + { \ + uint32_t row = i * 8 + tid / 16; \ + uint32_t col = k * 16 + tid % 16 * 2; \ + uint8_t x = 0; \ + if (tid % 16 < 8) \ + { \ + uint8_t lo = a[row * K + col] >> 2 & 0xf; \ + uint8_t hi = a[row * K + col + 1] >> 2 & 0xf; \ + x = lo | hi << 4; \ + } \ + a_smem[i * 256 + k * 128 + tid] = x; \ + } \ + } \ + a_desc = desc_encode(__cvta_generic_to_shared(a_smem)); \ + a_desc |= desc_encode(128ll) << 16; \ + a_desc |= desc_encode(256ll) << 32; \ + a_desc |= 1ll << 46; \ + } while (0) + +// (8x2) tiles, each with (8x16) elements, packed in 8 bytes +#define LOAD_A_M64K32_FP4() \ + do \ + { \ + for (uint32_t i = 0; i < 8; i++) \ + { \ + for (uint32_t k = 0; k < 2; k++) \ + { \ + uint32_t row = i * 8 + tid / 16; \ + uint32_t col = k * 16 + tid % 16 * 2; \ + uint8_t x = 0; \ + if (tid % 16 < 8) \ + { \ + uint8_t lo = a[row * K + col] >> 2 & 0xf; \ + uint8_t hi = a[row * K + col + 1] >> 2 & 0xf; \ + x = lo | hi << 4; \ + } \ + a_smem[i * 256 + k * 128 + tid] = x; \ + } \ + } \ + a_desc = desc_encode(__cvta_generic_to_shared(a_smem)); \ + a_desc |= desc_encode(128ll) << 16; \ + a_desc |= desc_encode(256ll) << 32; \ + a_desc |= 1ll << 46; \ + } while (0) + // (16x2) tiles, each with (8x32) elements #define LOAD_A_M128K64_FP4() \ do \ @@ -198,6 +304,53 @@ b_desc |= 1ll << 46; \ } while (0) +// (2x1) tiles, each with (16x8) elements, packed in 12 bytes +#define LOAD_B_N8K32_FP6() \ + do \ + { \ + for (uint32_t k = 0; k < 2; k++) \ + { \ + uint32_t row = k * 16 + tid % 16 * 8 / 6; \ + uint32_t col = tid / 16; \ + uint32_t shift = tid % 16 * 8 % 6; \ + uint8_t x = 0; \ + if (tid % 16 < 12) \ + { \ + uint8_t lo = b[col * K + row] & 0x3f; \ + uint8_t hi = b[col * K + row + 1] & 0x3f; \ + x = lo >> shift | hi << (6 - shift); \ + } \ + b_smem[k * 128 + tid] = x; \ + } \ + b_desc = desc_encode(__cvta_generic_to_shared(b_smem)); \ + b_desc |= desc_encode(128ll) << 16; \ + b_desc |= desc_encode(256ll) << 32; \ + b_desc |= 1ll << 46; \ + } while (0) + +// (2x1) tiles, each with (16x8) elements, packed in 8 bytes +#define LOAD_B_N8K32_FP4() \ + do \ + { \ + for (uint32_t k = 0; k < 2; k++) \ + { \ + uint32_t row = k * 16 + tid % 16 * 2; \ + uint32_t col = tid / 16; \ + uint8_t x = 0; \ + if (tid % 16 < 8) \ + { \ + uint8_t lo = b[col * K + row] >> 2 & 0xf; \ + uint8_t hi = b[col * K + row + 1] >> 2 & 0xf; \ + x = lo | hi << 4; \ + } \ + b_smem[k * 128 + tid] = x; \ + } \ + b_desc = desc_encode(__cvta_generic_to_shared(b_smem)); \ + b_desc |= desc_encode(128ll) << 16; \ + b_desc |= desc_encode(256ll) << 32; \ + b_desc |= 1ll << 46; \ + } while (0) + // (2x1) tiles, each with (32x8) elements #define LOAD_B_N8K64_FP4() \ do \ diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/tcgen05mma_sm100a.cu b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/tcgen05mma_sm100a.cu index 2a91e1e..1f06589 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/tcgen05mma_sm100a.cu +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/tcgen05mma_sm100a.cu @@ -195,6 +195,102 @@ extern "C" // fp8 m64n8k32 f32_output } } +extern "C" // fp6 and fp4 m64n8k32 f32_output +{ + __global__ void tcgen05mma_m64n8k32_f32_e3m2_e3m2_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_FP6(); + LOAD_B_N8K32_FP6(); + ALLOC_TMEM(d_tmem_addr); + __syncthreads(); + LOAD_D_M64N8(); + __syncthreads(); + i_desc |= (1 << 4) | (4 << 7) | (4 << 10); // f32_e3m2_e3m2 + MMA("f8f6f4"); + STORE_D_M64N8(); + DEALLOC_TMEM(d_tmem_addr); + RELINQUISH_TMEM(); + } + + __global__ void tcgen05mma_m64n8k32_f32_e2m3_e2m3_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_FP6(); + LOAD_B_N8K32_FP6(); + ALLOC_TMEM(d_tmem_addr); + __syncthreads(); + LOAD_D_M64N8(); + __syncthreads(); + i_desc |= (1 << 4) | (3 << 7) | (3 << 10); // f32_e2m3_e2m3 + MMA("f8f6f4"); + STORE_D_M64N8(); + DEALLOC_TMEM(d_tmem_addr); + RELINQUISH_TMEM(); + } + + __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_FP4(); + LOAD_B_N8K32_FP4(); + 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_e3m2_e3m2( + uint32_t *d, uint8_t *a, uint8_t *b) + { + tcgen05mma_m64n8k32_f32_e3m2_e3m2_kernel<<<1, 128>>>(d, a, b); + } + + void tcgen05mma_m64n8k32_f32_e2m3_e2m3( + uint32_t *d, uint8_t *a, uint8_t *b) + { + tcgen05mma_m64n8k32_f32_e2m3_e2m3_kernel<<<1, 128>>>(d, a, b); + } + + 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); + } +} + extern "C" // fp8 m64n8k32 f16_output { __global__ void tcgen05mma_m64n8k32_f16_e5m2_e5m2_kernel( @@ -260,6 +356,102 @@ extern "C" // fp8 m64n8k32 f16_output } } +extern "C" // fp6 and fp4 m64n8k32 f16_output +{ + __global__ void tcgen05mma_m64n8k32_f16_e3m2_e3m2_kernel( + uint16_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_FP6(); + LOAD_B_N8K32_FP6(); + ALLOC_TMEM(d_tmem_addr); + __syncthreads(); + LOAD_D_M64N8(); + __syncthreads(); + i_desc |= (4 << 7) | (4 << 10); // f16_e3m2_e3m2 + MMA("f8f6f4"); + STORE_D_M64N8(); + DEALLOC_TMEM(d_tmem_addr); + RELINQUISH_TMEM(); + } + + __global__ void tcgen05mma_m64n8k32_f16_e2m3_e2m3_kernel( + uint16_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_FP6(); + LOAD_B_N8K32_FP6(); + ALLOC_TMEM(d_tmem_addr); + __syncthreads(); + LOAD_D_M64N8(); + __syncthreads(); + i_desc |= (3 << 7) | (3 << 10); // f16_e2m3_e2m3 + MMA("f8f6f4"); + STORE_D_M64N8(); + DEALLOC_TMEM(d_tmem_addr); + RELINQUISH_TMEM(); + } + + __global__ void tcgen05mma_m64n8k32_f16_e2m1_e2m1_kernel( + uint16_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_FP4(); + LOAD_B_N8K32_FP4(); + ALLOC_TMEM(d_tmem_addr); + __syncthreads(); + LOAD_D_M64N8(); + __syncthreads(); + i_desc |= (5 << 7) | (5 << 10); // f16_e2m1_e2m1 + MMA("f8f6f4"); + STORE_D_M64N8(); + DEALLOC_TMEM(d_tmem_addr); + RELINQUISH_TMEM(); + } + + void tcgen05mma_m64n8k32_f16_e3m2_e3m2( + uint16_t *d, uint8_t *a, uint8_t *b) + { + tcgen05mma_m64n8k32_f16_e3m2_e3m2_kernel<<<1, 128>>>(d, a, b); + } + + void tcgen05mma_m64n8k32_f16_e2m3_e2m3( + uint16_t *d, uint8_t *a, uint8_t *b) + { + tcgen05mma_m64n8k32_f16_e2m3_e2m3_kernel<<<1, 128>>>(d, a, b); + } + + void tcgen05mma_m64n8k32_f16_e2m1_e2m1( + uint16_t *d, uint8_t *a, uint8_t *b) + { + tcgen05mma_m64n8k32_f16_e2m1_e2m1_kernel<<<1, 128>>>(d, a, b); + } +} + extern "C" // mxfp8 { __global__ void tcgen05mma_m128n8k32_block32_f32_e5m2_e5m2_ue8m0_kernel( @@ -337,6 +529,120 @@ extern "C" // mxfp8 } } +extern "C" // mxfp6 mxfp4 +{ + __global__ void tcgen05mma_m128n8k32_block32_f32_e3m2_e3m2_ue8m0_kernel( + uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) + { + const uint32_t M = 128, N = 8, K = 32; + uint32_t tid = threadIdx.x, warpid = tid / 32, laneid = tid % 32; + uint32_t i_desc = (N >> 3 << 17) | (M >> 7 << 27), 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, sfa_tmem_addr, sfb_tmem_addr; + __shared__ uint64_t mma_barrier; + + INIT_MBARRIER(); + LOAD_A_M128K32_FP6(); + LOAD_B_N8K32_FP6(); + ALLOC_TMEM(d_tmem_addr); + ALLOC_TMEM(sfa_tmem_addr); + ALLOC_TMEM(sfb_tmem_addr); + __syncthreads(); + LOAD_D_M128N8(); + LOAD_SFA_M128SFK1(); + LOAD_SFB_N8SFK1(); + __syncthreads(); + i_desc |= (4 << 7) | (4 << 10) | (1 << 23); // e3m2_e3m2, ue8m0_scale + MMA_WITH_BLOCK_SCALE("mxf8f6f4", "scale_vec::1X"); + STORE_D_M128N8(); + DEALLOC_TMEM(d_tmem_addr); + DEALLOC_TMEM(sfa_tmem_addr); + DEALLOC_TMEM(sfb_tmem_addr); + RELINQUISH_TMEM(); + } + + __global__ void tcgen05mma_m128n8k32_block32_f32_e2m3_e2m3_ue8m0_kernel( + uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) + { + const uint32_t M = 128, N = 8, K = 32; + uint32_t tid = threadIdx.x, warpid = tid / 32, laneid = tid % 32; + uint32_t i_desc = (N >> 3 << 17) | (M >> 7 << 27), 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, sfa_tmem_addr, sfb_tmem_addr; + __shared__ uint64_t mma_barrier; + + INIT_MBARRIER(); + LOAD_A_M128K32_FP6(); + LOAD_B_N8K32_FP6(); + ALLOC_TMEM(d_tmem_addr); + ALLOC_TMEM(sfa_tmem_addr); + ALLOC_TMEM(sfb_tmem_addr); + __syncthreads(); + LOAD_D_M128N8(); + LOAD_SFA_M128SFK1(); + LOAD_SFB_N8SFK1(); + __syncthreads(); + i_desc |= (3 << 7) | (3 << 10) | (1 << 23); // e2m3_e2m3, ue8m0_scale + MMA_WITH_BLOCK_SCALE("mxf8f6f4", "scale_vec::1X"); + STORE_D_M128N8(); + DEALLOC_TMEM(d_tmem_addr); + DEALLOC_TMEM(sfa_tmem_addr); + DEALLOC_TMEM(sfb_tmem_addr); + RELINQUISH_TMEM(); + } + + __global__ void tcgen05mma_m128n8k32_block32_f32_e2m1_e2m1_ue8m0_kernel( + uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) + { + const uint32_t M = 128, N = 8, K = 32; + uint32_t tid = threadIdx.x, warpid = tid / 32, laneid = tid % 32; + uint32_t i_desc = (N >> 3 << 17) | (M >> 7 << 27), 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, sfa_tmem_addr, sfb_tmem_addr; + __shared__ uint64_t mma_barrier; + + INIT_MBARRIER(); + LOAD_A_M128K32_FP4(); + LOAD_B_N8K32_FP4(); + ALLOC_TMEM(d_tmem_addr); + ALLOC_TMEM(sfa_tmem_addr); + ALLOC_TMEM(sfb_tmem_addr); + __syncthreads(); + LOAD_D_M128N8(); + LOAD_SFA_M128SFK1(); + LOAD_SFB_N8SFK1(); + __syncthreads(); + i_desc |= (5 << 7) | (5 << 10) | (1 << 23); // e2m1_e2m1, ue8m0_scale + MMA_WITH_BLOCK_SCALE("mxf8f6f4", "scale_vec::1X"); + STORE_D_M128N8(); + DEALLOC_TMEM(d_tmem_addr); + DEALLOC_TMEM(sfa_tmem_addr); + DEALLOC_TMEM(sfb_tmem_addr); + RELINQUISH_TMEM(); + } + + void tcgen05mma_m128n8k32_block32_f32_e3m2_e3m2_ue8m0( + uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) + { + tcgen05mma_m128n8k32_block32_f32_e3m2_e3m2_ue8m0_kernel<<<1, 128>>>(d, a, b, sfa, sfb); + } + + void tcgen05mma_m128n8k32_block32_f32_e2m3_e2m3_ue8m0( + uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) + { + tcgen05mma_m128n8k32_block32_f32_e2m3_e2m3_ue8m0_kernel<<<1, 128>>>(d, a, b, sfa, sfb); + } + + void tcgen05mma_m128n8k32_block32_f32_e2m1_e2m1_ue8m0( + uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) + { + tcgen05mma_m128n8k32_block32_f32_e2m1_e2m1_ue8m0_kernel<<<1, 128>>>(d, a, b, sfa, sfb); + } +} + extern "C" // mxfp4 nvfp4 { __global__ void tcgen05mma_m128n8k64_block32_f32_e2m1_e2m1_ue8m0_kernel( diff --git a/tests/directed/test_nv_directed.py b/tests/directed/test_nv_directed.py index d9a4008..0c6d55e 100644 --- a/tests/directed/test_nv_directed.py +++ b/tests/directed/test_nv_directed.py @@ -4,6 +4,8 @@ # Expected D bits are measured on H200 (Hopper), B300 and B200 (Blackwell) +storage_type = {4: torch.uint32, 2: torch.uint16} + def vector(dtype: torch.dtype, k: int, entries: dict[int, float]) -> torch.Tensor: x = torch.zeros(k, dtype=torch.float32) @@ -13,9 +15,7 @@ def vector(dtype: torch.dtype, k: int, entries: dict[int, float]) -> torch.Tenso def raw(x: torch.Tensor) -> int: - if x.dtype == torch.float16: - return int(x.view(torch.int16).item()) & 0xFFFF - return int(x.view(torch.int32).item()) & 0xFFFFFFFF + return x.view(storage_type[x.dtype.itemsize]).item() def directed_test(sim, a: dict, b: dict, c: float, expected: int): @@ -51,20 +51,16 @@ def test_zero_result_is_positive(): def test_e_zero_f16_accumulator(): # a zero f16 C anchors the fused sum at 2^-21 - hopper = [ - WGMMA("Hopper", "m64n8k16.f16.f16.f16"), - ] - for sim in hopper: - print(f"Testing {sim.arch} {sim.shape_and_type}") - # 2^-25 is a tie at the f16 subnormal step, 2^-q breaks it iff q <= 46 - for q, expected in [(46, 0x0001), (47, 0x0000)]: - a, b = {0: 2.0**-12, 1: 2.0**-23}, {0: 2.0**-13, 1: 2.0 ** (23 - q)} - directed_test(sim, a, b, 0.0, expected) - blackwell = [ + sim = WGMMA("Hopper", "m64n8k16.f16.f16.f16") + print(f"Testing {sim.arch} {sim.shape_and_type}") + # 2^-25 is a tie at the f16 subnormal step, 2^-q breaks it iff q <= 46 + for q, expected in [(46, 0x0001), (47, 0x0000)]: + a, b = {0: 2.0**-12, 1: 2.0**-23}, {0: 2.0**-13, 1: 2.0 ** (23 - q)} + directed_test(sim, a, b, 0.0, expected) + for sim in [ TCGen05MMA("Blackwell", "m64n8k16.f16.f16.f16"), MMA("Blackwell", "m16n8k16.f16.f16.f16.f16"), - ] - for sim in blackwell: + ]: print(f"Testing {sim.arch} {sim.shape_and_type}") # 1.5 * 2^-23 + 2^-25 is a tie, -2^-47 is below the 2^-46 grid a = {0: 1.5 * 2.0**-12, 1: 2.0**-13, 2: -(2.0**-24)} @@ -77,10 +73,19 @@ def test_e2m1_f8f6f4(): # a * b = 0.25 is 0.01b * 2^0, so C = -(0.25 - 2^-26) loses its last bit sim = TCGen05MMA("Blackwell", "m128n8k32.f32.e2m1.e2m1") print(f"Testing {sim.arch} {sim.shape_and_type}") - a = torch.zeros(sim.k, dtype=torch.uint8) - a[0] = 0x1 << 2 # e2m1 sits in bits 5:2 of the 8-bit container - d = sim.dpa(a, a, torch.tensor(-(0.25 - 2.0**-26), dtype=torch.float32)) - assert raw(d) == 0x33000000, f"D_sim = {raw(d):#x}, D_gpu = 0x33000000" + # e2m1 sits in bits 5:2 of the 8-bit container + a = b = {0: 0x1 << 2} + directed_test(sim, a, b, -(0.25 - 2.0**-26), 0x33000000) + print("Tests passed!") + + +def test_e3m2_f8f6f4(): + # a * b = 2^-8 is 0.0001b * 2^-4, so C = -(2^-8 - 2^-30) loses its last bit + sim = TCGen05MMA("Blackwell", "m64n8k32.f32.e3m2.e3m2") + print(f"Testing {sim.arch} {sim.shape_and_type}") + # e3m2 sits in bits 5:0 of the 8-bit container, bits 7:6 are ignored + a = b = {0: 0xC0 | 0x01} + directed_test(sim, a, b, -(2.0**-8 - 2.0**-30), 0x31000000) print("Tests passed!") @@ -88,3 +93,4 @@ def test_e2m1_f8f6f4(): test_zero_result_is_positive() test_e_zero_f16_accumulator() test_e2m1_f8f6f4() + test_e3m2_f8f6f4()