Skip to content

Commit 4dfd274

Browse files
gideon116gideon116
committed
* added nvfp4 mma samples and benchmark. Updated mma util functions.
Signed-off-by: Gideon Kassa <gkassa@nvidia.com>
1 parent 03193ae commit 4dfd274

11 files changed

Lines changed: 1208 additions & 271 deletions

‎samples/BlockScaledMatMul.py‎

Lines changed: 102 additions & 105 deletions
Large diffs are not rendered by default.

‎samples/NVFP4ScaledMatMul.py‎

Lines changed: 440 additions & 0 deletions
Large diffs are not rendered by default.

‎samples/templates/BlockScaledMatMul.py‎

Lines changed: 46 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -9,31 +9,32 @@
99
from cuda.tile._cext import get_compute_capability
1010
from cuda.tile._bytecode.version import BytecodeVersion
1111

12-
13-
from test.kernels.kernel_utils import block_quantize, swizzle_32_4_4, get_tileiras_version
12+
from test.kernels.kernel_utils import (block_quantize_f8e4m3fn_f8e8m0fnu,
13+
swizzle_32_4_4, get_tileiras_version)
1414
from test.kernels.scaled_matmul import block_scaled_matmul_kernel
1515

1616

1717
def cutile_block_scaled_matmul(A: torch.Tensor, A_scale: torch.Tensor,
18-
B: torch.Tensor, B_scale: torch.Tensor) -> torch.Tensor:
18+
B: torch.Tensor, B_scale: torch.Tensor,
19+
scaling_block_size: int) -> torch.Tensor:
1920

2021
"""
2122
Performs block-scaled matrix multiplication using a cuTile kernel.
2223
23-
This wrapper function handles input validation, determines appropriate
24-
tile sizes based on data type, calculates the necessary grid dimensions,
24+
This wrapper function handles input validation, calculates the necessary grid dimensions,
2525
and launches the `block_scaled_matmul_kernel`.
2626
2727
Args:
28-
A (torch.Tensor): The first input matrix (M x K). Must be on a CUDA device.
29-
B (torch.Tensor): The second input matrix (K x N). Must be on a CUDA device
30-
and have its K dimension match A's K dimension.
31-
A_scale (torch.Tensor): Either 2D scale with shape (M, K // scaling_block_size) or
32-
swizzled scale of (M // scaling_block_size // 4,
33-
K // scaling_block_size // 4, 32, 16).
34-
B_scale (torch.Tensor): Either 2D scale with shape (M, K // scaling_block_size) or
35-
swizzled scale of (M // scaling_block_size // 4,
36-
K // scaling_block_size // 4, 32, 16).
28+
A (torch.Tensor): The first input matrix (M x K). Must be on a CUDA device.
29+
A_scale (torch.Tensor): Either 2D scale with shape (M, K // scaling_block_size) or
30+
swizzled scale of (M // 32 // 4,
31+
K // scaling_block_size // 4, 32, 16).
32+
B (torch.Tensor): The second input matrix (N x K). Must be on a CUDA device
33+
and have its K dimension match A's K dimension.
34+
B_scale (torch.Tensor): Either 2D scale with shape (N, K // scaling_block_size) or
35+
swizzled scale of (N // 32 // 4,
36+
K // scaling_block_size // 4, 32, 16).
37+
scaling_block_size (int): The scaling block size.
3738
3839
Returns:
3940
torch.Tensor: The resulting matrix C (M x N) on the CUDA device.
@@ -43,21 +44,20 @@ def cutile_block_scaled_matmul(A: torch.Tensor, A_scale: torch.Tensor,
4344
or if they are not on a CUDA device.
4445
"""
4546
# --- Input Validation ---
46-
if A.shape[1] != B.shape[0]:
47-
raise ValueError(f"Incompatible matrices: K dimension of A ({A.shape[1]}) "
48-
f"must match K dimension of B ({B.shape[0]})")
47+
if A.shape[1] != B.shape[1]:
48+
raise ValueError("Incompatible matrices")
4949
if A.device != B.device or A.device != A_scale.device or A.device != B_scale.device:
5050
raise ValueError("Input tensors must be on the same device.")
5151
if not A.is_cuda or not A_scale.is_cuda or not B.is_cuda or not B_scale.is_cuda:
5252
raise ValueError("Input tensors must be on a CUDA device.")
5353
# Note: cuTile handles dtype compatibility within the kernel,
5454
# but inputs should generally match.
5555

