99 lines
4.0 KiB
Python
99 lines
4.0 KiB
Python
"""Diagnostic: print TMA partition tensor shapes for multi-tile K/V."""
|
|
import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils
|
|
from cutlass.cute.nvgpu import cpasync, tcgen05
|
|
from cutlass import Float32, BFloat16, Int32
|
|
import cutlass.torch as ct
|
|
import math
|
|
|
|
HEAD_DIM = 64
|
|
n = 256
|
|
|
|
q = torch.randn(128, HEAD_DIM, 1, dtype=torch.bfloat16, device='cuda')
|
|
k = torch.randn(n, HEAD_DIM, 1, dtype=torch.bfloat16, device='cuda')
|
|
v = torch.randn(n, HEAD_DIM, dtype=torch.bfloat16, device='cuda')
|
|
v_kernel = v.unsqueeze(-1)
|
|
|
|
mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
|
|
mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
|
|
mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
|
|
|
|
# Hardcode major modes since LayoutEnum.from_tensor needs JIT context
|
|
qk_mma = utils.sm100.make_trivial_tiled_mma(BFloat16, BFloat16, cute.nvgpu.OperandMajorMode.K, cute.nvgpu.OperandMajorMode.K, Float32, tcgen05.CtaGroup.ONE, (128,128), tcgen05.OperandSource.SMEM)
|
|
pv_mma = utils.sm100.make_trivial_tiled_mma(BFloat16, BFloat16, cute.nvgpu.OperandMajorMode.K, cute.nvgpu.OperandMajorMode.MN, Float32, tcgen05.CtaGroup.ONE, (128,HEAD_DIM), tcgen05.OperandSource.TMEM)
|
|
|
|
qk_ik = cute.size(qk_mma.shape_mnk, mode=[2])
|
|
qk_mma_tiler = (128, 128, qk_ik * 4)
|
|
pv_ik = cute.size(pv_mma.shape_mnk, mode=[2])
|
|
pv_mma_tiler = (128, HEAD_DIM, pv_ik * (128 // pv_ik))
|
|
cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,))
|
|
|
|
print(f'qk_mma_tiler: {qk_mma_tiler}')
|
|
print(f'pv_mma_tiler: {pv_mma_tiler}')
|
|
|
|
kv_stage = 2
|
|
k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, qk_mma_tiler, BFloat16, kv_stage)
|
|
v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, pv_mma_tiler, BFloat16, kv_stage)
|
|
|
|
k_s = cute.slice_(k_smem_s,(None,None,None,0))
|
|
v_s = cute.slice_(v_smem_s,(None,None,None,0))
|
|
|
|
print(f'k_s shape: {cute.shape(k_s)}')
|
|
print(f'v_s shape: {cute.shape(v_s)}')
|
|
|
|
tma_k, mK_tma = cute.nvgpu.make_tiled_tma_atom_B(
|
|
utils.sm100.cluster_shape_to_tma_atom_B(cluster_layout_vmnk.shape, qk_mma.thr_id),
|
|
mK, k_s, qk_mma_tiler, qk_mma, cluster_layout_vmnk.shape
|
|
)
|
|
|
|
gK = cute.local_tile(mK_tma, cute.slice_(qk_mma_tiler,(0,None,None)),(None,None,None))
|
|
print(f'mK_tma shape: {cute.shape(mK_tma)}')
|
|
print(f'gK shape: {cute.shape(gK)}')
|
|
print(f'n_kv_tiles: {cute.size(gK, mode=[3])}')
|
|
|
|
qk_thr = qk_mma.get_slice(0)
|
|
tCgK = qk_thr.partition_B(gK)
|
|
print(f'tCgK shape: {cute.shape(tCgK)}')
|
|
|
|
sK = cute.make_tensor(BFloat16, k_s)
|
|
b_lay = cute.make_layout(cute.slice_(cluster_layout_vmnk,(0,None,0,0)).shape)
|
|
|
|
tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3))
|
|
|
|
print(f'tBsK shape: {cute.shape(tBsK)}')
|
|
print(f'tBgK shape: {cute.shape(tBgK)}')
|
|
print(f'tBsK layout: {tBsK.layout}')
|
|
print(f'tBgK layout: {tBgK.layout}')
|
|
|
|
# Test slices
|
|
for desc, sl in [
|
|
("(None,0,None,0)", (None,0,None,0)), # Current (broken)
|
|
("(None,None,0,0)", (None,None,0,0)), # CUTLASS reference style
|
|
("(0,None,None,0)", (0,None,None,0)), # Alternative
|
|
]:
|
|
try:
|
|
result = tBgK[sl]
|
|
print(f'tBgK after {desc} shape: {cute.shape(result)}')
|
|
except Exception as e:
|
|
print(f'tBgK after {desc}: ERROR {type(e).__name__}: {e}')
|
|
|
|
# Also check V
|
|
tma_v, mV_tma = cute.nvgpu.make_tiled_tma_atom_B(
|
|
utils.sm100.cluster_shape_to_tma_atom_B(cluster_layout_vmnk.shape, pv_mma.thr_id),
|
|
mV, v_s, pv_mma_tiler, pv_mma, cluster_layout_vmnk.shape
|
|
)
|
|
gV = cute.local_tile(mV_tma, cute.slice_(pv_mma_tiler,(0,None,None)),(None,None,None))
|
|
pv_thr = pv_mma.get_slice(0)
|
|
tCgV = pv_thr.partition_B(gV)
|
|
sV = cute.make_tensor(BFloat16, v_s)
|
|
tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3))
|
|
print(f'tVgV shape: {cute.shape(tVgV)}')
|
|
for desc, sl in [
|
|
("(None,0,None,0)", (None,0,None,0)),
|
|
("(None,None,0,0)", (None,None,0,0)),
|
|
]:
|
|
try:
|
|
result = tVgV[sl]
|
|
print(f'tVgV after {desc} shape: {cute.shape(result)}')
|
|
except Exception as e:
|
|
print(f'tVgV after {desc}: ERROR {type(e).__name__}: {e}')
|