Fix zero-result sign, f16 zero anchor and e2m1 decoding on Hopper and Blackwell

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
 

src/mmasim/mmasim/arithmetic/fdpa.py

@@ -3,9 +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_fp6_fp4,
 )
 
 DoubleTensor = Annotated[torch.Tensor, "float64"]
@@ -48,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
@@ -57,7 +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)
+        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)
@@ -87,9 +89,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)
@@ -99,21 +103,45 @@ 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)
+        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 = 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):
             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,
+                self.e_subnormal_a,
+                self.e_subnormal_b,
             )
         return c
 
@@ -135,9 +163,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)
@@ -150,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,
@@ -164,6 +207,10 @@ def dpa(
         alpha: torch.Tensor,
         beta: torch.Tensor,
     ) -> torch.Tensor:
+        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 = 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):
@@ -176,6 +223,8 @@ def dpa(
                 self.F,
                 self.rho,
                 self.e_zero,
+                self.e_subnormal_a,
+                self.e_subnormal_b,
             )
         return c
 

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(
@@ -55,3 +57,14 @@ 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_fp6_fp4(x: torch.Tensor, x_type: str) -> torch.Tensor:
+    assert x.dtype == torch.uint8
+    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)

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
 }
 

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]

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]

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 == "Hopper":
+            else:  # Hopper, Blackwell, RTX Blackwell
                 e_zero = -21 if self.c_type == torch.float16 else -133
-            else:  # Blackwell, 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(
@@ -106,7 +111,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 +132,10 @@ 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, 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)
@@ -149,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(

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,