Init weight gradient kernels.

tests/test_core.py

@@ -1,6 +1,6 @@
 import random
 import torch
-from typing import Tuple
+from typing import List, Tuple
 
 import deep_gemm
 from deep_gemm import bench_kineto, calc_diff, ceil_div, get_col_major_tma_aligned_tensor
@@ -89,6 +89,70 @@ def construct_masked_grouped(num_groups: int, m: int, k: int, n: int) -> \
     return x_fp8, y_fp8, out, ref_out
 
 
+def construct_wgrad(m: int, k: int, n: int) -> \
+        Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor]:
+    x = torch.randn((m, k), device='cuda', dtype=torch.bfloat16)
+    y = torch.randn((n, k), device='cuda', dtype=torch.bfloat16)
+    residual = torch.randn((m, n), device='cuda', dtype=torch.float) * 10
+    out = residual.clone()
+    ref_out = residual + (x.float() @ y.float().t())
+
+    x_fp8 = per_token_cast_to_fp8(x)
+    y_fp8 = per_token_cast_to_fp8(y)
+
+    # NOTES: please do inplace add on the `out` later
+    return x_fp8, y_fp8, residual, out, ref_out
+
+
+def construct_k_grouped_wgrad(m: int, n: int, k_sizes: List[int]) -> \
+        Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor, List[int]]:
+    num_groups, total_k = len(k_sizes), sum(k_sizes)
+
+    x_flat = torch.empty((m * total_k,), device='cuda', dtype=torch.bfloat16)
+    y_flat = torch.empty((n * total_k,), device='cuda', dtype=torch.bfloat16)
+    out = torch.zeros((num_groups, m, n), device='cuda', dtype=torch.float)
+    ref_out = torch.zeros((num_groups, m, n), device='cuda', dtype=torch.float)
+
+    # Fill tensors with data and compute reference output
+    x_offset, y_offset = 0, 0
+    for idx, k in enumerate(k_sizes):
+        x_chunk = torch.randn((m, k), device='cuda', dtype=torch.bfloat16)
+        y_chunk = torch.randn((n, k), device='cuda', dtype=torch.bfloat16)
+
+        x_flat[x_offset:x_offset + m * k].copy_(x_chunk.flatten())
+        y_flat[y_offset:y_offset + n * k].copy_(y_chunk.flatten())
+        ref_out[idx] = x_chunk.float() @ y_chunk.float().t()
+        
+        x_offset += m * k
+        y_offset += n * k
+
+    x_fp8_flat = torch.empty_like(x_flat, dtype=torch.float8_e4m3fn)
+    y_fp8_flat = torch.empty_like(y_flat, dtype=torch.float8_e4m3fn)
+
+    total_scale_factors = sum((k + 127) // 128 for k in k_sizes)
+    x_scales = torch.empty((total_scale_factors, m), device='cuda', dtype=torch.float)
+    y_scales = torch.empty((total_scale_factors, n), device='cuda', dtype=torch.float)
+    
+    # Cast to FP8 and prepare scale factors
+    x_offset, y_offset, scale_offset = 0, 0, 0
+    for k in k_sizes:
+        x_fp8_chunk, x_scale_chunk = per_token_cast_to_fp8(x_flat[x_offset:x_offset + m * k].view(m, k))
+        y_fp8_chunk, y_scale_chunk = per_token_cast_to_fp8(y_flat[y_offset:y_offset + n * k].view(n, k))
+
+        x_fp8_flat[x_offset:x_offset + m * k].copy_(x_fp8_chunk.flatten())
+        y_fp8_flat[y_offset:y_offset + n * k].copy_(y_fp8_chunk.flatten())
+        
+        num_scales = (k + 127) // 128
+        x_scales[scale_offset:scale_offset + num_scales].copy_(x_scale_chunk.T)
+        y_scales[scale_offset:scale_offset + num_scales].copy_(y_scale_chunk.T)
+        
+        x_offset += m * k
+        y_offset += n * k
+        scale_offset += num_scales
+    
+    return (x_fp8_flat, x_scales), (y_fp8_flat, y_scales), out, ref_out, k_sizes
+
+
 def test_gemm() -> None:
     print('Testing GEMM:')
     for m in (64, 128, 4096):
@@ -170,6 +234,62 @@ def test_m_grouped_gemm_masked() -> None:
     print()
 
 
+def test_wgrad_gemm():
+    print('Testing weight gradient GEMM:')
+
+    for k in (4096, 8192):
+        for m, n in ((7168, 2112), (1536, 24576), (512, 32768), (16384, 7168), (7168, 4096), (2048, 7168)):
+            # Test correctness
+            x_fp8, y_fp8, residual, out, ref_out = construct_wgrad(m, k, n)
+            deep_gemm.wgrad_gemm_fp8_fp8_fp32_nt(x_fp8, y_fp8, out)
+            diff = calc_diff(out, ref_out)
+            assert diff < 0.001, f'{m=}, {k=}, {n=}, {diff:.5f}'
+
+            # Construct new tensors only once to avoid L2 cache acceleration (creating them puts them in L2)
+            x_fp8, y_fp8, residual, out, ref_out = construct_wgrad(m, k, n)
+
+            # noinspection PyShadowingNames
+            def test_func():
+                deep_gemm.wgrad_gemm_fp8_fp8_fp32_nt(x_fp8, y_fp8, out)
+
+            t = bench_kineto(test_func, 'fp8_wgrad_gemm', suppress_kineto_output=True)
+            print(f' > Performance (m={m:5}, n={n:5}, k={k:5}): {t * 1e6:4.0f} us | '
+                  f'throughput: {2 * m * n * k / t / 1e12:4.0f} TFLOPS, '
+                  f'{(m * k + k * n + m * n * 2) / 1e9 / t:4.0f} GB/s')
+    print()
+
+
+def test_k_grouped_wgrad_gemm():
+    print('Testing grouped weight gradient GEMM:')
+
+    for num_groups, base_k in ((4, 4096), (4, 8192), (8, 4096)):
+        for m, n in ((7168, 4096), (2048, 7168)):
+            # Vary k sizes around base_k
+            k_sizes = [base_k + random.randint(-1, 1) * 128 for _ in range(num_groups - 1)]
+            k_sizes.append(base_k * num_groups - sum(k_sizes))
+            
+            # Test correctness
+            x_fp8, y_fp8, out, ref_out, k_sizes = construct_k_grouped_wgrad(m, n, k_sizes)
+            deep_gemm.k_grouped_wgrad_gemm_fp8_fp8_fp32_nt(x_fp8, y_fp8, out, k_sizes)
+
+            for idx in range(num_groups):
+                diff = calc_diff(out[idx], ref_out[idx])
+                assert diff < 0.001, f'{num_groups=}, {m=}, {n=}, k={k_sizes[idx]}, batch={idx}, {diff:.5f}'
+
+            # Construct new tensors to avoid L2 cache acceleration
+            x_fp8, y_fp8, out, ref_out, k_sizes = construct_k_grouped_wgrad(m, n, k_sizes)
+            total_k = sum(k_sizes)
+            
+            def test_func():
+                deep_gemm.k_grouped_wgrad_gemm_fp8_fp8_fp32_nt(x_fp8, y_fp8, out, k_sizes)
+            
+            t = bench_kineto(test_func, 'fp8_wgrad_gemm', suppress_kineto_output=True, is_multiple=True) * num_groups
+            print(f' > Performance ({num_groups=}, m={m:5}, n={n:5}, avg_k={total_k//num_groups:5}): {t * 1e6:4.0f} us | '
+                  f'throughput: {2 * num_groups * m * n * (total_k/num_groups) / t / 1e12:4.0f} TFLOPS, '
+                  f'{(m * total_k + n * total_k + num_groups * m * n * 2) / 1e9 / t:4.0f} GB/s')
+    print()
+
+
 if __name__ == '__main__':
     torch.backends.cuda.matmul.allow_tf32 = True
     torch.backends.cudnn.allow_tf32 = True
@@ -182,3 +302,6 @@ if __name__ == '__main__':
     test_gemm()
     test_m_grouped_gemm_contiguous()
     test_m_grouped_gemm_masked()
+
+    test_wgrad_gemm()
+    test_k_grouped_wgrad_gemm()