[2026-06-01] pull origin main

Author: reubenconducts

.github/workflows/ci.yml

@@ -9,7 +9,10 @@ permissions:
 
 env:
   CI_WORK_DIR: ${{ vars.CI_WORK_DIR || format('/scratch/user/{0}', github.actor) }}
-  FA4_TEST_FILTER: "1024-1024-128-True-0-0.0-False-False-False-mha-dtype0 or 1024-1024-128-False-0-0.0-False-False-False-mha-dtype0"
+  FA4_TEST_FILTER: >-
+    1024-1024-128-True-0-0.0-False-False-False-mha-dtype0
+    or 1024-1024-128-False-0-0.0-False-False-False-mha-dtype0
+    or test_flash_attn_ex2_emu_decode_prefill_consistency
 
 jobs:
   lint:

CLAUDE.md

@@ -20,7 +20,7 @@ pip install flash-attn-4
 pip install -e "flash_attn/cute[dev]"
 ```
 
-Dependencies: `nvidia-cutlass-dsl>=4.4.1`, `torch`, `einops`, `apache-tvm-ffi`, `quack-kernels>=0.4.0`.
+Dependencies: `nvidia-cutlass-dsl>=4.5.2`, `torch`, `einops`, `apache-tvm-ffi`, `quack-kernels>=0.5.0`.
 
 ## Running Tests
 

benchmarks/benchmark_mla_paged_kv.py

@@ -0,0 +1,108 @@
+# Copyright (c) 2025, Johnsonms.
+
+# We recommend locking GPU clocks before running the benchmark to ensure consistent results.
+# This can be done using the following commands (2619 MHz is the max clock for B200):
+# sudo nvidia-smi -i 0 -pm 1
+# sudo nvidia-smi -i 0 --lock-gpu-clocks 2619,2619
+# See more here: https://github.com/triton-lang/triton/blob/d9f10ebdc5da53f73eb852fde73d8d7d80b679d1/python/triton/testing.py#L487
+
+import time
+import torch
+
+from triton.testing import do_bench
+
+from flash_attn.cute.interface import flash_attn_varlen_func
+
+
+device = "cuda"
+dtype = torch.bfloat16
+seqlen_q = 1
+nheads_q = 128
+nheads_kv = 1  # MQA-128
+headdim = 64
+headdim_v = 512
+causal = True
+
+batch_size = 128
+page_sizes = [None, 16, 64, 128]  # None = non-paged baseline
+
+torch.manual_seed(0)
+
+print(f"\nMLA paged KV, nheads_q = {nheads_q}, nheads_kv = {nheads_kv}, headdim = {headdim}, headdim_v = {headdim_v}, causal = {causal}")
+
+for seqlen in [s * 1024 for s in [1, 2, 4, 8, 16, 32, 64]]:
+    # Varlen format: (total_tokens, nheads, hdim)
+    total_q = batch_size * seqlen_q
+    total_k = batch_size * seqlen
+
+    try:
+        q = torch.randn(total_q, nheads_q, headdim, dtype=dtype, device=device)
+        k = torch.randn(total_k, nheads_kv, headdim, dtype=dtype, device=device)
+        v = torch.randn(total_k, nheads_kv, headdim_v, dtype=dtype, device=device)
+        qv = torch.randn(total_q, nheads_q, headdim_v, dtype=dtype, device=device)
+    except torch.OutOfMemoryError:
+        continue
+
+    cu_seqlens_q = torch.arange(0, total_q + seqlen_q, seqlen_q, dtype=torch.int32, device=device)
+    cu_seqlens_k = torch.arange(0, total_k + seqlen, seqlen, dtype=torch.int32, device=device)
+
+    # Mem I/O: KV read + Q/QV read + O write
+    total_seqlen = seqlen * batch_size
+    mem_io = (
+        total_seqlen * nheads_kv * (headdim + headdim_v) * 2  # K + V read
+        + q.numel() * 2 + qv.numel() * 2  # Q + QV read
+        + total_q * nheads_q * headdim_v * 2  # O write
+    )
+    # FLOPs: QK^T + PV (with qv, PV uses headdim_v)
+    flops = seqlen_q * total_seqlen * nheads_q * (headdim + headdim_v * 2) * 2
+
+    for page_size in page_sizes:
+        if page_size is None:
+            # Non-paged baseline
+            fn = lambda: flash_attn_varlen_func(  # noqa: E731
+                q, k, v, qv=qv,
+                cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=seqlen_q, max_seqlen_k=seqlen,
+                causal=causal,
+            )
+            label = "non-paged"
+        else:
+            # Create paged KV
+            num_pages_per_seq = (seqlen + page_size - 1) // page_size
+            total_pages = num_pages_per_seq * batch_size
+            k_paged = torch.zeros(total_pages, page_size, nheads_kv, headdim, device=device, dtype=dtype)
+            v_paged = torch.zeros(total_pages, page_size, nheads_kv, headdim_v, device=device, dtype=dtype)
+            page_table = torch.zeros(batch_size, num_pages_per_seq, dtype=torch.int32, device=device)
+            for b in range(batch_size):
+                for p in range(num_pages_per_seq):
+                    page_idx = b * num_pages_per_seq + p
+                    start = p * page_size
+                    end = min(start + page_size, seqlen)
+                    k_offset = b * seqlen
+                    if start < seqlen:
+                        k_paged[page_idx, :end - start] = k[k_offset + start:k_offset + end]
+                        v_paged[page_idx, :end - start] = v[k_offset + start:k_offset + end]
+                    page_table[b, p] = page_idx
+            seqused_k = torch.full((batch_size,), seqlen, dtype=torch.int32, device=device)
+
+            fn = lambda kp=k_paged, vp=v_paged, pt=page_table, su=seqused_k: flash_attn_varlen_func(  # noqa: E731
+                q, kp, vp, qv=qv,
+                cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=None,
+                max_seqlen_q=seqlen_q, max_seqlen_k=None,
+                seqused_k=su, page_table=pt,
+                causal=causal,
+            )
+            path = "TMA" if page_size == 128 else "cp.async"
+            label = f"paged-{page_size} ({path})"
+
+        fn()  # warmup / compile
+        time.sleep(1)  # avoid power throttling
+        t = do_bench(fn, warmup=5, rep=20)
+        print(
+            f"Seqlen = {seqlen}, {label}: {t * 1e3:.1f} us, "
+            f"{mem_io * 1e-9 / (t * 1e-3):.0f} GB/s, "
+            f"{flops * 1e-12 / (t * 1e-3):.0f} TFLOPS/s"
+        )
+
+    print(f"Arithmetic intensity: {flops / mem_io:.1f}")
+    print()

flash_attn/cute/block_sparse_utils.py

@@ -754,7 +754,7 @@ def handle_block_sparse_empty_tile_correction_sm100(
     sScale: cute.Tensor,
     stats: list,
     correction_epilogue: Callable,
-    thr_mma_pv: cute.core.ThrMma,
+    thr_mma_pv: cute.ThrMma,
     tOtO: cute.Tensor,
     sO: cute.Tensor,
     pipeline_sm_stats: cutlass.pipeline.PipelineAsync,

flash_attn/cute/cute_dsl_utils.py

@@ -61,6 +61,8 @@ def assume_tensor_aligned(t):
 
 def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False, enable_tvm_ffi=True):
     """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1."""
+    if t is None:
+        return None
     # NOTE: torch 2.9.1 doesn't support fp8 via DLPack but 2.11.0 nightly does
     # currently export raw bytes as uint8 and tell cutlass correct type
     # can directly export as fp8 when torch supports it

flash_attn/cute/flash_bwd.py

@@ -639,8 +639,8 @@ def kernel(
             thr_mma_dq = tiled_mma_dq.get_slice(tidx)
             acc_shape_dK = thr_mma_dkv.partition_shape_C((self.n_block_size, self.head_dim_padded))
             acc_shape_dV = thr_mma_dkv.partition_shape_C((self.n_block_size, self.head_dim_v_padded))
-            acc_dK = cute.make_fragment(acc_shape_dK, cutlass.Float32)
-            acc_dV = cute.make_fragment(acc_shape_dV, cutlass.Float32)
+            acc_dK = cute.make_rmem_tensor(acc_shape_dK, cutlass.Float32)
+            acc_dV = cute.make_rmem_tensor(acc_shape_dV, cutlass.Float32)
             acc_dK.fill(0.0)
             acc_dV.fill(0.0)
 
@@ -887,7 +887,7 @@ def load_dO_next():
         acc_shape_SdP = mma_params.thr_mma_sdp.partition_shape_C(
             (self.m_block_size, self.n_block_size) if cutlass.const_expr(not self.SdP_swapAB) else (self.n_block_size, self.m_block_size)
         )
-        acc_S = cute.make_fragment(acc_shape_SdP, cutlass.Float32)
+        acc_S = cute.make_rmem_tensor(acc_shape_SdP, cutlass.Float32)
         acc_S.fill(0.0)
         cute.arch.cp_async_wait_group(1 if cutlass.const_expr(self.num_stages_Q > 1) else 0)
         cute.arch.barrier()
@@ -925,7 +925,7 @@ def load_dO_next():
         # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(acc_S_mn)
 
         # MMA dP
-        acc_dP = cute.make_fragment(acc_shape_SdP, cutlass.Float32)
+        acc_dP = cute.make_rmem_tensor(acc_shape_SdP, cutlass.Float32)
         acc_dP.fill(0.0)
         cute.arch.cp_async_wait_group(1 if cutlass.const_expr(self.num_stages_dO > 1) else 0)
         cute.arch.barrier()
@@ -989,7 +989,7 @@ def dQ_mma(hook_fn):
             acc_shape_dQ = mma_params.thr_mma_dq.partition_shape_C(
                 (self.m_block_size, self.head_dim_padded) if cutlass.const_expr(not self.dQ_swapAB) else (self.head_dim_padded, self.m_block_size)
             )
-            acc_dQ = cute.make_fragment(acc_shape_dQ, cutlass.Float32)
+            acc_dQ = cute.make_rmem_tensor(acc_shape_dQ, cutlass.Float32)
             acc_dQ.fill(0.0)
             sm80_utils.gemm(
                 mma_params.thr_mma_dq, acc_dQ, mma_params.tdQrdS, mma_params.tdQrK,

flash_attn/cute/flash_bwd_postprocess.py

@@ -398,7 +398,7 @@ def kernel(
                 g2s_thr_copy = tiled_copy_accum.get_slice(tidx)
 
                 # S -> R
-                tdQrdQ_fp32 = cute.make_fragment(tdQrdQ.shape, cutlass.Float32)
+                tdQrdQ_fp32 = cute.make_rmem_tensor(tdQrdQ.shape, cutlass.Float32)
                 tdQrdQ_s2r = cute.make_tensor(tdQrdQ_fp32.iterator, tdQrdQ_fp32.shape)
 
                 smem_copy_atom = sm100_utils_basic.get_smem_store_op(
@@ -410,7 +410,7 @@ def kernel(
                     tiler_mn=tiled_tmem_ld.tiler_mn,
                 )
                 tdQsdQ_r2s = thr_tmem_ld.partition_D(thr_mma_dsk.partition_C(sdQ))
-                tdQrdQ_r2s = cute.make_fragment(tdQsdQ_r2s.shape, self.dtype)
+                tdQrdQ_r2s = cute.make_rmem_tensor(tdQsdQ_r2s.shape, self.dtype)
 
                 num_stages = cute.size(tdQrdQ_fp32, mode=[1])
                 stage_stride = self.dQ_reduce_ncol
@@ -510,7 +510,7 @@ def kernel(
                     acc_shape = tiled_mma.partition_shape_C(
                         tile_shape if const_expr(not dQ_swapAB) else tile_shape[::-1]
                     )
-                    acc = cute.make_fragment(acc_shape, cutlass.Float32)
+                    acc = cute.make_rmem_tensor(acc_shape, cutlass.Float32)
                     assert cute.size(acc) == cute.size(tdQsdQaccum)
                 else:
                     thr_mma = tiled_mma.get_slice(0)  # 1-CTA
@@ -526,7 +526,7 @@ def kernel(
                     tiled_copy_t2r = tcgen05.make_tmem_copy(tmem_load_atom, tdQtdQ)
                     thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
                     tdQrdQ_t2r_shape = thr_copy_t2r.partition_D(tdQcdQ).shape
-                    acc = cute.make_fragment(tdQrdQ_t2r_shape, Float32)
+                    acc = cute.make_rmem_tensor(tdQrdQ_t2r_shape, Float32)
                 tdQrdQaccum = cute.make_tensor(acc.iterator, cute.make_layout(tdQsdQaccum.shape))
                 cute.autovec_copy(tdQsdQaccum, tdQrdQaccum)
                 # Convert tdQrdQaccum from fp32 to fp16/bf16

flash_attn/cute/flash_bwd_sm100.py

@@ -1674,11 +1674,11 @@ def relay(
     @cute.jit
     def load(
         self,
-        thr_mma_S: cute.core.ThrMma,
-        thr_mma_dP: cute.core.ThrMma,
-        thr_mma_dV: cute.core.ThrMma,
-        thr_mma_dK: cute.core.ThrMma,
-        thr_mma_dQ: cute.core.ThrMma,
+        thr_mma_S: cute.ThrMma,
+        thr_mma_dP: cute.ThrMma,
+        thr_mma_dV: cute.ThrMma,
+        thr_mma_dK: cute.ThrMma,
+        thr_mma_dQ: cute.ThrMma,
         mQ: cute.Tensor,
         mK: cute.Tensor,
         mKt: Optional[cute.Tensor],
@@ -2821,10 +2821,10 @@ def apply_score_mod_bwd(
     @cute.jit
     def compute_loop(
         self,
-        thr_mma_S: cute.core.ThrMma,
-        thr_mma_dP: cute.core.ThrMma,
-        thr_mma_dV: cute.core.ThrMma,
-        thr_mma_dK: cute.core.ThrMma,
+        thr_mma_S: cute.ThrMma,
+        thr_mma_dP: cute.ThrMma,
+        thr_mma_dV: cute.ThrMma,
+        thr_mma_dK: cute.ThrMma,
         tStS: cute.Tensor,
         tdPtdP: cute.Tensor,
         tdVtdV: cute.Tensor,
@@ -3047,14 +3047,14 @@ def compute_loop(
                     m_block_oob = m_block >= m_block_max
                 # Prefetch 1 stage of LSE
                 pipeline_LSE.consumer_wait(consumer_state_LSE)
-                tSrLSE_s2r = cute.make_fragment(tScS_t2r[None, 0, 0, 0].shape, Float32)
+                tSrLSE_s2r = cute.make_rmem_tensor(tScS_t2r[None, 0, 0, 0].shape, Float32)
                 if const_expr(prefetch_LSE and not self.shuffle_LSE):
                     cute.autovec_copy(tSsLSE[None, 0, 0, 0, consumer_state_LSE.index], tSrLSE_s2r)
 
                 pipeline_S_P.consumer_wait(consumer_state_S_P_dP)
                 # pipeline_S_P.sync_object_full.wait(0, consumer_phase_S_P_dP)
                 #### TMEM->RMEM (Load S from TMEM)
-                tSrS_t2r = cute.make_fragment(tScS_t2r.shape, Float32)
+                tSrS_t2r = cute.make_rmem_tensor(tScS_t2r.shape, Float32)
                 cute.copy(thr_copy_t2r, tStS_t2r, tSrS_t2r)
 
                 if const_expr(self.tile_hdim == 192):
@@ -3105,7 +3105,7 @@ def compute_loop(
                 #### P = exp(S - LSE)
                 # ---------------------------------------------
                 lane_idx = cute.arch.lane_idx()
-                tSrP_r2t_f32 = cute.make_fragment(tScP_r2t.shape, Float32)  # 64
+                tSrP_r2t_f32 = cute.make_rmem_tensor(tScP_r2t.shape, Float32)  # 64
                 tSrP_r2t = cute.recast_tensor(tSrP_r2t_f32, self.q_dtype)
                 for stage in cutlass.range_constexpr(num_stages):
                     tSrS_cur = tSrS_t2r[None, stage, 0, 0]
@@ -3167,7 +3167,7 @@ def compute_loop(
 
                 ##### dS.T = P.T * (dP.T - Psum)
                 for stage in cutlass.range_constexpr(num_stages):
-                    tdPrdP_t2r = cute.make_fragment(tScS_t2r[None, 0, None, None].shape, Float32)
+                    tdPrdP_t2r = cute.make_rmem_tensor(tScS_t2r[None, 0, None, None].shape, Float32)
                     cute.copy(thr_copy_t2r, tdPtdP_t2r[None, stage, None, None], tdPrdP_t2r)
                     cute.arch.fence_view_async_tmem_load()
                     self.compute_sync_barrier.arrive_and_wait()
@@ -3459,7 +3459,7 @@ def dQacc_reduce(
         self,
         mdQaccum: cute.Tensor,
         sdQaccum: cute.Tensor,
-        thr_mma_dQ: cute.core.ThrMma,
+        thr_mma_dQ: cute.ThrMma,
         tdQtdQ: cute.Tensor,
         pipeline_dQ: PipelineAsync,
         dQaccum_empty_mbar_ptr: Optional[cute.Pointer],
@@ -3590,7 +3590,7 @@ def dQacc_reduce(
                     m_block_oob_upper = m_block >= m_block_max
                 pipeline_dQ.consumer_wait(dQ_consumer_state)
                 # TMEM -> RMEM
-                tdQrdQ_t2r = cute.make_fragment(tdQrdQ_t2r_shape, Float32)
+                tdQrdQ_t2r = cute.make_rmem_tensor(tdQrdQ_t2r_shape, Float32)
                 cute.copy(thr_copy_t2r, tdQtdQ_t2r, tdQrdQ_t2r)
                 cute.arch.fence_view_async_tmem_load()
                 cute.arch.sync_warp()
@@ -3721,8 +3721,8 @@ def epilogue_dKV(
         head_idx: Int32,
         n_block: Int32,
         seqlen,
-        thr_mma_dV: cute.core.ThrMma,
-        thr_mma_dK: cute.core.ThrMma,
+        thr_mma_dV: cute.ThrMma,
+        thr_mma_dK: cute.ThrMma,
         tdVtdV: cute.Tensor,
         tdKtdK: cute.Tensor,
         mdV: cute.Tensor,
@@ -3758,7 +3758,7 @@ def epilogue_dKV(
 
         tdVcdV_t2r_p = thr_tmem_ld_dV.partition_D(tdVcdV_tensor)
         tdVcdV_t2r = self.split_wg(tdVcdV_t2r_p, wg_idx, num_wg)
-        tdVrdV_t2r = cute.make_fragment(tdVcdV_t2r.shape, Float32)
+        tdVrdV_t2r = cute.make_rmem_tensor(tdVcdV_t2r.shape, Float32)
 
         cute.copy(thr_tmem_ld_dV, tdVtdV_t2r, tdVrdV_t2r)
         cute.arch.fence_view_async_tmem_load()
@@ -3775,7 +3775,7 @@ def epilogue_dKV(
             tiler_mn=tiled_tmem_ld_dV.tiler_mn,
         )
 
-        tdVrdV_r2s = cute.make_fragment(tdVrdV_t2r.shape, self.dv_dtype)
+        tdVrdV_r2s = cute.make_rmem_tensor(tdVrdV_t2r.shape, self.dv_dtype)
         for i in cutlass.range_constexpr(cute.size(tdVrdV_t2r, mode=[1])):
             dV_vec = tdVrdV_t2r[(None, i, 0, 0)].load()
             tdVrdV_r2s[(None, i, 0, 0)].store(dV_vec.to(self.dv_dtype))
@@ -3810,7 +3810,7 @@ def epilogue_dKV(
 
         tdKcdK_t2r_p = thr_tmem_ld_dK.partition_D(tdKcdK_tensor)
         tdKcdK_t2r = self.split_wg(tdKcdK_t2r_p, wg_idx, num_wg)
-        tdKrdK_t2r = cute.make_fragment(tdKcdK_t2r.shape, Float32)
+        tdKrdK_t2r = cute.make_rmem_tensor(tdKcdK_t2r.shape, Float32)
 
         cute.copy(tiled_tmem_ld_dK, tdKtdK_t2r, tdKrdK_t2r)
         cute.arch.fence_view_async_tmem_load()
@@ -3828,7 +3828,7 @@ def epilogue_dKV(
             tiler_mn=tiled_tmem_ld_dK.tiler_mn,
         )
 
-        tdKrdK_r2s = cute.make_fragment(tdKrdK_t2r.shape, self.dk_dtype)
+        tdKrdK_r2s = cute.make_rmem_tensor(tdKrdK_t2r.shape, self.dk_dtype)
 
         for i in cutlass.range_constexpr(cute.size(tdKrdK_t2r, mode=[1])):
             dK_vec = tdKrdK_t2r[(None, i, 0, 0)].load() * softmax_scale
@@ -3857,7 +3857,7 @@ def epilogue_dK_or_dV_tma(
         head_idx: Int32,
         n_block: Int32,
         seqlen,
-        thr_mma: cute.core.ThrMma,
+        thr_mma: cute.ThrMma,
         tdKVtdKV: cute.Tensor,
         mdKV: cute.Tensor,
         sdKV: cute.Tensor,
@@ -3976,7 +3976,7 @@ def epilogue_dK_or_dV_tma(
             if const_expr(num_epi_stages > 1):
                 tdKVcdKV_t2r = tdKVcdKV_t2r[None, epi_stage]
 
-            tdKVrdKV_t2r = cute.make_fragment(tdKVcdKV_t2r.shape, Float32)
+            tdKVrdKV_t2r = cute.make_rmem_tensor(tdKVcdKV_t2r.shape, Float32)
 
             assert cute.size(tdKVrdKV_t2r) == cute.size(tdKVtdKV_t2r) // cute.arch.WARP_SIZE, (
                 "RMEM<->TMEM fragment size mismatch"
@@ -3992,7 +3992,7 @@ def epilogue_dK_or_dV_tma(
                     tdKVrdKV_t2r[2 * i], tdKVrdKV_t2r[2 * i + 1] = cute.arch.mul_packed_f32x2(
                         (tdKVrdKV_t2r[2 * i], tdKVrdKV_t2r[2 * i + 1]), (scale, scale)
                     )
-            tdKVrdKV = cute.make_fragment(tdKVrdKV_t2r.shape, dtype)  # (32 columns)
+            tdKVrdKV = cute.make_rmem_tensor(tdKVrdKV_t2r.shape, dtype)  # (32 columns)
             tdKVrdKV.store(tdKVrdKV_t2r.load().to(dtype))
 
             # RMEM -> SMEM -- copy, fence and barrier

flash_attn/cute/flash_bwd_sm90.py

@@ -1017,7 +1017,7 @@ def load(
     def apply_score_mod(
         self,
         acc_S: cute.Tensor,
-        thr_mma_SdP: cute.core.ThrMma,
+        thr_mma_SdP: cute.ThrMma,
         batch_idx,
         head_idx,
         m_block,
@@ -1061,7 +1061,7 @@ def apply_score_mod_bwd(
         self,
         grad_tensor: cute.Tensor,
         score_tensor: cute.Tensor,
-        thr_mma_SdP: cute.core.ThrMma,
+        thr_mma_SdP: cute.ThrMma,
         batch_idx,
         head_idx,
         m_block,