From fd407ab4733f483065c4188cf31ad14c9ee74218 Mon Sep 17 00:00:00 2001 From: Peichen Xie Date: Fri, 16 Oct 2026 15:26:38 +0800 Subject: [PATCH 1/5] minor fix --- src/mmasim/mmasim/arithmetic/fdpa.py | 48 +++++++++++----------- src/mmasim/mmasim/arithmetic/helper.py | 6 +-- src/mmasim/mmasim/nv_ptx/isa_mma.py | 8 ++-- src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py | 8 ++-- src/mmasim/mmasim/nv_ptx/isa_wgmma.py | 2 + src/mmasim/mmasim/nv_ptx/sim.py | 16 ++++---- 6 files changed, 46 insertions(+), 42 deletions(-) diff --git a/src/mmasim/mmasim/arithmetic/fdpa.py b/src/mmasim/mmasim/arithmetic/fdpa.py index d4d2262..b956a7b 100644 --- a/src/mmasim/mmasim/arithmetic/fdpa.py +++ b/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): diff --git a/src/mmasim/mmasim/arithmetic/helper.py b/src/mmasim/mmasim/arithmetic/helper.py index 2dbc5ec..9a94a96 100644 --- a/src/mmasim/mmasim/arithmetic/helper.py +++ b/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) diff --git a/src/mmasim/mmasim/nv_ptx/isa_mma.py b/src/mmasim/mmasim/nv_ptx/isa_mma.py index c57a8b1..b31af42 100644 --- a/src/mmasim/mmasim/nv_ptx/isa_mma.py +++ b/src/mmasim/mmasim/nv_ptx/isa_mma.py @@ -136,8 +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.a_type_ptx = a_type + self.b_type_ptx = b_type self.c_type = nv_torch_dtype[c_type] self.d_type = nv_torch_dtype[d_type] @@ -163,8 +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.a_type_ptx = a_type + self.b_type_ptx = 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 5200375..c89e167 100644 --- a/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py +++ b/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py @@ -36,8 +36,8 @@ def __init__(self, arch: str, shape_and_type: str): 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.a_type_ptx = a_type + self.b_type_ptx = b_type self.c_type = nv_torch_dtype[d_type] self.d_type = nv_torch_dtype[d_type] @@ -73,8 +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.a_type_ptx = a_type + self.b_type_ptx = 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/isa_wgmma.py b/src/mmasim/mmasim/nv_ptx/isa_wgmma.py index 21e3416..6f4e9f9 100644 --- a/src/mmasim/mmasim/nv_ptx/isa_wgmma.py +++ b/src/mmasim/mmasim/nv_ptx/isa_wgmma.py @@ -32,5 +32,7 @@ def __init__(self, arch: str, shape_and_type: str): 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_type_ptx = a_type + self.b_type_ptx = b_type self.c_type = nv_torch_dtype[d_type] self.d_type = nv_torch_dtype[d_type] diff --git a/src/mmasim/mmasim/nv_ptx/sim.py b/src/mmasim/mmasim/nv_ptx/sim.py index 44d6b68..c66b8f2 100644 --- a/src/mmasim/mmasim/nv_ptx/sim.py +++ b/src/mmasim/mmasim/nv_ptx/sim.py @@ -45,7 +45,7 @@ def __init__(self, arch: str, shape_and_type: str): else: # Hopper, Blackwell, RTX Blackwell 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 + F, rho, L_max, e_zero, self.a_type_ptx, self.b_type_ptx ) def dpa(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: @@ -73,8 +73,8 @@ def __init__(self, arch: str, shape_and_type: str): rho="RZ-FP32", L_max=32, e_zero=-133, - a_ptx_type=self.a_ptx_type, - b_ptx_type=self.b_ptx_type, + a_type_ptx=self.a_type_ptx, + b_type_ptx=self.b_type_ptx, ) def dpa( @@ -112,7 +112,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_type_ptx, self.b_type_ptx + ) def dpa(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: return self.arithmetic_op.dpa(a, b, c) @@ -134,7 +136,7 @@ def __init__(self, arch: str, shape_and_type: str): 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.a_ptx_type, self.b_ptx_type + F, rho, L_max, e_zero, self.a_type_ptx, self.b_type_ptx ) def dpa(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: @@ -162,8 +164,8 @@ def __init__(self, arch: str, shape_and_type: str): rho="RZ-FP32", L_max=32, e_zero=-133, - a_ptx_type=self.a_ptx_type, - b_ptx_type=self.b_ptx_type, + a_type_ptx=self.a_type_ptx, + b_type_ptx=self.b_type_ptx, ) def dpa( From 1aa764ba39fad958fef55c9bcfa613708e440b1a Mon Sep 17 00:00:00 2001 From: Peichen Xie Date: Fri, 16 Oct 2026 15:27:31 +0800 Subject: [PATCH 2/5] remove extra test --- tests/directed/test_nv_directed.py | 96 ------------------------------ 1 file changed, 96 deletions(-) delete mode 100644 tests/directed/test_nv_directed.py diff --git a/tests/directed/test_nv_directed.py b/tests/directed/test_nv_directed.py deleted file mode 100644 index 0c6d55e..0000000 --- a/tests/directed/test_nv_directed.py +++ /dev/null @@ -1,96 +0,0 @@ -import torch - -from mmasim.nv_ptx import MMA, WGMMA, TCGen05MMA - -# 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) - for i, v in entries.items(): - x[i] = v - return x.to(dtype) - - -def raw(x: torch.Tensor) -> int: - return x.view(storage_type[x.dtype.itemsize]).item() - - -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) - print("Tests passed!") - - -def test_e_zero_f16_accumulator(): - # a zero f16 C anchors the fused sum at 2^-21 - 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"), - ]: - 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!") - - -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}") - # 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!") - - -if __name__ == "__main__": - test_zero_result_is_positive() - test_e_zero_f16_accumulator() - test_e2m1_f8f6f4() - test_e3m2_f8f6f4() From 9ca813c4722b3bec5b9606af9568bf1932c83557 Mon Sep 17 00:00:00 2001 From: Peichen Xie Date: Tue, 20 Oct 2026 19:15:04 +0800 Subject: [PATCH 3/5] refactor code and support more sm_90, sm_100, and sm_120 instructions --- src/mmasim/mmasim/nv_ptx/isa_common.py | 1 + src/mmasim/mmasim/nv_ptx/isa_mma.py | 203 ++-- src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py | 127 ++- src/mmasim/mmasim/nv_ptx/isa_wgmma.py | 51 +- .../mmasim_kernels/nv_ptx/ada_lovelace.py | 84 +- .../mmasim_kernels/nv_ptx/ampere.py | 42 +- .../mmasim_kernels/nv_ptx/blackwell.py | 150 +-- .../mmasim_kernels/nv_ptx/hopper.py | 96 +- .../mmasim_kernels/nv_ptx/kernels/mma.h | 27 + .../nv_ptx/kernels/mma_sm120a.cu | 242 ++++- .../nv_ptx/kernels/tcgen05mma.h | 869 ++++++----------- .../nv_ptx/kernels/tcgen05mma_sm100a.cu | 893 +++--------------- .../mmasim_kernels/nv_ptx/kernels/wgmma.h | 451 ++++++--- .../nv_ptx/kernels/wgmma_sm90a.cu | 447 ++------- .../mmasim_kernels/nv_ptx/rtx_blackwell.py | 126 +-- .../mmasim_kernels/nv_ptx/turing.py | 21 +- .../mmasim_kernels/nv_ptx/volta.py | 21 +- tests/equivalence/test_nv_rtx_blackwell.py | 8 +- 18 files changed, 1375 insertions(+), 2484 deletions(-) diff --git a/src/mmasim/mmasim/nv_ptx/isa_common.py b/src/mmasim/mmasim/nv_ptx/isa_common.py index 80d2864..012607c 100644 --- a/src/mmasim/mmasim/nv_ptx/isa_common.py +++ b/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]: diff --git a/src/mmasim/mmasim/nv_ptx/isa_mma.py b/src/mmasim/mmasim/nv_ptx/isa_mma.py index b31af42..a3e46ff 100644 --- a/src/mmasim/mmasim/nv_ptx/isa_mma.py +++ b/src/mmasim/mmasim/nv_ptx/isa_mma.py @@ -2,135 +2,134 @@ MMAOperation, MMABlockScaleOperation, ) -from .isa_common import nv_torch_dtype, nv_shape_to_mnk +from .isa_common import ( + nv_torch_dtype, + f8f6f4_types, + nv_shape_to_mnk, +) -volta_mma_shape_and_type = [ - # sm_70 f16 +mma_shape_and_type_from_sm_70 = [ + # fp16 "m8n8k4.f32.f16.f16.f32", "m8n8k4.f32.f16.f16.f16", "m8n8k4.f16.f16.f16.f16", ] -turing_mma_shape_and_type = [ - # sm_75 f16 +mma_shape_and_type_from_sm_75 = [ + # fp16 "m16n8k8.f32.f16.f16.f32", "m16n8k8.f16.f16.f16.f16", ] -ampere_mma_shape_and_type = [ - # sm_80 f64 +mma_shape_and_type_from_sm_80 = [ + # fp64 "m8n8k4.f64.f64.f64.f64", - # sm_80 tf32 + # tf32 "m16n8k8.f32.tf32.tf32.f32", "m16n8k4.f32.tf32.tf32.f32", - # sm_80 f16 + # fp16 "m16n8k16.f32.f16.f16.f32", "m16n8k16.f16.f16.f16.f16", - # sm_80 bf16 + # bf16 "m16n8k16.f32.bf16.bf16.f32", "m16n8k8.f32.bf16.bf16.f32", ] -adalovelace_mma_shape_and_type = [ - # sm_89 fp8 m16n8k32 f32_output - "m16n8k32.f32.e5m2.e5m2.f32", - "m16n8k32.f32.e5m2.e4m3.f32", - "m16n8k32.f32.e4m3.e5m2.f32", - "m16n8k32.f32.e4m3.e4m3.f32", - # sm_89 fp8 m16n8k16 f32_output - "m16n8k16.f32.e5m2.e5m2.f32", - "m16n8k16.f32.e5m2.e4m3.f32", - "m16n8k16.f32.e4m3.e5m2.f32", - "m16n8k16.f32.e4m3.e4m3.f32", - # sm_89 fp8 m16n8k32 f16_output - "m16n8k32.f16.e5m2.e5m2.f16", - "m16n8k32.f16.e5m2.e4m3.f16", - "m16n8k32.f16.e4m3.e5m2.f16", - "m16n8k32.f16.e4m3.e4m3.f16", - # sm_89 fp8 m16n8k16 f16_output - "m16n8k16.f16.e5m2.e5m2.f16", - "m16n8k16.f16.e5m2.e4m3.f16", - "m16n8k16.f16.e4m3.e5m2.f16", - "m16n8k16.f16.e4m3.e4m3.f16", +mma_shape_and_type_from_sm_89 = [ + # fp8 + f"m16n8k{k}.{d_type}.{a_type}.{b_type}.{d_type}" + for k in [32, 16] + for d_type in ["f32", "f16"] + for a_type in ["e5m2", "e4m3"] + for b_type in ["e5m2", "e4m3"] ] -hopper_mma_shape_and_type = [ - # sm_90a f64 +mma_shape_and_type_from_sm_90 = [ + # fp64 "m16n8k16.f64.f64.f64.f64", "m16n8k8.f64.f64.f64.f64", "m16n8k4.f64.f64.f64.f64", ] -rtx_blackwell_mma_shape_and_type = [ - # sm_120a f8f6f4 f32_output - # TODO: support e3m2 and e2m3 - "m16n8k32.f32.e5m2.e5m2.f32", - "m16n8k32.f32.e5m2.e4m3.f32", - "m16n8k32.f32.e5m2.e2m1.f32", - "m16n8k32.f32.e4m3.e5m2.f32", - "m16n8k32.f32.e4m3.e4m3.f32", - "m16n8k32.f32.e4m3.e2m1.f32", - "m16n8k32.f32.e2m1.e5m2.f32", - "m16n8k32.f32.e2m1.e4m3.f32", - "m16n8k32.f32.e2m1.e2m1.f32", - # sm_120a f8f6f4 f16_output - "m16n8k32.f16.e5m2.e5m2.f16", - "m16n8k32.f16.e5m2.e4m3.f16", - "m16n8k32.f16.e5m2.e2m1.f16", - "m16n8k32.f16.e4m3.e5m2.f16", - "m16n8k32.f16.e4m3.e4m3.f16", - "m16n8k32.f16.e4m3.e2m1.f16", - "m16n8k32.f16.e2m1.e5m2.f16", - "m16n8k32.f16.e2m1.e4m3.f16", - "m16n8k32.f16.e2m1.e2m1.f16", -] -rtx_blackwell_mma_block_scale_shape_and_type = [ - # sm_120a mxf8f6f4 - # TODO: support e3m2 and e2m3 - "m16n8k32.block32.f32.e5m2.e5m2.f32.ue8m0", - "m16n8k32.block32.f32.e5m2.e4m3.f32.ue8m0", - "m16n8k32.block32.f32.e5m2.e2m1.f32.ue8m0", - "m16n8k32.block32.f32.e4m3.e5m2.f32.ue8m0", - "m16n8k32.block32.f32.e4m3.e4m3.f32.ue8m0", - "m16n8k32.block32.f32.e4m3.e2m1.f32.ue8m0", - "m16n8k32.block32.f32.e2m1.e5m2.f32.ue8m0", - "m16n8k32.block32.f32.e2m1.e4m3.f32.ue8m0", - "m16n8k32.block32.f32.e2m1.e2m1.f32.ue8m0", - # sm_120a mxf4nvf4 - "m16n8k64.block32.f32.e2m1.e2m1.f32.ue8m0", - "m16n8k64.block16.f32.e2m1.e2m1.f32.ue8m0", - "m16n8k64.block32.f32.e2m1.e2m1.f32.ue4m3", - "m16n8k64.block16.f32.e2m1.e2m1.f32.ue4m3", +mma_shape_and_type_from_sm_120 = [ + f"m16n8k32.kind::f8f6f4.{d_type}.{a_type}.{b_type}.{d_type}" + for d_type in ["f32", "f16"] + for a_type in f8f6f4_types + for b_type in f8f6f4_types ] - -arch_mma_shape_and_type = { +volta_mma_shape_and_type = mma_shape_and_type_from_sm_70 +turing_mma_shape_and_type = ( + mma_shape_and_type_from_sm_75 + mma_shape_and_type_from_sm_70 +) +ampere_mma_shape_and_type = ( + mma_shape_and_type_from_sm_80 + mma_shape_and_type_from_sm_75 +) +ada_lovelace_mma_shape_and_type = ( + mma_shape_and_type_from_sm_89 + + mma_shape_and_type_from_sm_80 + + mma_shape_and_type_from_sm_75 +) +hopper_mma_shape_and_type = ( + mma_shape_and_type_from_sm_90 + + mma_shape_and_type_from_sm_80 + + mma_shape_and_type_from_sm_75 +) +blackwell_mma_shape_and_type = ( + mma_shape_and_type_from_sm_80 + mma_shape_and_type_from_sm_75 +) +rtx_blackwell_mma_shape_and_type = ( + mma_shape_and_type_from_sm_120 + + mma_shape_and_type_from_sm_89 + + mma_shape_and_type_from_sm_80 + + mma_shape_and_type_from_sm_75 +) +mma_shape_and_type = { "Volta": volta_mma_shape_and_type, - "Turing": turing_mma_shape_and_type + volta_mma_shape_and_type, - "Ampere": ampere_mma_shape_and_type + turing_mma_shape_and_type, - "Ada Lovelace": adalovelace_mma_shape_and_type - + ampere_mma_shape_and_type - + turing_mma_shape_and_type, - "Hopper": hopper_mma_shape_and_type - + ampere_mma_shape_and_type - + turing_mma_shape_and_type, - "Blackwell": ampere_mma_shape_and_type + turing_mma_shape_and_type, - "RTX Blackwell": rtx_blackwell_mma_shape_and_type - + adalovelace_mma_shape_and_type - + ampere_mma_shape_and_type - + turing_mma_shape_and_type, + "Turing": turing_mma_shape_and_type, + "Ampere": ampere_mma_shape_and_type, + "Ada Lovelace": ada_lovelace_mma_shape_and_type, + "Hopper": hopper_mma_shape_and_type, + "Blackwell": blackwell_mma_shape_and_type, + "RTX Blackwell": rtx_blackwell_mma_shape_and_type, } -arch_mma_block_scale_shape_and_type = { + +mma_block_scale_shape_and_type_from_sm_120 = ( + [ + f"m16n8k32.kind::mxf8f6f4.block32.f32.{a_type}.{b_type}.f32.ue8m0" + for a_type in f8f6f4_types + for b_type in f8f6f4_types + ] + + [ + "m16n8k64.kind::mxf4.block32.f32.e2m1.e2m1.f32.ue8m0", + ] + + [ + "m16n8k64.kind::mxf4nvf4.block32.f32.e2m1.e2m1.f32.ue8m0", + # "m16n8k64.kind::mxf4nvf4.block32.f32.e2m1.e2m1.f32.ue4m3", # invalid + "m16n8k64.kind::mxf4nvf4.block16.f32.e2m1.e2m1.f32.ue8m0", + "m16n8k64.kind::mxf4nvf4.block16.f32.e2m1.e2m1.f32.ue4m3", + ] +) +rtx_blackwell_mma_block_scale_shape_and_type = ( + mma_block_scale_shape_and_type_from_sm_120 +) +mma_block_scale_shape_and_type = { "RTX Blackwell": rtx_blackwell_mma_block_scale_shape_and_type, } class MMA(MMAOperation): def __init__(self, arch: str, shape_and_type: str): - assert arch in arch_mma_shape_and_type.keys(), ( + assert arch in mma_shape_and_type.keys(), ( f"Unsupported architecture {arch} for mma.\n" - f"Supported architectures: {list(arch_mma_shape_and_type.keys())}" + f"Supported architectures: {list(mma_shape_and_type.keys())}" ) - supported_shape_and_type = arch_mma_shape_and_type[arch] + supported_shape_and_type = mma_shape_and_type[arch] assert shape_and_type in supported_shape_and_type, ( f"Unsupported shape_and_type {shape_and_type} for mma on {arch} architecture.\n" f"Supported shape_and_type: {supported_shape_and_type}" ) - shape, d_type, a_type, b_type, c_type = shape_and_type.split(".") + qualifiers = shape_and_type.split(".") + self.kind = None + if len(qualifiers) == 6: + shape, kind, d_type, a_type, b_type, c_type = qualifiers + self.kind = kind.removeprefix("kind::") + else: + shape, d_type, a_type, b_type, c_type = qualifiers self.arch = arch self.shape_and_type = shape_and_type self.m, self.n, self.k = nv_shape_to_mnk(shape) @@ -144,23 +143,25 @@ def __init__(self, arch: str, shape_and_type: str): class MMABlockScale(MMABlockScaleOperation): def __init__(self, arch: str, shape_and_type: str): - assert arch in arch_mma_block_scale_shape_and_type.keys(), ( - f"Unsupported architecture {arch} for mma.block_scale.\n" - f"Supported architectures: {list(arch_mma_block_scale_shape_and_type.keys())}" + assert arch in mma_block_scale_shape_and_type.keys(), ( + f"Unsupported architecture {arch} for block scaled mma.\n" + f"Supported architectures: {list(mma_block_scale_shape_and_type.keys())}" ) - supported_shape_and_type = arch_mma_block_scale_shape_and_type[arch] + supported_shape_and_type = mma_block_scale_shape_and_type[arch] assert shape_and_type in supported_shape_and_type, ( - f"Unsupported shape_and_type {shape_and_type} for mma.block_scale on {arch} architecture.\n" + f"Unsupported shape_and_type {shape_and_type} for block scaled mma on {arch} architecture.\n" f"Supported shape_and_type: {supported_shape_and_type}" ) - shape, block_size, d_type, a_type, b_type, c_type, s_type = ( + shape, kind, block_size, d_type, a_type, b_type, c_type, s_type = ( shape_and_type.split(".") ) + kind = kind.removeprefix("kind::") self.arch = arch self.shape_and_type = shape_and_type self.m, self.n, self.k = nv_shape_to_mnk(shape) - self.block_size = int(block_size[-2:]) - self.packing = 2 if self.k == 64 else 1 + self.kind = kind + self.block_size = int(block_size.removeprefix("block")) + self.packing = 2 if kind in ("mxf4", "mxf4nvf4") else 1 self.a_type = nv_torch_dtype[a_type] self.b_type = nv_torch_dtype[b_type] self.a_type_ptx = a_type diff --git a/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py b/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py index c89e167..ac33ef9 100644 --- a/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py +++ b/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py @@ -2,38 +2,85 @@ MMAOperation, MMABlockScaleOperation, ) -from .isa_common import nv_torch_dtype, nv_shape_to_mnk +from .isa_common import ( + nv_torch_dtype, + f8f6f4_types, + nv_shape_to_mnk, +) # TODO: support tcgen05mma with .ws and .cta_group::2 +# TODO: support sm_103/sm_107 + +_tcgen05_mn = [(m, n) for m in (64, 128) for n in range(8, 257, 8)] +blackwell_tcgen05mma_shape_and_type = ( + [f"m{m}n{n}k8.kind::tf32.f32.tf32.tf32" for m, n in _tcgen05_mn] + + [ + f"m{m}n{n}k16.kind::f16.{d_type}.{a_type}.{b_type}" + for m, n in _tcgen05_mn + for d_type in ("f32", "f16") + for a_type in ("f16", "bf16") + for b_type in ("f16", "bf16") + if d_type == "f32" or a_type == b_type == "f16" + ] + + [ + f"m{m}n{n}k32.kind::f8f6f4.{d_type}.{a_type}.{b_type}" + for m, n in _tcgen05_mn + for d_type in ("f32", "f16") + for a_type in f8f6f4_types + for b_type in f8f6f4_types + ] +) +tcgen05mma_shape_and_type = { + "Blackwell": blackwell_tcgen05mma_shape_and_type, +} + +_tcgen05_block_scale_n = list(range(8, 257, 8)) +blackwell_tcgen05mma_block_scale_shape_and_type = ( + [ + f"m128n{n}k32.kind::mxf8f6f4.block32.f32.{a_type}.{b_type}.ue8m0" + for n in _tcgen05_block_scale_n + for a_type in f8f6f4_types + for b_type in f8f6f4_types + ] + + [ + f"m128n{n}k64.kind::mxf4.block32.f32.e2m1.e2m1.ue8m0" + for n in _tcgen05_block_scale_n + ] + + [ + f"m128n{n}k64.kind::mxf4nvf4.block32.f32.e2m1.e2m1.f32.ue8m0" + for n in _tcgen05_block_scale_n + ] + + [ + f"m128n{n}k64.kind::mxf4nvf4.block16.f32.e2m1.e2m1.f32.{s_type}" + for n in _tcgen05_block_scale_n + for s_type in ["ue8m0", "ue4m3"] + ] +) +tcgen05mma_block_scale_shape_and_type = { + "Blackwell": blackwell_tcgen05mma_block_scale_shape_and_type, +} class TCGen05MMA(MMAOperation): def __init__(self, arch: str, shape_and_type: str): - assert arch == "Blackwell" + assert arch in tcgen05mma_shape_and_type.keys(), ( + f"Unsupported architecture {arch} for tcgen05.mma.\n" + f"Supported architectures: {list(tcgen05mma_shape_and_type.keys())}" + ) + supported_shape_and_type = tcgen05mma_shape_and_type[arch] + assert shape_and_type in supported_shape_and_type, ( + f"Unsupported shape_and_type {shape_and_type} for tcgen05.mma on {arch} architecture.\n" + "Supported shape_and_type format: m{m}n{n}k{k}.kind::{kind}.{d_type}.{a_type}.{b_type}\n" + "Refer to the PTX documentation for valid combinations." + ) qualifiers = shape_and_type.split(".") - assert len(qualifiers) == 4 - shape, d_type, a_type, b_type = qualifiers + shape, kind, d_type, a_type, b_type = qualifiers m, n, k = nv_shape_to_mnk(shape) - assert m in [64, 128] - assert n in list(range(8, 256 + 1, 8)) - assert k in [8, 16, 32] - if k == 8: # kind::tf32 - assert a_type == b_type == "tf32" - assert d_type == "f32" - elif k == 16: # kind::f16 - assert d_type in ["f32", "f16"] - if d_type == "f32": - assert a_type in ["f16", "bf16"] - assert b_type in ["f16", "bf16"] - else: # d_type == "f16" - assert a_type == b_type == "f16" - else: # k == 32, kind::f8f6f4 - assert a_type in ["e5m2", "e4m3", "e3m2", "e2m3", "e2m1"] - assert b_type in ["e5m2", "e4m3", "e3m2", "e2m3", "e2m1"] - assert d_type in ["f32", "f16"] + kind = kind.removeprefix("kind::") self.arch = arch self.shape_and_type = shape_and_type self.m, self.n, self.k = m, n, k + self.kind = kind self.a_type = nv_torch_dtype[a_type] self.b_type = nv_torch_dtype[b_type] self.a_type_ptx = a_type @@ -44,33 +91,27 @@ def __init__(self, arch: str, shape_and_type: str): class TCGen05MMABlockScale(MMABlockScaleOperation): def __init__(self, arch: str, shape_and_type: str): - assert arch == "Blackwell" + assert arch in tcgen05mma_block_scale_shape_and_type.keys(), ( + f"Unsupported architecture {arch} for block scaled tcgen05.mma.\n" + f"Supported architectures: {list(tcgen05mma_block_scale_shape_and_type.keys())}" + ) + supported_shape_and_type = tcgen05mma_block_scale_shape_and_type[arch] + assert shape_and_type in supported_shape_and_type, ( + f"Unsupported shape_and_type {shape_and_type} for block scaled tcgen05.mma on {arch} architecture.\n" + "Supported shape_and_type format: m{m}n{n}k{k}.kind::{kind}.block{block_size}.{d_type}.{a_type}.{b_type}.{s_type}\n" + "Refer to the PTX documentation for valid combinations." + ) qualifiers = shape_and_type.split(".") - assert len(qualifiers) == 6 - shape, block_size, d_type, a_type, b_type, s_type = qualifiers + shape, kind, block_size, d_type, a_type, b_type, s_type = qualifiers m, n, k = nv_shape_to_mnk(shape) - assert m == 128 - assert n in list(range(8, 256 + 1, 8)) - assert k in [32, 64] - assert d_type == "f32" - if k == 32: # kind::mxf8f6f4 - 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 - assert a_type == b_type == "e2m1" - assert d_type == "f32" - assert s_type in ["ue8m0", "ue4m3"] - if s_type == "ue8m0": - assert block_size in ["block32", "block16"] - else: # s_type == "ue4m3" - assert block_size == "block16" + kind = kind.removeprefix("kind::") + block = int(block_size.removeprefix("block")) self.arch = arch self.shape_and_type = shape_and_type self.m, self.n, self.k = m, n, k - self.block_size = int(block_size[-2:]) - self.packing = 2 if k == 64 else 1 + self.kind = kind + self.block_size = block + self.packing = 2 if kind in ("mxf4", "mxf4nvf4") else 1 self.a_type = nv_torch_dtype[a_type] self.b_type = nv_torch_dtype[b_type] self.a_type_ptx = a_type diff --git a/src/mmasim/mmasim/nv_ptx/isa_wgmma.py b/src/mmasim/mmasim/nv_ptx/isa_wgmma.py index 6f4e9f9..b1554b2 100644 --- a/src/mmasim/mmasim/nv_ptx/isa_wgmma.py +++ b/src/mmasim/mmasim/nv_ptx/isa_wgmma.py @@ -1,32 +1,39 @@ from .. import MMAOperation from .isa_common import nv_torch_dtype, nv_shape_to_mnk +_wgmma_n = list(range(8, 257, 8)) +hopper_wgmma_shape_and_type = ( + [f"m64n{n}k8.f32.tf32.tf32" for n in _wgmma_n] + + [f"m64n{n}k16.f32.f16.f16" for n in _wgmma_n] + + [f"m64n{n}k16.f16.f16.f16" for n in _wgmma_n] + + [f"m64n{n}k16.f32.bf16.bf16" for n in _wgmma_n] + + [ + f"m64n{n}k32.{d_type}.{a_type}.{b_type}" + for n in _wgmma_n + for d_type in ("f32", "f16") + for a_type in ("e5m2", "e4m3") + for b_type in ("e5m2", "e4m3") + ] +) +wgmma_shape_and_type = { + "Hopper": hopper_wgmma_shape_and_type, +} + class WGMMA(MMAOperation): def __init__(self, arch: str, shape_and_type: str): - assert arch == "Hopper" - qualifiers = shape_and_type.split(".") - assert len(qualifiers) == 4 - shape, d_type, a_type, b_type = qualifiers + assert arch in wgmma_shape_and_type, ( + f"Unsupported architecture {arch} for wgmma.\n" + f"Supported architectures: {list(wgmma_shape_and_type.keys())}" + ) + supported_shape_and_type = wgmma_shape_and_type[arch] + assert shape_and_type in supported_shape_and_type, ( + f"Unsupported shape_and_type {shape_and_type} for wgmma on {arch} architecture.\n" + "Supported shape_and_type format: m{m}n{n}k{k}.{d_type}.{a_type}.{b_type}\n" + "Refer to the PTX documentation for valid combinations." + ) + shape, d_type, a_type, b_type = shape_and_type.split(".") m, n, k = nv_shape_to_mnk(shape) - assert m == 64 - assert n in list(range(8, 256 + 1, 8)) - assert k in [8, 16, 32] - if k == 8: - assert a_type == b_type == "tf32" - assert d_type == "f32" - elif k == 16: - assert a_type in ["f16", "bf16"] - if a_type == "f16": - assert b_type == "f16" - assert d_type in ["f32", "f16"] - else: # a_type == "bf16" - assert b_type == "bf16" - assert d_type == "f32" - else: # k == 32 - assert a_type in ["e5m2", "e4m3"] - assert b_type in ["e5m2", "e4m3"] - 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 diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/ada_lovelace.py b/src/mmasim_kernels/mmasim_kernels/nv_ptx/ada_lovelace.py index 43a7f91..af5ec81 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/ada_lovelace.py +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/ada_lovelace.py @@ -2,6 +2,7 @@ from pathlib import Path from torch.utils.cpp_extension import load +from mmasim.nv_ptx.isa_mma import ada_lovelace_mma_shape_and_type from . import MMA dir = Path(__file__).with_name("kernels") @@ -14,81 +15,8 @@ ) lib = ctypes.CDLL(kernel_path) -# sm_89 fp8 m16n8k32 f32_output -lib.mma_m16n8k32_f32_e5m2_e5m2_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f32_e5m2_e4m3_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f32_e4m3_e5m2_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f32_e4m3_e4m3_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_89 fp8 m16n8k16 f32_output -lib.mma_m16n8k16_f32_e5m2_e5m2_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f32_e5m2_e4m3_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f32_e4m3_e5m2_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f32_e4m3_e4m3_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_89 fp8 m16n8k32 f16_output -lib.mma_m16n8k32_f16_e5m2_e5m2_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f16_e5m2_e4m3_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f16_e4m3_e5m2_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f16_e4m3_e4m3_f16.argtypes = [ctypes.c_void_p] * 4 -# sm_89 fp8 m16n8k16 f16_output -lib.mma_m16n8k16_f16_e5m2_e5m2_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f16_e5m2_e4m3_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f16_e4m3_e5m2_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f16_e4m3_e4m3_f16.argtypes = [ctypes.c_void_p] * 4 -# sm_80 f64 -lib.mma_m8n8k4_f64_f64_f64_f64.argtypes = [ctypes.c_void_p] * 4 -# sm_80 tf32 -lib.mma_m16n8k8_f32_tf32_tf32_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k4_f32_tf32_tf32_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_80 f16 -lib.mma_m16n8k16_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 -# sm_80 bf16 -lib.mma_m16n8k16_f32_bf16_bf16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f32_bf16_bf16_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_75 f16 -lib.mma_m16n8k8_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 - -mma_cuda_kernels = { - # sm_89 fp8 m16n8k32 f32_output - "m16n8k32.f32.e5m2.e5m2.f32": lib.mma_m16n8k32_f32_e5m2_e5m2_f32, - "m16n8k32.f32.e5m2.e4m3.f32": lib.mma_m16n8k32_f32_e5m2_e4m3_f32, - "m16n8k32.f32.e4m3.e5m2.f32": lib.mma_m16n8k32_f32_e4m3_e5m2_f32, - "m16n8k32.f32.e4m3.e4m3.f32": lib.mma_m16n8k32_f32_e4m3_e4m3_f32, - # sm_89 fp8 m16n8k16 f32_output - "m16n8k16.f32.e5m2.e5m2.f32": lib.mma_m16n8k16_f32_e5m2_e5m2_f32, - "m16n8k16.f32.e5m2.e4m3.f32": lib.mma_m16n8k16_f32_e5m2_e4m3_f32, - "m16n8k16.f32.e4m3.e5m2.f32": lib.mma_m16n8k16_f32_e4m3_e5m2_f32, - "m16n8k16.f32.e4m3.e4m3.f32": lib.mma_m16n8k16_f32_e4m3_e4m3_f32, - # sm_89 fp8 m16n8k32 f16_output - "m16n8k32.f16.e5m2.e5m2.f16": lib.mma_m16n8k32_f16_e5m2_e5m2_f16, - "m16n8k32.f16.e5m2.e4m3.f16": lib.mma_m16n8k32_f16_e5m2_e4m3_f16, - "m16n8k32.f16.e4m3.e5m2.f16": lib.mma_m16n8k32_f16_e4m3_e5m2_f16, - "m16n8k32.f16.e4m3.e4m3.f16": lib.mma_m16n8k32_f16_e4m3_e4m3_f16, - # sm_89 fp8 m16n8k16 f16_output - "m16n8k16.f16.e5m2.e5m2.f16": lib.mma_m16n8k16_f16_e5m2_e5m2_f16, - "m16n8k16.f16.e5m2.e4m3.f16": lib.mma_m16n8k16_f16_e5m2_e4m3_f16, - "m16n8k16.f16.e4m3.e5m2.f16": lib.mma_m16n8k16_f16_e4m3_e5m2_f16, - "m16n8k16.f16.e4m3.e4m3.f16": lib.mma_m16n8k16_f16_e4m3_e4m3_f16, - # sm_80 f64 - "m8n8k4.f64.f64.f64.f64": lib.mma_m8n8k4_f64_f64_f64_f64, - # sm_80 tf32 - "m16n8k8.f32.tf32.tf32.f32": lib.mma_m16n8k8_f32_tf32_tf32_f32, - "m16n8k4.f32.tf32.tf32.f32": lib.mma_m16n8k4_f32_tf32_tf32_f32, - # sm_80 f16 - "m16n8k16.f32.f16.f16.f32": lib.mma_m16n8k16_f32_f16_f16_f32, - "m16n8k16.f16.f16.f16.f16": lib.mma_m16n8k16_f16_f16_f16_f16, - # sm_80 bf16 - "m16n8k16.f32.bf16.bf16.f32": lib.mma_m16n8k16_f32_bf16_bf16_f32, - "m16n8k8.f32.bf16.bf16.f32": lib.mma_m16n8k8_f32_bf16_bf16_f32, - # sm_75 f16 - "m16n8k8.f32.f16.f16.f32": lib.mma_m16n8k8_f32_f16_f16_f32, - "m16n8k8.f16.f16.f16.f16": lib.mma_m16n8k8_f16_f16_f16_f16, -} - -mma_kernels = { - shape_and_type: MMA( - "Ada Lovelace", shape_and_type, mma_cuda_kernels[shape_and_type] - ) - for shape_and_type in mma_cuda_kernels -} +mma_kernels = {} +for shape_and_type in ada_lovelace_mma_shape_and_type: + kernel = lib.__getattr__(f"mma_{shape_and_type.replace('.', '_')}") + kernel.argtypes = [ctypes.c_void_p] * 4 + mma_kernels[shape_and_type] = MMA("Ada Lovelace", shape_and_type, kernel) diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/ampere.py b/src/mmasim_kernels/mmasim_kernels/nv_ptx/ampere.py index 84148bd..19d9b66 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/ampere.py +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/ampere.py @@ -2,6 +2,7 @@ from pathlib import Path from torch.utils.cpp_extension import load +from mmasim.nv_ptx.isa_mma import ampere_mma_shape_and_type from . import MMA dir = Path(__file__).with_name("kernels") @@ -14,39 +15,8 @@ ) lib = ctypes.CDLL(kernel_path) -# sm_80 f64 -lib.mma_m8n8k4_f64_f64_f64_f64.argtypes = [ctypes.c_void_p] * 4 -# sm_80 tf32 -lib.mma_m16n8k8_f32_tf32_tf32_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k4_f32_tf32_tf32_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_80 f16 -lib.mma_m16n8k16_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 -# sm_80 bf16 -lib.mma_m16n8k16_f32_bf16_bf16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f32_bf16_bf16_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_75 f16 -lib.mma_m16n8k8_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 - -mma_cuda_kernels = { - # sm_80 f64 - "m8n8k4.f64.f64.f64.f64": lib.mma_m8n8k4_f64_f64_f64_f64, - # sm_80 tf32 - "m16n8k8.f32.tf32.tf32.f32": lib.mma_m16n8k8_f32_tf32_tf32_f32, - "m16n8k4.f32.tf32.tf32.f32": lib.mma_m16n8k4_f32_tf32_tf32_f32, - # sm_80 f16 - "m16n8k16.f32.f16.f16.f32": lib.mma_m16n8k16_f32_f16_f16_f32, - "m16n8k16.f16.f16.f16.f16": lib.mma_m16n8k16_f16_f16_f16_f16, - # sm_80 bf16 - "m16n8k16.f32.bf16.bf16.f32": lib.mma_m16n8k16_f32_bf16_bf16_f32, - "m16n8k8.f32.bf16.bf16.f32": lib.mma_m16n8k8_f32_bf16_bf16_f32, - # sm_75 f16 - "m16n8k8.f32.f16.f16.f32": lib.mma_m16n8k8_f32_f16_f16_f32, - "m16n8k8.f16.f16.f16.f16": lib.mma_m16n8k8_f16_f16_f16_f16, -} - -mma_kernels = { - shape_and_type: MMA("Ampere", shape_and_type, mma_cuda_kernels[shape_and_type]) - for shape_and_type in mma_cuda_kernels -} +mma_kernels = {} +for shape_and_type in ampere_mma_shape_and_type: + kernel = lib.__getattr__(f"mma_{shape_and_type.replace('.', '_')}") + kernel.argtypes = [ctypes.c_void_p] * 4 + mma_kernels[shape_and_type] = MMA("Ampere", shape_and_type, kernel) diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/blackwell.py b/src/mmasim_kernels/mmasim_kernels/nv_ptx/blackwell.py index 0bb22b6..f777a1c 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/blackwell.py +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/blackwell.py @@ -2,6 +2,12 @@ from pathlib import Path from torch.utils.cpp_extension import load +from mmasim.nv_ptx.isa_common import nv_shape_to_mnk +from mmasim.nv_ptx.isa_mma import blackwell_mma_shape_and_type +from mmasim.nv_ptx.isa_tcgen05mma import ( + blackwell_tcgen05mma_shape_and_type, + blackwell_tcgen05mma_block_scale_shape_and_type, +) from . import MMA, TCGen05MMA, TCGen05MMABlockScale dir = Path(__file__).with_name("kernels") @@ -14,119 +20,41 @@ ) lib = ctypes.CDLL(kernel_path) -# sm_80 f64 -lib.mma_m8n8k4_f64_f64_f64_f64.argtypes = [ctypes.c_void_p] * 4 -# sm_80 tf32 -lib.mma_m16n8k8_f32_tf32_tf32_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k4_f32_tf32_tf32_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_80 f16 -lib.mma_m16n8k16_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 -# sm_80 bf16 -lib.mma_m16n8k16_f32_bf16_bf16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f32_bf16_bf16_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_75 f16 -lib.mma_m16n8k8_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 +mma_kernels = {} +for shape_and_type in blackwell_mma_shape_and_type: + kernel = lib.__getattr__( + f"mma_{shape_and_type.replace('.', '_').replace('::', '_')}" + ) + kernel.argtypes = [ctypes.c_void_p] * 4 + mma_kernels[shape_and_type] = MMA("Blackwell", shape_and_type, kernel) -# sm_100a tf32 -lib.tcgen05mma_m64n8k8_f32_tf32_tf32.argtypes = [ctypes.c_void_p] * 3 -# sm_100a f16 and bf16 -lib.tcgen05mma_m64n8k16_f32_f16_f16.argtypes = [ctypes.c_void_p] * 3 -lib.tcgen05mma_m64n8k16_f32_bf16_bf16.argtypes = [ctypes.c_void_p] * 3 -lib.tcgen05mma_m64n8k16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 3 -# 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 -lib.tcgen05mma_m128n8k64_block16_f32_e2m1_e2m1_ue4m3.argtypes = [ctypes.c_void_p] * 5 +def _bind_shape(kernel, m, n, pointer_count): + """Bind runtime descriptor dimensions while keeping the common kernel ABI.""" + kernel.argtypes = [ctypes.c_void_p] * pointer_count + [ctypes.c_uint32] * 2 + kernel.restype = None -mma_cuda_kernels = { - # sm_80 f64 - "m8n8k4.f64.f64.f64.f64": lib.mma_m8n8k4_f64_f64_f64_f64, - # sm_80 tf32 - "m16n8k8.f32.tf32.tf32.f32": lib.mma_m16n8k8_f32_tf32_tf32_f32, - "m16n8k4.f32.tf32.tf32.f32": lib.mma_m16n8k4_f32_tf32_tf32_f32, - # sm_80 f16 - "m16n8k16.f32.f16.f16.f32": lib.mma_m16n8k16_f32_f16_f16_f32, - "m16n8k16.f16.f16.f16.f16": lib.mma_m16n8k16_f16_f16_f16_f16, - # sm_80 bf16 - "m16n8k16.f32.bf16.bf16.f32": lib.mma_m16n8k16_f32_bf16_bf16_f32, - "m16n8k8.f32.bf16.bf16.f32": lib.mma_m16n8k8_f32_bf16_bf16_f32, - # sm_75 f16 - "m16n8k8.f32.f16.f16.f32": lib.mma_m16n8k8_f32_f16_f16_f32, - "m16n8k8.f16.f16.f16.f16": lib.mma_m16n8k8_f16_f16_f16_f16, -} -tcgen05mma_cuda_kernels = { - # sm_100a tf32 - "m64n8k8.f32.tf32.tf32": lib.tcgen05mma_m64n8k8_f32_tf32_tf32, - # sm_100a f16 and bf16 - "m64n8k16.f32.f16.f16": lib.tcgen05mma_m64n8k16_f32_f16_f16, - "m64n8k16.f32.bf16.bf16": lib.tcgen05mma_m64n8k16_f32_bf16_bf16, - "m64n8k16.f16.f16.f16": lib.tcgen05mma_m64n8k16_f16_f16_f16, - # 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, - "m128n8k64.block16.f32.e2m1.e2m1.ue4m3": lib.tcgen05mma_m128n8k64_block16_f32_e2m1_e2m1_ue4m3, -} + def bound(*pointers): + return kernel(*pointers, m, n) -mma_kernels = { - shape_and_type: MMA("Blackwell", shape_and_type, mma_cuda_kernels[shape_and_type]) - for shape_and_type in mma_cuda_kernels -} -tcgen05mma_kernels = { - shape_and_type: TCGen05MMA( - "Blackwell", shape_and_type, tcgen05mma_cuda_kernels[shape_and_type] - ) - for shape_and_type in tcgen05mma_cuda_kernels -} -tcgen05mma_block_scale_kernels = { - shape_and_type: TCGen05MMABlockScale( - "Blackwell", shape_and_type, tcgen05mma_block_scale_cuda_kernels[shape_and_type] + bound.argtypes = kernel.argtypes[:pointer_count] + return bound + + +tcgen05mma_kernels = {} +for shape_and_type in blackwell_tcgen05mma_shape_and_type: + shape, type = shape_and_type.split(".", 1) + m, n, _ = nv_shape_to_mnk(shape) + kernel = lib.__getattr__(f"tcgen05mma_{type.replace('.', '_').replace('::', '_')}") + kernel = _bind_shape(kernel, m, n, 3) + tcgen05mma_kernels[shape_and_type] = TCGen05MMA("Blackwell", shape_and_type, kernel) + +tcgen05mma_block_scale_kernels = {} +for shape_and_type in blackwell_tcgen05mma_block_scale_shape_and_type: + shape, type = shape_and_type.split(".", 1) + m, n, _ = nv_shape_to_mnk(shape) + kernel = lib.__getattr__(f"tcgen05mma_{type.replace('.', '_').replace('::', '_')}") + kernel = _bind_shape(kernel, m, n, 5) + tcgen05mma_block_scale_kernels[shape_and_type] = TCGen05MMABlockScale( + "Blackwell", shape_and_type, kernel ) - for shape_and_type in tcgen05mma_block_scale_cuda_kernels -} diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/hopper.py b/src/mmasim_kernels/mmasim_kernels/nv_ptx/hopper.py index 16af37d..efa6e7e 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/hopper.py +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/hopper.py @@ -2,6 +2,8 @@ from pathlib import Path from torch.utils.cpp_extension import load +from mmasim.nv_ptx.isa_mma import hopper_mma_shape_and_type +from mmasim.nv_ptx.isa_wgmma import hopper_wgmma_shape_and_type from . import MMA, WGMMA dir = Path(__file__).with_name("kernels") @@ -16,88 +18,14 @@ ) lib = ctypes.CDLL(kernel_path) -# sm_90 f64 -lib.mma_m16n8k16_f64_f64_f64_f64.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f64_f64_f64_f64.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k4_f64_f64_f64_f64.argtypes = [ctypes.c_void_p] * 4 -# sm_80 f64 -lib.mma_m8n8k4_f64_f64_f64_f64.argtypes = [ctypes.c_void_p] * 4 -# sm_80 tf32 -lib.mma_m16n8k8_f32_tf32_tf32_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k4_f32_tf32_tf32_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_80 f16 -lib.mma_m16n8k16_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 -# sm_80 bf16 -lib.mma_m16n8k16_f32_bf16_bf16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f32_bf16_bf16_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_75 f16 -lib.mma_m16n8k8_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 +mma_kernels = {} +for shape_and_type in hopper_mma_shape_and_type: + kernel = lib.__getattr__(f"mma_{shape_and_type.replace('.', '_')}") + kernel.argtypes = [ctypes.c_void_p] * 4 + mma_kernels[shape_and_type] = MMA("Hopper", shape_and_type, kernel) -# sm_90a tf32 -lib.wgmma_m64n8k8_f32_tf32_tf32.argtypes = [ctypes.c_void_p] * 3 -# sm_90a f16 -lib.wgmma_m64n8k16_f32_f16_f16.argtypes = [ctypes.c_void_p] * 3 -lib.wgmma_m64n8k16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 3 -# sm_90a bf16 -lib.wgmma_m64n8k16_f32_bf16_bf16.argtypes = [ctypes.c_void_p] * 3 -# sm_90a fp8 f32_output -lib.wgmma_m64n8k32_f32_e5m2_e5m2.argtypes = [ctypes.c_void_p] * 3 -lib.wgmma_m64n8k32_f32_e5m2_e4m3.argtypes = [ctypes.c_void_p] * 3 -lib.wgmma_m64n8k32_f32_e4m3_e5m2.argtypes = [ctypes.c_void_p] * 3 -lib.wgmma_m64n8k32_f32_e4m3_e4m3.argtypes = [ctypes.c_void_p] * 3 -# sm_90a fp8 f16_output -lib.wgmma_m64n8k32_f16_e5m2_e5m2.argtypes = [ctypes.c_void_p] * 3 -lib.wgmma_m64n8k32_f16_e5m2_e4m3.argtypes = [ctypes.c_void_p] * 3 -lib.wgmma_m64n8k32_f16_e4m3_e5m2.argtypes = [ctypes.c_void_p] * 3 -lib.wgmma_m64n8k32_f16_e4m3_e4m3.argtypes = [ctypes.c_void_p] * 3 - -mma_cuda_kernels = { - # sm_90 f64 - "m16n8k16.f64.f64.f64.f64": lib.mma_m16n8k16_f64_f64_f64_f64, - "m16n8k8.f64.f64.f64.f64": lib.mma_m16n8k8_f64_f64_f64_f64, - "m16n8k4.f64.f64.f64.f64": lib.mma_m16n8k4_f64_f64_f64_f64, - # sm_80 f64 - "m8n8k4.f64.f64.f64.f64": lib.mma_m8n8k4_f64_f64_f64_f64, - # sm_80 tf32 - "m16n8k8.f32.tf32.tf32.f32": lib.mma_m16n8k8_f32_tf32_tf32_f32, - "m16n8k4.f32.tf32.tf32.f32": lib.mma_m16n8k4_f32_tf32_tf32_f32, - # sm_80 f16 - "m16n8k16.f32.f16.f16.f32": lib.mma_m16n8k16_f32_f16_f16_f32, - "m16n8k16.f16.f16.f16.f16": lib.mma_m16n8k16_f16_f16_f16_f16, - # sm_80 bf16 - "m16n8k16.f32.bf16.bf16.f32": lib.mma_m16n8k16_f32_bf16_bf16_f32, - "m16n8k8.f32.bf16.bf16.f32": lib.mma_m16n8k8_f32_bf16_bf16_f32, - # sm_75 f16 - "m16n8k8.f32.f16.f16.f32": lib.mma_m16n8k8_f32_f16_f16_f32, - "m16n8k8.f16.f16.f16.f16": lib.mma_m16n8k8_f16_f16_f16_f16, -} -wgmma_cuda_kernels = { - # sm_90a tf32 - "m64n8k8.f32.tf32.tf32": lib.wgmma_m64n8k8_f32_tf32_tf32, - # sm_90a f16 - "m64n8k16.f32.f16.f16": lib.wgmma_m64n8k16_f32_f16_f16, - "m64n8k16.f16.f16.f16": lib.wgmma_m64n8k16_f16_f16_f16, - # sm_90a bf16 - "m64n8k16.f32.bf16.bf16": lib.wgmma_m64n8k16_f32_bf16_bf16, - # sm_90a fp8 f32_output - "m64n8k32.f32.e5m2.e5m2": lib.wgmma_m64n8k32_f32_e5m2_e5m2, - "m64n8k32.f32.e5m2.e4m3": lib.wgmma_m64n8k32_f32_e5m2_e4m3, - "m64n8k32.f32.e4m3.e5m2": lib.wgmma_m64n8k32_f32_e4m3_e5m2, - "m64n8k32.f32.e4m3.e4m3": lib.wgmma_m64n8k32_f32_e4m3_e4m3, - # sm_90a fp8 f16_output - "m64n8k32.f16.e5m2.e5m2": lib.wgmma_m64n8k32_f16_e5m2_e5m2, - "m64n8k32.f16.e5m2.e4m3": lib.wgmma_m64n8k32_f16_e5m2_e4m3, - "m64n8k32.f16.e4m3.e5m2": lib.wgmma_m64n8k32_f16_e4m3_e5m2, - "m64n8k32.f16.e4m3.e4m3": lib.wgmma_m64n8k32_f16_e4m3_e4m3, -} - -mma_kernels = { - shape_and_type: MMA("Hopper", shape_and_type, mma_cuda_kernels[shape_and_type]) - for shape_and_type in mma_cuda_kernels -} -wgmma_kernels = { - shape_and_type: WGMMA("Hopper", shape_and_type, wgmma_cuda_kernels[shape_and_type]) - for shape_and_type in wgmma_cuda_kernels -} +wgmma_kernels = {} +for shape_and_type in hopper_wgmma_shape_and_type: + kernel = lib.__getattr__(f"wgmma_{shape_and_type.replace('.', '_')}") + kernel.argtypes = [ctypes.c_void_p] * 3 + wgmma_kernels[shape_and_type] = WGMMA("Hopper", shape_and_type, kernel) diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/mma.h b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/mma.h index 002d85a..60115e7 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/mma.h +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/mma.h @@ -107,6 +107,23 @@ } \ } while (0) +// Nonzero low FP4 padding bits can be consumed as additional mantissa bits by the hardware. +#define F8F6F4_MASK_e5m2 0xffffffffu +#define F8F6F4_MASK_e4m3 0xffffffffu +#define F8F6F4_MASK_e3m2 0x3f3f3f3fu +#define F8F6F4_MASK_e2m3 0x3f3f3f3fu +#define F8F6F4_MASK_e2m1 0x3c3c3c3cu + +#define LOAD_A_M16K32_F8F6F4(mask) \ + do \ + { \ + LOAD_A_M16K32_FP8(); \ + for (uint32_t i = 0; i < 4; i++) \ + { \ + a_frag[i] &= (mask); \ + } \ + } while (0) + #define LOAD_A_M16K16_FP8() \ do \ { \ @@ -219,6 +236,16 @@ } \ } while (0) +#define LOAD_B_N8K32_F8F6F4(mask) \ + do \ + { \ + LOAD_B_N8K32_FP8(); \ + for (uint32_t i = 0; i < 2; i++) \ + { \ + b_frag[i] &= (mask); \ + } \ + } while (0) + #define LOAD_B_N8K16_FP8() \ do \ { \ diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/mma_sm120a.cu b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/mma_sm120a.cu index efde656..d827312 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/mma_sm120a.cu +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/mma_sm120a.cu @@ -1,28 +1,147 @@ #include "mma.h" -extern "C" // f8f6f4 +extern "C" int mmasim_cuda_version() { return CUDART_VERSION; } + +extern "C" // kind::f8f6f4 { - // .kind::f8f6f4 seems to behave the same as mma without .kind for fp8 +#define DEFINE_MMA_F8F6F4_F32(atype, btype) \ + __global__ void mma_m16n8k32_kind_f8f6f4_f32_##atype##_##btype##_f32_kernel( \ + float *d, uint8_t *a, uint8_t *b, float *c) \ + { \ + const uint32_t N = 8, K = 32; \ + uint32_t laneid = threadIdx.x; \ + uint32_t a_frag[4], b_frag[2]; \ + float c_frag[4], d_frag[4]; \ + LOAD_A_M16K32_F8F6F4(F8F6F4_MASK_##atype); \ + LOAD_B_N8K32_F8F6F4(F8F6F4_MASK_##btype); \ + LOAD_C_M16N8_F32(); \ + asm volatile( \ + "mma.sync.aligned.m16n8k32.row.col.kind::f8f6f4.f32." #atype "." #btype ".f32 " \ + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};" \ + : "=f"(d_frag[0]), "=f"(d_frag[1]), "=f"(d_frag[2]), "=f"(d_frag[3]) \ + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), \ + "r"(b_frag[0]), "r"(b_frag[1]), \ + "f"(c_frag[0]), "f"(c_frag[1]), "f"(c_frag[2]), "f"(c_frag[3])); \ + STORE_D_M16N8_F32(); \ + } \ + void mma_m16n8k32_kind_f8f6f4_f32_##atype##_##btype##_f32( \ + float *d, uint8_t *a, uint8_t *b, float *c) \ + { \ + mma_m16n8k32_kind_f8f6f4_f32_##atype##_##btype##_f32_kernel<<<1, 32>>>(d, a, b, c); \ + } + +#define DEFINE_MMA_F8F6F4_F16(atype, btype) \ + __global__ void mma_m16n8k32_kind_f8f6f4_f16_##atype##_##btype##_f16_kernel( \ + uint16_t *d, uint8_t *a, uint8_t *b, uint16_t *c) \ + { \ + const uint32_t N = 8, K = 32; \ + uint32_t laneid = threadIdx.x; \ + uint32_t a_frag[4], b_frag[2]; \ + uint32_t c_frag[2], d_frag[2]; \ + LOAD_A_M16K32_F8F6F4(F8F6F4_MASK_##atype); \ + LOAD_B_N8K32_F8F6F4(F8F6F4_MASK_##btype); \ + LOAD_C_M16N8_F16(); \ + asm volatile( \ + "mma.sync.aligned.m16n8k32.row.col.kind::f8f6f4.f16." #atype "." #btype ".f16 " \ + "{%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%8, %9};" \ + : "=r"(d_frag[0]), "=r"(d_frag[1]) \ + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), \ + "r"(b_frag[0]), "r"(b_frag[1]), \ + "r"(c_frag[0]), "r"(c_frag[1])); \ + STORE_D_M16N8_F16(); \ + } \ + void mma_m16n8k32_kind_f8f6f4_f16_##atype##_##btype##_f16( \ + uint16_t *d, uint8_t *a, uint8_t *b, uint16_t *c) \ + { \ + mma_m16n8k32_kind_f8f6f4_f16_##atype##_##btype##_f16_kernel<<<1, 32>>>(d, a, b, c); \ + } + +#define DEFINE_MMA_F8F6F4(atype, btype) \ + DEFINE_MMA_F8F6F4_F32(atype, btype) \ + DEFINE_MMA_F8F6F4_F16(atype, btype) + +#define DEFINE_MMA_F8F6F4_FOR_A(atype) \ + DEFINE_MMA_F8F6F4(atype, e5m2) \ + DEFINE_MMA_F8F6F4(atype, e4m3) \ + DEFINE_MMA_F8F6F4(atype, e3m2) \ + DEFINE_MMA_F8F6F4(atype, e2m3) \ + DEFINE_MMA_F8F6F4(atype, e2m1) + + DEFINE_MMA_F8F6F4_FOR_A(e5m2) + DEFINE_MMA_F8F6F4_FOR_A(e4m3) + DEFINE_MMA_F8F6F4_FOR_A(e3m2) + DEFINE_MMA_F8F6F4_FOR_A(e2m3) + DEFINE_MMA_F8F6F4_FOR_A(e2m1) + +#define DEFINE_MMA_MXF8F6F4(atype, btype) \ + __global__ void mma_m16n8k32_kind_mxf8f6f4_block32_f32_##atype##_##btype##_f32_ue8m0_kernel( \ + float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) \ + { \ + const uint32_t N = 8, K = 32; \ + const uint16_t selector = 0; \ + uint32_t laneid = threadIdx.x; \ + uint32_t a_frag[4], b_frag[2], sfa_frag, sfb_frag; \ + float c_frag[4], d_frag[4]; \ + LOAD_A_M16K32_F8F6F4(F8F6F4_MASK_##atype); \ + LOAD_B_N8K32_F8F6F4(F8F6F4_MASK_##btype); \ + LOAD_C_M16N8_F32(); \ + LOAD_SFA_M16SFK1(); \ + LOAD_SFB_N8SFK1(); \ + asm volatile( \ + "mma.sync.aligned.m16n8k32.row.col.kind::mxf8f6f4.block_scale.scale_vec::1X.f32." #atype "." #btype ".f32.ue8m0 " \ + "{%0, %1, %2, %3}, " \ + "{%4, %5, %6, %7}, " \ + "{%8, %9}, " \ + "{%10, %11, %12, %13}, " \ + "%14, {%15, %16}, " \ + "%17, {%18, %19};" \ + : "=f"(d_frag[0]), "=f"(d_frag[1]), "=f"(d_frag[2]), "=f"(d_frag[3]) \ + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), \ + "r"(b_frag[0]), "r"(b_frag[1]), \ + "f"(c_frag[0]), "f"(c_frag[1]), "f"(c_frag[2]), "f"(c_frag[3]), \ + "r"(sfa_frag), "h"(selector), "h"(selector), \ + "r"(sfb_frag), "h"(selector), "h"(selector)); \ + STORE_D_M16N8_F32(); \ + } \ + void mma_m16n8k32_kind_mxf8f6f4_block32_f32_##atype##_##btype##_f32_ue8m0( \ + float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) \ + { \ + mma_m16n8k32_kind_mxf8f6f4_block32_f32_##atype##_##btype##_f32_ue8m0_kernel<<<1, 32>>>( \ + d, a, b, c, sfa, sfb); \ + } + +#define DEFINE_MMA_MXF8F6F4_FOR_A(atype) \ + DEFINE_MMA_MXF8F6F4(atype, e5m2) \ + DEFINE_MMA_MXF8F6F4(atype, e4m3) \ + DEFINE_MMA_MXF8F6F4(atype, e3m2) \ + DEFINE_MMA_MXF8F6F4(atype, e2m3) \ + DEFINE_MMA_MXF8F6F4(atype, e2m1) + + DEFINE_MMA_MXF8F6F4_FOR_A(e5m2) + DEFINE_MMA_MXF8F6F4_FOR_A(e4m3) + DEFINE_MMA_MXF8F6F4_FOR_A(e3m2) + DEFINE_MMA_MXF8F6F4_FOR_A(e2m3) + DEFINE_MMA_MXF8F6F4_FOR_A(e2m1) } -extern "C" // mxfp8 +extern "C" // kind::mxf4 { - __global__ void mma_m16n8k32_block32_f32_e5m2_e5m2_f32_ue8m0_kernel( + __global__ void mma_m16n8k64_kind_mxf4_block32_f32_e2m1_e2m1_f32_ue8m0_kernel( float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) { - const uint32_t N = 8, K = 32; + const uint32_t N = 8, K = 64; const uint16_t selector = 0; uint32_t laneid = threadIdx.x; uint32_t a_frag[4], b_frag[2], sfa_frag, sfb_frag; float c_frag[4], d_frag[4]; - LOAD_A_M16K32_FP8(); - LOAD_B_N8K32_FP8(); + LOAD_A_M16K64_FP4(); + LOAD_B_N8K64_FP4(); LOAD_C_M16N8_F32(); - LOAD_SFA_M16SFK1(); - LOAD_SFB_N8SFK1(); + LOAD_SFA_M16SFK2(); + LOAD_SFB_N8SFK2(); asm volatile( - "mma.sync.aligned.m16n8k32.row.col.kind::mxf8f6f4.block_scale.scale_vec::1X.f32.e5m2.e5m2.f32.ue8m0 " + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4.block_scale.scale_vec::2X.f32.e2m1.e2m1.f32.ue8m0 " "{%0, %1, %2, %3}, " "{%4, %5, %6, %7}, " "{%8, %9}, " @@ -38,22 +157,31 @@ extern "C" // mxfp8 STORE_D_M16N8_F32(); } - __global__ void mma_m16n8k32_block32_f32_e4m3_e4m3_f32_ue8m0_kernel( + void mma_m16n8k64_kind_mxf4_block32_f32_e2m1_e2m1_f32_ue8m0( + float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) + { + mma_m16n8k64_kind_mxf4_block32_f32_e2m1_e2m1_f32_ue8m0_kernel<<<1, 32>>>(d, a, b, c, sfa, sfb); + } +} + +extern "C" // kind::mxf4nvf4 +{ + __global__ void mma_m16n8k64_kind_mxf4nvf4_block32_f32_e2m1_e2m1_f32_ue8m0_kernel( float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) { - const uint32_t N = 8, K = 32; + const uint32_t N = 8, K = 64; const uint16_t selector = 0; uint32_t laneid = threadIdx.x; uint32_t a_frag[4], b_frag[2], sfa_frag, sfb_frag; float c_frag[4], d_frag[4]; - LOAD_A_M16K32_FP8(); - LOAD_B_N8K32_FP8(); + LOAD_A_M16K64_FP4(); + LOAD_B_N8K64_FP4(); LOAD_C_M16N8_F32(); - LOAD_SFA_M16SFK1(); - LOAD_SFB_N8SFK1(); + LOAD_SFA_M16SFK2(); + LOAD_SFB_N8SFK2(); asm volatile( - "mma.sync.aligned.m16n8k32.row.col.kind::mxf8f6f4.block_scale.scale_vec::1X.f32.e4m3.e4m3.f32.ue8m0 " + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::2X.f32.e2m1.e2m1.f32.ue8m0 " "{%0, %1, %2, %3}, " "{%4, %5, %6, %7}, " "{%8, %9}, " @@ -69,22 +197,51 @@ extern "C" // mxfp8 STORE_D_M16N8_F32(); } - void mma_m16n8k32_block32_f32_e5m2_e5m2_f32_ue8m0( + void mma_m16n8k64_kind_mxf4nvf4_block32_f32_e2m1_e2m1_f32_ue8m0( float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) { - mma_m16n8k32_block32_f32_e5m2_e5m2_f32_ue8m0_kernel<<<1, 32>>>(d, a, b, c, sfa, sfb); + mma_m16n8k64_kind_mxf4nvf4_block32_f32_e2m1_e2m1_f32_ue8m0_kernel<<<1, 32>>>(d, a, b, c, sfa, sfb); } - void mma_m16n8k32_block32_f32_e4m3_e4m3_f32_ue8m0( - float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) - { - mma_m16n8k32_block32_f32_e4m3_e4m3_f32_ue8m0_kernel<<<1, 32>>>(d, a, b, c, sfa, sfb); - } -} + // invalid + // __global__ void mma_m16n8k64_kind_mxf4nvf4_block32_f32_e2m1_e2m1_f32_ue4m3_kernel( + // float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) + // { + // const uint32_t N = 8, K = 64; + // const uint16_t selector = 0; + // uint32_t laneid = threadIdx.x; + // uint32_t a_frag[4], b_frag[2], sfa_frag, sfb_frag; + // float c_frag[4], d_frag[4]; -extern "C" // mxfp4 nvfp4 -{ - __global__ void mma_m16n8k64_block32_f32_e2m1_e2m1_f32_ue8m0_kernel( + // LOAD_A_M16K64_FP4(); + // LOAD_B_N8K64_FP4(); + // LOAD_C_M16N8_F32(); + // LOAD_SFA_M16SFK2(); + // LOAD_SFB_N8SFK2(); + // asm volatile( + // "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::2X.f32.e2m1.e2m1.f32.ue4m3 " + // "{%0, %1, %2, %3}, " + // "{%4, %5, %6, %7}, " + // "{%8, %9}, " + // "{%10, %11, %12, %13}, " + // "%14, {%15, %16}, " + // "%17, {%18, %19};" + // : "=f"(d_frag[0]), "=f"(d_frag[1]), "=f"(d_frag[2]), "=f"(d_frag[3]) + // : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), + // "r"(b_frag[0]), "r"(b_frag[1]), + // "f"(c_frag[0]), "f"(c_frag[1]), "f"(c_frag[2]), "f"(c_frag[3]), + // "r"(sfa_frag), "h"(selector), "h"(selector), + // "r"(sfb_frag), "h"(selector), "h"(selector)); + // STORE_D_M16N8_F32(); + // } + + // void mma_m16n8k64_kind_mxf4nvf4_block32_f32_e2m1_e2m1_f32_ue4m3( + // float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) + // { + // mma_m16n8k64_kind_mxf4nvf4_block32_f32_e2m1_e2m1_f32_ue4m3_kernel<<<1, 32>>>(d, a, b, c, sfa, sfb); + // } + + __global__ void mma_m16n8k64_kind_mxf4nvf4_block16_f32_e2m1_e2m1_f32_ue4m3_kernel( float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) { const uint32_t N = 8, K = 64; @@ -96,11 +253,10 @@ extern "C" // mxfp4 nvfp4 LOAD_A_M16K64_FP4(); LOAD_B_N8K64_FP4(); LOAD_C_M16N8_F32(); - LOAD_SFA_M16SFK2(); - LOAD_SFB_N8SFK2(); - // .kind::mxf4 seems to behave the same as .kind::mxf4nvf4 for mxfp4 + LOAD_SFA_M16SFK4(); + LOAD_SFB_N8SFK4(); asm volatile( - "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::2X.f32.e2m1.e2m1.f32.ue8m0 " + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X.f32.e2m1.e2m1.f32.ue4m3 " "{%0, %1, %2, %3}, " "{%4, %5, %6, %7}, " "{%8, %9}, " @@ -116,9 +272,16 @@ extern "C" // mxfp4 nvfp4 STORE_D_M16N8_F32(); } - __global__ void mma_m16n8k64_block16_f32_e2m1_e2m1_f32_ue4m3_kernel( + void mma_m16n8k64_kind_mxf4nvf4_block16_f32_e2m1_e2m1_f32_ue4m3( + float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) + { + mma_m16n8k64_kind_mxf4nvf4_block16_f32_e2m1_e2m1_f32_ue4m3_kernel<<<1, 32>>>(d, a, b, c, sfa, sfb); + } + + __global__ void mma_m16n8k64_kind_mxf4nvf4_block16_f32_e2m1_e2m1_f32_ue8m0_kernel( float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) { +#if CUDART_VERSION >= 13010 const uint32_t N = 8, K = 64; const uint16_t selector = 0; uint32_t laneid = threadIdx.x; @@ -131,7 +294,7 @@ extern "C" // mxfp4 nvfp4 LOAD_SFA_M16SFK4(); LOAD_SFB_N8SFK4(); asm volatile( - "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X.f32.e2m1.e2m1.f32.ue4m3 " + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X.f32.e2m1.e2m1.f32.ue8m0 " "{%0, %1, %2, %3}, " "{%4, %5, %6, %7}, " "{%8, %9}, " @@ -145,17 +308,12 @@ extern "C" // mxfp4 nvfp4 "r"(sfa_frag), "h"(selector), "h"(selector), "r"(sfb_frag), "h"(selector), "h"(selector)); STORE_D_M16N8_F32(); +#endif } - void mma_m16n8k64_block32_f32_e2m1_e2m1_f32_ue8m0( - float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) - { - mma_m16n8k64_block32_f32_e2m1_e2m1_f32_ue8m0_kernel<<<1, 32>>>(d, a, b, c, sfa, sfb); - } - - void mma_m16n8k64_block16_f32_e2m1_e2m1_f32_ue4m3( + void mma_m16n8k64_kind_mxf4nvf4_block16_f32_e2m1_e2m1_f32_ue8m0( float *d, uint8_t *a, uint8_t *b, float *c, uint8_t *sfa, uint8_t *sfb) { - mma_m16n8k64_block16_f32_e2m1_e2m1_f32_ue4m3_kernel<<<1, 32>>>(d, a, b, c, sfa, sfb); + mma_m16n8k64_kind_mxf4nvf4_block16_f32_e2m1_e2m1_f32_ue8m0_kernel<<<1, 32>>>(d, a, b, c, sfa, sfb); } } 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 f0075cf..a93a538 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/tcgen05mma.h +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/tcgen05mma.h @@ -1,593 +1,278 @@ +#pragma once #include - -#define desc_encode(x) (((x) & 0x3ffff) >> 4) - -#define ALLOC_TMEM(tmem_addr) \ - do \ - { \ - if (warpid == 0) \ - { \ - asm volatile( \ - "tcgen05.alloc.cta_group::1.sync.aligned.b32 [%0], 32;" \ - : \ - : "r"((uint32_t)__cvta_generic_to_shared(&tmem_addr))); \ - } \ - } while (0) - -#define DEALLOC_TMEM(tmem_addr) \ - do \ - { \ - if (warpid == 0) \ - { \ - asm volatile( \ - "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, 32;" \ - : \ - : "r"(tmem_addr)); \ - } \ - } while (0) - -#define RELINQUISH_TMEM() \ - do \ - { \ - if (warpid == 0) \ - { \ - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); \ - } \ - } while (0) - -#define INIT_MBARRIER() \ - do \ - { \ - if (tid == 0) \ - { \ - asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" \ - : \ - : "r"((uint32_t)__cvta_generic_to_shared(&mma_barrier))); \ - } \ - } while (0) - -#define WAIT_MBARRIER() \ - do \ - { \ - asm volatile( \ - "{\n" \ - ".reg .pred P1;\n" \ - "LAB_WAIT:\n" \ - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" \ - "@P1 bra DONE;\n" \ - "bra LAB_WAIT;\n" \ - "DONE:\n" \ - "}\n" \ - : \ - : "r"((uint32_t)__cvta_generic_to_shared(&mma_barrier)), \ - "r"(mma_barrier_phase_bit)); \ - mma_barrier_phase_bit ^= 1; \ - } while (0) - -// (8x2) tiles, each with (8x4) elements -#define LOAD_A_M64K8_TF32() \ - do \ - { \ - for (uint32_t i = 0; i < 8; i += 2) \ - { \ - uint32_t row = i * 8 + warpid / 2 * 8 + laneid / 4; \ - uint32_t col = warpid % 2 * 4 + laneid % 4; \ - a_smem[i * 64 + tid] = a[row * K + col]; \ - } \ - 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 (8x8) elements -#define LOAD_A_M64K16_F16BF16() \ - do \ - { \ - for (uint32_t i = 0; i < 8; i++) \ - { \ - uint32_t row = i * 8 + tid % 64 / 8; \ - uint32_t col = tid / 64 * 8 + tid % 8; \ - a_smem[i * 128 + tid] = a[row * K + col]; \ - } \ - 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 -#define LOAD_A_M128K32_FP8() \ - 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; \ - a_smem[i * 256 + k * 128 + tid] = a[row * K + col]; \ - } \ - } \ - 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 -#define LOAD_A_M64K32_FP8() \ - 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; \ - a_smem[i * 256 + k * 128 + tid] = a[row * K + col]; \ - } \ - } \ - 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 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 \ - { \ - 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 * 32 + tid % 16 * 2; \ - a_smem[i * 256 + k * 128 + tid] = a[row * K / 2 + col / 2]; \ - } \ - } \ - 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) - -// (2x1) tiles, each with (4x8) elements -#define LOAD_B_N8K8_TF32() \ - do \ - { \ - if (tid < 64) \ - { \ - uint32_t row = warpid * 4 + laneid % 4; \ - uint32_t col = laneid / 4; \ - b_smem[tid] = b[col * K + row]; \ - } \ - 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 (8x8) elements -#define LOAD_B_N8K16_F16BF16() \ - do \ - { \ - uint32_t row = tid / 64 * 8 + tid % 8; \ - uint32_t col = tid % 64 / 8; \ - b_smem[tid] = b[col * K + row]; \ - 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 -#define LOAD_B_N8K32_FP8() \ - do \ - { \ - for (uint32_t k = 0; k < 2; k++) \ - { \ - uint32_t row = k * 16 + tid % 16; \ - uint32_t col = tid / 16; \ - b_smem[k * 128 + tid] = b[col * K + row]; \ - } \ - 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 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 \ - { \ - for (uint32_t k = 0; k < 2; k++) \ - { \ - uint32_t row = k * 32 + tid % 16 * 2; \ - uint32_t col = tid / 16; \ - b_smem[k * 128 + tid] = b[col * K / 2 + row / 2]; \ - } \ - 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) - -#define LOAD_D_M128N8() \ - do \ - { \ - uint32_t d_frag[8]; \ - for (uint32_t i = 0; i < 8; i++) \ - { \ - uint32_t row = tid; \ - uint32_t col = i; \ - d_frag[i] = d[row * N + col]; \ - } \ - asm volatile("tcgen05.st.sync.aligned.32x32b.x8.b32 " \ - "[%0], {%1, %2, %3, %4, %5, %6, %7, %8};" \ - : \ - : "r"(d_tmem_addr + ((warpid * 32) << 16)), \ - "r"(d_frag[0]), "r"(d_frag[1]), "r"(d_frag[2]), "r"(d_frag[3]), \ - "r"(d_frag[4]), "r"(d_frag[5]), "r"(d_frag[6]), "r"(d_frag[7])); \ - asm volatile("tcgen05.wait::st.sync.aligned;"); \ - } while (0) - -#define LOAD_D_M64N8() \ - do \ - { \ - uint32_t d_frag[4]; \ - for (uint32_t i = 0; i < 4; i++) \ - { \ - uint32_t row = warpid * 16 + i / 2 * 8 + laneid / 4; \ - uint32_t col = laneid % 4 * 2 + i % 2; \ - d_frag[i] = d[row * N + col]; \ - } \ - asm volatile("tcgen05.st.sync.aligned.16x256b.x1.b32 " \ - "[%0], {%1, %2, %3, %4};" \ - : \ - : "r"(d_tmem_addr + ((warpid * 32) << 16)), \ - "r"(d_frag[0]), "r"(d_frag[1]), \ - "r"(d_frag[2]), "r"(d_frag[3])); \ - asm volatile("tcgen05.wait::st.sync.aligned;"); \ - } while (0) - -#define STORE_D_M128N8() \ - do \ - { \ - uint32_t d_frag[8]; \ - asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 " \ - "{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" \ - : "=r"(d_frag[0]), "=r"(d_frag[1]), "=r"(d_frag[2]), "=r"(d_frag[3]), \ - "=r"(d_frag[4]), "=r"(d_frag[5]), "=r"(d_frag[6]), "=r"(d_frag[7]) \ - : "r"(d_tmem_addr + ((warpid * 32) << 16))); \ - asm volatile("tcgen05.wait::ld.sync.aligned;"); \ - for (uint32_t i = 0; i < 8; i++) \ - { \ - uint32_t row = tid; \ - uint32_t col = i; \ - d[row * N + col] = d_frag[i]; \ - } \ - } while (0) - -#define STORE_D_M64N8() \ - do \ - { \ - uint32_t d_frag[4]; \ - asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 " \ - "{%0, %1, %2, %3}, [%4];" \ - : "=r"(d_frag[0]), "=r"(d_frag[1]), \ - "=r"(d_frag[2]), "=r"(d_frag[3]) \ - : "r"(d_tmem_addr + ((warpid * 32) << 16))); \ - asm volatile("tcgen05.wait::ld.sync.aligned;"); \ - for (uint32_t i = 0; i < 4; i++) \ - { \ - uint32_t row = warpid * 16 + i / 2 * 8 + laneid / 4; \ - uint32_t col = laneid % 4 * 2 + i % 2; \ - d[row * N + col] = d_frag[i]; \ - } \ - __syncthreads(); \ - } while (0) - -#define LOAD_SFA_M128SFK1() \ - do \ - { \ - uint32_t sfa_frag[4]; \ - for (uint32_t i = 0; i < 4; i++) \ - { \ - uint32_t row = i * 32 + laneid; \ - sfa_frag[i] = sfa[row]; \ - } \ - asm volatile("tcgen05.st.sync.aligned.32x32b.x4.b32 [%0], {%1, %2, %3, %4};" \ - : \ - : "r"(sfa_tmem_addr + ((warpid * 32) << 16)), \ - "r"(sfa_frag[0]), "r"(sfa_frag[1]), \ - "r"(sfa_frag[2]), "r"(sfa_frag[3])); \ - asm volatile("tcgen05.wait::st.sync.aligned;"); \ - } while (0) - -#define LOAD_SFA_M128SFK2() \ - do \ - { \ - uint32_t sfa_frag[4]; \ - for (uint32_t i = 0; i < 4; i++) \ - { \ - uint32_t row = i * 32 + laneid; \ - sfa_frag[i] = *(uint16_t *)(sfa + row * 2); \ - } \ - asm volatile("tcgen05.st.sync.aligned.32x32b.x4.b32 [%0], {%1, %2, %3, %4};" \ - : \ - : "r"(sfa_tmem_addr + ((warpid * 32) << 16)), \ - "r"(sfa_frag[0]), "r"(sfa_frag[1]), \ - "r"(sfa_frag[2]), "r"(sfa_frag[3])); \ - asm volatile("tcgen05.wait::st.sync.aligned;"); \ - } while (0) - -#define LOAD_SFA_M128SFK4() \ - do \ - { \ - uint32_t sfa_frag[4]; \ - for (uint32_t i = 0; i < 4; i++) \ - { \ - uint32_t row = i * 32 + laneid; \ - sfa_frag[i] = *(uint32_t *)(sfa + row * 4); \ - } \ - asm volatile("tcgen05.st.sync.aligned.32x32b.x4.b32 [%0], {%1, %2, %3, %4};" \ - : \ - : "r"(sfa_tmem_addr + ((warpid * 32) << 16)), \ - "r"(sfa_frag[0]), "r"(sfa_frag[1]), \ - "r"(sfa_frag[2]), "r"(sfa_frag[3])); \ - asm volatile("tcgen05.wait::st.sync.aligned;"); \ - } while (0) - -#define LOAD_SFB_N8SFK1() \ - do \ - { \ - uint32_t sfb_frag[1]; \ - uint32_t col = laneid % 8; \ - sfb_frag[0] = sfb[col]; \ - asm volatile("tcgen05.st.sync.aligned.32x32b.x1.b32 [%0], {%1};" \ - : \ - : "r"(sfb_tmem_addr + ((warpid * 32) << 16)), \ - "r"(sfb_frag[0])); \ - asm volatile("tcgen05.wait::st.sync.aligned;"); \ - } while (0) - -#define LOAD_SFB_N8SFK2() \ - do \ - { \ - uint32_t sfb_frag[1]; \ - uint32_t col = laneid % 8; \ - sfb_frag[0] = *(uint16_t *)(sfb + col * 2); \ - asm volatile("tcgen05.st.sync.aligned.32x32b.x1.b32 [%0], {%1};" \ - : \ - : "r"(sfb_tmem_addr + ((warpid * 32) << 16)), \ - "r"(sfb_frag[0])); \ - asm volatile("tcgen05.wait::st.sync.aligned;"); \ - } while (0) - -#define LOAD_SFB_N8SFK4() \ - do \ - { \ - uint32_t sfb_frag[1]; \ - uint32_t col = laneid % 8; \ - sfb_frag[0] = *(uint32_t *)(sfb + col * 4); \ - asm volatile("tcgen05.st.sync.aligned.32x32b.x1.b32 [%0], {%1};" \ - : \ - : "r"(sfb_tmem_addr + ((warpid * 32) << 16)), \ - "r"(sfb_frag[0])); \ - asm volatile("tcgen05.wait::st.sync.aligned;"); \ - } while (0) - -#define MMA(kind) \ - do \ - { \ - if (tid == 0) \ - { \ - const uint32_t mask[4] = {0, 0, 0, 0}; \ - asm volatile( \ - "tcgen05.mma.cta_group::1.kind::" kind " [%0], %1, %2, %3, {%4, %5, %6, %7}, 1;" \ - : \ - : "r"(d_tmem_addr), "l"(a_desc), "l"(b_desc), "r"(i_desc), \ - "r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3])); \ - asm volatile( \ - "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" \ - : \ - : "r"((uint32_t)__cvta_generic_to_shared(&mma_barrier))); \ - } \ - asm volatile( \ - "{\n" \ - ".reg .pred P1;\n" \ - "LAB_WAIT:\n" \ - " mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" \ - "@P1 bra DONE;\n" \ - " bra LAB_WAIT;\n" \ - "DONE:\n" \ - "}\n" \ - : \ - : "r"((uint32_t)__cvta_generic_to_shared(&mma_barrier)), "r"(mma_barrier_phase_bit)); \ - mma_barrier_phase_bit ^= 1; \ - } while (0) - -#define MMA_WITH_BLOCK_SCALE(kind, block_size) \ - do \ - { \ - if (tid == 0) \ - { \ - asm volatile( \ - "tcgen05.mma.cta_group::1.kind::" kind ".block_scale." block_size " " \ - "[%0], %1, %2, %3, [%4], [%5], 1;" \ - : \ - : "r"(d_tmem_addr), "l"(a_desc), "l"(b_desc), "r"(i_desc), \ - "r"(sfa_tmem_addr), "r"(sfb_tmem_addr)); \ - asm volatile( \ - "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" \ - : \ - : "r"((uint32_t)__cvta_generic_to_shared(&mma_barrier))); \ - } \ - asm volatile( \ - "{\n" \ - ".reg .pred P1;\n" \ - "LAB_WAIT:\n" \ - " mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" \ - "@P1 bra DONE;\n" \ - " bra LAB_WAIT;\n" \ - "DONE:\n" \ - "}\n" \ - : \ - : "r"((uint32_t)__cvta_generic_to_shared(&mma_barrier)), "r"(mma_barrier_phase_bit)); \ - mma_barrier_phase_bit ^= 1; \ - } while (0) +#include + +enum class TCGen05Kind +{ + tf32, + f16, + f8f6f4, + mxf8f6f4, + mxf4, + mxf4nvf4 +}; +enum class TCGen05Type +{ + tf32, + f16, + bf16, + e4m3, + e5m2, + e2m3, + e3m2, + e2m1 +}; +enum class TCGen05ScaleType +{ + ue8m0, + ue4m3 +}; + +template +__device__ constexpr uint32_t tcgen05_type_code() +{ + if constexpr (Kind == TCGen05Kind::mxf4 || Kind == TCGen05Kind::mxf4nvf4) + return 1; + if constexpr (Type == TCGen05Type::f16 || Type == TCGen05Type::e4m3) + return 0; + if constexpr (Type == TCGen05Type::bf16 || Type == TCGen05Type::e5m2) + return 1; + if constexpr (Type == TCGen05Type::tf32) + return 2; + if constexpr (Type == TCGen05Type::e2m3) + return 3; + if constexpr (Type == TCGen05Type::e3m2) + return 4; + return 5; +} + +__device__ inline uint32_t tcgen05_desc_encode(uint64_t value) +{ + return (value & 0x3ffff) >> 4; +} + +__device__ inline uint32_t tcgen05_tmem_columns(uint32_t n) +{ + uint32_t cols = 32; + while (cols < n) + cols *= 2; + return cols; +} + +__device__ inline void tcgen05_alloc_tmem(uint32_t &addr, uint32_t cols) +{ + if (threadIdx.x < 32) + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.b32 [%0], %1;" ::"r"((uint32_t)__cvta_generic_to_shared(&addr)), "r"(cols)); +} + +__device__ inline void tcgen05_dealloc_tmem(uint32_t addr, uint32_t cols) +{ + if (threadIdx.x < 32) + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(addr), "r"(cols)); +} + +__device__ inline void tcgen05_relinquish_tmem() +{ + if (threadIdx.x < 32) + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); +} + +__device__ inline void tcgen05_init_barrier(uint64_t &barrier) +{ + if (threadIdx.x == 0) + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"((uint32_t)__cvta_generic_to_shared(&barrier))); +} + +// Each kernel issues one MMA, so the newly initialized barrier uses phase 0. +__device__ inline void tcgen05_commit_and_wait(uint64_t &barrier) +{ + uint32_t addr = (uint32_t)__cvta_generic_to_shared(&barrier); + if (threadIdx.x == 0) + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" ::"r"(addr)); + asm volatile( + "{\n" + " .reg .pred ready;\n" + "wait_mma:\n" + " mbarrier.try_wait.parity.shared::cta.b64 ready, [%0], 0;\n" + " @!ready bra wait_mma;\n" + "}\n" ::"r"(addr) : "memory"); +} + +#define TCGEN05_MMA_ASM(kind) \ + asm volatile( \ + "tcgen05.mma.cta_group::1.kind::" kind " [%0], %1, %2, %3, {%4, %5, %6, %7}, 1;" ::"r"(d_addr), "l"(a_desc), "l"(b_desc), "r"(i_desc), \ + "r"(0u), "r"(0u), "r"(0u), "r"(0u)) + +template +__device__ void tcgen05_mma(uint32_t d_addr, uint64_t a_desc, uint64_t b_desc, uint32_t i_desc) +{ + if (threadIdx.x == 0) + { + if constexpr (Kind == TCGen05Kind::tf32) + { + TCGEN05_MMA_ASM("tf32"); + } + if constexpr (Kind == TCGen05Kind::f16) + { + TCGEN05_MMA_ASM("f16"); + } + if constexpr (Kind == TCGen05Kind::f8f6f4) + { + TCGEN05_MMA_ASM("f8f6f4"); + } + } +} + +#define TCGEN05_MMA_BLOCK_SCALE_ASM(kind, scale_vec) \ + asm volatile( \ + "tcgen05.mma.cta_group::1.kind::" kind ".block_scale.scale_vec::" scale_vec " " \ + "[%0], %1, %2, %3, [%4], [%5], 1;" ::"r"(d_addr), \ + "l"(a_desc), "l"(b_desc), "r"(i_desc), "r"(sfa_addr), "r"(sfb_addr)) + +template +__device__ void tcgen05_mma_block_scale(uint32_t d_addr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, uint32_t sfa_addr, uint32_t sfb_addr) +{ + if (threadIdx.x == 0) + { + if constexpr (Kind == TCGen05Kind::mxf8f6f4) + { + TCGEN05_MMA_BLOCK_SCALE_ASM("mxf8f6f4", "1X"); + } + if constexpr (Kind == TCGen05Kind::mxf4) + { + TCGEN05_MMA_BLOCK_SCALE_ASM("mxf4", "2X"); + } + if constexpr (Kind == TCGen05Kind::mxf4nvf4 && BlockSize == 32) + { + TCGEN05_MMA_BLOCK_SCALE_ASM("mxf4nvf4", "2X"); + } + if constexpr (Kind == TCGen05Kind::mxf4nvf4 && BlockSize == 16) + { + TCGEN05_MMA_BLOCK_SCALE_ASM("mxf4nvf4", "4X"); + } + } +} + +// Descriptor for K-major 8-row x 16-byte tiles, two K tiles per row group. +__device__ inline uint64_t tcgen05_smem_desc(const void *ptr) +{ + return tcgen05_desc_encode(__cvta_generic_to_shared(ptr)) | (uint64_t(tcgen05_desc_encode(128)) << 16) | (uint64_t(tcgen05_desc_encode(256)) << 32) | (1ull << 46); +} + +template +__device__ void tcgen05_pack_operand(uint8_t *dst, const void *input, + uint32_t rows, uint32_t tid) +{ + const uint8_t *src = static_cast(input); + for (uint32_t index = tid; index < rows * 32; index += 128) + { + uint32_t row = index / 32, tile = index % 32 / 16, byte = index % 16; + uint8_t value; + if constexpr ((Kind == TCGen05Kind::f8f6f4 || Kind == TCGen05Kind::mxf8f6f4) && (Type == TCGen05Type::e2m3 || Type == TCGen05Type::e3m2)) + { + // Sixteen 6-bit values occupy 12 bytes; the last 4 bytes are padding. + value = 0; + if (byte < 12) + { + uint32_t col = tile * 16 + byte * 8 / 6; + uint32_t shift = byte * 8 % 6; + uint32_t bits = src[row * 32 + col] & 0x3f; + if (col % 16 != 15) + bits |= (src[row * 32 + col + 1] & 0x3f) << 6; + value = bits >> shift; + } + } + else if constexpr ((Kind == TCGen05Kind::f8f6f4 || Kind == TCGen05Kind::mxf8f6f4) && Type == TCGen05Type::e2m1) + { + // Sixteen FP4 values occupy 8 bytes; the last 8 bytes are padding. + value = 0; + if (byte < 8) + { + uint32_t col = tile * 16 + byte * 2; + value = ((src[row * 32 + col] >> 2) & 15) | (((src[row * 32 + col + 1] >> 2) & 15) << 4); + } + } + else + { + // TF32/F16/BF16, FP8, and already packed MXFP4/NVFP4. + value = src[index]; + } + dst[(row / 8) * 256 + tile * 128 + (row % 8) * 16 + byte] = value; + } +} + +// PTX requires each D element (including F16) in a separate 32-bit word. +// M=64 uses only 16 datapath lanes in each 32-lane partition (layout F). +template +__device__ void tcgen05_transfer_d(D *d, uint32_t addr, uint32_t m, uint32_t n, + uint32_t tid) +{ + uint32_t warp = tid / 32, lane = tid % 32; + uint32_t base = addr + ((warp * 32) << 16); + if (m == 128) + { + for (uint32_t col = 0; col < n; ++col) + { + uint32_t word; + if constexpr (Store) + { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x1.b32 {%0}, [%1];" + : "=r"(word) : "r"(base + col)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + d[tid * n + col] = D(word); + } + else + { + word = uint32_t(d[tid * n + col]); + asm volatile("tcgen05.st.sync.aligned.32x32b.x1.b32 [%0], {%1};" ::"r"(base + col), "r"(word)); + asm volatile("tcgen05.wait::st.sync.aligned;"); + } + } + } + else + { + for (uint32_t col = 0; col < n; col += 8) + { + uint32_t r[4]; + if constexpr (Store) + { + asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(base + col)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + } + for (uint32_t i = 0; i < 4; ++i) + { + uint32_t row = warp * 16 + i / 2 * 8 + lane / 4; + uint32_t offset = row * n + col + lane % 4 * 2 + i % 2; + if constexpr (Store) + d[offset] = D(r[i]); + else + r[i] = uint32_t(d[offset]); + } + if constexpr (!Store) + { + asm volatile("tcgen05.st.sync.aligned.16x256b.x1.b32 [%0], {%1, %2, %3, %4};" ::"r"(base + col), "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3])); + asm volatile("tcgen05.wait::st.sync.aligned;"); + } + } + } +} + +template +__device__ void tcgen05_store_scales(uint32_t addr, const uint8_t *src, + uint32_t rows, uint32_t tid) +{ + for (uint32_t col = 0; col < (rows + 31) / 32; ++col) + { + uint32_t row = col * 32 + tid % 32; + uint32_t word = 0; + if (row < rows) + for (uint32_t k = 0; k < SF; ++k) + word |= uint32_t(src[row * SF + k]) << (k * 8); + uint32_t dst = addr + ((tid / 32 * 32) << 16) + col; + asm volatile("tcgen05.st.sync.aligned.32x32b.x1.b32 [%0], {%1};" ::"r"(dst), "r"(word)); + } + asm volatile("tcgen05.wait::st.sync.aligned;"); +} 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 1f06589..6afc03f 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 @@ -1,759 +1,150 @@ #include "tcgen05mma.h" -extern "C" // tf32 +template +__global__ void tcgen05mma_kernel(D *d, const void *a, const void *b, uint32_t M, uint32_t N) { - __global__ void tcgen05mma_m64n8k8_f32_tf32_tf32_kernel( - uint32_t *d, uint32_t *a, uint32_t *b) - { - const uint32_t M = 64, N = 8, K = 8; - 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__ uint32_t a_smem[M * K], b_smem[N * K]; - __shared__ uint32_t d_tmem_addr; - __shared__ uint64_t mma_barrier; - - INIT_MBARRIER(); - LOAD_A_M64K8_TF32(); - LOAD_B_N8K8_TF32(); - ALLOC_TMEM(d_tmem_addr); - __syncthreads(); - LOAD_D_M64N8(); - __syncthreads(); - i_desc |= (1 << 4) | (2 << 7) | (2 << 10); // f32_tf32_tf32 - MMA("tf32"); - STORE_D_M64N8(); - DEALLOC_TMEM(d_tmem_addr); - RELINQUISH_TMEM(); - } - - void tcgen05mma_m64n8k8_f32_tf32_tf32( - uint32_t *d, uint32_t *a, uint32_t *b) - { - tcgen05mma_m64n8k8_f32_tf32_tf32_kernel<<<1, 128>>>(d, a, b); - } + uint32_t tid = threadIdx.x; + uint32_t cols = tcgen05_tmem_columns(N); + __shared__ __align__(16) uint8_t a_smem[128 * 32], b_smem[256 * 32]; + __shared__ uint32_t d_tmem_addr; + __shared__ uint64_t mma_barrier; + + tcgen05_init_barrier(mma_barrier); + tcgen05_pack_operand(a_smem, a, M, tid); + tcgen05_pack_operand(b_smem, b, N, tid); + uint64_t a_desc = tcgen05_smem_desc(a_smem); + uint64_t b_desc = tcgen05_smem_desc(b_smem); + tcgen05_alloc_tmem(d_tmem_addr, cols); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncthreads(); + + tcgen05_transfer_d(d, d_tmem_addr, M, N, tid); + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + constexpr uint32_t d_code = std::is_same::value ? 0 : 1; + uint32_t i_desc = (N >> 3 << 17) | (M >> 4 << 24) | (d_code << 4) | (tcgen05_type_code() << 7) | (tcgen05_type_code() << 10); + tcgen05_mma(d_tmem_addr, a_desc, b_desc, i_desc); + tcgen05_commit_and_wait(mma_barrier); + + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + tcgen05_transfer_d(d, d_tmem_addr, M, N, tid); + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); + __syncthreads(); + tcgen05_dealloc_tmem(d_tmem_addr, cols); + tcgen05_relinquish_tmem(); } -extern "C" // f16 and bf16 +template +__global__ void tcgen05mma_block_scale_kernel(uint32_t *d, const void *a, const void *b, + const uint8_t *sfa, const uint8_t *sfb, + uint32_t M, uint32_t N) { - __global__ void tcgen05mma_m64n8k16_f32_f16_f16_kernel( - uint32_t *d, uint16_t *a, uint16_t *b) - { - const uint32_t M = 64, N = 8, K = 16; - 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__ uint16_t a_smem[M * K], b_smem[N * K]; - __shared__ uint32_t d_tmem_addr; - __shared__ uint64_t mma_barrier; - - INIT_MBARRIER(); - LOAD_A_M64K16_F16BF16(); - LOAD_B_N8K16_F16BF16(); - ALLOC_TMEM(d_tmem_addr); - __syncthreads(); - LOAD_D_M64N8(); - __syncthreads(); - i_desc |= 1 << 4; // f32_f16_f16 - MMA("f16"); - STORE_D_M64N8(); - DEALLOC_TMEM(d_tmem_addr); - RELINQUISH_TMEM(); - } - - __global__ void tcgen05mma_m64n8k16_f32_bf16_bf16_kernel( - uint32_t *d, uint16_t *a, uint16_t *b) - { - const uint32_t M = 64, N = 8, K = 16; - 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__ uint16_t a_smem[M * K], b_smem[N * K]; - __shared__ uint32_t d_tmem_addr; - __shared__ uint64_t mma_barrier; - - INIT_MBARRIER(); - LOAD_A_M64K16_F16BF16(); - LOAD_B_N8K16_F16BF16(); - ALLOC_TMEM(d_tmem_addr); - __syncthreads(); - LOAD_D_M64N8(); - __syncthreads(); - i_desc |= (1 << 4) | (1 << 7) | (1 << 10); // f32_bf16_bf16 - MMA("f16"); - STORE_D_M64N8(); - DEALLOC_TMEM(d_tmem_addr); - RELINQUISH_TMEM(); - } - - __global__ void tcgen05mma_m64n8k16_f16_f16_f16_kernel( - uint16_t *d, uint16_t *a, uint16_t *b) - { - const uint32_t M = 64, N = 8, K = 16; - 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__ uint16_t a_smem[M * K], b_smem[N * K]; - __shared__ uint32_t d_tmem_addr; - __shared__ uint64_t mma_barrier; - - INIT_MBARRIER(); - LOAD_A_M64K16_F16BF16(); - LOAD_B_N8K16_F16BF16(); - ALLOC_TMEM(d_tmem_addr); - __syncthreads(); - LOAD_D_M64N8(); - __syncthreads(); - i_desc |= 0; // f16_f16_f16 - MMA("f16"); - STORE_D_M64N8(); - DEALLOC_TMEM(d_tmem_addr); - RELINQUISH_TMEM(); - } - - void tcgen05mma_m64n8k16_f32_f16_f16( - uint32_t *d, uint16_t *a, uint16_t *b) - { - tcgen05mma_m64n8k16_f32_f16_f16_kernel<<<1, 128>>>(d, a, b); - } - - void tcgen05mma_m64n8k16_f32_bf16_bf16( - uint32_t *d, uint16_t *a, uint16_t *b) - { - tcgen05mma_m64n8k16_f32_bf16_bf16_kernel<<<1, 128>>>(d, a, b); - } - - void tcgen05mma_m64n8k16_f16_f16_f16( - uint16_t *d, uint16_t *a, uint16_t *b) - { - tcgen05mma_m64n8k16_f16_f16_f16_kernel<<<1, 128>>>(d, a, b); - } + constexpr int scale_vectors = Kind == TCGen05Kind::mxf8f6f4 ? 1 : 64 / BlockSize; + constexpr uint32_t scale_code = SType == TCGen05ScaleType::ue8m0 ? 1 : 0; + uint32_t tid = threadIdx.x; + uint32_t cols = tcgen05_tmem_columns(N); + __shared__ __align__(16) uint8_t a_smem[128 * 32], b_smem[256 * 32]; + __shared__ uint32_t d_tmem_addr, sfa_tmem_addr, sfb_tmem_addr; + __shared__ uint64_t mma_barrier; + + tcgen05_init_barrier(mma_barrier); + tcgen05_pack_operand(a_smem, a, M, tid); + tcgen05_pack_operand(b_smem, b, N, tid); + uint64_t a_desc = tcgen05_smem_desc(a_smem); + uint64_t b_desc = tcgen05_smem_desc(b_smem); + tcgen05_alloc_tmem(d_tmem_addr, cols); + tcgen05_alloc_tmem(sfa_tmem_addr, 32); + tcgen05_alloc_tmem(sfb_tmem_addr, 32); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncthreads(); + + tcgen05_transfer_d(d, d_tmem_addr, M, N, tid); + tcgen05_store_scales(sfa_tmem_addr, sfa, M, tid); + tcgen05_store_scales(sfb_tmem_addr, sfb, N, tid); + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + uint32_t i_desc = (N >> 3 << 17) | (M >> 4 << 24) | (scale_code << 23) | (tcgen05_type_code() << 7) | (tcgen05_type_code() << 10); + tcgen05_mma_block_scale( + d_tmem_addr, a_desc, b_desc, i_desc, sfa_tmem_addr, sfb_tmem_addr); + tcgen05_commit_and_wait(mma_barrier); + + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + tcgen05_transfer_d(d, d_tmem_addr, M, N, tid); + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); + __syncthreads(); + tcgen05_dealloc_tmem(d_tmem_addr, cols); + tcgen05_dealloc_tmem(sfa_tmem_addr, 32); + tcgen05_dealloc_tmem(sfb_tmem_addr, 32); + tcgen05_relinquish_tmem(); } -extern "C" // fp8 m64n8k32 f32_output -{ - __global__ void tcgen05mma_m64n8k32_f32_e5m2_e5m2_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) | (1 << 7) | (1 << 10); // f32_e5m2_e5m2 - MMA("f8f6f4"); - STORE_D_M64N8(); - DEALLOC_TMEM(d_tmem_addr); - RELINQUISH_TMEM(); - } - - __global__ void tcgen05mma_m64n8k32_f32_e4m3_e4m3_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; // f32_e4m3_e4m3 - MMA("f8f6f4"); - STORE_D_M64N8(); - DEALLOC_TMEM(d_tmem_addr); - RELINQUISH_TMEM(); - } - - void tcgen05mma_m64n8k32_f32_e5m2_e5m2( - uint32_t *d, uint8_t *a, uint8_t *b) - { - tcgen05mma_m64n8k32_f32_e5m2_e5m2_kernel<<<1, 128>>>(d, a, b); - } - - void tcgen05mma_m64n8k32_f32_e4m3_e4m3( - uint32_t *d, uint8_t *a, uint8_t *b) - { - tcgen05mma_m64n8k32_f32_e4m3_e4m3_kernel<<<1, 128>>>(d, a, b); - } -} - -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( - 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_FP8(); - LOAD_B_N8K32_FP8(); - ALLOC_TMEM(d_tmem_addr); - __syncthreads(); - LOAD_D_M64N8(); - __syncthreads(); - i_desc |= (1 << 7) | (1 << 10); // f16_e5m2_e5m2 - MMA("f8f6f4"); - STORE_D_M64N8(); - DEALLOC_TMEM(d_tmem_addr); - RELINQUISH_TMEM(); - } - - __global__ void tcgen05mma_m64n8k32_f16_e4m3_e4m3_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_FP8(); - LOAD_B_N8K32_FP8(); - ALLOC_TMEM(d_tmem_addr); - __syncthreads(); - LOAD_D_M64N8(); - __syncthreads(); - i_desc |= 0; // f16_e4m3_e4m3 - MMA("f8f6f4"); - STORE_D_M64N8(); - DEALLOC_TMEM(d_tmem_addr); - RELINQUISH_TMEM(); - } - - void tcgen05mma_m64n8k32_f16_e5m2_e5m2( - uint16_t *d, uint8_t *a, uint8_t *b) - { - tcgen05mma_m64n8k32_f16_e5m2_e5m2_kernel<<<1, 128>>>(d, a, b); - } - - void tcgen05mma_m64n8k32_f16_e4m3_e4m3( - uint16_t *d, uint8_t *a, uint8_t *b) - { - tcgen05mma_m64n8k32_f16_e4m3_e4m3_kernel<<<1, 128>>>(d, a, b); - } -} - -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( - 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_FP8(); - LOAD_B_N8K32_FP8(); - 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 |= (1 << 7) | (1 << 10) | (1 << 23); // e5m2_e5m2, 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_e4m3_e4m3_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_FP8(); - LOAD_B_N8K32_FP8(); - 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 |= 1 << 23; // e4m3_e4m3, 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_e5m2_e5m2_ue8m0( - uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) - { - tcgen05mma_m128n8k32_block32_f32_e5m2_e5m2_ue8m0_kernel<<<1, 128>>>(d, a, b, sfa, sfb); - } - - void tcgen05mma_m128n8k32_block32_f32_e4m3_e4m3_ue8m0( - uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) - { - tcgen05mma_m128n8k32_block32_f32_e4m3_e4m3_ue8m0_kernel<<<1, 128>>>(d, a, b, sfa, sfb); - } -} - -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( - uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) - { - const uint32_t M = 128, N = 8, K = 64; - 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 / 2], b_smem[N * K / 2]; - __shared__ uint32_t d_tmem_addr, sfa_tmem_addr, sfb_tmem_addr; - __shared__ uint64_t mma_barrier; - - INIT_MBARRIER(); - LOAD_A_M128K64_FP4(); - LOAD_B_N8K64_FP4(); - ALLOC_TMEM(d_tmem_addr); - ALLOC_TMEM(sfa_tmem_addr); - ALLOC_TMEM(sfb_tmem_addr); - __syncthreads(); - LOAD_D_M128N8(); - LOAD_SFA_M128SFK2(); - LOAD_SFB_N8SFK2(); - __syncthreads(); - i_desc |= (1 << 7) | (1 << 10) | (1 << 23); // e2m1, ue8m0_scale - // .kind::mxf4 seems to behave the same as .kind::mxf4nvf4 for mxfp4 - MMA_WITH_BLOCK_SCALE("mxf4nvf4", "scale_vec::2X"); - STORE_D_M128N8(); - DEALLOC_TMEM(d_tmem_addr); - DEALLOC_TMEM(sfa_tmem_addr); - DEALLOC_TMEM(sfb_tmem_addr); - RELINQUISH_TMEM(); - } - - __global__ void tcgen05mma_m128n8k64_block16_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 = 64; - 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 / 2], b_smem[N * K / 2]; - __shared__ uint32_t d_tmem_addr, sfa_tmem_addr, sfb_tmem_addr; - __shared__ uint64_t mma_barrier; - - INIT_MBARRIER(); - LOAD_A_M128K64_FP4(); - LOAD_B_N8K64_FP4(); - ALLOC_TMEM(d_tmem_addr); - ALLOC_TMEM(sfa_tmem_addr); - ALLOC_TMEM(sfb_tmem_addr); - __syncthreads(); - LOAD_D_M128N8(); - LOAD_SFA_M128SFK4(); - LOAD_SFB_N8SFK4(); - __syncthreads(); - i_desc |= (1 << 7) | (1 << 10) | (1 << 23); // e2m1, ue8m0_scale - MMA_WITH_BLOCK_SCALE("mxf4nvf4", "scale_vec::4X"); - STORE_D_M128N8(); - DEALLOC_TMEM(d_tmem_addr); - DEALLOC_TMEM(sfa_tmem_addr); - DEALLOC_TMEM(sfb_tmem_addr); - RELINQUISH_TMEM(); - } - - __global__ void tcgen05mma_m128n8k64_block16_f32_e2m1_e2m1_ue4m3_kernel( - uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) - { - const uint32_t M = 128, N = 8, K = 64; - 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 / 2], b_smem[N * K / 2]; - __shared__ uint32_t d_tmem_addr, sfa_tmem_addr, sfb_tmem_addr; - __shared__ uint64_t mma_barrier; - - INIT_MBARRIER(); - LOAD_A_M128K64_FP4(); - LOAD_B_N8K64_FP4(); - ALLOC_TMEM(d_tmem_addr); - ALLOC_TMEM(sfa_tmem_addr); - ALLOC_TMEM(sfb_tmem_addr); - __syncthreads(); - LOAD_D_M128N8(); - LOAD_SFA_M128SFK4(); - LOAD_SFB_N8SFK4(); - __syncthreads(); - i_desc |= (1 << 7) | (1 << 10); // e2m1, ue4m3_scale - MMA_WITH_BLOCK_SCALE("mxf4nvf4", "scale_vec::4X"); - STORE_D_M128N8(); - DEALLOC_TMEM(d_tmem_addr); - DEALLOC_TMEM(sfa_tmem_addr); - DEALLOC_TMEM(sfb_tmem_addr); - RELINQUISH_TMEM(); - } - - void tcgen05mma_m128n8k64_block32_f32_e2m1_e2m1_ue8m0( - uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) - { - tcgen05mma_m128n8k64_block32_f32_e2m1_e2m1_ue8m0_kernel<<<1, 128>>>(d, a, b, sfa, sfb); - } - - void tcgen05mma_m128n8k64_block16_f32_e2m1_e2m1_ue8m0( - uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) - { - tcgen05mma_m128n8k64_block16_f32_e2m1_e2m1_ue8m0_kernel<<<1, 128>>>(d, a, b, sfa, sfb); - } - - void tcgen05mma_m128n8k64_block16_f32_e2m1_e2m1_ue4m3( - uint32_t *d, uint8_t *a, uint8_t *b, uint8_t *sfa, uint8_t *sfb) - { - tcgen05mma_m128n8k64_block16_f32_e2m1_e2m1_ue4m3_kernel<<<1, 128>>>(d, a, b, sfa, sfb); - } -} +#define TCGEN05_DTYPE_f32 uint32_t +#define TCGEN05_DTYPE_f16 uint16_t + +// One entry point per kind/type combination; Python binds the M/N dimensions. +#define DEFINE_TCGEN05(kind, dtype, atype, btype) \ + extern "C" void tcgen05mma_kind_##kind##_##dtype##_##atype##_##btype( \ + TCGEN05_DTYPE_##dtype *d, void *a, void *b, uint32_t m, uint32_t n) \ + { \ + tcgen05mma_kernel \ + <<<1, 128>>>(d, a, b, m, n); \ + } + +#define DEFINE_TCGEN05_BLOCK_SCALE(kind, block, stype, atype, btype) \ + extern "C" void tcgen05mma_kind_##kind##_block##block##_f32_##atype##_##btype##_##stype( \ + uint32_t *d, void *a, void *b, uint8_t *sfa, uint8_t *sfb, uint32_t m, uint32_t n) \ + { \ + tcgen05mma_block_scale_kernel \ + <<<1, 128>>>(d, a, b, sfa, sfb, m, n); \ + } + +DEFINE_TCGEN05(tf32, f32, tf32, tf32) + +DEFINE_TCGEN05(f16, f32, f16, f16) +DEFINE_TCGEN05(f16, f32, f16, bf16) +DEFINE_TCGEN05(f16, f32, bf16, f16) +DEFINE_TCGEN05(f16, f32, bf16, bf16) +DEFINE_TCGEN05(f16, f16, f16, f16) + +#define DEFINE_TCGEN05_F8F6F4(atype, btype) \ + DEFINE_TCGEN05(f8f6f4, f32, atype, btype) \ + DEFINE_TCGEN05(f8f6f4, f16, atype, btype) + +#define DEFINE_TCGEN05_F8F6F4_FOR_A(atype) \ + DEFINE_TCGEN05_F8F6F4(atype, e5m2) \ + DEFINE_TCGEN05_F8F6F4(atype, e4m3) \ + DEFINE_TCGEN05_F8F6F4(atype, e3m2) \ + DEFINE_TCGEN05_F8F6F4(atype, e2m3) \ + DEFINE_TCGEN05_F8F6F4(atype, e2m1) + +DEFINE_TCGEN05_F8F6F4_FOR_A(e5m2) +DEFINE_TCGEN05_F8F6F4_FOR_A(e4m3) +DEFINE_TCGEN05_F8F6F4_FOR_A(e3m2) +DEFINE_TCGEN05_F8F6F4_FOR_A(e2m3) +DEFINE_TCGEN05_F8F6F4_FOR_A(e2m1) + +#define DEFINE_TCGEN05_MXF8F6F4_FOR_A(atype) \ + DEFINE_TCGEN05_BLOCK_SCALE(mxf8f6f4, 32, ue8m0, atype, e5m2) \ + DEFINE_TCGEN05_BLOCK_SCALE(mxf8f6f4, 32, ue8m0, atype, e4m3) \ + DEFINE_TCGEN05_BLOCK_SCALE(mxf8f6f4, 32, ue8m0, atype, e3m2) \ + DEFINE_TCGEN05_BLOCK_SCALE(mxf8f6f4, 32, ue8m0, atype, e2m3) \ + DEFINE_TCGEN05_BLOCK_SCALE(mxf8f6f4, 32, ue8m0, atype, e2m1) + +DEFINE_TCGEN05_MXF8F6F4_FOR_A(e5m2) +DEFINE_TCGEN05_MXF8F6F4_FOR_A(e4m3) +DEFINE_TCGEN05_MXF8F6F4_FOR_A(e3m2) +DEFINE_TCGEN05_MXF8F6F4_FOR_A(e2m3) +DEFINE_TCGEN05_MXF8F6F4_FOR_A(e2m1) + +DEFINE_TCGEN05_BLOCK_SCALE(mxf4, 32, ue8m0, e2m1, e2m1) + +DEFINE_TCGEN05_BLOCK_SCALE(mxf4nvf4, 32, ue8m0, e2m1, e2m1) +DEFINE_TCGEN05_BLOCK_SCALE(mxf4nvf4, 16, ue8m0, e2m1, e2m1) +// DEFINE_TCGEN05_BLOCK_SCALE(mxf4nvf4, 32, ue4m3, e2m1, e2m1) // requires sm_107 +DEFINE_TCGEN05_BLOCK_SCALE(mxf4nvf4, 16, ue4m3, e2m1, e2m1) diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/wgmma.h b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/wgmma.h index be6831a..cee66c7 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/wgmma.h +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/wgmma.h @@ -1,137 +1,338 @@ +#pragma once #include +#include -#define desc_encode(x) (((x) & 0x3ffff) >> 4) +// All supported floating types use 32 bytes per row. Arrange two K tiles +// of eight rows x 16 bytes, matching the unswizzled K-major descriptor. +__device__ inline uint64_t wgmma_pack_operand(uint8_t *dst, const void *input, + uint32_t rows, uint32_t tid) +{ + const uint8_t *src = static_cast(input); + for (uint32_t index = tid; index < rows * 32; index += 128) + { + uint32_t row = index / 32, tile = index % 32 / 16, byte = index % 16; + dst[(row / 8) * 256 + tile * 128 + (row % 8) * 16 + byte] = src[index]; + } + return ((__cvta_generic_to_shared(dst) & 0x3ffffull) >> 4) + | (8ull << 16) | (16ull << 32); +} -// (8x2) tiles, each with (8x4) elements -#define LOAD_A_M64K8_TF32() \ - do \ - { \ - for (uint32_t i = 0; i < 8; i += 2) \ - { \ - uint32_t row = i * 8 + warpid / 2 * 8 + laneid / 4; \ - uint32_t col = warpid % 2 * 4 + laneid % 4; \ - a_smem[i * 64 + tid] = a[row * K + col]; \ - } \ - a_desc = desc_encode(__cvta_generic_to_shared(a_smem)); \ - a_desc |= desc_encode(128ll) << 16; \ - a_desc |= desc_encode(256ll) << 32; \ - } while (0) +template +__device__ void wgmma_transfer_f32(float *d, float *frag, uint32_t tid) +{ + uint32_t warp = tid / 32, lane = tid % 32; + #pragma unroll + for (uint32_t i = 0; i < N / 2; ++i) + { + uint32_t row = warp * 16 + (i % 4) / 2 * 8 + lane / 4; + uint32_t col = i / 4 * 8 + lane % 4 * 2 + i % 2; + if constexpr (Store) d[row * N + col] = frag[i]; + else frag[i] = d[row * N + col]; + } +} -// (8x2) tiles, each with (8x8) elements -#define LOAD_A_M64K16_F16BF16() \ - do \ - { \ - for (uint32_t i = 0; i < 8; i++) \ - { \ - uint32_t row = i * 8 + tid % 64 / 8; \ - uint32_t col = tid / 64 * 8 + tid % 8; \ - a_smem[i * 128 + tid] = a[row * K + col]; \ - } \ - a_desc = desc_encode(__cvta_generic_to_shared(a_smem)); \ - a_desc |= desc_encode(128ll) << 16; \ - a_desc |= desc_encode(256ll) << 32; \ - } while (0) +template +__device__ void wgmma_transfer_f16(uint16_t *d, uint32_t *frag, uint32_t tid) +{ + uint32_t warp = tid / 32, lane = tid % 32; + #pragma unroll + for (uint32_t i = 0; i < N / 4; ++i) + { + uint32_t row = warp * 16 + i % 2 * 8 + lane / 4; + uint32_t col = i / 2 * 8 + lane % 4 * 2; + uint32_t *word = reinterpret_cast(d + row * N + col); + if constexpr (Store) *word = frag[i]; + else frag[i] = *word; + } +} -// (8x2) tiles, each with (8x16) elements -#define LOAD_A_M64K32_FP8() \ - 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; \ - a_smem[i * 256 + k * 128 + tid] = a[row * K + col]; \ - } \ - } \ - a_desc = desc_encode(__cvta_generic_to_shared(a_smem)); \ - a_desc |= desc_encode(128ll) << 16; \ - a_desc |= desc_encode(256ll) << 32; \ - } while (0) +// The output element type and the type of one accumulator register. +enum class WGMMAType { f32, f16, bf16, tf32, e4m3, e5m2 }; -// (2x1) tiles, each with (4x8) elements -#define LOAD_B_N8K8_TF32() \ - do \ - { \ - if (tid < 64) \ - { \ - uint32_t row = warpid * 4 + laneid % 4; \ - uint32_t col = laneid / 4; \ - b_smem[tid] = b[col * K + row]; \ - } \ - b_desc = desc_encode(__cvta_generic_to_shared(b_smem)); \ - b_desc |= desc_encode(128ll) << 16; \ - b_desc |= desc_encode(256ll) << 32; \ - } while (0) +template struct WGMMAAccumulator; -// (2x1) tiles, each with (8x8) elements -#define LOAD_B_N8K16_F16BF16() \ - do \ - { \ - uint32_t row = tid / 64 * 8 + tid % 8; \ - uint32_t col = tid % 64 / 8; \ - b_smem[tid] = b[col * K + row]; \ - b_desc = desc_encode(__cvta_generic_to_shared(b_smem)); \ - b_desc |= desc_encode(128ll) << 16; \ - b_desc |= desc_encode(256ll) << 32; \ - } while (0) +template <> struct WGMMAAccumulator +{ + using Element = float; + using Register = float; + static constexpr int registers_per_8_columns = 4; +}; -// (2x1) tiles, each with (16x8) elements -#define LOAD_B_N8K32_FP8() \ - do \ - { \ - for (uint32_t k = 0; k < 2; k++) \ - { \ - uint32_t row = k * 16 + tid % 16; \ - uint32_t col = tid / 16; \ - b_smem[k * 128 + tid] = b[col * K + row]; \ - } \ - b_desc = desc_encode(__cvta_generic_to_shared(b_smem)); \ - b_desc |= desc_encode(128ll) << 16; \ - b_desc |= desc_encode(256ll) << 32; \ - } while (0) +template <> struct WGMMAAccumulator +{ + using Element = uint16_t; + using Register = uint32_t; + static constexpr int registers_per_8_columns = 2; +}; -#define LOAD_D_M64N8_F32() \ - do \ - { \ - for (uint32_t i = 0; i < 4; i++) \ - { \ - uint32_t row = warpid * 16 + i / 2 * 8 + laneid / 4; \ - uint32_t col = laneid % 4 * 2 + i % 2; \ - d_frag[i] = d[row * N + col]; \ - } \ - } while (0) +template +__device__ void wgmma_transfer_d(Element *d, Register *frag, uint32_t tid) +{ + if constexpr (std::is_same::value) + wgmma_transfer_f32(d, frag, tid); + else + wgmma_transfer_f16(d, frag, tid); +} -#define LOAD_D_M64N8_F16() \ - do \ - { \ - for (uint32_t i = 0; i < 4; i += 2) \ - { \ - uint32_t row = warpid * 16 + i / 2 * 8 + laneid / 4; \ - uint32_t col = laneid % 4 * 2; \ - d_frag[i / 2] = *(uint32_t *)(d + row * N + col); \ - } \ - } while (0) +// Inline PTX needs an explicit register list for each N. Keep this mechanical +// operand-list expansion here; wgmma_sm90a.cu uses ordinary C++ templates. +#define WGMMA_REGS_2 "%0, %1" +#define WGMMA_ARGS_2(arg) arg(0), arg(1) +#define WGMMA_REGS_4 WGMMA_REGS_2 ", %2, %3" +#define WGMMA_ARGS_4(arg) WGMMA_ARGS_2(arg), arg(2), arg(3) +#define WGMMA_REGS_6 WGMMA_REGS_4 ", %4, %5" +#define WGMMA_ARGS_6(arg) WGMMA_ARGS_4(arg), arg(4), arg(5) +#define WGMMA_REGS_8 WGMMA_REGS_6 ", %6, %7" +#define WGMMA_ARGS_8(arg) WGMMA_ARGS_6(arg), arg(6), arg(7) +#define WGMMA_REGS_10 WGMMA_REGS_8 ", %8, %9" +#define WGMMA_ARGS_10(arg) WGMMA_ARGS_8(arg), arg(8), arg(9) +#define WGMMA_REGS_12 WGMMA_REGS_10 ", %10, %11" +#define WGMMA_ARGS_12(arg) WGMMA_ARGS_10(arg), arg(10), arg(11) +#define WGMMA_REGS_14 WGMMA_REGS_12 ", %12, %13" +#define WGMMA_ARGS_14(arg) WGMMA_ARGS_12(arg), arg(12), arg(13) +#define WGMMA_REGS_16 WGMMA_REGS_14 ", %14, %15" +#define WGMMA_ARGS_16(arg) WGMMA_ARGS_14(arg), arg(14), arg(15) +#define WGMMA_REGS_18 WGMMA_REGS_16 ", %16, %17" +#define WGMMA_ARGS_18(arg) WGMMA_ARGS_16(arg), arg(16), arg(17) +#define WGMMA_REGS_20 WGMMA_REGS_18 ", %18, %19" +#define WGMMA_ARGS_20(arg) WGMMA_ARGS_18(arg), arg(18), arg(19) +#define WGMMA_REGS_22 WGMMA_REGS_20 ", %20, %21" +#define WGMMA_ARGS_22(arg) WGMMA_ARGS_20(arg), arg(20), arg(21) +#define WGMMA_REGS_24 WGMMA_REGS_22 ", %22, %23" +#define WGMMA_ARGS_24(arg) WGMMA_ARGS_22(arg), arg(22), arg(23) +#define WGMMA_REGS_26 WGMMA_REGS_24 ", %24, %25" +#define WGMMA_ARGS_26(arg) WGMMA_ARGS_24(arg), arg(24), arg(25) +#define WGMMA_REGS_28 WGMMA_REGS_26 ", %26, %27" +#define WGMMA_ARGS_28(arg) WGMMA_ARGS_26(arg), arg(26), arg(27) +#define WGMMA_REGS_30 WGMMA_REGS_28 ", %28, %29" +#define WGMMA_ARGS_30(arg) WGMMA_ARGS_28(arg), arg(28), arg(29) +#define WGMMA_REGS_32 WGMMA_REGS_30 ", %30, %31" +#define WGMMA_ARGS_32(arg) WGMMA_ARGS_30(arg), arg(30), arg(31) +#define WGMMA_REGS_34 WGMMA_REGS_32 ", %32, %33" +#define WGMMA_ARGS_34(arg) WGMMA_ARGS_32(arg), arg(32), arg(33) +#define WGMMA_REGS_36 WGMMA_REGS_34 ", %34, %35" +#define WGMMA_ARGS_36(arg) WGMMA_ARGS_34(arg), arg(34), arg(35) +#define WGMMA_REGS_38 WGMMA_REGS_36 ", %36, %37" +#define WGMMA_ARGS_38(arg) WGMMA_ARGS_36(arg), arg(36), arg(37) +#define WGMMA_REGS_40 WGMMA_REGS_38 ", %38, %39" +#define WGMMA_ARGS_40(arg) WGMMA_ARGS_38(arg), arg(38), arg(39) +#define WGMMA_REGS_42 WGMMA_REGS_40 ", %40, %41" +#define WGMMA_ARGS_42(arg) WGMMA_ARGS_40(arg), arg(40), arg(41) +#define WGMMA_REGS_44 WGMMA_REGS_42 ", %42, %43" +#define WGMMA_ARGS_44(arg) WGMMA_ARGS_42(arg), arg(42), arg(43) +#define WGMMA_REGS_46 WGMMA_REGS_44 ", %44, %45" +#define WGMMA_ARGS_46(arg) WGMMA_ARGS_44(arg), arg(44), arg(45) +#define WGMMA_REGS_48 WGMMA_REGS_46 ", %46, %47" +#define WGMMA_ARGS_48(arg) WGMMA_ARGS_46(arg), arg(46), arg(47) +#define WGMMA_REGS_50 WGMMA_REGS_48 ", %48, %49" +#define WGMMA_ARGS_50(arg) WGMMA_ARGS_48(arg), arg(48), arg(49) +#define WGMMA_REGS_52 WGMMA_REGS_50 ", %50, %51" +#define WGMMA_ARGS_52(arg) WGMMA_ARGS_50(arg), arg(50), arg(51) +#define WGMMA_REGS_54 WGMMA_REGS_52 ", %52, %53" +#define WGMMA_ARGS_54(arg) WGMMA_ARGS_52(arg), arg(52), arg(53) +#define WGMMA_REGS_56 WGMMA_REGS_54 ", %54, %55" +#define WGMMA_ARGS_56(arg) WGMMA_ARGS_54(arg), arg(54), arg(55) +#define WGMMA_REGS_58 WGMMA_REGS_56 ", %56, %57" +#define WGMMA_ARGS_58(arg) WGMMA_ARGS_56(arg), arg(56), arg(57) +#define WGMMA_REGS_60 WGMMA_REGS_58 ", %58, %59" +#define WGMMA_ARGS_60(arg) WGMMA_ARGS_58(arg), arg(58), arg(59) +#define WGMMA_REGS_62 WGMMA_REGS_60 ", %60, %61" +#define WGMMA_ARGS_62(arg) WGMMA_ARGS_60(arg), arg(60), arg(61) +#define WGMMA_REGS_64 WGMMA_REGS_62 ", %62, %63" +#define WGMMA_ARGS_64(arg) WGMMA_ARGS_62(arg), arg(62), arg(63) +#define WGMMA_REGS_66 WGMMA_REGS_64 ", %64, %65" +#define WGMMA_ARGS_66(arg) WGMMA_ARGS_64(arg), arg(64), arg(65) +#define WGMMA_REGS_68 WGMMA_REGS_66 ", %66, %67" +#define WGMMA_ARGS_68(arg) WGMMA_ARGS_66(arg), arg(66), arg(67) +#define WGMMA_REGS_70 WGMMA_REGS_68 ", %68, %69" +#define WGMMA_ARGS_70(arg) WGMMA_ARGS_68(arg), arg(68), arg(69) +#define WGMMA_REGS_72 WGMMA_REGS_70 ", %70, %71" +#define WGMMA_ARGS_72(arg) WGMMA_ARGS_70(arg), arg(70), arg(71) +#define WGMMA_REGS_74 WGMMA_REGS_72 ", %72, %73" +#define WGMMA_ARGS_74(arg) WGMMA_ARGS_72(arg), arg(72), arg(73) +#define WGMMA_REGS_76 WGMMA_REGS_74 ", %74, %75" +#define WGMMA_ARGS_76(arg) WGMMA_ARGS_74(arg), arg(74), arg(75) +#define WGMMA_REGS_78 WGMMA_REGS_76 ", %76, %77" +#define WGMMA_ARGS_78(arg) WGMMA_ARGS_76(arg), arg(76), arg(77) +#define WGMMA_REGS_80 WGMMA_REGS_78 ", %78, %79" +#define WGMMA_ARGS_80(arg) WGMMA_ARGS_78(arg), arg(78), arg(79) +#define WGMMA_REGS_82 WGMMA_REGS_80 ", %80, %81" +#define WGMMA_ARGS_82(arg) WGMMA_ARGS_80(arg), arg(80), arg(81) +#define WGMMA_REGS_84 WGMMA_REGS_82 ", %82, %83" +#define WGMMA_ARGS_84(arg) WGMMA_ARGS_82(arg), arg(82), arg(83) +#define WGMMA_REGS_86 WGMMA_REGS_84 ", %84, %85" +#define WGMMA_ARGS_86(arg) WGMMA_ARGS_84(arg), arg(84), arg(85) +#define WGMMA_REGS_88 WGMMA_REGS_86 ", %86, %87" +#define WGMMA_ARGS_88(arg) WGMMA_ARGS_86(arg), arg(86), arg(87) +#define WGMMA_REGS_90 WGMMA_REGS_88 ", %88, %89" +#define WGMMA_ARGS_90(arg) WGMMA_ARGS_88(arg), arg(88), arg(89) +#define WGMMA_REGS_92 WGMMA_REGS_90 ", %90, %91" +#define WGMMA_ARGS_92(arg) WGMMA_ARGS_90(arg), arg(90), arg(91) +#define WGMMA_REGS_94 WGMMA_REGS_92 ", %92, %93" +#define WGMMA_ARGS_94(arg) WGMMA_ARGS_92(arg), arg(92), arg(93) +#define WGMMA_REGS_96 WGMMA_REGS_94 ", %94, %95" +#define WGMMA_ARGS_96(arg) WGMMA_ARGS_94(arg), arg(94), arg(95) +#define WGMMA_REGS_98 WGMMA_REGS_96 ", %96, %97" +#define WGMMA_ARGS_98(arg) WGMMA_ARGS_96(arg), arg(96), arg(97) +#define WGMMA_REGS_100 WGMMA_REGS_98 ", %98, %99" +#define WGMMA_ARGS_100(arg) WGMMA_ARGS_98(arg), arg(98), arg(99) +#define WGMMA_REGS_102 WGMMA_REGS_100 ", %100, %101" +#define WGMMA_ARGS_102(arg) WGMMA_ARGS_100(arg), arg(100), arg(101) +#define WGMMA_REGS_104 WGMMA_REGS_102 ", %102, %103" +#define WGMMA_ARGS_104(arg) WGMMA_ARGS_102(arg), arg(102), arg(103) +#define WGMMA_REGS_106 WGMMA_REGS_104 ", %104, %105" +#define WGMMA_ARGS_106(arg) WGMMA_ARGS_104(arg), arg(104), arg(105) +#define WGMMA_REGS_108 WGMMA_REGS_106 ", %106, %107" +#define WGMMA_ARGS_108(arg) WGMMA_ARGS_106(arg), arg(106), arg(107) +#define WGMMA_REGS_110 WGMMA_REGS_108 ", %108, %109" +#define WGMMA_ARGS_110(arg) WGMMA_ARGS_108(arg), arg(108), arg(109) +#define WGMMA_REGS_112 WGMMA_REGS_110 ", %110, %111" +#define WGMMA_ARGS_112(arg) WGMMA_ARGS_110(arg), arg(110), arg(111) +#define WGMMA_REGS_114 WGMMA_REGS_112 ", %112, %113" +#define WGMMA_ARGS_114(arg) WGMMA_ARGS_112(arg), arg(112), arg(113) +#define WGMMA_REGS_116 WGMMA_REGS_114 ", %114, %115" +#define WGMMA_ARGS_116(arg) WGMMA_ARGS_114(arg), arg(114), arg(115) +#define WGMMA_REGS_118 WGMMA_REGS_116 ", %116, %117" +#define WGMMA_ARGS_118(arg) WGMMA_ARGS_116(arg), arg(116), arg(117) +#define WGMMA_REGS_120 WGMMA_REGS_118 ", %118, %119" +#define WGMMA_ARGS_120(arg) WGMMA_ARGS_118(arg), arg(118), arg(119) +#define WGMMA_REGS_122 WGMMA_REGS_120 ", %120, %121" +#define WGMMA_ARGS_122(arg) WGMMA_ARGS_120(arg), arg(120), arg(121) +#define WGMMA_REGS_124 WGMMA_REGS_122 ", %122, %123" +#define WGMMA_ARGS_124(arg) WGMMA_ARGS_122(arg), arg(122), arg(123) +#define WGMMA_REGS_126 WGMMA_REGS_124 ", %124, %125" +#define WGMMA_ARGS_126(arg) WGMMA_ARGS_124(arg), arg(124), arg(125) +#define WGMMA_REGS_128 WGMMA_REGS_126 ", %126, %127" +#define WGMMA_ARGS_128(arg) WGMMA_ARGS_126(arg), arg(126), arg(127) +#define WGMMA_INPUTS_2 "%2, %3" +#define WGMMA_INPUTS_4 "%4, %5" +#define WGMMA_INPUTS_6 "%6, %7" +#define WGMMA_INPUTS_8 "%8, %9" +#define WGMMA_INPUTS_10 "%10, %11" +#define WGMMA_INPUTS_12 "%12, %13" +#define WGMMA_INPUTS_14 "%14, %15" +#define WGMMA_INPUTS_16 "%16, %17" +#define WGMMA_INPUTS_18 "%18, %19" +#define WGMMA_INPUTS_20 "%20, %21" +#define WGMMA_INPUTS_22 "%22, %23" +#define WGMMA_INPUTS_24 "%24, %25" +#define WGMMA_INPUTS_26 "%26, %27" +#define WGMMA_INPUTS_28 "%28, %29" +#define WGMMA_INPUTS_30 "%30, %31" +#define WGMMA_INPUTS_32 "%32, %33" +#define WGMMA_INPUTS_34 "%34, %35" +#define WGMMA_INPUTS_36 "%36, %37" +#define WGMMA_INPUTS_38 "%38, %39" +#define WGMMA_INPUTS_40 "%40, %41" +#define WGMMA_INPUTS_42 "%42, %43" +#define WGMMA_INPUTS_44 "%44, %45" +#define WGMMA_INPUTS_46 "%46, %47" +#define WGMMA_INPUTS_48 "%48, %49" +#define WGMMA_INPUTS_50 "%50, %51" +#define WGMMA_INPUTS_52 "%52, %53" +#define WGMMA_INPUTS_54 "%54, %55" +#define WGMMA_INPUTS_56 "%56, %57" +#define WGMMA_INPUTS_58 "%58, %59" +#define WGMMA_INPUTS_60 "%60, %61" +#define WGMMA_INPUTS_62 "%62, %63" +#define WGMMA_INPUTS_64 "%64, %65" +#define WGMMA_INPUTS_66 "%66, %67" +#define WGMMA_INPUTS_68 "%68, %69" +#define WGMMA_INPUTS_70 "%70, %71" +#define WGMMA_INPUTS_72 "%72, %73" +#define WGMMA_INPUTS_74 "%74, %75" +#define WGMMA_INPUTS_76 "%76, %77" +#define WGMMA_INPUTS_78 "%78, %79" +#define WGMMA_INPUTS_80 "%80, %81" +#define WGMMA_INPUTS_82 "%82, %83" +#define WGMMA_INPUTS_84 "%84, %85" +#define WGMMA_INPUTS_86 "%86, %87" +#define WGMMA_INPUTS_88 "%88, %89" +#define WGMMA_INPUTS_90 "%90, %91" +#define WGMMA_INPUTS_92 "%92, %93" +#define WGMMA_INPUTS_94 "%94, %95" +#define WGMMA_INPUTS_96 "%96, %97" +#define WGMMA_INPUTS_98 "%98, %99" +#define WGMMA_INPUTS_100 "%100, %101" +#define WGMMA_INPUTS_102 "%102, %103" +#define WGMMA_INPUTS_104 "%104, %105" +#define WGMMA_INPUTS_106 "%106, %107" +#define WGMMA_INPUTS_108 "%108, %109" +#define WGMMA_INPUTS_110 "%110, %111" +#define WGMMA_INPUTS_112 "%112, %113" +#define WGMMA_INPUTS_114 "%114, %115" +#define WGMMA_INPUTS_116 "%116, %117" +#define WGMMA_INPUTS_118 "%118, %119" +#define WGMMA_INPUTS_120 "%120, %121" +#define WGMMA_INPUTS_122 "%122, %123" +#define WGMMA_INPUTS_124 "%124, %125" +#define WGMMA_INPUTS_126 "%126, %127" +#define WGMMA_INPUTS_128 "%128, %129" -#define STORE_D_M64N8_F32() \ - do \ - { \ - for (uint32_t i = 0; i < 4; i++) \ - { \ - uint32_t row = warpid * 16 + i / 2 * 8 + laneid / 4; \ - uint32_t col = laneid % 4 * 2 + i % 2; \ - d[row * N + col] = d_frag[i]; \ - } \ - } while (0) +#define WGMMA_ARG_f32(i) "+f"(d_frag[i]) +#define WGMMA_ARG_f16(i) "+r"(d_frag[i]) -#define STORE_D_M64N8_F16() \ - do \ - { \ - for (uint32_t i = 0; i < 4; i += 2) \ - { \ - uint32_t row = warpid * 16 + i / 2 * 8 + laneid / 4; \ - uint32_t col = laneid % 4 * 2; \ - *(uint32_t *)(d + row * N + col) = d_frag[i / 2]; \ - } \ - } while (0) +template +struct WGMMAInstruction; + +#define DEFINE_WGMMA_INSTRUCTION(n, ksize, dtype, atype, btype, regs, extra) \ +template <> \ +struct WGMMAInstruction \ +{ \ + using Register = WGMMAAccumulator::Register; \ + __device__ static void run(Register *d_frag, uint64_t a_desc, uint64_t b_desc) \ + { \ + asm volatile( \ + "wgmma.mma_async.sync.aligned.m64n" #n "k" #ksize "." #dtype "." #atype "." #btype " " \ + "{" WGMMA_REGS_##regs "}, " WGMMA_INPUTS_##regs ", 1, 1, 1" extra ";" \ + : WGMMA_ARGS_##regs(WGMMA_ARG_##dtype) \ + : "l"(a_desc), "l"(b_desc)); \ + } \ +}; + +#define DEFINE_WGMMA_INSTRUCTIONS(n, fregs, hregs) \ + DEFINE_WGMMA_INSTRUCTION(n, 8, f32, tf32, tf32, fregs, "") \ + DEFINE_WGMMA_INSTRUCTION(n, 16, f32, f16, f16, fregs, ", 0, 0") \ + DEFINE_WGMMA_INSTRUCTION(n, 16, f16, f16, f16, hregs, ", 0, 0") \ + DEFINE_WGMMA_INSTRUCTION(n, 16, f32, bf16, bf16, fregs, ", 0, 0") \ + DEFINE_WGMMA_INSTRUCTION(n, 32, f32, e5m2, e5m2, fregs, "") \ + DEFINE_WGMMA_INSTRUCTION(n, 32, f32, e5m2, e4m3, fregs, "") \ + DEFINE_WGMMA_INSTRUCTION(n, 32, f32, e4m3, e5m2, fregs, "") \ + DEFINE_WGMMA_INSTRUCTION(n, 32, f32, e4m3, e4m3, fregs, "") \ + DEFINE_WGMMA_INSTRUCTION(n, 32, f16, e5m2, e5m2, hregs, "") \ + DEFINE_WGMMA_INSTRUCTION(n, 32, f16, e5m2, e4m3, hregs, "") \ + DEFINE_WGMMA_INSTRUCTION(n, 32, f16, e4m3, e5m2, hregs, "") \ + DEFINE_WGMMA_INSTRUCTION(n, 32, f16, e4m3, e4m3, hregs, "") + +DEFINE_WGMMA_INSTRUCTIONS(8, 4, 2) +DEFINE_WGMMA_INSTRUCTIONS(16, 8, 4) +DEFINE_WGMMA_INSTRUCTIONS(24, 12, 6) +DEFINE_WGMMA_INSTRUCTIONS(32, 16, 8) +DEFINE_WGMMA_INSTRUCTIONS(40, 20, 10) +DEFINE_WGMMA_INSTRUCTIONS(48, 24, 12) +DEFINE_WGMMA_INSTRUCTIONS(56, 28, 14) +DEFINE_WGMMA_INSTRUCTIONS(64, 32, 16) +DEFINE_WGMMA_INSTRUCTIONS(72, 36, 18) +DEFINE_WGMMA_INSTRUCTIONS(80, 40, 20) +DEFINE_WGMMA_INSTRUCTIONS(88, 44, 22) +DEFINE_WGMMA_INSTRUCTIONS(96, 48, 24) +DEFINE_WGMMA_INSTRUCTIONS(104, 52, 26) +DEFINE_WGMMA_INSTRUCTIONS(112, 56, 28) +DEFINE_WGMMA_INSTRUCTIONS(120, 60, 30) +DEFINE_WGMMA_INSTRUCTIONS(128, 64, 32) +DEFINE_WGMMA_INSTRUCTIONS(136, 68, 34) +DEFINE_WGMMA_INSTRUCTIONS(144, 72, 36) +DEFINE_WGMMA_INSTRUCTIONS(152, 76, 38) +DEFINE_WGMMA_INSTRUCTIONS(160, 80, 40) +DEFINE_WGMMA_INSTRUCTIONS(168, 84, 42) +DEFINE_WGMMA_INSTRUCTIONS(176, 88, 44) +DEFINE_WGMMA_INSTRUCTIONS(184, 92, 46) +DEFINE_WGMMA_INSTRUCTIONS(192, 96, 48) +DEFINE_WGMMA_INSTRUCTIONS(200, 100, 50) +DEFINE_WGMMA_INSTRUCTIONS(208, 104, 52) +DEFINE_WGMMA_INSTRUCTIONS(216, 108, 54) +DEFINE_WGMMA_INSTRUCTIONS(224, 112, 56) +DEFINE_WGMMA_INSTRUCTIONS(232, 116, 58) +DEFINE_WGMMA_INSTRUCTIONS(240, 120, 60) +DEFINE_WGMMA_INSTRUCTIONS(248, 124, 62) +DEFINE_WGMMA_INSTRUCTIONS(256, 128, 64) diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/wgmma_sm90a.cu b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/wgmma_sm90a.cu index df80e63..080d509 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/wgmma_sm90a.cu +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/kernels/wgmma_sm90a.cu @@ -1,380 +1,81 @@ #include "wgmma.h" -extern "C" // tf32 +// One full m64nNkK instruction per kernel, including the complete N-column +// accumulator. The matrix is never decomposed into smaller WGMMA shapes. +template +__global__ void wgmma_kernel(typename WGMMAAccumulator::Element *d, + const void *a, const void *b) { - __global__ void wgmma_m64n8k8_f32_tf32_tf32_kernel( - float *d, uint32_t *a, uint32_t *b) - { - const uint32_t M = 64, N = 8, K = 8; - uint32_t tid = threadIdx.x, warpid = tid / 32, laneid = tid % 32; - uint64_t a_desc, b_desc; - __shared__ uint32_t a_smem[M * K], b_smem[N * K]; - float d_frag[4]; - - LOAD_A_M64K8_TF32(); - LOAD_B_N8K8_TF32(); - LOAD_D_M64N8_F32(); - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k8.f32.tf32.tf32 " - "{%0, %1, %2, %3}, %4, %5, 1, 1, 1;" - : "+f"(d_frag[0]), "+f"(d_frag[1]), "+f"(d_frag[2]), "+f"(d_frag[3]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F32(); - } - - void wgmma_m64n8k8_f32_tf32_tf32( - float *d, uint32_t *a, uint32_t *b) - { - wgmma_m64n8k8_f32_tf32_tf32_kernel<<<1, 128>>>(d, a, b); - } + using Accumulator = WGMMAAccumulator; + constexpr int registers = N / 8 * Accumulator::registers_per_8_columns; + typename Accumulator::Register d_frag[registers]; + uint32_t tid = threadIdx.x; + __shared__ __align__(16) uint8_t a_smem[64 * 32], b_smem[N * 32]; + + uint64_t a_desc = wgmma_pack_operand(a_smem, a, 64, tid); + uint64_t b_desc = wgmma_pack_operand(b_smem, b, N, tid); + wgmma_transfer_d(d, d_frag, tid); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncthreads(); + + asm volatile("wgmma.fence.sync.aligned;" ::: "memory"); + WGMMAInstruction::run(d_frag, a_desc, b_desc); + asm volatile("wgmma.commit_group.sync.aligned;" ::: "memory"); + asm volatile("wgmma.wait_group.sync.aligned 0;" ::: "memory"); + wgmma_transfer_d(d, d_frag, tid); } -extern "C" // f16 -{ - __global__ void wgmma_m64n8k16_f32_f16_f16_kernel( - float *d, uint16_t *a, uint16_t *b) - { - const uint32_t M = 64, N = 8, K = 16; - uint32_t tid = threadIdx.x, warpid = tid / 32, laneid = tid % 32; - uint64_t a_desc, b_desc; - __shared__ uint16_t a_smem[M * K], b_smem[N * K]; - float d_frag[4]; - - LOAD_A_M64K16_F16BF16(); - LOAD_B_N8K16_F16BF16(); - LOAD_D_M64N8_F32(); - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " - "{%0, %1, %2, %3}, %4, %5, 1, 1, 1, 0, 0;" - : "+f"(d_frag[0]), "+f"(d_frag[1]), "+f"(d_frag[2]), "+f"(d_frag[3]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F32(); - } - - __global__ void wgmma_m64n8k16_f16_f16_f16_kernel( - uint16_t *d, uint16_t *a, uint16_t *b) - { - const uint32_t M = 64, N = 8, K = 16; - uint32_t tid = threadIdx.x, warpid = tid / 32, laneid = tid % 32; - uint64_t a_desc, b_desc; - __shared__ uint16_t a_smem[M * K], b_smem[N * K]; - uint32_t d_frag[2]; - - LOAD_A_M64K16_F16BF16(); - LOAD_B_N8K16_F16BF16(); - LOAD_D_M64N8_F16(); - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k16.f16.f16.f16 " - "{%0, %1}, %2, %3, 1, 1, 1, 0, 0;" - : "+r"(d_frag[0]), "+r"(d_frag[1]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F16(); - } - - void wgmma_m64n8k16_f32_f16_f16( - float *d, uint16_t *a, uint16_t *b) - { - wgmma_m64n8k16_f32_f16_f16_kernel<<<1, 128>>>(d, a, b); - } - - void wgmma_m64n8k16_f16_f16_f16( - uint16_t *d, uint16_t *a, uint16_t *b) - { - wgmma_m64n8k16_f16_f16_f16_kernel<<<1, 128>>>(d, a, b); - } +#define DEFINE_WGMMA(n, ksize, dtype, atype, btype) \ +extern "C" void wgmma_m64n##n##k##ksize##_##dtype##_##atype##_##btype( \ + WGMMAAccumulator::Element *d, void *a, void *b) \ +{ \ + wgmma_kernel \ + <<<1, 128>>>(d, a, b); \ } -extern "C" // bf16 -{ - __global__ void wgmma_m64n8k16_f32_bf16_bf16_kernel( - float *d, uint16_t *a, uint16_t *b) - { - const uint32_t M = 64, N = 8, K = 16; - uint32_t tid = threadIdx.x, warpid = tid / 32, laneid = tid % 32; - uint64_t a_desc, b_desc; - __shared__ uint16_t a_smem[M * K], b_smem[N * K]; - float d_frag[4]; - - LOAD_A_M64K16_F16BF16(); - LOAD_B_N8K16_F16BF16(); - LOAD_D_M64N8_F32(); - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k16.f32.bf16.bf16 " - "{%0, %1, %2, %3}, %4, %5, 1, 1, 1, 0, 0;" - : "+f"(d_frag[0]), "+f"(d_frag[1]), "+f"(d_frag[2]), "+f"(d_frag[3]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F32(); - } - - void wgmma_m64n8k16_f32_bf16_bf16( - float *d, uint16_t *a, uint16_t *b) - { - wgmma_m64n8k16_f32_bf16_bf16_kernel<<<1, 128>>>(d, a, b); - } -} - -extern "C" // fp8 m64n8k32 f32_output -{ - - __global__ void wgmma_m64n8k32_f32_e5m2_e5m2_kernel( - float *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; - uint64_t a_desc, b_desc; - __shared__ uint8_t a_smem[M * K], b_smem[N * K]; - float d_frag[4]; - - LOAD_A_M64K32_FP8(); - LOAD_B_N8K32_FP8(); - LOAD_D_M64N8_F32(); - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k32.f32.e5m2.e5m2 " - "{%0, %1, %2, %3}, %4, %5, 1, 1, 1;" - : "+f"(d_frag[0]), "+f"(d_frag[1]), "+f"(d_frag[2]), "+f"(d_frag[3]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F32(); - } - - __global__ void wgmma_m64n8k32_f32_e5m2_e4m3_kernel( - float *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; - uint64_t a_desc, b_desc; - __shared__ uint8_t a_smem[M * K], b_smem[N * K]; - float d_frag[4]; - - LOAD_A_M64K32_FP8(); - LOAD_B_N8K32_FP8(); - LOAD_D_M64N8_F32(); - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k32.f32.e5m2.e4m3 " - "{%0, %1, %2, %3}, %4, %5, 1, 1, 1;" - : "+f"(d_frag[0]), "+f"(d_frag[1]), "+f"(d_frag[2]), "+f"(d_frag[3]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F32(); - } - - __global__ void wgmma_m64n8k32_f32_e4m3_e5m2_kernel( - float *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; - uint64_t a_desc, b_desc; - __shared__ uint8_t a_smem[M * K], b_smem[N * K]; - float d_frag[4]; - - LOAD_A_M64K32_FP8(); - LOAD_B_N8K32_FP8(); - LOAD_D_M64N8_F32(); - __syncthreads(); - // mma - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k32.f32.e4m3.e5m2 " - "{%0, %1, %2, %3}, %4, %5, 1, 1, 1;" - : "+f"(d_frag[0]), "+f"(d_frag[1]), "+f"(d_frag[2]), "+f"(d_frag[3]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F32(); - } - - __global__ void wgmma_m64n8k32_f32_e4m3_e4m3_kernel( - float *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; - uint64_t a_desc, b_desc; - __shared__ uint8_t a_smem[M * K], b_smem[N * K]; - float d_frag[4]; - - LOAD_A_M64K32_FP8(); - LOAD_B_N8K32_FP8(); - LOAD_D_M64N8_F32(); - __syncthreads(); - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k32.f32.e4m3.e4m3 " - "{%0, %1, %2, %3}, %4, %5, 1, 1, 1;" - : "+f"(d_frag[0]), "+f"(d_frag[1]), "+f"(d_frag[2]), "+f"(d_frag[3]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F32(); - } - - void wgmma_m64n8k32_f32_e5m2_e5m2( - float *d, uint8_t *a, uint8_t *b) - { - wgmma_m64n8k32_f32_e5m2_e5m2_kernel<<<1, 128>>>(d, a, b); - } - - void wgmma_m64n8k32_f32_e5m2_e4m3( - float *d, uint8_t *a, uint8_t *b) - { - wgmma_m64n8k32_f32_e5m2_e4m3_kernel<<<1, 128>>>(d, a, b); - } - - void wgmma_m64n8k32_f32_e4m3_e5m2( - float *d, uint8_t *a, uint8_t *b) - { - wgmma_m64n8k32_f32_e4m3_e5m2_kernel<<<1, 128>>>(d, a, b); - } - - void wgmma_m64n8k32_f32_e4m3_e4m3( - float *d, uint8_t *a, uint8_t *b) - { - wgmma_m64n8k32_f32_e4m3_e4m3_kernel<<<1, 128>>>(d, a, b); - } -} - -extern "C" // fp8 m64n8k32 f16_output -{ - __global__ void wgmma_m64n8k32_f16_e5m2_e5m2_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; - uint64_t a_desc, b_desc; - __shared__ uint8_t a_smem[M * K], b_smem[N * K]; - uint32_t d_frag[2]; - - LOAD_A_M64K32_FP8(); - LOAD_B_N8K32_FP8(); - LOAD_D_M64N8_F16(); - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k32.f16.e5m2.e5m2 " - "{%0, %1}, %2, %3, 1, 1, 1;" - : "+r"(d_frag[0]), "+r"(d_frag[1]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F16(); - } - - __global__ void wgmma_m64n8k32_f16_e5m2_e4m3_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; - uint64_t a_desc, b_desc; - __shared__ uint8_t a_smem[M * K], b_smem[N * K]; - uint32_t d_frag[2]; - - LOAD_A_M64K32_FP8(); - LOAD_B_N8K32_FP8(); - LOAD_D_M64N8_F16(); - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k32.f16.e5m2.e4m3 " - "{%0, %1}, %2, %3, 1, 1, 1;" - : "+r"(d_frag[0]), "+r"(d_frag[1]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F16(); - } - - __global__ void wgmma_m64n8k32_f16_e4m3_e5m2_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; - uint64_t a_desc, b_desc; - __shared__ uint8_t a_smem[M * K], b_smem[N * K]; - uint32_t d_frag[2]; - - LOAD_A_M64K32_FP8(); - LOAD_B_N8K32_FP8(); - LOAD_D_M64N8_F16(); - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k32.f16.e4m3.e5m2 " - "{%0, %1}, %2, %3, 1, 1, 1;" - : "+r"(d_frag[0]), "+r"(d_frag[1]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F16(); - } - - __global__ void wgmma_m64n8k32_f16_e4m3_e4m3_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; - uint64_t a_desc, b_desc; - __shared__ uint8_t a_smem[M * K], b_smem[N * K]; - uint32_t d_frag[2]; - - LOAD_A_M64K32_FP8(); - LOAD_B_N8K32_FP8(); - LOAD_D_M64N8_F16(); - __syncthreads(); - asm volatile("wgmma.fence.sync.aligned;"); - asm volatile( - "wgmma.mma_async.sync.aligned.m64n8k32.f16.e4m3.e4m3 " - "{%0, %1}, %2, %3, 1, 1, 1;" - : "+r"(d_frag[0]), "+r"(d_frag[1]) - : "l"(a_desc), "l"(b_desc)); - asm volatile("wgmma.commit_group.sync.aligned;"); - asm volatile("wgmma.wait_group.sync.aligned 0;"); - STORE_D_M64N8_F16(); - } - - void wgmma_m64n8k32_f16_e5m2_e5m2( - uint16_t *d, uint8_t *a, uint8_t *b) - { - wgmma_m64n8k32_f16_e5m2_e5m2_kernel<<<1, 128>>>(d, a, b); - } - - void wgmma_m64n8k32_f16_e5m2_e4m3( - uint16_t *d, uint8_t *a, uint8_t *b) - { - wgmma_m64n8k32_f16_e5m2_e4m3_kernel<<<1, 128>>>(d, a, b); - } - - void wgmma_m64n8k32_f16_e4m3_e5m2( - uint16_t *d, uint8_t *a, uint8_t *b) - { - wgmma_m64n8k32_f16_e4m3_e5m2_kernel<<<1, 128>>>(d, a, b); - } - - void wgmma_m64n8k32_f16_e4m3_e4m3( - uint16_t *d, uint8_t *a, uint8_t *b) - { - wgmma_m64n8k32_f16_e4m3_e4m3_kernel<<<1, 128>>>(d, a, b); - } -} +#define DEFINE_WGMMA_SHAPE(n) \ + DEFINE_WGMMA(n, 8, f32, tf32, tf32) \ + DEFINE_WGMMA(n, 16, f32, f16, f16) \ + DEFINE_WGMMA(n, 16, f16, f16, f16) \ + DEFINE_WGMMA(n, 16, f32, bf16, bf16) \ + DEFINE_WGMMA(n, 32, f32, e5m2, e5m2) \ + DEFINE_WGMMA(n, 32, f32, e5m2, e4m3) \ + DEFINE_WGMMA(n, 32, f32, e4m3, e5m2) \ + DEFINE_WGMMA(n, 32, f32, e4m3, e4m3) \ + DEFINE_WGMMA(n, 32, f16, e5m2, e5m2) \ + DEFINE_WGMMA(n, 32, f16, e5m2, e4m3) \ + DEFINE_WGMMA(n, 32, f16, e4m3, e5m2) \ + DEFINE_WGMMA(n, 32, f16, e4m3, e4m3) + +DEFINE_WGMMA_SHAPE(8) +DEFINE_WGMMA_SHAPE(16) +DEFINE_WGMMA_SHAPE(24) +DEFINE_WGMMA_SHAPE(32) +DEFINE_WGMMA_SHAPE(40) +DEFINE_WGMMA_SHAPE(48) +DEFINE_WGMMA_SHAPE(56) +DEFINE_WGMMA_SHAPE(64) +DEFINE_WGMMA_SHAPE(72) +DEFINE_WGMMA_SHAPE(80) +DEFINE_WGMMA_SHAPE(88) +DEFINE_WGMMA_SHAPE(96) +DEFINE_WGMMA_SHAPE(104) +DEFINE_WGMMA_SHAPE(112) +DEFINE_WGMMA_SHAPE(120) +DEFINE_WGMMA_SHAPE(128) +DEFINE_WGMMA_SHAPE(136) +DEFINE_WGMMA_SHAPE(144) +DEFINE_WGMMA_SHAPE(152) +DEFINE_WGMMA_SHAPE(160) +DEFINE_WGMMA_SHAPE(168) +DEFINE_WGMMA_SHAPE(176) +DEFINE_WGMMA_SHAPE(184) +DEFINE_WGMMA_SHAPE(192) +DEFINE_WGMMA_SHAPE(200) +DEFINE_WGMMA_SHAPE(208) +DEFINE_WGMMA_SHAPE(216) +DEFINE_WGMMA_SHAPE(224) +DEFINE_WGMMA_SHAPE(232) +DEFINE_WGMMA_SHAPE(240) +DEFINE_WGMMA_SHAPE(248) +DEFINE_WGMMA_SHAPE(256) diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/rtx_blackwell.py b/src/mmasim_kernels/mmasim_kernels/nv_ptx/rtx_blackwell.py index 125923d..ff9de16 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/rtx_blackwell.py +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/rtx_blackwell.py @@ -2,6 +2,10 @@ from pathlib import Path from torch.utils.cpp_extension import load +from mmasim.nv_ptx.isa_mma import ( + rtx_blackwell_mma_shape_and_type, + rtx_blackwell_mma_block_scale_shape_and_type, +) from . import MMA, MMABlockScale dir = Path(__file__).with_name("kernels") @@ -15,103 +19,33 @@ is_python_module=False, ) lib = ctypes.CDLL(kernel_path) +lib.mmasim_cuda_version.argtypes = [] +lib.mmasim_cuda_version.restype = ctypes.c_int -# sm_89 fp8 m16n8k32 f32_output -lib.mma_m16n8k32_f32_e5m2_e5m2_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f32_e5m2_e4m3_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f32_e4m3_e5m2_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f32_e4m3_e4m3_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_89 fp8 m16n8k16 f32_output -lib.mma_m16n8k16_f32_e5m2_e5m2_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f32_e5m2_e4m3_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f32_e4m3_e5m2_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f32_e4m3_e4m3_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_89 fp8 m16n8k32 f16_output -lib.mma_m16n8k32_f16_e5m2_e5m2_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f16_e5m2_e4m3_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f16_e4m3_e5m2_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k32_f16_e4m3_e4m3_f16.argtypes = [ctypes.c_void_p] * 4 -# sm_89 fp8 m16n8k16 f16_output -lib.mma_m16n8k16_f16_e5m2_e5m2_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f16_e5m2_e4m3_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f16_e4m3_e5m2_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f16_e4m3_e4m3_f16.argtypes = [ctypes.c_void_p] * 4 -# sm_80 f64 -lib.mma_m8n8k4_f64_f64_f64_f64.argtypes = [ctypes.c_void_p] * 4 -# sm_80 tf32 -lib.mma_m16n8k8_f32_tf32_tf32_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k4_f32_tf32_tf32_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_80 f16 -lib.mma_m16n8k16_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k16_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 -# sm_80 bf16 -lib.mma_m16n8k16_f32_bf16_bf16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f32_bf16_bf16_f32.argtypes = [ctypes.c_void_p] * 4 -# sm_75 f16 -lib.mma_m16n8k8_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 - -# sm_120a mxf8f6f4 -lib.mma_m16n8k32_block32_f32_e5m2_e5m2_f32_ue8m0.argtypes = [ctypes.c_void_p] * 6 -lib.mma_m16n8k32_block32_f32_e4m3_e4m3_f32_ue8m0.argtypes = [ctypes.c_void_p] * 6 -# sm_120a mxf4 and nvf4 -lib.mma_m16n8k64_block32_f32_e2m1_e2m1_f32_ue8m0.argtypes = [ctypes.c_void_p] * 6 -lib.mma_m16n8k64_block16_f32_e2m1_e2m1_f32_ue4m3.argtypes = [ctypes.c_void_p] * 6 - -mma_cuda_kernels = { - # sm_89 fp8 m16n8k32 f32_output - "m16n8k32.f32.e5m2.e5m2.f32": lib.mma_m16n8k32_f32_e5m2_e5m2_f32, - "m16n8k32.f32.e5m2.e4m3.f32": lib.mma_m16n8k32_f32_e5m2_e4m3_f32, - "m16n8k32.f32.e4m3.e5m2.f32": lib.mma_m16n8k32_f32_e4m3_e5m2_f32, - "m16n8k32.f32.e4m3.e4m3.f32": lib.mma_m16n8k32_f32_e4m3_e4m3_f32, - # sm_89 fp8 m16n8k16 f32_output - "m16n8k16.f32.e5m2.e5m2.f32": lib.mma_m16n8k16_f32_e5m2_e5m2_f32, - "m16n8k16.f32.e5m2.e4m3.f32": lib.mma_m16n8k16_f32_e5m2_e4m3_f32, - "m16n8k16.f32.e4m3.e5m2.f32": lib.mma_m16n8k16_f32_e4m3_e5m2_f32, - "m16n8k16.f32.e4m3.e4m3.f32": lib.mma_m16n8k16_f32_e4m3_e4m3_f32, - # sm_89 fp8 m16n8k32 f16_output - "m16n8k32.f16.e5m2.e5m2.f16": lib.mma_m16n8k32_f16_e5m2_e5m2_f16, - "m16n8k32.f16.e5m2.e4m3.f16": lib.mma_m16n8k32_f16_e5m2_e4m3_f16, - "m16n8k32.f16.e4m3.e5m2.f16": lib.mma_m16n8k32_f16_e4m3_e5m2_f16, - "m16n8k32.f16.e4m3.e4m3.f16": lib.mma_m16n8k32_f16_e4m3_e4m3_f16, - # sm_89 fp8 m16n8k16 f16_output - "m16n8k16.f16.e5m2.e5m2.f16": lib.mma_m16n8k16_f16_e5m2_e5m2_f16, - "m16n8k16.f16.e5m2.e4m3.f16": lib.mma_m16n8k16_f16_e5m2_e4m3_f16, - "m16n8k16.f16.e4m3.e5m2.f16": lib.mma_m16n8k16_f16_e4m3_e5m2_f16, - "m16n8k16.f16.e4m3.e4m3.f16": lib.mma_m16n8k16_f16_e4m3_e4m3_f16, - # sm_80 f64 - "m8n8k4.f64.f64.f64.f64": lib.mma_m8n8k4_f64_f64_f64_f64, - # sm_80 tf32 - "m16n8k8.f32.tf32.tf32.f32": lib.mma_m16n8k8_f32_tf32_tf32_f32, - "m16n8k4.f32.tf32.tf32.f32": lib.mma_m16n8k4_f32_tf32_tf32_f32, - # sm_80 f16 - "m16n8k16.f32.f16.f16.f32": lib.mma_m16n8k16_f32_f16_f16_f32, - "m16n8k16.f16.f16.f16.f16": lib.mma_m16n8k16_f16_f16_f16_f16, - # sm_80 bf16 - "m16n8k16.f32.bf16.bf16.f32": lib.mma_m16n8k16_f32_bf16_bf16_f32, - "m16n8k8.f32.bf16.bf16.f32": lib.mma_m16n8k8_f32_bf16_bf16_f32, - # sm_75 f16 - "m16n8k8.f32.f16.f16.f32": lib.mma_m16n8k8_f32_f16_f16_f32, - "m16n8k8.f16.f16.f16.f16": lib.mma_m16n8k8_f16_f16_f16_f16, -} -mma_block_scale_cuda_kernels = { - # sm_120a mxf8f6f4 - "m16n8k32.block32.f32.e5m2.e5m2.f32.ue8m0": lib.mma_m16n8k32_block32_f32_e5m2_e5m2_f32_ue8m0, - "m16n8k32.block32.f32.e4m3.e4m3.f32.ue8m0": lib.mma_m16n8k32_block32_f32_e4m3_e4m3_f32_ue8m0, - # sm_120a mxf4 and nvf4 - "m16n8k64.block32.f32.e2m1.e2m1.f32.ue8m0": lib.mma_m16n8k64_block32_f32_e2m1_e2m1_f32_ue8m0, - "m16n8k64.block16.f32.e2m1.e2m1.f32.ue4m3": lib.mma_m16n8k64_block16_f32_e2m1_e2m1_f32_ue4m3, -} +mma_kernels = {} +for shape_and_type in rtx_blackwell_mma_shape_and_type: + kernel = lib.__getattr__( + f"mma_{shape_and_type.replace('.', '_').replace('::', '_')}" + ) + kernel.argtypes = [ctypes.c_void_p] * 4 + mma_kernels[shape_and_type] = MMA("RTX Blackwell", shape_and_type, kernel) -mma_kernels = { - shape_and_type: MMA( - "RTX Blackwell", shape_and_type, mma_cuda_kernels[shape_and_type] +mma_block_scale_kernels = {} +unavailable_mma_block_scale_kernels = {} +for shape_and_type in rtx_blackwell_mma_block_scale_shape_and_type: + # PTX ISA 9.1 introduced UE8M0 with scale_vec::4X for this instruction. + if ( + shape_and_type == "m16n8k64.kind::mxf4nvf4.block16.f32.e2m1.e2m1.f32.ue8m0" + and lib.mmasim_cuda_version() < 13010 + ): + unavailable_mma_block_scale_kernels[shape_and_type] = ( + "requires PTX ISA 9.1 / CUDA 13.1+" + ) + continue + kernel = lib.__getattr__( + f"mma_{shape_and_type.replace('.', '_').replace('::', '_')}" ) - for shape_and_type in mma_cuda_kernels -} -mma_block_scale_kernels = { - shape_and_type: MMABlockScale( - "RTX Blackwell", shape_and_type, mma_block_scale_cuda_kernels[shape_and_type] + kernel.argtypes = [ctypes.c_void_p] * 6 + mma_block_scale_kernels[shape_and_type] = MMABlockScale( + "RTX Blackwell", shape_and_type, kernel ) - for shape_and_type in mma_block_scale_cuda_kernels -} diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/turing.py b/src/mmasim_kernels/mmasim_kernels/nv_ptx/turing.py index 587d257..e436e46 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/turing.py +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/turing.py @@ -2,10 +2,11 @@ from pathlib import Path from torch.utils.cpp_extension import load +from mmasim.nv_ptx.isa_mma import turing_mma_shape_and_type from . import MMA dir = Path(__file__).with_name("kernels") -files = [dir / f for f in ["mma_sm75.cu"]] +files = [dir / f for f in ["mma_sm75.cu", "mma_sm70.cu"]] kernel_path = load( "mmasim_kernels_nv_turing", [str(f) for f in files], @@ -14,16 +15,8 @@ ) lib = ctypes.CDLL(kernel_path) -# f16 -lib.mma_m16n8k8_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m16n8k8_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 - -mma_cuda_kernels = { - # f16 - "m16n8k8.f32.f16.f16.f32": lib.mma_m16n8k8_f32_f16_f16_f32, - "m16n8k8.f16.f16.f16.f16": lib.mma_m16n8k8_f16_f16_f16_f16, -} -mma_kernels = { - shape_and_type: MMA("Turing", shape_and_type, mma_cuda_kernels[shape_and_type]) - for shape_and_type in mma_cuda_kernels -} +mma_kernels = {} +for shape_and_type in turing_mma_shape_and_type: + kernel = lib.__getattr__(f"mma_{shape_and_type.replace('.', '_')}") + kernel.argtypes = [ctypes.c_void_p] * 4 + mma_kernels[shape_and_type] = MMA("Turing", shape_and_type, kernel) diff --git a/src/mmasim_kernels/mmasim_kernels/nv_ptx/volta.py b/src/mmasim_kernels/mmasim_kernels/nv_ptx/volta.py index 40f470e..9cf9c0f 100644 --- a/src/mmasim_kernels/mmasim_kernels/nv_ptx/volta.py +++ b/src/mmasim_kernels/mmasim_kernels/nv_ptx/volta.py @@ -2,6 +2,7 @@ from pathlib import Path from torch.utils.cpp_extension import load +from mmasim.nv_ptx.isa_mma import volta_mma_shape_and_type from . import MMA dir = Path(__file__).with_name("kernels") @@ -14,18 +15,8 @@ ) lib = ctypes.CDLL(kernel_path) -# f16 -lib.mma_m8n8k4_f32_f16_f16_f32.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m8n8k4_f32_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 -lib.mma_m8n8k4_f16_f16_f16_f16.argtypes = [ctypes.c_void_p] * 4 - -mma_cuda_kernels = { - # f16 - "m8n8k4.f32.f16.f16.f32": lib.mma_m8n8k4_f32_f16_f16_f32, - "m8n8k4.f32.f16.f16.f16": lib.mma_m8n8k4_f32_f16_f16_f16, - "m8n8k4.f16.f16.f16.f16": lib.mma_m8n8k4_f16_f16_f16_f16, -} -mma_kernels = { - shape_and_type: MMA("Volta", shape_and_type, mma_cuda_kernels[shape_and_type]) - for shape_and_type in mma_cuda_kernels -} +mma_kernels = {} +for shape_and_type in volta_mma_shape_and_type: + kernel = lib.__getattr__(f"mma_{shape_and_type.replace('.', '_')}") + kernel.argtypes = [ctypes.c_void_p] * 4 + mma_kernels[shape_and_type] = MMA("Volta", shape_and_type, kernel) diff --git a/tests/equivalence/test_nv_rtx_blackwell.py b/tests/equivalence/test_nv_rtx_blackwell.py index b1b4874..557bf0e 100644 --- a/tests/equivalence/test_nv_rtx_blackwell.py +++ b/tests/equivalence/test_nv_rtx_blackwell.py @@ -1,10 +1,16 @@ -from mmasim_kernels.nv_ptx.rtx_blackwell import mma_kernels, mma_block_scale_kernels +from mmasim_kernels.nv_ptx.rtx_blackwell import ( + mma_kernels, + mma_block_scale_kernels, + unavailable_mma_block_scale_kernels, +) from mmasim.nv_ptx import MMA, MMABlockScale from helper import random_test def test_nv_rtx_blackwell(trials: int = 100): + for signature, reason in unavailable_mma_block_scale_kernels.items(): + print(f"Skipping RTX Blackwell instruction mma.{signature}: {reason}") for shape_and_type, kernel in mma_kernels.items(): print(f"Testing RTX Blackwell instruction mma.{shape_and_type}") random_test( From 4aab4a71820be6c837c9b54119dcd99f7e0fad0f Mon Sep 17 00:00:00 2001 From: Peichen Xie Date: Wed, 21 Oct 2026 10:56:14 +0800 Subject: [PATCH 4/5] fix invalid tcgen05 instructions --- src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py | 16 +++++++--------- .../nv_ptx/kernels/tcgen05mma_sm100a.cu | 2 +- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py b/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py index ac33ef9..f018eed 100644 --- a/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py +++ b/src/mmasim/mmasim/nv_ptx/isa_tcgen05mma.py @@ -11,21 +11,19 @@ # TODO: support tcgen05mma with .ws and .cta_group::2 # TODO: support sm_103/sm_107 -_tcgen05_mn = [(m, n) for m in (64, 128) for n in range(8, 257, 8)] +_tcgen05_mn = [(m, n) for m in [64, 128] for n in range(8, 257, 8)] blackwell_tcgen05mma_shape_and_type = ( [f"m{m}n{n}k8.kind::tf32.f32.tf32.tf32" for m, n in _tcgen05_mn] + [ - f"m{m}n{n}k16.kind::f16.{d_type}.{a_type}.{b_type}" + f"m{m}n{n}k16.kind::f16.f32.{ab_type}.{ab_type}" for m, n in _tcgen05_mn - for d_type in ("f32", "f16") - for a_type in ("f16", "bf16") - for b_type in ("f16", "bf16") - if d_type == "f32" or a_type == b_type == "f16" + for ab_type in ["f16", "bf16"] ] + + [f"m{m}n{n}k16.kind::f16.f16.f16.f16" for m, n in _tcgen05_mn] + [ f"m{m}n{n}k32.kind::f8f6f4.{d_type}.{a_type}.{b_type}" for m, n in _tcgen05_mn - for d_type in ("f32", "f16") + for d_type in ["f32", "f16"] for a_type in f8f6f4_types for b_type in f8f6f4_types ] @@ -47,11 +45,11 @@ for n in _tcgen05_block_scale_n ] + [ - f"m128n{n}k64.kind::mxf4nvf4.block32.f32.e2m1.e2m1.f32.ue8m0" + f"m128n{n}k64.kind::mxf4nvf4.block32.f32.e2m1.e2m1.ue8m0" for n in _tcgen05_block_scale_n ] + [ - f"m128n{n}k64.kind::mxf4nvf4.block16.f32.e2m1.e2m1.f32.{s_type}" + f"m128n{n}k64.kind::mxf4nvf4.block16.f32.e2m1.e2m1.{s_type}" for n in _tcgen05_block_scale_n for s_type in ["ue8m0", "ue4m3"] ] 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 6afc03f..1b325ad 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 @@ -145,6 +145,6 @@ DEFINE_TCGEN05_MXF8F6F4_FOR_A(e2m1) DEFINE_TCGEN05_BLOCK_SCALE(mxf4, 32, ue8m0, e2m1, e2m1) DEFINE_TCGEN05_BLOCK_SCALE(mxf4nvf4, 32, ue8m0, e2m1, e2m1) -DEFINE_TCGEN05_BLOCK_SCALE(mxf4nvf4, 16, ue8m0, e2m1, e2m1) // DEFINE_TCGEN05_BLOCK_SCALE(mxf4nvf4, 32, ue4m3, e2m1, e2m1) // requires sm_107 +DEFINE_TCGEN05_BLOCK_SCALE(mxf4nvf4, 16, ue8m0, e2m1, e2m1) DEFINE_TCGEN05_BLOCK_SCALE(mxf4nvf4, 16, ue4m3, e2m1, e2m1) From 85a37c4252ed51fee933576737c9be01df0b5a91 Mon Sep 17 00:00:00 2001 From: Peichen Xie Date: Wed, 21 Oct 2026 11:30:32 +0800 Subject: [PATCH 5/5] fix bug and update version number --- src/mmasim/mmasim/nv_ptx/sim.py | 2 ++ src/mmasim/pyproject.toml | 4 ++-- src/mmasim_kernels/pyproject.toml | 4 ++-- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/mmasim/mmasim/nv_ptx/sim.py b/src/mmasim/mmasim/nv_ptx/sim.py index c66b8f2..ed66655 100644 --- a/src/mmasim/mmasim/nv_ptx/sim.py +++ b/src/mmasim/mmasim/nv_ptx/sim.py @@ -27,6 +27,8 @@ def __init__(self, arch: str, shape_and_type: str): "RTX Blackwell": 25, } F = F_table[arch] + if shape_and_type in isa_mma.volta_mma_shape_and_type: + F = 23 # special case on Turing rho = "RNE-FP16" if self.d_type == torch.float16 else "RZ-FP32" L_max_bytes = { "Volta": 4 * 2, diff --git a/src/mmasim/pyproject.toml b/src/mmasim/pyproject.toml index 45ee3e7..2f2e041 100644 --- a/src/mmasim/pyproject.toml +++ b/src/mmasim/pyproject.toml @@ -4,10 +4,10 @@ build-backend = "setuptools.build_meta" [project] name = "mmasim" -version = "2.0.3" +version = "2.1.0" authors = [{ name="Peichen Xie", email="peichenxie@microsoft.com" }] requires-python = ">=3.10" -dependencies = ["torch>=2.9", "ninja"] +dependencies = ["torch>=2.9", "ninja", "setuptools"] [tool.setuptools.packages.find] where = ["."] diff --git a/src/mmasim_kernels/pyproject.toml b/src/mmasim_kernels/pyproject.toml index f6820a6..ee76457 100644 --- a/src/mmasim_kernels/pyproject.toml +++ b/src/mmasim_kernels/pyproject.toml @@ -4,10 +4,10 @@ build-backend = "setuptools.build_meta" [project] name = "mmasim-kernels" -version = "2.0.3" +version = "2.1.0" authors = [{ name="Peichen Xie", email="peichenxie@microsoft.com" }] requires-python = ">=3.10" -dependencies = ["mmasim>=2.0.3"] +dependencies = ["mmasim>=2.1.0"] [tool.setuptools.packages.find] where = ["."]