Quick test: working v3 with n=256 multi-tile

tests/quick_v3_multitile.py

@@ -0,0 +1,53 @@
+"""Quick test: run the working test_fmha_v3.py with n=256 to check multi-tile."""
+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
+import math
+
+HEAD_DIM = 64
+
+# Import the working kernel class
+import sys
+sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel')
+from tests.unit.test_fmha_v3 import FmhaV3
+
+for n in [128, 256]:
+    torch.manual_seed(42)
+    m, hd = 128, HEAD_DIM
+    q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
+    k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda')
+    v = torch.ones(n, hd, dtype=torch.bfloat16, device='cuda')  # V=ones like the working test
+    v_kernel = v.unsqueeze(-1)
+    c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
+
+    qf = q[:, :, 0].float()
+    kf = k[:, :, 0].float()
+    scale = 1.0 / math.sqrt(hd)
+    attn = qf @ kf.T * scale
+    attn = torch.softmax(attn, dim=-1)
+    ref = attn @ v.float()
+
+    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_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
+    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 = FmhaV3(s_k=n)
+    print(f'n={n}: Compiling...', flush=True)
+    compiled = cute.compile(kernel, mQ, mK, mV, mC, stream)
+    compiled(mQ, mK, mV, mC, stream)
+    torch.cuda.synchronize()
+
+    out = c[:, :, 0].float()
+    cos = torch.nn.functional.cosine_similarity(
+        out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)
+    ).item()
+    print(f'FMHA v3 n={n}: cos {cos:.6f}  {"PASS" if cos >= 0.99 else "FAIL"}')
+    if cos < 0.99:
+        print(f'  out[0,:4]={out[0,:4].tolist()}')
+        print(f'  ref[0,:4]={ref[0,:4].tolist()}')