biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

D1.3: Fix cotiled diagnostic - use proper MMA construction

Browse Source
This commit is contained in:
biondizzle
2026-05-24 00:06:50 +00:00
parent 63e3ed0fed
commit d9f3fcd71d
1 changed files with 294 additions and 276 deletions
Download Patch File Download Diff File Expand all files Collapse all files

570
tests/unit/test_d1_3_cotiled.py
View File

@@ -1,15 +1,11 @@
"""
D1.3 SMEM-P: Diagnostic for make_cotiled_copy approach.
Goal: Build a custom R→S tiled copy that maps softmax thread registers
(TMEM-load ownership) to sP (PV A-operand SMEM with swizzle).
Steps:
1. Print tiled_tmem_load TV layout shapes
2. Print tTMEM_LOADcS coordinate partition
3. Build atom_layout_tv: (tid, vid) -> sP address
4. Create make_cotiled_copy and print partition shapes
5. Test: write P to sP via cotiled copy, read back via PV MMA, verify
Runs a minimal kernel that:
1. Prints tiled_tmem_load TV layout shapes
2. Prints tTMEM_LOADcS coordinate partition shapes
3. Prints sP layout shapes
4. Attempts to build make_cotiled_copy and prints partition shapes
"""
import torch, math
import cutlass, cutlass.cute as cute, cutlass.utils as utils
@@ -21,277 +17,299 @@ import cutlass.torch as ct
import cuda.bindings.driver as cuda
def test_cotiled_copy_diag():
"""Print TV layout shapes from TMEM load and sP to understand the mapping."""
class SmemPDiag:
def __init__(self, head_dim=64, s_k=128):
self.head_dim = head_dim
self.s_k = s_k
self.pv_n_tile = min(head_dim, 256)
self.use_smem_p = True # We're diagnosing SMEM-P
self.normalize = False
self.qk_mma_tiler = (128, 128, 128 * 4)
self.pv_mma_tiler = (128, self.pv_n_tile, 128)
self.acc_dtype = Float32
self.qk_acc_dtype = Float32
self.q_dtype = BFloat16
self.o_dtype = BFloat16
self.use_2cta_instrs = False
self.cta_group = tcgen05.CtaGroup.ONE
self.cluster_shape_mn = (1, 1)
self.epilogue_warp_id = (0, 1, 2, 3)
self.mma_warp_id = 4
self.tma_warp_id = 5
self.threads_per_cta = 192
self.num_acc_stage = 1
self.scale_softmax_log2 = 1.0 / math.sqrt(self.head_dim) * math.log2(math.e)
self.kv_stage = 2
self.q_stage = 1
self.num_c_stage = 2
self.c_layout = LayoutEnum.ROW_MAJOR
@cute.jit
def __call__(self, q, k, v, c, stream):
self.a_major = LayoutEnum.from_tensor(q).mma_major_mode()
self.b_major = LayoutEnum.from_tensor(k).mma_major_mode()
v_fmha = cute.make_tensor(
v.iterator,
cute.make_layout(
(self.pv_n_tile, self.s_k, 1),
stride=(1, self.pv_n_tile, self.pv_n_tile * self.s_k),
),
)
self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode()
qk_mma = utils.sm100.make_trivial_tiled_mma(
self.q_dtype, self.q_dtype, self.a_major, self.b_major,
self.qk_acc_dtype, self.cta_group, (128, 128), tcgen05.OperandSource.SMEM
)
pv_a_major = cute.nvgpu.OperandMajorMode.K # SMEM-P uses K major
pv_mma = utils.sm100.make_trivial_tiled_mma(
self.q_dtype, self.q_dtype, pv_a_major, self.v_major,
self.qk_acc_dtype, self.cta_group, (128, self.pv_n_tile),
tcgen05.OperandSource.SMEM
)
# Setup TMEM offsets (simplified for SMEM-P)
qk_thr = qk_mma.get_slice(0)
qk_as = qk_thr.partition_shape_C(self.qk_mma_tiler[:2])
tStS = qk_thr.make_fragment_C(qk_as)
pv_thr = pv_mma.get_slice(0)
pv_as = pv_thr.partition_shape_C(self.pv_mma_tiler[:2])
tOtO = pv_thr.make_fragment_C(pv_as)
self.tmem_s0_offset = 0
self.tmem_p0_offset = -1 # SMEM-P
self.tmem_o0_offset = 0
s_cols = self.qk_mma_tiler[1]
o_cols = find_tmem_tensor_col_offset(tOtO)
total = max(s_cols, o_cols)
self.num_tmem_alloc_cols = 1
while self.num_tmem_alloc_cols < total:
self.num_tmem_alloc_cols *= 2
self.tOrP0_offset = 0 # SMEM-P
# Build SMEM layouts
p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1)
q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage)
k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, self.q_dtype, self.kv_stage)
v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, self.kv_stage)
c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, (128, self.pv_n_tile), 2)
# TMA atoms
q_s = cute.slice_(q_smem_s, (None, None, None, 0))
k_s = cute.slice_(k_smem_s, (None, None, None, 0))
v_s = cute.slice_(v_smem_s, (None, None, None, 0))
cta = cute.size(qk_mma.thr_id.shape)
self.q_tx_bytes = cute.size_in_bytes(self.q_dtype, q_s) * cta
self.kv_tx_bytes = (cute.size_in_bytes(self.q_dtype, k_s) + cute.size_in_bytes(self.q_dtype, v_s)) * cta
tma_q, mQ = cute.nvgpu.make_tiled_tma_atom_A(
utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id),
q, q_s, self.qk_mma_tiler, qk_mma, (1, 1, 1, 1)
)
tma_k, mK = cute.nvgpu.make_tiled_tma_atom_B(
utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id),
k, k_s, self.qk_mma_tiler, qk_mma, (1, 1, 1, 1)
)
tma_v, mV = cute.nvgpu.make_tiled_tma_atom_B(
utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id),
v_fmha, v_s, self.pv_mma_tiler, pv_mma, (1, 1, 1, 1)
)
epi_s = cute.select(c_smem_s, mode=[0, 1])
tma_c, mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_s, (128, self.pv_n_tile))
# ===== KERNEL =====
self._kernel(qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC,
p_smem_s, q_smem_s, k_smem_s, v_smem_s, c_smem_s,
tStS, tOtO, qk_thr, pv_thr).launch(
grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream
)
@cute.kernel
def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC,
p_smem_s, q_smem_s, k_smem_s, v_smem_s, c_smem_s,
tStS, tOtO, qk_thr, pv_thr):
tidx, _, _ = cute.arch.thread_idx()
warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
@cute.struct
class SS:
q_bar: cute.struct.MemRange[cutlass.Int64, self.q_stage * 2]
kv_bar: cute.struct.MemRange[cutlass.Int64, self.kv_stage * 2]
s_bar: cute.struct.MemRange[cutlass.Int64, 2]
acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
tmem_dealloc: cutlass.Int64
holding: cutlass.Int32
smem = utils.SmemAllocator()
st = smem.allocate(SS)
sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=q_smem_s.outer, byte_alignment=128, swizzle=q_smem_s.inner)
sK = smem.allocate_tensor(element_type=self.q_dtype, layout=k_smem_s.outer, byte_alignment=128, swizzle=k_smem_s.inner)
sV = smem.allocate_tensor(element_type=self.q_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner)
sP = smem.allocate_tensor(element_type=self.q_dtype, layout=p_smem_s.outer, byte_alignment=128, swizzle=p_smem_s.inner)
# ===== PRINT DIAGNOSTICS (only from warp 0, thread 0) =====
if warp_idx < 4: # softmax warps
# TMEM load atoms
tmem_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS)
sfw_idx = tidx % (32 * len(self.epilogue_warp_id))
thr_load = tiled_tmem_load.get_slice(sfw_idx)
tTMEM_LOADtS = thr_load.partition_S(tStS)
cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1]))
tScS = qk_thr.partition_C(cS)
tTMEM_LOADcS = thr_load.partition_D(tScS)
if sfw_idx == 0:
print(f"=== TMEM Load Diagnostics (sfw_idx=0) ===")
print(f"tiled_tmem_load shape: {cute.shape(tiled_tmem_load)}")
print(f"tTMEM_LOADtS shape: {cute.shape(tTMEM_LOADtS)}")
print(f"tTMEM_LOADtS layout: {tTMEM_LOADtS.layout}")
print(f"tTMEM_LOADcS shape: {cute.shape(tTMEM_LOADcS)}")
print(f"tTMEM_LOADcS layout: {tTMEM_LOADcS.layout}")
# Print TV layout
tv = tiled_tmem_load.layout_dst_tv_tiled
print(f"TV layout: {tv}")
print(f"TV layout shape: {cute.shape(tv)}")
# Print sP layout
sP_stage = sP[(None, None, None, 0)]
print(f"sP shape: {cute.shape(sP)}")
print(f"sP layout (outer): {sP.layout}")
print(f"sP_stage shape: {cute.shape(sP_stage)}")
print(f"sP_stage layout: {sP_stage.layout}")
# Print tStS layout
print(f"tStS shape: {cute.shape(tStS)}")
print(f"tStS layout: {tStS.layout}")
# Try to build make_cotiled_copy
# atom_layout_tv: (tid, vid) -> sP flat address
# We need to map from the TV layout's codomain (tStS addresses)
# to sP addresses.
#
# tStS layout: ((128,128),1,1):((65536,1),0,0)
# So tStS_addr(m,k) = m*65536 + k
# sP layout: ((128,16),1,(4,2),1):(((64,1),0,((16,8192),0)
# So sP_addr((m,k0),0,(k1,k2),0) = m*64 + k0 + k1*16 + k2*8192
# where k0=k%16, k1=(k//16)%4, k2=k//64
#
# We can't directly compose because the codomains are different.
# But we CAN build a new Layout that maps (tid, vid) -> sP addr.
#
# The key: enumerate all (tid, vid) pairs, find their (m, k) from
# tTMEM_LOADcS, then compute sP_addr.
# But we can't do this at Python trace time inside @cute.kernel
# because we don't have the actual coordinate values.
#
# OR: we can try composition. Let's see if
# cute.composition(sP_stage.layout, tv) works.
# This requires the codomain of tv to be compatible with
# the domain of sP_stage.layout.
print(f"\n=== Attempting composition ===")
print(f"TV codomain size: {cute.size(tv)}")
print(f"sP_stage domain size: {cute.size(sP_stage.layout)}")
# Try: make a layout that maps (m, k) -> sP_addr
# then compose with the inverse of tStS.layout
# Actually, tStS has 3 modes. Let me flatten it.
# The 2D logical P is (128, 128) = (M, K)
# tStS has layout ((128,128),1,1):((65536,1),0,0)
# The 2D part is (128,128) with strides (65536, 1)
# So the 2D layout is: (128, 128) : (65536, 1)
# sP_stage has shape ((128,16),1,(4,2),1)
# The logical P is 128 x 128 where:
# mode0: (128, 16) with strides (64, 1)
# mode2: (4, 2) with strides (16, 8192)
# So the 2D logical layout (grouping modes 0 and 2) is:
# (128*4, 16*2) = (512, 32) with... complex strides.
# Actually, sP is 4D, not 2D.
# Let me try the group_modes approach:
sP_2d = cute.group_modes(sP_stage, 0, 4) # flatten to 1D? No.
# group_modes(sP, 0, 4) makes it (N,) where N = total elements
# Actually group_modes(sP, 0, 4) groups the first 4 modes into 1
# and keeps the rest. sP_stage has 4 modes -> becomes 1 mode.
# That's just a 1D tensor.
print(f"sP_2d (grouped) shape: {cute.shape(sP_2d)}")
print(f"sP_2d layout: {sP_2d.layout}")
# Now try to compose the TV layout with the identity tensor's layout
# to get (tid, vid) -> (m, k), then use that to index into sP.
# Actually, the right approach per the CUTLASS LLM:
# 1. Build a custom atom_layout_tv from scratch
# 2. Each thread's (m, k) pairs come from tTMEM_LOADcS
# 3. Convert (m, k) to sP address
# 4. Build a Layout((n_threads, n_vals_per_thread), (sP_addr_strides,))
#
# But building this at trace time requires knowing all (m, k) values.
# The identity tensor partition tTMEM_LOADcS has these values
# but they're runtime values in CuTeDSL, not Python values.
#
# HOWEVER: since tTMEM_LOADcS is an identity tensor partition,
# the coordinates are STATIC (known at compile/trace time).
# We should be able to extract them at Python trace time.
#
# Let me check: can we iterate over tTMEM_LOADcS at trace time?
# tTMEM_LOADcS has shape ((32,1),4,1,1)
# Each element is a (m, k) pair.
# At trace time, cute.shape gives us the static shape.
# But the actual coordinate VALUES might be dynamic.
#
# In CuTe, identity tensors have static values.
# cute.make_identity_tensor((M, K)) creates a tensor where
# the value at each coordinate IS the coordinate.
# So tTMEM_LOADcS[j0, j1, 0, 0] should give a static (m, k) pair.
# At trace time, we should be able to extract these.
#
# But wait - at trace time, the identity tensor is being traced.
# The values are MLIR constants, not Python ints.
# We can't iterate and build a Python Layout from them.
#
# The alternative: define atom_layout_tv analytically.
# The Ld32x32bOp with Repetition(32) partitions 128 threads
# over a 128x128 matrix. We need to figure out the exact
# (tid, vid) -> (m, k) mapping analytically.
# Let me print what we know about the TMEM load's thread mapping.
# The layout_dst_tv_tiled is a Layout with shape (n_threads, n_vals).
print(f"\n=== Analyzing TV layout for analytical construction ===")
print(f"TV layout type: {type(tv)}")
# Try to decompose it
if hasattr(tv, 'shape'):
print(f" shape: {tv.shape}")
if hasattr(tv, 'stride'):
print(f" stride: {tv.stride}")
def test_cotiled_diag():
print("=== make_cotiled_copy Diagnostic ===\n")
hd = 64
q = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda')
k = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda')
v = torch.randn(128, min(hd, 256), dtype=torch.bfloat16, device='cuda')
c = torch.zeros(128, min(hd, 256), 1, dtype=torch.bfloat16, device='cuda')
head_dim = 64 # Start with proven hd=64 (TMEM-P works at cos 0.973)
s_k = 128
pv_n_tile = min(head_dim, 256)
qk_mma_tiler = (128, 128, 128 * 4)
pv_mma_tiler = (128, pv_n_tile, 128)
diag = SmemPDiag(head_dim=hd, s_k=128)
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
# Build MMA objects
a_major = LayoutEnum.ROW_MAJOR
b_major = LayoutEnum.ROW_MAJOR
v_major = LayoutEnum.ROW_MAJOR
v_tile = 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))
qk_mma = utils.sm100.make_trivial_tiled_mma(
BFloat16, BFloat16, a_major, b_major, Float32,
tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM
)
pv_mma = utils.sm100.make_trivial_tiled_mma(
BFloat16, BFloat16, a_major, v_major, Float32,
tcgen05.CtaGroup.ONE, (128, pv_n_tile), tcgen05.OperandSource.SMEM
)
# Build SMEM layouts
p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1)
# Build QK C-fragment and TMEM load partition
qk_thr = qk_mma.get_slice(0)
qk_as = qk_thr.partition_shape_C(qk_mma_tiler[:2])
tStS = qk_thr.make_fragment_C(qk_as)
# TMEM load atoms
tmem_load_atom = cute.make_copy_atom(
tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), Float32
)
tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS)
# Print the TV layout of the TMEM load
print(f"tiled_tmem_load shape: {cute.shape(tiled_tmem_load)}")
tv_layout = tiled_tmem_load.layout_dst_tv_tiled
print(f"TV layout: {tv_layout}")
print(f"TV layout shape: {cute.shape(tv_layout)}")
# Print per-thread coordinate partition
sfw_idx = 0 # first softmax thread
thr_load = tiled_tmem_load.get_slice(sfw_idx)
tTMEM_LOADtS = thr_load.partition_S(tStS)
print(f"\ntTMEM_LOADtS shape (thread 0): {cute.shape(tTMEM_LOADtS)}")
print(f"tTMEM_LOADtS layout: {tTMEM_LOADtS.layout}")
cS = cute.make_identity_tensor((qk_mma_tiler[0], qk_mma_tiler[1]))
tScS = qk_thr.partition_C(cS)
tTMEM_LOADcS = thr_load.partition_D(tScS)
print(f"tTMEM_LOADcS shape (thread 0): {cute.shape(tTMEM_LOADcS)}")
print(f"tTMEM_LOADcS layout: {tTMEM_LOADcS.layout}")
# Print sP layout
sP_shape = p_smem_s.outer
print(f"\nsP outer shape: {cute.shape(sP_shape)}")
print(f"sP outer layout: {sP_shape}")
print(f"sP inner (swizzle): {p_smem_s.inner}")
# Print what each softmax thread's coordinates look like
# For 128 threads, each with ((32,1),4,1,1) = 128 elements
print(f"\n=== Coordinate analysis ===")
print(f"Total softmax threads: 128 (warps 0-3)")
print(f"Per thread: {cute.size(tTMEM_LOADcS)} coordinate pairs")
print(f"Total coordinates: 128 * {cute.size(tTMEM_LOADcS)} = {128 * cute.size(tTMEM_LOADcS)}")
print(f"Expected: 128*128 = {128*128}")
# The key question: can we build an atom_layout_tv that maps
# (tid, vid) -> sP address?
#
# For make_cotiled_copy:
# atom_layout_tv: (tid, vid) -> data address in sP's codomain
# data_layout: data coord -> data address (this is sP's layout)
#
# The TMEM load's TV layout maps (tid, vid) -> tStS0 coordinates.
# tStS0 is a TMEM tensor with layout ((128,128),1,1):((65536,1),0,0)
# So (tid, vid) -> flat TMEM address.
# But we need (tid, vid) -> flat sP address.
#
# The connection: tStS0 stores S with logical (m, k) = (128, 128).
# sP stores P with the same logical (m, k) but in SMEM layout.
# So the mapping is:
# (tid, vid) -> tStS0 address -> (m, k) via inverse of tStS0.layout
# (m, k) -> sP address via sP.layout
#
# But computing the inverse of tStS0 layout at Python time is the challenge.
# Let's try a different approach: build the atom_layout_tv directly.
# The TMEM load has 128 threads and 128 values per thread.
# Total values = 128 * 128 = 16384 = 128 * 128 ✓
#
# We need: atom_layout_tv such that for thread tid and value vid,
# the output is the flat address in sP's layout.
#
# The TMEM load's layout_dst_tv_tiled already maps (tid, vid) -> tStS0 flat addr.
# But tStS0 flat addr = m * 65536 + k (from layout ((128,128):((65536,1)))
# So we can extract (m, k) from the address: m = addr // 65536, k = addr % 65536
# Wait, that's not right. The layout is ((128,128),1,1):((65536,1),0,0)
# So the address for coordinate (m, k) is m * 65536 + k * 1
# Meaning: addr = m * 65536 + k
# So: m = addr // 65536, k = addr % 1... no, stride of k is 1.
# Actually: addr = m * 65536 + k. So m = addr // 65536, k = addr % 65536
# But k ranges from 0 to 127, so k = addr % 65536 (which should be < 128).
# And m ranges from 0 to 127, so m = addr // 65536 (which should be < 128).
# Wait, that doesn't work because addr could be huge.
# Let me reconsider.
# tStS layout: ((128,128),1,1):((65536,1),0,0)
# This is a 3-mode tensor. For coordinate ((m, k), i, j):
# addr = m * 65536 + k * 1 + i * 0 + j * 0
# So effectively: addr = m * 65536 + k
#
# sP layout (from p_smem_s.outer): ((128,16),1,(4,2),1):(((64,1),0,((16,8192),0)
# For coordinate ((m, k0), i, (k1, k2), j):
# addr = m * 64 + k0 * 1 + i * 0 + k1 * 16 + k2 * 8192 + j * 0
# addr = m * 64 + k0 + k1 * 16 + k2 * 8192
#
# Given (m, k) from TMEM load coords:
# k0 = k % 16, k1 = (k // 16) % 4, k2 = k // 64
# sP_addr = m * 64 + (k % 16) + ((k // 16) % 4) * 16 + (k // 64) * 8192
#
# But we need to account for the swizzle! The swizzle is applied by
# CuTe automatically when you index into sP. So the sP.layout already
# includes the swizzle in the address computation.
#
# Actually, the swizzle is in p_smem_s.inner, not in p_smem_s.outer.
# The outer layout is the logical layout (no swizzle).
# When we allocate sP with swizzle=p_smem_s.inner, the indexing
# automatically applies the swizzle XOR.
#
# So for make_cotiled_copy, data_layout should be the COMPOSED layout
# (outer with swizzle applied). Let me check how CUTLASS handles this.
#
# Actually, in CuTe, when you do cute.make_tensor(ptr, layout, swizzle),
# the swizzle is applied during tensor creation. The layout is the
# pre-swizzle layout, and the swizzle XOR is applied on top.
#
# For make_cotiled_copy, we need to think about what "address" means.
# The atom_layout_tv maps (tid, vid) to data addresses.
# If data_layout includes the swizzle, then the addresses are post-swizzle.
# If not, they're pre-swizzle.
#
# The safest approach: use the 2D sP without the stage dimension,
# and let CuTe handle swizzle via the tensor's layout+swizzle.
#
# Actually, let me look at what make_cotiled_copy expects.
# The docstring says: "atom_layout_tv: (tid, vid) -> data addr"
# and "data_layout: data coord -> data addr"
# So the atom_layout_tv's codomain should match data_layout's codomain.
#
# For sP with swizzle, the "data addr" is the swizzled address.
# So we need to compose: (tid, vid) -> (m, k) -> swizzled sP addr.
#
# This is getting complex. Let me try the simpler approach first:
# use the current coordinate-indexed write but verify it works.
# Then come back and optimize with make_cotiled_copy.
print("\n=== Attempting make_cotiled_copy ===")
# Step 1: Build atom_layout_tv from TMEM load's TV layout.
# The TMEM load's layout_dst_tv_tiled maps (tid, vid) -> tStS0 flat address.
# We need to transform this to map (tid, vid) -> sP flat address.
#
# The transformation is:
# tStS0_addr = m * 65536 + k (from tStS0 layout)
# sP_addr = m * 64 + (k % 16) + ((k // 16) % 4) * 16 + (k // 64) * 8192
#
# This is NOT a simple layout composition because the (m, k) decomposition
# from tStS0_addr is not trivial (stride 65536 for m).
#
# Alternative: Use make_cotiled_copy with the TMEM load's TV layout
# but change the data_layout to something that maps tStS0 addresses to
# the same codomain as sP addresses.
#
# Actually, the cleanest approach is to build atom_layout_tv from scratch
# using the coordinate information from tTMEM_LOADcS.
#
# Each softmax thread owns 128 (m, k) pairs.
# We can enumerate all 128 * 128 = 16384 (tid, (m, k)) pairs
# and compute the sP address for each.
# But in CuTeDSL, we can't iterate and build layouts at Python time
# inside @cute.jit. We need to build the layout at Python (trace) time.
# Let me try a different approach: use make_tiled_copy_tv.
# thr_layout: maps (TileM, TileN) -> tid
# val_layout: maps (ValueM, ValueN) -> vid
#
# The softmax threads are indexed 0..127 (4 warps × 32 threads).
# Each thread owns a (32, 4) sub-tile of the P matrix (from tTMEM_LOADcS shape ((32,1),4)).
# So thr_layout should map (32, 4) tiles to 128 threads.
# And val_layout should map (32, 1) values per tile position.
#
# Wait, the TMEM load's coordinate partition is ((32,1),4,1,1).
# This means each thread has 32 × 4 = 128 coordinate pairs.
# The first mode (32,1) is the "row" within a fragment.
# The second mode 4 is the "fragment" index.
#
# For make_tiled_copy_tv:
# - thr_layout: how threads tile the (M, K) = (128, 128) P matrix
# - val_layout: how values are arranged within a thread's tile
#
# From the TMEM load partition, each thread owns:
# - 32 M-values (not necessarily contiguous)
# - 4 K-fragments
# The thread layout is determined by the Ld32x32bOp atom's thread mapping.
# Let me just print the actual TV layout to understand it.
print(f"tiled_tmem_load layout_dst_tv shape: {cute.shape(tv_layout)}")
if hasattr(tv_layout, 'a'):
print(f" a (thread): {tv_layout.a}")
if hasattr(tv_layout, 'b'):
print(f" b (value): {tv_layout.b}")
# Print sP composed layout (with swizzle)
# The key insight from the CUTLASS LLM: we need atom_layout_tv
# such that atom_layout_tv(tid, vid) gives a coordinate in sP's codomain.
#
# But actually, for make_cotiled_copy, the atom_layout_tv maps to
# "data addr" which is the same codomain as data_layout.
# data_layout maps data coordinates to addresses.
# So atom_layout_tv(tid, vid) should produce an address.
#
# The TV layout from tiled_tmem_load maps (tid, vid) to tStS0 addresses.
# If we could remap tStS0 addresses to sP addresses, we'd have it.
#
# tStS0 addresses: m * 65536 + k (m in [0,128), k in [0,128))
# sP addresses: m * 64 + (k % 16) + ((k // 16) % 4) * 16 + (k // 64) * 8192
# + swizzle XOR
#
# The mapping is: tStS0_addr -> (m, k) -> sP_coord -> sP_addr
# tStS0_addr = m * 65536 + k
# m = tStS0_addr // 65536
# k = tStS0_addr % 65536 (but should be < 128)
# Actually since k < 128 and stride is 1, k = tStS0_addr % 65536 is correct
# and since 128 < 65536, this gives k directly.
# And m = tStS0_addr // 65536 (since k < 65536, integer division works).
# The problem: computing m = addr // 65536 and k = addr % 65536
# from a Layout object is not straightforward. Layouts are affine maps.
# Division/modulo by 65536 is not affine in general.
#
# BUT: the TMEM load's TV layout already encodes the (tid, vid) -> addr map
# as a Layout. We need to transform this Layout to produce sP addresses
# instead of tStS0 addresses.
#
# Let me try yet another approach: just print what we have and think.
print("\n=== Checking if we can extract (m,k) from TV layout ===")
# The TV layout has shape (128_threads, 128_values) mapping to addresses.
# If we can reshape this to (128, 128) and interpret the output as (m, k)
# coordinates, then compose with sP's layout, we get what we need.
# Actually, let me try the simplest possible thing:
# Print the existing TV layout and see if it's compatible with sP
# in any way.
try:
sP_stage = p_smem_s.outer
# Try to see what the layout composition would look like
print(f"sP_stage layout: {sP_stage}")
print(f"TV layout: {tv_layout}")
# Can we compose tv_layout with the inverse of tStS.layout, then with sP.layout?
# cute.composition(sP_stage, tv_layout) might work if the codomains match
print(f"tStS layout: {tStS.layout}")
except Exception as e:
print(f"Error: {e}")
print('Compiling diagnostic kernel...', flush=True)
compiled = cute.compile(diag, mQ, mK, mV, mC, stream)
print('Running...', flush=True)
compiled(mQ, mK, mV, mC, stream)
torch.cuda.synchronize()
print('Done.')
if __name__ == '__main__':
test_cotiled_copy_diag()
test_cotiled_diag()