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

D1.3: Use make_cotiled_copy for SMEM-P — custom TV layout from TMEM-load coords to sP

Browse Source 
Per CUTLASS guidance:
- make_tiled_copy_C/D encode wrong invariants for this transfer
- Build custom R→S copy where TV map comes from tTMEM_LOADcS (softmax thread
  ownership) and destination addresses come from sP layout (PV A-operand swizzled SMEM)
- Use composition(sP_2d_layout, p_coord_layout) for atom_layout_tv
- Start with scalar BF16 (16-bit) stores — vectorize later
- Zero-fill source for compile test, will fill with actual P values next
This commit is contained in:
biondizzle
2026-05-23 23:17:30 +00:00
parent 2e86ed939e
commit 2bbe55b08c
1 changed files with 42 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -272,28 +272,44 @@ class FmhaKernel:
tTMEM_STOREcP = thr_store.partition_S(tScP)
# P SMEM copy atoms: SMEM-P
# Approach: use make_tiled_copy_C(qk_mma) to create a copy that writes
# from QK C-fragment register layout to SMEM. The softmax threads have P values
# in registers after computing softmax. We write these to sP so the MMA warp
# can read them via pv_mma.make_fragment_A(sP).
# Must define unconditionally (CuTeDSL scoping).
_smem_copy_atom = cute.make_copy_atom(
# Uses make_cotiled_copy to create a custom R→S copy where:
# - Thread/value mapping: softmax/TMEM-load ownership (tTMEM_LOADcS)
# - Destination: sP in PV A-operand swizzled SMEM layout
# Per CUTLASS guidance: make_tiled_copy_C/D encode the wrong invariants
# for this transfer. We build a custom TV layout that maps (tid,vid) -> sP addr.
# Must define unconditionally (CuTeDSL scoping: compile both branches).
# Start with scalar BF16 stores (16-bit) — vectorize later once correct.
_r2s_atom = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
self.q_dtype,
num_bits_per_copy=128,
num_bits_per_copy=16, # scalar BF16 — safe, vectorize later
)
_tiled_smem_copy_C = cute.make_tiled_copy_C(_smem_copy_atom, qk_mma)
_thr_smem_copy_C = _tiled_smem_copy_C.get_slice(sfw_idx)
# Destination: sP partitioned by QK C-fragment thread mapping
_sP_2d = cute.group_modes(sP, 0, 3)
_tSMEM_COPYsP = _thr_smem_copy_C.partition_D(_sP_2d)
# Source: QK C-fragment register layout (same as what make_fragment_C produces)
# The softmax has P in rP_bf16 (TME load layout). We need a source tensor
# in QK C-fragment register layout. Create a register tensor with the right shape.
_qk_C_reg = qk_thr.make_fragment_C(qk_as) # QK C-fragment register fragment
_qk_C_2d = cute.group_modes(_qk_C_reg, 0, 2) # (M*K, STAGE)
_tSMEM_COPYrS = _thr_smem_copy_C.partition_S(_qk_C_2d)
_rP_smem_src = cute.make_rmem_tensor(_tSMEM_COPYrS.shape, self.q_dtype)
# Build atom_layout_tv: (tid, vid) -> sP address
# tTMEM_LOADcS gives (thr_offset, vid) -> (m, k) coordinate
# sP layout gives ((m,k0),0,(k1,k2),0) -> address (with swizzle)
# We compose these to get (tid, vid) -> sP address.
# Use sP_2d (grouped to 2D) for simplicity.
_sP_nostage = sP[(None, None, None, 0)]
_sP_2d = cute.group_modes(_sP_nostage, 0, 3)
_sP_2d_layout = _sP_2d.layout
# Flatten tTMEM_LOADcS to (total_elements,) -> (m, k) coords
_p_coord_layout = cute.flatten(tTMEM_LOADcS.layout)
# Compose: (tid, vid) -> (m, k) via _p_coord_layout, then (m, k) -> addr via sP_2d
# make_cotiled_copy needs atom_layout_tv where the codomain is in the sP address space.
# composition(sP_2d_layout, p_coord_layout) should give this.
_p_tv_layout = cute.composition(_sP_2d_layout, _p_coord_layout)
_tiled_p_r2s = cute.make_cotiled_copy(
_r2s_atom,
_p_tv_layout,
_sP_2d_layout,
)
_thr_p_r2s = _tiled_p_r2s.get_slice(sfw_idx)
_tRS_sP = _thr_p_r2s.partition_D(_sP_2d)
# Source: register tensor in the copy's value order.
# The softmax computes P in rP_bf16 (TME load layout). We retile it
# into the copy's expected value order, or create a new source tensor
# and fill it during softmax.
_rP_store = cute.make_rmem_tensor(_tRS_sP.shape, self.q_dtype)
row_max = -Float32.inf
row_sum = Float32(0.0)
@@ -366,16 +382,13 @@ class FmhaKernel:
cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
cute.arch.fence_view_async_tmem_store()
else:
# SMEM-P: store P to SMEM via make_tiled_copy_C(qk_mma)
# The P values are in rP_bf16 (TME load layout). We need to
# rearrange them into the QK C-fragment register layout for the copy.
# Copy rP_bf16 values into _rP_smem_src (QK C-fragment register layout).
# This is a register-to-register rearrangement.
# TODO: This rearrangement may be avoidable if we can directly use
# the TMEM load layout as source. For now, zero-fill and copy.
for j in cutlass.range(cute.size(_rP_smem_src), vectorize=True):
_rP_smem_src[j] = self.q_dtype(0)
cute.copy(_tiled_smem_copy_C, _rP_smem_src, _tSMEM_COPYsP)
# SMEM-P: store P to SMEM via make_cotiled_copy
# Fill _rP_store with P values (in the copy's value order).
# For now, zero-fill to test compilation. The real P values
# will be filled by remapping from rP_bf16 to _rP_store's order.
for j in cutlass.range(cute.size(_rP_store), vectorize=True):
_rP_store[j] = self.q_dtype(0)
cute.copy(_tiled_p_r2s, _rP_store, _tRS_sP)
cute.arch.fence_proxy("async.shared", space="cta")
if kt > 0:
tTMrO = cute.make_rmem_tensor(