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

D1: SMEM-P using make_tiled_copy_C + retile (correct CUTLASS pattern)

Browse Source
This commit is contained in:
biondizzle
2026-05-24 03:32:03 +00:00
parent d0aec403e4
commit 5e154770de
2 changed files with 47 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
dsv4/kernels/attention/fmha.py
View File

@@ -366,16 +366,19 @@ class FmhaKernel:
cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
cute.arch.fence_view_async_tmem_store()
else:
# SMEM-P: write P to sP using coordinate-indexed store.
for j0 in range(32):
for j1 in range(4):
coord = tTMEM_LOADcS[(j0, 0), j1, 0, 0]
m_coord = coord[0]
k_coord = coord[1]
k0 = k_coord % 16
k1 = (k_coord // 16) % 4
k2 = k_coord // 64
_sP_nostage[(m_coord, k0), 0, (k1, k2)] = rP_bf16[(j0, 0), j1, 0, 0]
# SMEM-P: write P to sP using make_tiled_copy_C + retile.
# The retile() call transforms rP_bf16 from QK C-fragment layout
# to the SMEM copy's source layout, matching partition_D(sP).
_smem_p_store_atom = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
self.q_dtype,
num_bits_per_copy=16,
)
_tiled_smem_p = cute.make_tiled_copy_C(_smem_p_store_atom, qk_mma)
_thr_smem_p = _tiled_smem_p.get_slice(sfw_idx)
_tRS_sP = _thr_smem_p.partition_D(_sP_nostage)
_tRS_rP = _tiled_smem_p.retile(rP_bf16)
cute.copy(_tiled_smem_p, _tRS_rP, _tRS_sP)
cute.arch.fence_proxy("async.shared", space="cta")
if kt > 0:
for i in range(n_corr_tiles):

62
tests/unit/test_d1_layout_diag.py
View File

@@ -1,35 +1,41 @@
"""D1: Print SMEM-P layout diagnostics for hd=128."""
import torch, math, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.torch as ct
"""D1: Print SMEM-P layout diagnostics inside the kernel."""
import torch, math
import cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.torch as ct
import cuda.bindings.driver as cuda
from cutlass.cute.nvgpu import tcgen05
from cutlass import Float32, BFloat16
from cutlass.utils import LayoutEnum
from dsv4.kernels.attention.fmha import FmhaKernel
for hd in [64, 128, 256]:
q = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda')
k = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda')
a_major = LayoutEnum.from_tensor(ct.from_dlpack(q)).mma_major_mode()
b_major = LayoutEnum.from_tensor(ct.from_dlpack(k)).mma_major_mode()
qk_mma = utils.sm100.make_trivial_tiled_mma(BFloat16, BFloat16, a_major, b_major, Float32, tcgen05.CtaGroup.ONE, (128,128), tcgen05.OperandSource.SMEM)
pv_n_tile = min(hd, 256)
pv_a_major = a_major # SMEM-P path uses a_major
pv_mma = utils.sm100.make_trivial_tiled_mma(BFloat16, BFloat16, pv_a_major, b_major, Float32, tcgen05.CtaGroup.ONE, (128,pv_n_tile), tcgen05.OperandSource.SMEM)
def test_layout(hd):
m = 128; n_kv = 1 # minimal
torch.manual_seed(42)
q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
k = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda')
v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda')
c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
qk_mma_tiler = (128, 128, cute.size(qk_mma.shape_mnk, mode=[2]) * 4)
pv_mma_tiler = (128, pv_n_tile, cute.size(pv_mma.shape_mnk, mode=[2]) * (128 // cute.size(pv_mma.shape_mnk, mode=[2])))
kernel = FmhaKernel(head_dim=hd, s_k=n_kv, use_smem_p=False)
pv_n_tile = kernel.pv_n_tile
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1)
p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1)
v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1)
c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
# QK C-fragment layout
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)
# The _setup method prints are visible at compile time
print(f'--- hd={hd} ---', flush=True)
print(f' pv_n_tile={pv_n_tile}, use_smem_p={kernel.use_smem_p}', flush=True)
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE)
print(f' tmem_p0_offset={kernel.tmem_p0_offset}', flush=True)
print(f' tmem_o0_offset={kernel.tmem_o0_offset}', flush=True)
print(f' tOrP0_offset={kernel.tOrP0_offset}', flush=True)
print(f' num_tmem_alloc_cols={kernel.num_tmem_alloc_cols}', flush=True)
print(f'--- hd={hd} ---')
print(f' p_tmem_s.outer = {p_tmem_s.outer}')
print(f' p_smem_s.outer = {p_smem_s.outer}')
print(f' tStS.layout = {tStS.layout}')
print(f' pv_mma_tiler = {pv_mma_tiler}')
print()
if __name__ == '__main__':
for hd in [64, 128, 256]:
test_layout(hd)