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

SMEM-P: implement simple test pattern instead of coord lookup

Browse Source
This commit is contained in:
biondizzle
2026-05-23 19:21:31 +00:00
parent fc3e9bf0ae
commit 5ddaea6dad
1 changed files with 25 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -341,33 +341,38 @@ class FmhaKernel:
print(f"[SMEM-P MANUAL] Starting manual P write to SMEM")
# Get thread's logical position in QK C-fragment partition
# tStS0 is QK C-fragment tensor for this thread
thread_qk_coord = cute.coord(tStS0) # Logical coordinates in QK C-fragment space
print(f"[SMEM-P MANUAL] Thread QK coord: {thread_qk_coord}")
# Debug layouts
print(f"[SMEM-P MANUAL] tStS0 layout: {tStS0.layout}")
print(f"[SMEM-P MANUAL] sP layout: {sP.layout}")
# Get the shape of P values this thread owns
# tStS0 has shape ((128, 128), 1, 1) - total 128×128 P matrix
# Each thread owns a subtile of this
qk_fragment_shape = cute.shape(tStS0)
print(f"[SMEM-P MANUAL] QK fragment shape: {qk_fragment_shape}")
# rP_bf16 contains P values in TMEM layout
# We need to copy them to sP (PV A-operand SMEM layout)
# For now, implement simple test: write thread ID to first element
# This is WRONG but helps debug
smem_test_offset = Int32(0)
sP[smem_test_offset] = BFloat16(float(sfw_idx) * 0.01)
print(f"[SMEM-P MANUAL] Wrote thread {sfw_idx} value to SMEM offset 0")
# Approach: Get P values from rP_bf16 and write to sP
# rP_bf16 has shape ((32, 1), 4, 1, 1) in TMEM layout
# Need to map to sP shape ((128, 16), 1, (4, 2), 1)
# For now, simple test: each thread writes one value
# Compute linear thread index in softmax warp
thread_in_warp = sfw_idx % 32 # Assuming 32 threads per warp
smem_offset = thread_in_warp * 1024 # Arbitrary stride
if smem_offset < cute.size(sP):
sP[smem_offset] = BFloat16(float(thread_in_warp) * 0.01)
print(f"[SMEM-P MANUAL] Thread {thread_in_warp} wrote to SMEM offset {smem_offset}")
else:
print(f"[SMEM-P MANUAL] Thread {thread_in_warp} offset {smem_offset} out of bounds")
# TODO: Implement proper mapping
# 1. Determine which P indices this thread owns
# 2. For each P index (i,j), compute SMEM address in PV A-operand layout
# 3. Write P value to that address
# Need to understand:
# 1. Which P values does this thread own in rP_bf16?
# 2. Where in sP should they go?
# For now, zero out sP as fallback (produces wrong results but compiles)
# For now, zero out rest of sP
for j in cutlass.range(cute.size(sP), vectorize=True):
sP[j] = BFloat16(0.0)
if j != smem_offset:
sP[j] = BFloat16(0.0)
print(f"[SMEM-P MANUAL] Used zero-fallback (TODO: implement proper mapping)")
print(f"[SMEM-P MANUAL] Used test pattern + zeros (TODO: implement proper mapping)")
cute.arch.fence_proxy("async.shared", space="cta")
softmax_done_bar.arrive() # Per-tile O rescale (hand-constructed atoms with logical_divide layout)
if kt > 0: