ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: bypass L1 for attention tile loads

Browse Source
This commit is contained in:
ViperEkura
2026-08-15 21:23:20 +08:00
parent f8d9ab344d
commit a01c1fd427
2 changed files with 13 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_mma_utils.cuh
+10 -10
View File
@@ -99,22 +99,22 @@ __device__ __forceinline__ int swiz_col(int d, int r, int mask = 7) {
return ((d >> 3) ^ (r & mask)) << 3 | (d & 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. // 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, __device__ __forceinline__ void cp_async_16_pred(bf16* smem_ptr,
const void* gmem_ptr, const void* gmem_ptr,
bool pred) { bool pred) {
unsigned smem_addr = __cvta_generic_to_shared(smem_ptr); unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
int src_size = pred ? 16 : 0; int src_size = pred ? 16 : 0;
asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;" if constexpr (BypassL1) {
:: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size)); 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() { __device__ __forceinline__ void cp_async_commit() {
csrc/tests/attn_paged_test.cu
+3 -3
View File
@@ -667,8 +667,8 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
// ====================================================================== // ======================================================================
template <int HEAD_DIM> template <int HEAD_DIM>
static void bench_decode(int B, int Hq, int Hkv, int seq_len) { static void bench_decode(int B, int Hq, int Hkv, int seq_len) {
int max_ctx = seq_len + 16; int max_ctx = max(16384, seq_len + 16);
int pool_size = B * max_ctx; int pool_size = B * (seq_len + 16);
int num_reqs = B; int num_reqs = B;
size_t sz_q = (size_t)B * Hq * HEAD_DIM * sizeof(bf16); 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, 1024);
bench_decode<128>(1, 32, 4, 2048); bench_decode<128>(1, 32, 4, 2048);
bench_decode<128>(1, 32, 4, 4096); bench_decode<128>(1, 32, 4, 4096);
bench_decode<128>(1, 32, 4, 16384);
bench_decode<128>(4, 32, 4, 2048); bench_decode<128>(4, 32, 4, 2048);
bench_decode<128>(16, 32, 4, 2048); bench_decode<128>(16, 32, 4, 2048);
bench_decode<128>(32, 32, 4, 1024);
printf("\n===== PAGED PREFILL BENCH =====\n"); printf("\n===== PAGED PREFILL BENCH =====\n");
print_bench_header(); print_bench_header();