diff --git a/README.md b/README.md index 4c744727..6316d0ef 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ CuTeDSL kernels for DeepSeek-V4 (Blackwell B200, SM100). All kernels use `cutlass.cute` (CuTeDSL) with Blackwell tensor cores. -## Status (May 21, 2026 — 04:10 UTC) +## Status (May 21, 2026 — 04:35 UTC) ### ✅ Stage A: Bare Q@K^T via tcgen05.mma → TMEM → GMEM — COMPLETE @@ -13,7 +13,9 @@ CuTeDSL kernels for DeepSeek-V4 (Blackwell B200, SM100). All kernels use `cutlas **Pipeline deadlock: FIXED. Kernel runs without deadlock.** **Bug 1 (V MN-major): Fix applied.** -**Bug 2 (softmax packing): Fix applied, but PV output is garbage.** +**Bug 2 (softmax packing): Confirmed correct (V=I test: cosine 1.0).** +**Bug 3 (ACCUMULATE): Fix applied.** +**Bug 4 (non-square PV): PV works for (128,128) output, broken for (128,64) output.** #### Bug 1: V B-Operand Must Be MN-Major — ✅ FIX APPLIED @@ -22,18 +24,38 @@ PV MMA uses `v_major` (OperandMajorMode.MN) instead of `b_major` (K). V must use `as_strided` — default PyTorch (64,128) gives strides (128,1) which is K-major. -#### Bug 2 (Packing): C-Fragment Composition Store — ✅ APPLIED, ❌ PV OUTPUT WRONG +#### Bug 2 (Packing): C-Fragment Composition Store — ✅ CONFIRMED CORRECT -FP32→BF16 packing via C-fragment composition store (FMHA pattern) runs without error. -The softmax packing overwrites part of S in TMEM (P at tmem_p0_offset=32 overlaps S at offset 0). -This is intentional — S is no longer needed after softmax. +FP32→BF16 packing via C-fragment composition store (FMHA pattern) is correct. +Proven by V=I test (cosine 1.0) and random V 128x128 test (cosine 0.999999). ⛔ **FOOTGUN**: `St32x32bOp` MUST use Float32, NOT BFloat16. ⚠️ The recast view for P packing uses the LOAD layout (128 BF16 elements), not the store composition shape. -#### Bug 3 (NEW): PV MMA Output Is Garbage — 🔨 INVESTIGATING +#### Bug 3 (ACCUMULATE): First PV Must Use ACCUMULATE=False — ✅ FIX APPLIED -The PV MMA produces cosine ~0.01 against the reference. Suspected cause: TMEM layout mismatch between the softmax P store (C-fragment composition layout) and the PV MMA A-fragment read (`p_tmem_s` layout from `make_smem_layout_a`). These should alias the same physical TMEM columns by the sequential-flattening property, but the specific layout functions may compute different shapes/strides. +If ACCUMULATE=True on the first PV, `O = P@V + old_O` adds uninitialized TMEM. Always ACCUMULATE=False for first PV, then True for subsequent tiles. + +#### Bug 4 (CURRENT): PV MMA Broken for Non-Square Output — 🔨 ROOT CAUSE UNKNOWN + +**What works:** +- PV with (128,128) output, V=I: cosine 1.0 ✅ +- PV with (128,128) output, random V: cosine 0.999999 ✅ + +**What doesn't work:** +- PV with (128,64) output, V MN-major (64,128): cosine ~0.01 ❌ + +**Possible causes:** +1. `make_trivial_tiled_mma` with (128,64) produces different A-fragment layout — alias with softmax P may break +2. V TMA load wrong for (128,64) PV — SMEM layout, TMA descriptor, or partitioning incorrect +3. Epilogue/gC mismatch — output c is (128,64) but epilogue may write (128,128) tile +4. PV mma_tiler_mn doesn't affect the MMA atom (which is always (128,128,16)) + +**Diagnostic findings:** +- Pointer arithmetic correct: softmax P and PV A-fragment address same TMEM location +- Layout aliasing correct: C-fragment composition and A-fragment produce same physical addresses +- Pipeline ordering correct: softmax completes before PV starts +- Softmax packing correct: proven by V=I test ### 🔨 Stage C: Online Softmax — AFTER B @@ -98,6 +120,23 @@ Passing `cta_layout_vmnk` to the mma_si PipelineUmmaAsync causes deadlock. FMHA The tmem allocation barrier only includes MMA + epilogue warps. The TMA warp is excluded. Calling `wait_for_alloc()` from the TMA warp corrupts the barrier. +### 7. PV MMA ACCUMULATE Must Be False on First Tile + +If ACCUMULATE=True on the first PV MMA, `O = P@V + old_O` adds uninitialized TMEM to the result. Always set ACCUMULATE=False for the first PV, then True for subsequent tiles. FMHA: `pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, kphase_idx != 0)`. + +### 8. TMEM Pointer Arithmetic: Offset Units Depend on Pointer Type + +When computing PV A-fragment offset from the softmax P offset: +```python +# Softmax store: FP32 pointer + tmem_p0_offset (in FP32 elements) +tStS_P = cute.make_tensor(tStS.iterator + tmem_p0_offset, tStS_P_layout) + +# PV A-fragment: BF16 pointer + scaled offset (in BF16 elements) +p_offset = acc_dtype.width // q_dtype.width * tmem_p0_offset # 2 * 32 = 64 +tOrP0 = cute.make_tensor(tOrP.iterator + p_offset, tOrP.layout) +``` +Both must address the same physical TMEM column. The 2× scaling accounts for FP32→BF16 element size difference. + --- ## Architecture: Per-Tile Flow @@ -128,7 +167,9 @@ After all tiles: epilogue warps tcgen05.ld tmem_output, divide by row_sum, cast | `test_stage_a_v2.py` | Q@K^T only | 0.999999 | ✅ PASS | | `test_mma_si_only.py` | Q@K^T + mma_si pipeline (no PV) | 0.999999 | ✅ PASS | | `test_softmax_only.py` | Q@K^T + softmax packing, output S | 0.52 | ❌ S overwritten by P (expected) | -| `test_mma_si_pv.py` | Q@K^T + softmax + P@V (V MN-major) | 0.01 | ❌ PV output garbage | +| `test_mma_si_pv.py` | Q@K^T + softmax + P@V (V MN-major, 128x64) | 0.01 | ❌ PV output garbage | +| `test_pv_diag.py` | Q@K^T + softmax + P@V (V=I/random, 128x128) | 1.0 / 0.999999 | ✅ PASS | +| `test_layout_compare.py` | Print TMEM layouts for QK S and PV A-fragment | N/A | ℹ️ layout inspection | | `test_stage_b_v7.py` | Q@K^T + C-fragment softmax (V=K, wrong major) | -0.02 | ❌ wrong major + P packing | | `test_stage_b_v20.py` | Q@K^T + softmax (V=K, PipelineTmaStore bug) | N/A | ❌ compile error | @@ -147,6 +188,8 @@ pv_mma_tiler = (qk_mma_tiler[0], qk_mma_tiler[2], qk_mma_tiler[1]) # = (M, head_dim, QK_N) = (128, 64, 128) for head_dim=64 ``` +FMHA passes `pv_mma_tiler[:2] = (128, head_dim)` to `make_trivial_tiled_mma`, NOT the QK tiler `(128, 128)`. + ### make_trivial_tiled_mma — Use New Overload ```python make_trivial_tiled_mma(a_dtype, b_dtype, a_leading_mode, b_leading_mode, diff --git a/tests/test_inspect_types.py b/tests/test_inspect_types.py new file mode 100644 index 00000000..7ef47b12 --- /dev/null +++ b/tests/test_inspect_types.py @@ -0,0 +1,73 @@ +import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils +from cutlass.cute.nvgpu import tcgen05 +from cutlass import Float32, BFloat16 +from cutlass.utils import LayoutEnum +from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset +import cuda.bindings.driver as cuda +import cutlass.torch as ct + +m, n, head_dim = 128, 128, 64 +q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') +k = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda') +v_base = torch.randn(head_dim, n, dtype=torch.bfloat16, device='cuda') +v = v_base.as_strided((head_dim, n), (1, head_dim)).unsqueeze(-1) +c = torch.zeros(m, head_dim, 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).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) +mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) +stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + +class InspectKernel: + def __init__(self): + self.q_dtype = BFloat16; self.acc_dtype = Float32 + + @cute.jit + def __call__(self, q, k, v, c, stream): + a_major = LayoutEnum.from_tensor(q).mma_major_mode() + b_major = LayoutEnum.from_tensor(k).mma_major_mode() + v_major = LayoutEnum.from_tensor(v).mma_major_mode() + + qk_mma = utils.sm100.make_trivial_tiled_mma( + self.q_dtype, self.q_dtype, a_major, b_major, + self.acc_dtype, tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM) + pv_mma = utils.sm100.make_trivial_tiled_mma( + self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, v_major, + self.acc_dtype, tcgen05.CtaGroup.ONE, (128, 64), tcgen05.OperandSource.TMEM) + + qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) + qk_mma_tiler = (128, 128, qk_inst_k * 4) + pv_mma_tiler = (qk_mma_tiler[0], qk_mma_tiler[2], qk_mma_tiler[1]) + + qk_thr = qk_mma.get_slice(0) + pv_thr = pv_mma.get_slice(0) + + qk_acc_shape = qk_thr.partition_shape_C(qk_mma_tiler[:2]) + tStS = qk_thr.make_fragment_C(qk_acc_shape) + + p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, self.q_dtype, 1) + tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) + tOrP_base = pv_thr.make_fragment_A(tP) + tOrP = tOrP_base[(None, None, None, 0)] + + tmem_p0_offset = 32 + p_offset = self.acc_dtype.width // self.q_dtype.width * tmem_p0_offset + + # Print pointer values - if tOrP inherits FP32 type, +64 adds 256 bytes + # If BF16 type, +64 adds 128 bytes (correct, matches tStS+32 FP32 = 128 bytes) + cute.printf("tStS ptr value: %d", tStS.iterator) + cute.printf("tStS_P ptr (tStS+32): %d", tStS.iterator + tmem_p0_offset) + cute.printf("tOrP ptr value: %d", tOrP.iterator) + cute.printf("tOrP0 ptr (tOrP+64): %d", tOrP.iterator + p_offset) + cute.printf("p_offset: %d", p_offset) + + pv_acc_shape = pv_thr.partition_shape_C(pv_mma_tiler[:2]) + tOtO = pv_thr.make_fragment_C(pv_acc_shape) + o_cols = find_tmem_tensor_col_offset(tOtO) + cute.printf("o_cols: %d", o_cols) + +kernel = InspectKernel() +compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) +compiled(mQ, mK, mV, mC, stream) +torch.cuda.synchronize() diff --git a/tests/test_layout_compare.py b/tests/test_layout_compare.py new file mode 100644 index 00000000..71a94aee --- /dev/null +++ b/tests/test_layout_compare.py @@ -0,0 +1,95 @@ +"""Compare C-fragment composition layout vs A-fragment layout for PV P operand.""" +import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils +from cutlass.cute.nvgpu import tcgen05 +from cutlass import Float32, BFloat16 +from cutlass.utils import LayoutEnum +from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset +import cuda.bindings.driver as cuda +import cutlass.torch as ct + + +class LayoutCompareKernel: + def __init__(self): + self.acc_dtype = Float32; self.qk_acc_dtype = Float32 + self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 + self.mma_tiler_mn = (128, 128) + self.cta_group = tcgen05.CtaGroup.ONE + self.threads_per_cta = 64 # minimal + + @cute.jit + def __call__(self, q, k, v, c, stream): + a_major = LayoutEnum.from_tensor(q).mma_major_mode() + b_major = LayoutEnum.from_tensor(k).mma_major_mode() + v_major = LayoutEnum.from_tensor(v).mma_major_mode() + + qk_mma = utils.sm100.make_trivial_tiled_mma( + self.q_dtype, self.q_dtype, a_major, b_major, + self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) + pv_mma = utils.sm100.make_trivial_tiled_mma( + self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, v_major, + self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) + + qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) + qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) + pv_mma_tiler = (qk_mma_tiler[0], qk_mma_tiler[2], qk_mma_tiler[1]) + + qk_thr = qk_mma.get_slice(0) + pv_thr = pv_mma.get_slice(0) + + qk_acc_shape = qk_thr.partition_shape_C(qk_mma_tiler[:2]) + tStS = qk_thr.make_fragment_C(qk_acc_shape) + tStS0 = cute.make_tensor(tStS.iterator, tStS.layout) + + pv_acc_shape = pv_thr.partition_shape_C(pv_mma_tiler[:2]) + tOtO = pv_thr.make_fragment_C(pv_acc_shape) + + # P A-fragment + p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, self.q_dtype, 1) + tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) + tOrP_base = pv_thr.make_fragment_A(tP) + tOrP = tOrP_base[(None, None, None, 0)] + + # C-fragment composition layout + tilePlikeFP32 = qk_mma_tiler[1] // Float32.width * self.o_dtype.width + tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) + tStS_P = cute.make_tensor(tStS.iterator + 32, tStS_P_layout) # offset 32 FP32 columns + + # With scaled offset for A-fragment + p_offset_in_a_elements = self.qk_acc_dtype.width // self.q_dtype.width * 32 # = 64 + tOrP0 = cute.make_tensor(tOrP.iterator + p_offset_in_a_elements, tOrP.layout) + + # Print layouts + cute.printf("tStS layout: {}", tStS.layout) + cute.printf("tOrP layout: {}", tOrP.layout) + cute.printf("tStS_P layout: {}", tStS_P_layout) + cute.printf("tOrP0 layout: {}", tOrP0.layout) + cute.printf("tOrP shape: {}", tOrP.shape) + cute.printf("tStS_P shape: {}", tStS_P.shape) + cute.printf("tOtO layout: {}", tOtO.layout) + cute.printf("pv_mma_tiler: {}", pv_mma_tiler) + cute.printf("qk_mma_tiler: {}", qk_mma_tiler) + + +def test(): + m, n, head_dim = 128, 128, 64 + q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda') + k = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda') + v_base = torch.randn(head_dim, n, dtype=torch.bfloat16, device='cuda') + v = v_base.as_strided((head_dim, n), (1, head_dim)).unsqueeze(-1) + c = torch.zeros(m, head_dim, 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).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) + mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + kernel = LayoutCompareKernel() + print('Compiling...', flush=True) + compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) + print('Running...', flush=True) + compiled(mQ, mK, mV, mC, stream) + torch.cuda.synchronize() + + +if __name__ == '__main__': + test() diff --git a/tests/test_mma_si_pv.py b/tests/test_mma_si_pv.py index 26118bec..182fd69f 100644 --- a/tests/test_mma_si_pv.py +++ b/tests/test_mma_si_pv.py @@ -11,7 +11,8 @@ import cuda.bindings.driver as cuda class MmaSiPvTest: - def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True): + def __init__(self, mma_tiler_mn, head_dim, use_2cta_instrs=False, use_tma_store=True): + self.head_dim = head_dim self.acc_dtype = Float32; self.qk_acc_dtype = Float32 self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store @@ -79,12 +80,19 @@ class MmaSiPvTest: self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() self.c_layout = LayoutEnum.from_tensor(c) + # Compute PV tiler: swap N and K from QK tiler (FMHA convention) + # QK: (M=128, N=128, K=64) -> PV: (M=128, N=64, K=128) + # PV mma_tiler_mn is (M, N_pv) = (128, head_dim=64), NOT (128, 128) 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, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) + # BUG FIX: pv_mma_tiler_mn must be (M, head_dim), not (M, N_qk) + # Passing mma_tiler_mn=(128,128) creates a (128,128) MMA that expects 128-column output + # but PV output is (128,64). This caused cosine ~0.01. + pv_mma_tiler_mn = (self.mma_tiler_mn[0], self.head_dim) # (128, 64) pv_mma = utils.sm100.make_trivial_tiled_mma( self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, - self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) + self.qk_acc_dtype, self.cta_group, pv_mma_tiler_mn, tcgen05.OperandSource.TMEM) self._setup(qk_mma, pv_mma) q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) @@ -164,12 +172,12 @@ class MmaSiPvTest: gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) - gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) + gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None,0,None)), (None,None,None)) # Use PV tiler for output k_cnt = cute.size(gQ, mode=[3]) qk_thr = qk_mma.get_slice(0) pv_thr = pv_mma.get_slice(0) - tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) + tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = pv_thr.partition_C(gC) # PV output: partition with pv_thr a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape) @@ -239,11 +247,13 @@ class MmaSiPvTest: s0_handle = mma_si_prod.acquire_and_advance() # wait for softmax done # PV MMA - pv_mma.set(tcgen05.Field.ACCUMULATE, True) + # FMHA pattern: first PV overwrites (ACCUMULATE=False), then accumulates + pv_mma.set(tcgen05.Field.ACCUMULATE, False) tCrV_s = tCrV[(None, None, None, 0)] nblk_pv = cute.size(tOrP0, mode=[2]) for kb in cutlass.range(nblk_pv, unroll_full=True): cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV_s[(None,None,kb)], tOtO0) + pv_mma.set(tcgen05.Field.ACCUMULATE, True) acc_pipe.producer_commit(acc_prod_st) acc_prod_st.advance() @@ -330,7 +340,7 @@ def test(): mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v)) mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c)) stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - kernel = MmaSiPvTest(mma_tiler_mn=(128, 128)) + kernel = MmaSiPvTest(mma_tiler_mn=(128, 128), head_dim=head_dim) print('Compiling...', flush=True) compiled = cute.compile(kernel, mQ, mK, mV, mC, stream) print('Running...', flush=True) diff --git a/tests/test_pv_diag.py b/tests/test_pv_diag.py new file mode 100644 index 00000000..80184627 --- /dev/null +++ b/tests/test_pv_diag.py @@ -0,0 +1,382 @@ +""" +Minimal PV-only test: Load P from GMEM to TMEM via QK-style MMA, then PV from TMEM. +Step 1: QK MMA writes FP32 S to TMEM (we know this works) +Step 2: Softmax packing writes BF16 P to TMEM (test this) +Step 3: PV MMA reads BF16 P from TMEM and V from SMEM, produces O + +But to isolate the bug, let me test just the PV MMA in isolation. +I'll write known BF16 values to TMEM using the softmax packing path, +then immediately read them back using the PV A-fragment path, +and compare. + +Actually, the simplest isolation test: +1. Do QK MMA to get S in TMEM (cosine 0.999999 verified) +2. Do softmax packing: S → P in TMEM (at offset 32) +3. Skip PV entirely — read P from TMEM using the C-fragment composition LOAD path +4. Output P to GMEM and compare against S.to(BF16) + +This tests whether the softmax packing writes P correctly to the same TMEM +that the PV would read from. + +But we can't easily read P from TMEM using the standard epilogue path +because the epilogue expects FP32 accumulator data. + +Alternative: Use the PV MMA with V=I (identity). If P is correct, +then P @ I = P. But V needs to be MN-major and (128, 128), not (128, 64). +The output would be (128, 128) which doesn't match our (128, 64) c tensor. + +Let me use V that selects the first 64 columns: V[k, n] = delta(k, n) for k in [0,63]. +This gives P @ V = P[:, :64], and the output is (128, 64). +But V is (128, 128) in the MMA K,N dims. V[k, n] for k in [0,127], n in [0,63]. +Hmm, this is getting complicated. Let me just do the identity approach with a (128, 128) output. +""" +import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass import Float32, BFloat16, Int32, Boolean, const_expr +from cutlass.utils import LayoutEnum +from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset +import cuda.bindings.driver as cuda +import cutlass.torch as ct + + +class PvDiagKernel: + """QK + softmax packing + PV with V=I to isolate PV MMA correctness. + Output should be P = S.to(BF16), i.e. (Q@K^T).bfloat16() + With V=I, O = P @ I = P. + But V is (K=128, N=128) in the MMA. We need a 128x128 identity in MN-major. + Output tensor is (128, 128). + """ + def __init__(self, mma_tiler_mn): + self.acc_dtype = Float32; self.qk_acc_dtype = Float32 + self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 + self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1) + self.use_2cta_instrs = False # needed by epilogue_tma_store + self.epilog_sync_bar_id = 1 # needed by epilogue_tma_store + self.cluster_shape_mn = (1, 1) + self.cta_group = tcgen05.CtaGroup.ONE + 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_c_stage = 2 + + def _setup(self, qk_mma, pv_mma): + qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2]) + self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4) + # PV with V=I: output is (128, 128), same as QK + self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[1]) + # pv_mma_tiler = (128, 128, 128) since V is 128x128 + self.mma_tiler = self.qk_mma_tiler + + self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,)) + self.cta_tile_shape_mnk = ( + self.qk_mma_tiler[0] // cute.size(qk_mma.thr_id.shape), + self.qk_mma_tiler[1], self.qk_mma_tiler[2]) + self.c_layout = LayoutEnum.ROW_MAJOR + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, False, self.c_layout, self.o_dtype) + self.num_ab_stage = 1; self.num_acc_stage = 1 + + self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1) + self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.q_dtype, 1) + self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) + self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1) + self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2) + + qk_thr = qk_mma.get_slice(0) + qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) + tStS = qk_thr.make_fragment_C(qk_acc_shape) + s_cols = find_tmem_tensor_col_offset(tStS) + pv_thr = pv_mma.get_slice(0) + pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) + tOtO = pv_thr.make_fragment_C(pv_acc_shape) + o_cols = find_tmem_tensor_col_offset(tOtO) + + self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width + self.tmem_s0_offset = 0 + self.tmem_p0_offset = 32 + self.tmem_o0_offset = s_cols + + tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) + tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100") + + a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) + b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) + v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) + self.num_tma_load_bytes = ( + cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.q_dtype, b_smem) + + cute.size_in_bytes(self.q_dtype, v_smem) + ) * cute.size(qk_mma.thr_id.shape) + + @cute.jit + def __call__(self, q, k, v, c, stream): + self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype + self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() + self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() + self.v_major = LayoutEnum.from_tensor(v).mma_major_mode() + self.c_layout = LayoutEnum.from_tensor(c) + + 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, self.mma_tiler_mn, tcgen05.OperandSource.SMEM) + # PV with 128x128 output (V=I) + pv_mma = utils.sm100.make_trivial_tiled_mma( + self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, + self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.TMEM) + self._setup(qk_mma, pv_mma) + + q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0)) + k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0)) + v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0)) + tma_q, tma_tq = 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_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) + tma_k, tma_tk = 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_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape) + tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B( + utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id), + v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape) + epi_smem = cute.select(self.c_smem_s, mode=[0, 1]) + tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile) + + self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv, + tma_c, tma_tc, self.cluster_layout_vmnk, + self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile + ).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, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + use_2cta = cute.size(qk_mma.thr_id.shape) == 2 + + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k) + cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c) + + @cute.struct + class SS: + ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + mma_si_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) + + ab_p, ab_c = pipeline.PipelineTmaUmma.create( + barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True + ).make_participants() + + mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create( + barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), + ).make_participants() + + acc_pipe = pipeline.PipelineUmmaAsync.create( + barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)), + cta_layout_vmnk=cl_vmnk, defer_sync=True) + + tmem_bar = pipeline.NamedBarrier(barrier_id=2, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id))) + tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar, + allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta, + two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr) + + pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True) + + sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner) + sK = smem.allocate_tensor(element_type=self.q_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_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) + sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner) + + gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None)) + gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None)) + gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None)) + k_cnt = cute.size(gQ, mode=[3]) + + qk_thr = qk_mma.get_slice(0) + pv_thr = pv_mma.get_slice(0) + tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK); tCgC = qk_thr.partition_C(gC) + a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape) + tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3)) + b_lay = cute.make_layout(cute.slice_(cl_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)) + tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)] + + gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None)) + tCgV = pv_thr.partition_B(gV) + tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3)) + tVgV = tVgV[(None,0,None,0)] + + tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) + tCrV = pv_mma.make_fragment_B(sV) + + qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2]) + tStS = qk_thr.make_fragment_C(qk_acc_shape) + tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) + + pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2]) + tOtO = pv_thr.make_fragment_C(pv_acc_shape) + tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) + + tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer) + tOrP_base = pv_thr.make_fragment_A(tP) + tOrP = tOrP_base[(None, None, None, 0)] + tOrP0 = cute.make_tensor( + tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, + tOrP.layout) + + tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage)) + tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage)) + + pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) + + # ═══ TMA LOAD WARP ═══ + if warp_idx == self.tma_warp_id: + ab_p.reset(); peek = ab_p.try_acquire() + for kt in cutlass.range(k_cnt, unroll=1): + h = ab_p.acquire_and_advance(peek) + cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier) + cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier) + cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier) + peek = cutlass.Boolean(1) + if h.count+1= 0.99 else 'FAIL')) + + +if __name__ == '__main__': + test()