56-
tm, tn, tk, scaling_block_size = 256, 256, 128, 32
56+
tm, tn, tk = 256, 256, 128
5757

5858
# --- Get Matrix Dimensions ---
5959
m, _ = A.shape
60-
_, n = B.shape
60+
n, _ = B.shape
6161

6262
# --- Calculate Grid Dimensions for Kernel Launch (1D Grid) ---
6363
# The grid defines how many CUDA blocks (CTAs) will be launched.
@@ -72,15 +72,17 @@ def cutile_block_scaled_matmul(A: torch.Tensor, A_scale: torch.Tensor,
7272

7373
# --- Create Output Tensor C ---
7474
# The output tensor `C` is initialized with the correct dimensions (M x N),
75-
# on the same device, and with the same data type as the input matrices.
75+
# on the same device, and with a datatype of float32.
7676
C = torch.empty((m, n), device=A.device, dtype=torch.float32)
7777

7878
# --- Launch the cuTile Kernel ---
7979
# The `block_scaled_matmul_kernel` is launched with the calculated grid dimensions.
8080
# `tm`, `tn`, and `tk` are passed as Constant integers to the kernel.
8181
kernel = block_scaled_matmul_kernel
8282
ct.launch(torch.cuda.current_stream(), grid, kernel, (
83-
A, A_scale, B, B_scale, C, tm, tn, tk, scaling_block_size))
83+
A, A_scale,
84+
B, B_scale,
85+
C, tm, tn, tk, scaling_block_size))
8486

8587
return C
8688

@@ -101,7 +103,7 @@ def cutile_block_scaled_matmul(A: torch.Tensor, A_scale: torch.Tensor,
101103

102104
if get_tileiras_version() < BytecodeVersion.V_13_3:
103105
print("Skipped test: NOT Running cuTile Block Scaled Matrix Multiplication Examples "
104-
"tileiras versiom 13.3 required.")
106+
"tileiras version 13.3 required.")
105107
sys.exit(0)
106108

