add CuTe descriptor printing script
This commit is contained in:
87
scripts/compute_swizzle.py
Normal file
87
scripts/compute_swizzle.py
Normal file
@@ -0,0 +1,87 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Print UMMA SMEM layout offsets for FMHA decode.
|
||||
Directly compute the swizzle pattern without CuTeDSL imports.
|
||||
"""
|
||||
import sys
|
||||
|
||||
# SWIZZLE_128B atom layout for BF16 MN-major:
|
||||
# Atom shape: (1024, 8) BF16, stride (1, 1024)
|
||||
# tile_to_shape for (128, 64):
|
||||
# m is always < 1024, so m/1024 = 0
|
||||
# n needs 64/8 = 8 atoms
|
||||
# offset = m + (n % 8) * 1024 + (n / 8) * 8192
|
||||
#
|
||||
# Swizzle<3,4,3>:
|
||||
# swizzled = offset ^ ((((offset >> 3) ^ (offset >> 7)) & 0x7) << 3)
|
||||
|
||||
def swizzle_3_4_3(offset):
|
||||
"""Apply Swizzle<3,4,3> to an element offset."""
|
||||
return offset ^ ((((offset >> 3) ^ (offset >> 7)) & 0x7) << 3)
|
||||
|
||||
def mn_sw128_offset(row, col, hd=64):
|
||||
"""Compute swizzled SMEM offset for MN-major (128, HD) BF16 matrix."""
|
||||
# Logical offset in BF16 elements
|
||||
logical = row + (col % 8) * 1024 + (col // 8) * 8192
|
||||
# Apply swizzle
|
||||
return swizzle_3_4_3(logical)
|
||||
|
||||
def k_sw128_offset(row, col, hd=64):
|
||||
"""Compute swizzled SMEM offset for K-major (128, HD) BF16 matrix.
|
||||
K_SW128 atom: (8, 1024) BF16, stride (1, 8)
|
||||
For (128, 64):
|
||||
k ranges 0..127, mn ranges 0..63
|
||||
logical = (k % 8) + (mn % 1024) * 8 + (k / 8) * 8 * 1024
|
||||
= (k % 8) + mn * 8 + (k / 8) * 8192
|
||||
"""
|
||||
logical = (row % 8) + col * 8 + (row // 8) * 8192
|
||||
return swizzle_3_4_3(logical)
|
||||
|
||||
# Print Q (MN-major) swizzle pattern
|
||||
print("=== MN_SW128 Q layout (row, col) -> offset (BF16 elements) ===")
|
||||
max_offset = 0
|
||||
for row in [0, 1, 2, 3, 127]:
|
||||
for col in [0, 1, 7, 8, 15, 16, 63]:
|
||||
offset = mn_sw128_offset(row, col)
|
||||
max_offset = max(max_offset, offset)
|
||||
print(f" ({row:3d},{col:2d}) -> {offset:6d}", end="")
|
||||
print()
|
||||
|
||||
print(f"\nMax Q offset: {max_offset} BF16 = {max_offset * 2} bytes = {max_offset * 2 / 1024:.1f} KB")
|
||||
|
||||
# Print K (K-major) swizzle pattern
|
||||
print("\n=== K_SW128 K layout (row, col) -> offset (BF16 elements) ===")
|
||||
max_offset_k = 0
|
||||
for row in [0, 1, 7, 8, 15, 127]:
|
||||
for col in [0, 1, 7, 8, 63]:
|
||||
offset = k_sw128_offset(row, col)
|
||||
max_offset_k = max(max_offset_k, offset)
|
||||
print(f" ({row:3d},{col:2d}) -> {offset:6d}", end="")
|
||||
print()
|
||||
|
||||
print(f"\nMax K offset: {max_offset_k} BF16 = {max_offset_k * 2} bytes = {max_offset_k * 2 / 1024:.1f} KB")
|
||||
|
||||
# Verify: does the swizzle stay within the matrix size?
|
||||
matrix_size = 128 * 64 # 8192 BF16
|
||||
print(f"\nMatrix size: {matrix_size} BF16 = {matrix_size * 2} bytes = {matrix_size * 2 / 1024:.1f} KB")
|
||||
print(f"Q layout footprint: {max_offset + 1} BF16 ({(max_offset + 1) * 2 / 1024:.1f} KB)")
|
||||
print(f"K layout footprint: {max_offset_k + 1} BF16 ({(max_offset_k + 1) * 2 / 1024:.1f} KB)")
|
||||
|
||||
# Check if any offsets exceed the simple row-major layout
|
||||
row_major_max = 128 * 64 - 1 # 8191
|
||||
q_overflow = any(mn_sw128_offset(r, c) > row_major_max
|
||||
for r in range(128) for c in range(64))
|
||||
k_overflow = any(k_sw128_offset(r, c) > row_major_max
|
||||
for r in range(128) for c in range(64))
|
||||
print(f"\nQ offsets exceed row-major bounds: {q_overflow}")
|
||||
print(f"K offsets exceed row-major bounds: {k_overflow}")
|
||||
|
||||
# Count unique offsets (verify no collisions)
|
||||
q_offsets = set()
|
||||
k_offsets = set()
|
||||
for r in range(128):
|
||||
for c in range(64):
|
||||
q_offsets.add(mn_sw128_offset(r, c))
|
||||
k_offsets.add(k_sw128_offset(r, c))
|
||||
print(f"Q unique offsets: {len(q_offsets)} (expected {128*64})")
|
||||
print(f"K unique offsets: {len(k_offsets)} (expected {128*64})")
|
||||
53
scripts/print_umma_desc_cute.py
Normal file
53
scripts/print_umma_desc_cute.py
Normal file
@@ -0,0 +1,53 @@
|
||||
"""
|
||||
Minimal CuTeDSL script to print UMMA descriptors for FMHA Q and K.
|
||||
Run on B200 with: python3 print_umma_desc_cute.py
|
||||
"""
|
||||
import sys
|
||||
sys.path.insert(0, '/root/cutlass/python/CuTeDSL')
|
||||
|
||||
import torch
|
||||
import cutlass
|
||||
import cute
|
||||
|
||||
# Configuration matching our FMHA decode
|
||||
HEAD_DIM = 64
|
||||
M = 128
|
||||
N = 128 # KV tile
|
||||
|
||||
# Create a simple kernel that prints the UMMA descriptors
|
||||
@cute.jit
|
||||
def print_umma_descriptors():
|
||||
# Construct SMEM layouts the same way the FMHA does
|
||||
# Using the SM100 utils
|
||||
from cutlass.utils.blackwell_helpers import make_smem_layout_a, make_smem_layout_b, OperandMajorMode
|
||||
|
||||
# Q: MN-major, (M, HEAD_DIM) BF16
|
||||
q_layout = make_smem_layout_a(
|
||||
major_mode=OperandMajorMode.MN_MAJOR,
|
||||
smem_tile_shape=cute.make_shape(M, HEAD_DIM),
|
||||
element_type=cutlass.float16, # BF16
|
||||
stage=1,
|
||||
)
|
||||
|
||||
# K: K-major, (N, HEAD_DIM) BF16
|
||||
k_layout = make_smem_layout_b(
|
||||
major_mode=OperandMajorMode.K_MAJOR,
|
||||
smem_tile_shape=cute.make_shape(N, HEAD_DIM),
|
||||
element_type=cutlass.float16,
|
||||
stage=1,
|
||||
)
|
||||
|
||||
cute.printf("Q layout shape: {}", cute.shape(q_layout))
|
||||
cute.printf("Q layout size: {}", cute.size(q_layout))
|
||||
cute.printf("K layout shape: {}", cute.shape(k_layout))
|
||||
cute.printf("K layout size: {}", cute.size(k_layout))
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Just trigger the JIT compilation to see the prints
|
||||
print("Compiling...")
|
||||
try:
|
||||
print_umma_descriptors()
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
Reference in New Issue
Block a user