perf: bypass L1 for attention tile loads

csrc/kernels/attn_mma_utils.cuh

@@ -99,23 +99,23 @@ __device__ __forceinline__ int swiz_col(int d, int r, int mask = 7) {
     return ((d >> 3) ^ (r & mask)) << 3 | (d & 7);
 }
 
-// cp.async: copy 16 bytes (8 bf16) from global to shared memory directly.
-__device__ __forceinline__ void cp_async_16(bf16* smem_ptr, const void* gmem_ptr) {
-    unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
-    asm volatile("cp.async.ca.shared.global [%0], [%1], 16;"
-                 :: "r"(smem_addr), "l"(gmem_ptr));
-}
-
 // Predicated cp.async: copy 16 bytes when `pred`, otherwise zero-fill.
-// src_size=0 → no bytes read from src, so out-of-bounds src address is safe.
+// BypassL1 defaults to .cg (L2 only); false selects .ca (L1 + L2).
+// src_size=0 means no bytes are read, so an out-of-bounds address is safe.
+template <bool BypassL1 = true>
 __device__ __forceinline__ void cp_async_16_pred(bf16* smem_ptr,
                                                   const void* gmem_ptr,
                                                   bool pred) {
     unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
     int src_size = pred ? 16 : 0;
+    if constexpr (BypassL1) {
+        asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"
+                     :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
+    } else {
         asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"
                      :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
     }
+}
 
 __device__ __forceinline__ void cp_async_commit() {
     asm volatile("cp.async.commit_group;");

csrc/tests/attn_paged_test.cu

@@ -667,8 +667,8 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
 // ======================================================================
 template <int HEAD_DIM>
 static void bench_decode(int B, int Hq, int Hkv, int seq_len) {
-    int max_ctx = seq_len + 16;
-    int pool_size = B * max_ctx;
+    int max_ctx = max(16384, seq_len + 16);
+    int pool_size = B * (seq_len + 16);
     int num_reqs = B;
 
     size_t sz_q  = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
@@ -933,9 +933,9 @@ int main() {
     bench_decode<128>(1, 32, 4, 1024);
     bench_decode<128>(1, 32, 4, 2048);
     bench_decode<128>(1, 32, 4, 4096);
+    bench_decode<128>(1, 32, 4, 16384);
     bench_decode<128>(4, 32, 4, 2048);
     bench_decode<128>(16, 32, 4, 2048);
-    bench_decode<128>(32, 32, 4, 1024);
 
     printf("\n===== PAGED PREFILL BENCH =====\n");
     print_bench_header();