107109
# --- Running cuTile Block Scaled Matrix Multiplication Examples ---
@@ -115,41 +117,49 @@ def cutile_block_scaled_matmul(A: torch.Tensor, A_scale: torch.Tensor,
115117
scaling_block_size = 32
116118
KS_dim = K_dim // scaling_block_size
117119

118-
print("\n--- Test Case: Block Scaled Matrix Multiplication with M = 512, N = 512, "
119-
"K = 768, Scaling Block Size = 32 ---")
120+
print(f"\n--- Test Case: Block Scaled Matrix Multiplication with M = {M_dim}, N = {N_dim}, "
121+
f"K = {K_dim}, Scaling Block Size = {scaling_block_size} ---")
120122

121123
A = torch.rand((M_dim, K_dim), device='cuda:0')
122124
B = torch.rand((N_dim, K_dim), device='cuda:0')
123125

124-
A, A_scale = block_quantize(A, scaling_block_size)
125-
B, B_scale = block_quantize(B, scaling_block_size)
126+
uncompressed_ref = A @ B.T
126127

127-
B = B.T
128-
B_scale = B_scale.T
128+
A, A_scale = block_quantize_f8e4m3fn_f8e8m0fnu(A, scaling_block_size)
129+
B, B_scale = block_quantize_f8e4m3fn_f8e8m0fnu(B, scaling_block_size)
129130

130131
k = A.shape[-1]
131132
ks = k // scaling_block_size
132133

133134
A_s_swizzled = swizzle_32_4_4(A_scale)
134-
B_s_swizzled = swizzle_32_4_4(B_scale.T.contiguous())
135+
B_s_swizzled = swizzle_32_4_4(B_scale)
135136

136137
print(f"Input A shape: {A.shape}, dtype: {A.dtype}")
137138
print(f"Input B shape: {B.shape}, dtype: {B.dtype}")
138139

139-
atol, rtol = 1e-4, 1e-3
140+
kernel_atol, kernel_rtol = 1e-4, 1e-3
141+
compression_atol, compression_rtol = 0.5, 0.05
140142

141143
# Perform matrix multiplication using the cuTile wrapper function.
142-
C_cutile = cutile_block_scaled_matmul(A, A_scale, B, B_scale)
143-
C_cutile_swizzled = cutile_block_scaled_matmul(A, A_s_swizzled, B, B_s_swizzled)
144+
C_cutile = cutile_block_scaled_matmul(A, A_scale, B, B_scale, scaling_block_size)
145+
torch.cuda.synchronize()
146+
C_cutile_swizzled = cutile_block_scaled_matmul(A, A_s_swizzled, B, B_s_swizzled,
147+
scaling_block_size)
148+
torch.cuda.synchronize()
144149
print(f"cuTile Output C shape: {C_cutile.shape}, dtype: {C_cutile.dtype}")
145150

146151
if args.correctness_check:
147152
ref_A_scale = torch.repeat_interleave(A_scale, scaling_block_size, dim=1).to(torch.float32)
148-
ref_B_scale = torch.repeat_interleave(B_scale, scaling_block_size, dim=0).to(torch.float32)
149-
ref = (A.to(torch.float32) * ref_A_scale) @ (B.to(torch.float32) * ref_B_scale)
153+
ref_B_scale = torch.repeat_interleave(B_scale, scaling_block_size, dim=1).to(torch.float32)
154+
ref = (A.to(torch.float32) * ref_A_scale) @ (B.T.to(torch.float32) * ref_B_scale.T)
155+
156+
torch.testing.assert_close(C_cutile, ref, atol=kernel_atol, rtol=kernel_rtol)
157+
torch.testing.assert_close(C_cutile_swizzled, ref, atol=kernel_atol, rtol=kernel_rtol)
150158

151-
torch.testing.assert_close(C_cutile, ref, atol=atol, rtol=rtol)
152-
torch.testing.assert_close(C_cutile_swizzled, ref, atol=atol, rtol=rtol)
159+
torch.testing.assert_close(C_cutile, uncompressed_ref,
160+
atol=compression_atol, rtol=compression_rtol)
161+
torch.testing.assert_close(C_cutile_swizzled, uncompressed_ref,
162+
atol=compression_atol, rtol=compression_rtol)
153163
print("Correctness check passed")
154164
else:
155165
print("Correctness check disabled")

‎samples/templates/NVFP4ScaledMatMul.py‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
#
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
import argparse
6+
import cuda.tile as ct
7+
import torch
8+
import sys
9+
from cuda.tile._cext import get_compute_capability
10+
from cuda.tile._bytecode.version import BytecodeVersion
11+
12+
from test.kernels.kernel_utils import swizzle_32_4_4, get_tileiras_version, \
13+
unpack_e2m1_bytes_to_float, block_quantize_f4e2m1fn_f8e4m3fn
14+
from test.kernels.scaled_matmul import nvfp4_block_scaled_matmul_kernel
15+
16+
17+
def cutile_nvfp4_matmul(A: torch.Tensor, A_scale: torch.Tensor, A_global: torch.Tensor,
18+
B: torch.Tensor, B_scale: torch.Tensor, B_global: torch.Tensor,
19+
scaling_block_size: int) -> torch.Tensor:
20+
21+
"""
22+
Performs NVFP4 matrix multiplication using a cuTile kernel.
23+
24+
This wrapper function handles input validation, calculates the necessary grid dimensions,
25+
and launches the `nvfp4_block_scaled_matmul_kernel`.
26+
27+
Args:
28+
A (torch.Tensor): Packed input matrix with physical shape (M, K // 2).
29+
Must be on a CUDA device.
30+
A_scale (torch.Tensor): Either 2D scale with shape (M, K // NVFP4_BLOCK_SIZE) or swizzled
31+
scale of (M // 32 // 4, K // NVFP4_BLOCK_SIZE // 4, 32, 16).
32+
A_global (torch.Tensor): One-element float32 tensor containing A's tensor-wide scale.
33+
B (torch.Tensor): Packed input matrix with physical shape (N, K // 2).
34+
Must be on a CUDA device and have its 2nd dimension match
35+
A's 2nd dimension.
36+
B_scale (torch.Tensor): Either 2D scale with shape (N, K // NVFP4_BLOCK_SIZE) or swizzled
37+
scale of (N // 32 // 4, K // NVFP4_BLOCK_SIZE // 4, 32, 16).
38+
B_global (torch.Tensor): One-element float32 tensor containing B's tensor-wide scale.
39+
scaling_block_size (int): The scaling block size.
40+
41+
Returns:
42+
torch.Tensor: The resulting matrix C (M x N) on the CUDA device.
43+
44+
Raises:
45+
ValueError: If matrices are incompatible (K dimensions don't match),
46+
or if they are not on a CUDA device.
47+
"""
48+
# --- Input Validation ---
49+
if A.shape[1] != B.shape[1]:
50+
raise ValueError("Incompatible matrices")
51+
if not (A.device == A_scale.device == A_global.device ==
52+
B.device == B_scale.device == B_global.device):
53+
raise ValueError("Input tensors must be on the same device.")
54+
if not (A.is_cuda and A_scale.is_cuda and A_global.is_cuda and
55+
B.is_cuda and B_scale.is_cuda and B_global.is_cuda):
56+
raise ValueError("Input tensors must be on a CUDA device.")
57+
58+
A = A.view(torch.uint8)
59+
B = B.view(torch.uint8)
60+
61+
# Note: cuTile handles dtype compatibility within the kernel,
62+
# but inputs should generally match.
63+
tm, tn, tk = 256, 256, 256
64+
65+
# --- Get Matrix Dimensions ---
66+
m, _ = A.shape
67+
n, _ = B.shape
68+
69+
# --- Calculate Grid Dimensions for Kernel Launch (1D Grid) ---
70+
# The grid defines how many CUDA blocks (CTAs) will be launched.
71+
# Each block computes one (tm x tn) output tile of matrix C.
72+
# `ct.cdiv(total_dim, tile_dim)` ensures enough blocks are launched to cover
73+
# the entire matrix, even if dimensions are not perfect multiples of tile sizes.
74+
grid_x = ct.cdiv(m, tm) # Number of blocks needed along the M dimension (rows of C)
75+
grid_y = ct.cdiv(n, tn) # Number of blocks needed along the N dimension (columns of C)
76+
grid_size = grid_x * grid_y
77+
78+
grid = (grid_size, 1, 1)
79+
80+
# --- Create Output Tensor C ---
81+
# The output tensor `C` is initialized with the correct dimensions (M x N),
82+
# on the same device, and with a datatype of float32.
83+
C = torch.empty((m, n), device=A.device, dtype=torch.float32)
84+
85+
# --- Launch the cuTile Kernel ---
86+
# The `nvfp4_block_scaled_matmul_kernel` is launched with the calculated grid dimensions.
87+
# `tm`, `tn`, and `tk` are passed as Constant integers to the kernel.
88+
kernel = nvfp4_block_scaled_matmul_kernel
89+
ct.launch(torch.cuda.current_stream(), grid, kernel, (
90+
A, A_scale, A_global, B, B_scale, B_global, C, tm, tn, tk, scaling_block_size))
91+
return C
92+
93+
94+
if __name__ == "__main__":
95+
parser = argparse.ArgumentParser()
96+
parser.add_argument(
97+
"--correctness-check",
98+
action="store_true",
99+
help="Check the correctness of the results",
100+
)
101+
args = parser.parse_args()
102+
103+
if get_compute_capability()[0] < 10:
104+
print("Skipped test: NOT Running cuTile NVFP4 Matrix Multiplication Examples "
105+
"Blackwell or newer required.")
106+
sys.exit(0)
107+
108+
if get_tileiras_version() < BytecodeVersion.V_13_4:
109+
print("Skipped test: NOT Running cuTile NVFP4 Matrix Multiplication Examples "
110+
"tileiras version 13.4 required.")
111+
sys.exit(0)
112+
113+
# --- Running cuTile NVFP4 Matrix Multiplication Examples ---
114+
print("--- Running cuTile NVFP4 Matrix Multiplication Examples ---")
115+
116+
# Define common matrix dimensions for the examples
117+
M_dim = 512
118+
N_dim = 512
119+
K_dim = 768
120+
121+
scaling_block_size = 16
122+
123+
print(f"\n--- Test Case: NVFP4 Matrix Multiplication with M = {M_dim}, N = {N_dim}, "
124+
f"K = {K_dim}, Scaling Block Size = {scaling_block_size} ---")
125+
126+
A = torch.rand((M_dim, K_dim), device='cuda:0')
127+
B = torch.rand((N_dim, K_dim), device='cuda:0')
128+
129+
uncompressed_ref = A @ B.T
130+
131+
A, A_scale, A_global = block_quantize_f4e2m1fn_f8e4m3fn(A, scaling_block_size)
132+
B, B_scale, B_global = block_quantize_f4e2m1fn_f8e4m3fn(B, scaling_block_size)
133+
134+
A_s_swizzled = swizzle_32_4_4(A_scale)
135+
B_s_swizzled = swizzle_32_4_4(B_scale)
136+
137+
print(f"Input A packed shape: {A.shape}, dtype: {A.dtype}")
138+
print(f"Input B packed shape: {B.shape}, dtype: {B.dtype}")
139+
140+
kernel_atol, kernel_rtol = 1e-4, 1e-3
141+
compression_atol, compression_rtol = 0.5, 0.05
142+
143+
# Perform NVFP4 matrix multiplication using the cuTile wrapper function.
144+
C_cutile = cutile_nvfp4_matmul(A, A_scale, A_global, B, B_scale, B_global, scaling_block_size)
145+
torch.cuda.synchronize()
146+
C_cutile_swizzled = cutile_nvfp4_matmul(A, A_s_swizzled, A_global, B, B_s_swizzled, B_global,
147+
scaling_block_size)
148+
torch.cuda.synchronize()
149+
print(f"cuTile Output C shape: {C_cutile.shape}, dtype: {C_cutile.dtype}")
150+
151+
if args.correctness_check:
152+
ref_A_scale = torch.repeat_interleave(A_scale, scaling_block_size, dim=1).to(torch.float32)
153+
ref_B_scale = torch.repeat_interleave(B_scale, scaling_block_size, dim=1).to(torch.float32)
154+
ref_A = unpack_e2m1_bytes_to_float(A.view(torch.uint8))
155+
ref_B = unpack_e2m1_bytes_to_float(B.view(torch.uint8))
156+
ref = (ref_A * ref_A_scale * A_global) @ (ref_B.T * ref_B_scale.T * B_global)
157+
158+
torch.testing.assert_close(C_cutile, ref, atol=kernel_atol, rtol=kernel_rtol)
159+
torch.testing.assert_close(C_cutile_swizzled, ref, atol=kernel_atol, rtol=kernel_rtol)
160+
161+
torch.testing.assert_close(C_cutile, uncompressed_ref,
162+
atol=compression_atol, rtol=compression_rtol)
163+
torch.testing.assert_close(C_cutile_swizzled, uncompressed_ref,
164+
atol=compression_atol, rtol=compression_rtol)
165+
print("Correctness check passed")
166+
else:
167+
print("Correctness check disabled")
168+
169+
print("\n--- cuTile NVFP4 matrix multiplication example completed. ---")

‎samples/test_samples.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,8 @@
2828
"LayerNorm.py",
2929
"MoE.py",
3030
"AllGatherMatmul.py",
31-
"BlockScaledMatMul.py"
31+
"BlockScaledMatMul.py",
32+
"NVFP4ScaledMatMul.py",
3233
]
3334

3435
# Get the absolute path of the current directory to ensure the script

0 commit comments

Comments
 (0)