Release v2.1.0

src/mmasim/mmasim/arithmetic/fdpa.py

@@ -3,7 +3,7 @@
 from .. import MMAOperation, MMABlockScaleOperation
 from .helper import (
     dtype_subnormal_exponent,
-    uint8_subnormal_exponent,
+    fp6_fp4_subnormal_exponent,
     truncate_fp32_to_tf32,
     truncate_e4m3_to_ue4m3,
     unpack_uint8_to_fp4,
@@ -50,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) + 0.0).to(torch.float16)  # -0.0 + 0.0 = +0.0
+        res = (s * pow2(e)).to(torch.float16)
     elif rho == "RNE-FP32":
         res = (s * pow2(e)).to(torch.float32)
     else:  # RZ
@@ -59,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) + 0.0).to(torch.float32)  # -0.0 + 0.0 = +0.0
+        res = (s * pow2(e)).to(torch.float32)
     nans = res.isnan()
     if res.dtype == torch.float16:
         res = res.view(torch.int16)
@@ -89,7 +89,7 @@ def t_fdpa(
     F: int,
     rho: str,
     e_zero: int,
-    e_subnormal_a: int | None = None,
+    e_subnormal_a: int | None = None,  # for FP6 and FP4
     e_subnormal_b: int | None = None,
 ) -> torch.Tensor:  # [...]
     s_a, e_a = frexp_and_normalize(a, e_subnormal_a)
@@ -99,7 +99,7 @@ def t_fdpa(
     e = torch.cat([e_a + e_b, e_c.unsqueeze(-1)], dim=-1)
     e[s == 0.0] = e_zero  # handle zero
     sum, e_max = truncated_fused_sum(s, e, F)
-    return ldexp_and_normalize(sum, e_max, rho)
+    return ldexp_and_normalize(sum, e_max, rho) + 0.0  # -0.0 to +0.0
 
 
 class MMA_T_FDPA(MMAOperation):
@@ -109,27 +109,27 @@ def __init__(
         rho: str,
         L_max: int,
         e_zero: int,
-        a_ptx_type: str | None = None,
-        b_ptx_type: str | None = None,
+        a_type_ptx: str,
+        b_type_ptx: str,
     ):
         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)
+        self.a_type_ptx = a_type_ptx
+        self.b_type_ptx = b_type_ptx
+        self.e_subnormal_a = fp6_fp4_subnormal_exponent.get(a_type_ptx)
+        self.e_subnormal_b = fp6_fp4_subnormal_exponent.get(b_type_ptx)
 
     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)
+            a = decode_uint8_to_fp6_fp4(a, self.a_type_ptx)
         if b.dtype == torch.uint8:
-            b = decode_uint8_to_fp6_fp4(b, self.b_ptx_type)
+            b = decode_uint8_to_fp6_fp4(b, self.b_type_ptx)
         K = a.shape[-1]
         L = min(K, self.L_max)
         for i in range(0, K, L):
@@ -163,7 +163,7 @@ def st_fdpa(
     F: int,
     rho: str,
     e_zero: int,
-    e_subnormal_a: int | None = None,
+    e_subnormal_a: int | None = None,  # for FP6 and FP4
     e_subnormal_b: int | None = None,
 ) -> torch.Tensor:  # [...]
     s_a, e_a = frexp_and_normalize(a, e_subnormal_a)
@@ -176,7 +176,7 @@ def st_fdpa(
     e = torch.cat([e_a + e_b + e_alpha + e_beta, e_c.unsqueeze(-1)], dim=-1)
     e[s == 0.0] = e_zero  # handle zero
     sum, e_max = truncated_fused_sum(s, e, F)
-    return ldexp_and_normalize(sum, e_max, rho)
+    return ldexp_and_normalize(sum, e_max, rho) + 0.0  # -0.0 to +0.0
 
 
 class MMA_ST_FDPA(MMABlockScaleOperation):
@@ -186,18 +186,18 @@ def __init__(
         rho: str,
         L_max: int,
         e_zero: int,
-        a_ptx_type: str | None = None,
-        b_ptx_type: str | None = None,
+        a_type_ptx: str,
+        b_type_ptx: str,
     ):
         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)
+        self.a_type_ptx = a_type_ptx
+        self.b_type_ptx = b_type_ptx
+        self.e_subnormal_a = fp6_fp4_subnormal_exponent.get(a_type_ptx)
+        self.e_subnormal_b = fp6_fp4_subnormal_exponent.get(b_type_ptx)
 
     def dpa(
         self,
@@ -208,9 +208,9 @@ def dpa(
         beta: torch.Tensor,
     ) -> torch.Tensor:
         if a.dtype == torch.uint8:  # fp6 or fp4
-            a = decode_uint8_to_fp6_fp4(a, self.a_ptx_type)
+            a = decode_uint8_to_fp6_fp4(a, self.a_type_ptx)
         if b.dtype == torch.uint8:
-            b = decode_uint8_to_fp6_fp4(b, self.b_ptx_type)
+            b = decode_uint8_to_fp6_fp4(b, self.b_type_ptx)
         K = a.shape[-1]
         L = min(K, self.L_max)
         for i in range(0, K, L):
@@ -273,7 +273,7 @@ def gst_fdpa(
     s = torch.cat([p * s_scale, s_c.unsqueeze(-1)], dim=-1)
     e = torch.cat([e_scale, e_c.unsqueeze(-1)], dim=-1)
     sum, e_max = truncated_fused_sum(s, e, F)
-    return ldexp_and_normalize(sum, e_max, rho)
+    return ldexp_and_normalize(sum, e_max, rho) + 0.0  # -0.0 to +0.0
 
 
 class MMA_GST_FDPA(MMABlockScaleOperation):

src/mmasim/mmasim/arithmetic/helper.py

@@ -11,8 +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)
+fp6_fp4_subnormal_exponent = {"e3m2": -2, "e2m3": 0, "e2m1": 0}
+fp6_fp4_bits = {"e3m2": (3, 2), "e2m3": (2, 3), "e2m1": (2, 1)}  # (E, M)
 
 
 fp4_value_table = torch.tensor(
@@ -61,7 +61,7 @@ def unpack_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
-    E, M = uint8_type_bits[x_type]
+    E, M = fp6_fp4_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)

src/mmasim/mmasim/nv_ptx/isa_common.py

@@ -14,6 +14,7 @@
     "e2m3": torch.uint8,  # torch has no fp6 dtype
     "e2m1": torch.uint8,  # torch.float4_e2m1fn_x2 is not well-implemented
 }
+f8f6f4_types = ["e5m2", "e4m3", "e3m2", "e2m3", "e2m1"]
 
 
 def nv_shape_to_mnk(shape: str) -> tuple[int, int, int]: