From 1779830c01f653ff546cf3c9b87f8a6535765407 Mon Sep 17 00:00:00 2001 From: waveringana Date: Thu, 8 Oct 2026 23:20:04 -0400 Subject: [PATCH] valefar: bonsai prefill +8-13% (staged PTQ1_0 MMQ main loop, unpadded GQA q8_0 prefill attention), no core dumps on restart --- hosts/valefar/bonsai.nix | 4 +- .../llama-cpp-bonsai-fattn-q8-prefill.patch | 427 +++++++++++++++ .../pkgs/llama-cpp-bonsai-mmq-staged.patch | 508 ++++++++++++++++++ pkgs-set/pkgs/llama-cpp-bonsai.nix | 8 + 4 files changed, 946 insertions(+), 1 deletion(-) create mode 100644 pkgs-set/pkgs/llama-cpp-bonsai-fattn-q8-prefill.patch create mode 100644 pkgs-set/pkgs/llama-cpp-bonsai-mmq-staged.patch diff --git a/hosts/valefar/bonsai.nix b/hosts/valefar/bonsai.nix index d706bf5..7d4e16c 100644 --- a/hosts/valefar/bonsai.nix +++ b/hosts/valefar/bonsai.nix @@ -101,7 +101,9 @@ in (toString port) ]; Restart = "on-failure"; - RestartSec = "5s"; + RestartSec = "15s"; + # a crash dump is the whole mapped model (~130 MB each); a restart loop filled 4 GB with them + LimitCORE = 0; TimeoutStartSec = "infinity"; # stopping writes the slots and the RAM prompt cache to the disk tier TimeoutStopSec = "5min"; diff --git a/pkgs-set/pkgs/llama-cpp-bonsai-fattn-q8-prefill.patch b/pkgs-set/pkgs/llama-cpp-bonsai-fattn-q8-prefill.patch new file mode 100644 index 0000000..864a40f --- /dev/null +++ b/pkgs-set/pkgs/llama-cpp-bonsai-fattn-q8-prefill.patch @@ -0,0 +1,427 @@ +diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh +index a95dc4dc..2a558bce 100644 +--- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh ++++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh +@@ -480,32 +480,83 @@ static __device__ __forceinline__ void flash_attn_ext_f16_store_q4_0_block( + } + + // Dequantize one q8_0 block (32 weights) given its scale and the 32 int8 bytes as 8 words, store 16 half2. ++// Each byte pair becomes the half2 1024 + (q + 128) by placing it under the exponent byte 0x64, minus 1152 that is ++// exactly q. d*q is exact in float (11 x 8 significant bits), so the correctly rounded half product equals the ++// rounding of d*float(q) that dequantize_q8_0 -> half does: the tile is bit-identical, at a fraction of the cost. + static __device__ __forceinline__ void flash_attn_ext_f16_store_q8_0_block( +- half2 * const __restrict__ dst, const float d, const uint32_t * const __restrict__ q) { ++ half2 * const __restrict__ dst, const half d, const uint32_t * const __restrict__ q) { ++ const half2 d2 = __half2half2(d); ++ const half2 bias = make_half2(1152.0f, 1152.0f); + #pragma unroll + for (int m = 0; m < 8; ++m) { +- const uint32_t v = q[m]; +- dst[2*m ] = __floats2half2_rn(d*float(int(int8_t(v & 0xFF))), d*float(int(int8_t((v >> 8) & 0xFF)))); +- dst[2*m + 1] = __floats2half2_rn(d*float(int(int8_t((v >> 16) & 0xFF))), d*float(int(int8_t((v >> 24) & 0xFF)))); ++ const uint32_t u = q[m] ^ 0x80808080u; // int8 -> uint8 offset by 128 ++ const uint32_t lo = __byte_perm(u, 0x64646464u, 0x4140); ++ const uint32_t hi = __byte_perm(u, 0x64646464u, 0x4342); ++ dst[2*m ] = __hmul2(__hsub2(*((const half2 *) &lo), bias), d2); ++ dst[2*m + 1] = __hmul2(__hsub2(*((const half2 *) &hi), bias), d2); ++ } ++} ++ ++// A q4_0 / q8_0 K or V row is read as block pairs (64 weights): a q4_0 pair is 36 bytes and a q8_0 pair 68 bytes, ++// both 4-byte aligned at even block offsets, so the pair is read as aligned 32-bit words and the odd block's bytes ++// are realigned with byte permutes. ++template ++struct flash_attn_ext_f16_q_pair { ++ static constexpr int bs = 32; // weights per block, both types ++ static constexpr int ts = type_KV == GGML_TYPE_Q4_0 ? 18 : 34; // bytes per block ++ static constexpr int nwords = 2*ts/4; // 32-bit words per block pair: 9 or 17 ++ static_assert((2*ts) % 4 == 0, "block pair is not word aligned"); ++}; ++ ++template ++static __device__ __forceinline__ void flash_attn_ext_f16_load_q_pair( ++ const char * const __restrict__ src, uint32_t * const __restrict__ w) { ++#pragma unroll ++ for (int k = 0; k < flash_attn_ext_f16_q_pair::nwords; ++k) { ++ w[k] = __ldg((const uint32_t *) src + k); ++ } ++} ++ ++// Dequantize a block pair from its words into 32 half2 (64 weights). ++template ++static __device__ __forceinline__ void flash_attn_ext_f16_dequantize_q_pair( ++ const uint32_t * const __restrict__ w, half2 * const __restrict__ dst) { ++ static_assert(type_KV == GGML_TYPE_Q4_0 || type_KV == GGML_TYPE_Q8_0, "unsupported quantized K/V type"); ++ constexpr int bs = flash_attn_ext_f16_q_pair::bs; ++ if constexpr (type_KV == GGML_TYPE_Q4_0) { ++ // block 0: d = bytes 0-1, qs = bytes 2-17. block 1: d = bytes 18-19, qs = bytes 20-35 (= w5..w8). ++ const float d0 = __half2float(__ushort_as_half((unsigned short) (w[0] & 0xFFFF))); ++ const float d1 = __half2float(__ushort_as_half((unsigned short) (w[4] >> 16))); ++ flash_attn_ext_f16_store_q4_0_block(dst, d0, ++ __byte_perm(w[0], w[1], 0x5432), __byte_perm(w[1], w[2], 0x5432), ++ __byte_perm(w[2], w[3], 0x5432), __byte_perm(w[3], w[4], 0x5432)); ++ flash_attn_ext_f16_store_q4_0_block(dst + bs/2, d1, w[5], w[6], w[7], w[8]); ++ } else { ++ // block 0: d = bytes 0-1, qs = bytes 2-33. block 1: d = bytes 34-35, qs = bytes 36-67 (= w9..w16). ++ const half d0 = __ushort_as_half((unsigned short) (w[0] & 0xFFFF)); ++ const half d1 = __ushort_as_half((unsigned short) (w[8] >> 16)); ++ uint32_t q0[8]; ++#pragma unroll ++ for (int k = 0; k < 8; ++k) { ++ q0[k] = __byte_perm(w[k], w[k + 1], 0x5432); ++ } ++ flash_attn_ext_f16_store_q8_0_block(dst, d0, q0); ++ flash_attn_ext_f16_store_q8_0_block(dst + bs/2, d1, w + 9); + } + } + + // Loads a [nbatch_fa x ne] tile of a q4_0 / q8_0 K or V cache and dequantizes it into the F16 shared tile, in the + // same layout the F16 loader produces. KV points at byte 0 of tile row 0, e0 / ne are the element offset and count +-// within a row (multiples of 64), stride_KV is the row stride in bytes. Each thread handles two adjacent blocks +-// (64 weights): a q4_0 pair is 36 bytes and a q8_0 pair 68 bytes, both 4-byte aligned at even block offsets, so the +-// pair is read as aligned 32-bit words and the odd block's bytes are realigned with byte permutes. ++// within a row (multiples of 64), stride_KV is the row stride in bytes. Each thread handles one block pair at a time. + template + static __device__ __forceinline__ void flash_attn_ext_f16_load_tile_q( + const char * const __restrict__ KV, half2 * const __restrict__ tile_KV, + const int e0, const int ne, const int stride_KV, const int i_sup) { +- static_assert(type_KV == GGML_TYPE_Q4_0 || type_KV == GGML_TYPE_Q8_0, "unsupported quantized K/V type"); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nthreads = nwarps*warp_size; +- constexpr int bs = 32; // weights per block, both types +- constexpr int ts = type_KV == GGML_TYPE_Q4_0 ? 18 : 34; // bytes per block +- constexpr int nwords = 2*ts/4; // 32-bit words per block pair: 9 or 17 +- static_assert((2*ts) % 4 == 0, "block pair is not word aligned"); ++ constexpr int bs = flash_attn_ext_f16_q_pair::bs; ++ constexpr int ts = flash_attn_ext_f16_q_pair::ts; ++ constexpr int nwords = flash_attn_ext_f16_q_pair::nwords; + + const int pairs_per_row = ne / (2*bs); + const int ntotal = nbatch_fa * pairs_per_row; +@@ -523,32 +574,89 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile_q( + continue; + } + +- const uint32_t * src = (const uint32_t *) (KV + int64_t(i)*stride_KV + ((e0 + p*2*bs)/bs)*ts); + uint32_t w[nwords]; +-#pragma unroll +- for (int k = 0; k < nwords; ++k) { +- w[k] = __ldg(src + k); ++ flash_attn_ext_f16_load_q_pair(KV + int64_t(i)*stride_KV + ((e0 + p*2*bs)/bs)*ts, w); ++ flash_attn_ext_f16_dequantize_q_pair(w, dst); ++ } ++} ++ ++// Prompt-sized tiles over a quantized cache are software pipelined: the raw rows of the next K or V tile and the next ++// mask tile are copied into shared memory with cp.async while the tensor cores work on the current tile, then ++// dequantized shared -> shared. Needs whole 16-byte aligned rows (head size 256: 272 bytes of q8_0, 144 of q4_0) ++// and no OOB rows (ncols2 > 1 implies a padded cache); the staging buffers fit in the Q tile's shared memory. ++static constexpr __host__ __device__ int ggml_cuda_fattn_mma_q_row_bytes(const ggml_type type_KV, const int D) { ++ return D/32 * (type_KV == GGML_TYPE_Q4_0 ? 18 : 34); ++} ++ ++static constexpr __host__ __device__ bool ggml_cuda_fattn_mma_kv_prefetch_shape( ++ const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2) { ++ return ggml_cuda_fattn_mma_kv_native(type_KV) && ncols1*ncols2 >= 64 && ncols2 > 1 && ++ ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ) % 16 == 0 && ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV) % 16 == 0; ++} ++ ++static __host__ bool ggml_cuda_fattn_mma_kv_prefetch( ++ const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2, const int cc) { ++ return cp_async_available(cc) && ggml_cuda_fattn_mma_kv_prefetch_shape(type_KV, DKQ, DV, ncols1, ncols2); ++} ++ ++static constexpr __device__ bool ggml_cuda_fattn_mma_kv_prefetch( ++ const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2) { ++#ifdef CP_ASYNC_AVAILABLE ++ return ggml_cuda_fattn_mma_kv_prefetch_shape(type_KV, DKQ, DV, ncols1, ncols2); ++#else ++ GGML_UNUSED_VARS(type_KV, DKQ, DV, ncols1, ncols2); ++ return false; ++#endif // CP_ASYNC_AVAILABLE ++} ++ ++// Starts copying the raw rows of a [nbatch_fa x row_bytes] quantized K/V tile into shared memory, row after row. ++template ++static __device__ __forceinline__ void flash_attn_ext_f16_load_raw_q( ++ const char * const __restrict__ KV, char * const __restrict__ raw, const int stride_KV) { ++ constexpr int warp_size = ggml_cuda_get_physical_warp_size(); ++ constexpr int nthreads = nwarps*warp_size; ++ constexpr int cpr = row_bytes/16; // 16-byte chunks per row ++ constexpr int nchunks = nbatch_fa*cpr; ++ static_assert(row_bytes % 16 == 0, "rows must be whole 16-byte chunks"); ++ ++ const unsigned int raw_32 = ggml_cuda_cvta_generic_to_shared(raw); ++ // Not unrolled: the 64-bit addresses of all chunks would spill the Q and VKQ registers at head size 256. ++#pragma unroll 1 ++ for (int c0 = 0; c0 < nchunks; c0 += nthreads) { ++ const int c = c0 + threadIdx.y*warp_size + threadIdx.x; ++ if (nchunks % nthreads == 0 || c < nchunks) { ++ const int i = c / cpr; ++ const int k = c % cpr; ++ cp_async_cg_16<0>(raw_32 + i*row_bytes + 16*k, KV + int64_t(i)*stride_KV + 16*k); + } ++ } ++} + +- if constexpr (type_KV == GGML_TYPE_Q4_0) { +- // block 0: d = bytes 0-1, qs = bytes 2-17. block 1: d = bytes 18-19, qs = bytes 20-35 (= w5..w8). +- const float d0 = __half2float(__ushort_as_half((unsigned short) (w[0] & 0xFFFF))); +- const float d1 = __half2float(__ushort_as_half((unsigned short) (w[4] >> 16))); +- flash_attn_ext_f16_store_q4_0_block(dst, d0, +- __byte_perm(w[0], w[1], 0x5432), __byte_perm(w[1], w[2], 0x5432), +- __byte_perm(w[2], w[3], 0x5432), __byte_perm(w[3], w[4], 0x5432)); +- flash_attn_ext_f16_store_q4_0_block(dst + bs/2, d1, w[5], w[6], w[7], w[8]); +- } else { +- // block 0: d = bytes 0-1, qs = bytes 2-33. block 1: d = bytes 34-35, qs = bytes 36-67 (= w9..w16). +- const float d0 = __half2float(__ushort_as_half((unsigned short) (w[0] & 0xFFFF))); +- const float d1 = __half2float(__ushort_as_half((unsigned short) (w[8] >> 16))); +- uint32_t q0[8]; +-#pragma unroll +- for (int k = 0; k < 8; ++k) { +- q0[k] = __byte_perm(w[k], w[k + 1], 0x5432); ++// Dequantizes a raw [nbatch_fa x D] tile from shared memory into the F16 tile, block pair by block pair. ++template ++static __device__ __forceinline__ void flash_attn_ext_f16_dequantize_raw_q( ++ const char * const __restrict__ raw, half2 * const __restrict__ tile_KV) { ++ constexpr int warp_size = ggml_cuda_get_physical_warp_size(); ++ constexpr int nthreads = nwarps*warp_size; ++ constexpr int row_bytes = ggml_cuda_fattn_mma_q_row_bytes(type_KV, D); ++ constexpr int ts = flash_attn_ext_f16_q_pair::ts; ++ constexpr int nwords = flash_attn_ext_f16_q_pair::nwords; ++ constexpr int ppr = D/64; // block pairs per row ++ constexpr int ntotal = nbatch_fa*ppr; ++ ++#pragma unroll ++ for (int t0 = 0; t0 < ntotal; t0 += nthreads) { ++ const int t = t0 + threadIdx.y*warp_size + threadIdx.x; ++ if (ntotal % nthreads == 0 || t < ntotal) { ++ const int i = t / ppr; ++ const int p = t % ppr; ++ const uint32_t * src = (const uint32_t *) (raw + i*row_bytes + p*(2*ts)); ++ uint32_t w[nwords]; ++#pragma unroll ++ for (int k = 0; k < nwords; ++k) { ++ w[k] = src[k]; + } +- flash_attn_ext_f16_store_q8_0_block(dst, d0, q0); +- flash_attn_ext_f16_store_q8_0_block(dst + bs/2, d1, w + 9); ++ flash_attn_ext_f16_dequantize_q_pair(w, tile_KV + i*stride_tile + p*32); + } + } + } +@@ -683,6 +791,15 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( + static_assert(type_KV == GGML_TYPE_F16 || ((2*nbatch_K2) % 64 == 0 && (2*nbatch_V2) % 64 == 0 && !V_is_K_view), + "quantized K/V loading needs whole block pairs per pass"); + ++ constexpr bool kv_prefetch = ggml_cuda_fattn_mma_kv_prefetch(type_KV, DKQ, DV, ncols1, ncols2); ++ static_assert(!kv_prefetch || (nstages == 0 && nbatch_K2 == DKQ/2 && nbatch_V2 == DV/2 && !oob_check), ++ "K/V prefetch needs whole rows per pass and no OOB rows"); ++ constexpr int row_bytes_K = ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ); ++ constexpr int row_bytes_V = ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV); ++ // Raw staging tiles behind the mask tile, see flash_attn_ext_f16_process_tile. ++ char * raw_K = (char *) (tile_mask + ncols1*(nbatch_fa + 8)); ++ char * raw_V = raw_K + nbatch_fa*row_bytes_K; ++ + const int k_VKQ_0 = kb0 * nbatch_fa; + #if defined(TURING_MMA_AVAILABLE) + T_C_KQ KQ_C[nbatch_fa/(np*(cols_per_warp == 8 ? T_C_KQ::I : T_C_KQ::J))]; +@@ -701,6 +818,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( + __syncthreads(); + flash_attn_ext_f16_load_tile + (V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup); ++ } else if constexpr (kv_prefetch) { ++ // Raw K and the mask of this tile were started by the previous iteration (or before the first one); ++ // the barrier also retires every warp's VKQ reads of tile_V before tile_K (the same memory) is written. ++ cp_async_wait_all(); ++ __syncthreads(); ++ flash_attn_ext_f16_load_raw_q ++ ((const char *) V_h2 + int64_t(k_VKQ_0)*stride_V, raw_V, stride_V); + } else { + constexpr bool use_cp_async = nstages == 1; + if (ncols2 > 1 || mask_h) { +@@ -724,6 +848,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( + if (use_cp_async) { + cp_async_wait_all(); + } ++ } else if constexpr (kv_prefetch) { ++ flash_attn_ext_f16_dequantize_raw_q(raw_K, tile_K); + } else { + flash_attn_ext_f16_load_tile_q + ((const char *) K_h2 + int64_t(k_VKQ_0)*stride_K, tile_K, 2*k0_start, 2*k0_diff, stride_K, k_VKQ_sup); +@@ -783,7 +909,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( + } + } + +- if constexpr (nstages <= 1) { ++ if constexpr (nstages <= 1 && !kv_prefetch) { + __syncthreads(); // Only needed if tile_K == tile_V. + } + } +@@ -1081,6 +1207,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( + if (use_cp_async) { + cp_async_wait_all(); + } ++ } else if constexpr (kv_prefetch) { ++ // Every warp is past KQ and the softmax, so raw K and the mask tile can take the next tile. ++ cp_async_wait_all(); ++ __syncthreads(); ++ if constexpr (!last_iter) { ++ flash_attn_ext_f16_load_mask ++ (mask_h + k_VKQ_0 + nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); ++ flash_attn_ext_f16_load_raw_q ++ ((const char *) K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, raw_K, stride_K); ++ } ++ flash_attn_ext_f16_dequantize_raw_q(raw_V, tile_V); + } else { + flash_attn_ext_f16_load_tile_q + ((const char *) V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, i0_start, i0_diff, stride_V, k_VKQ_sup); +@@ -1131,7 +1268,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( + } + #endif // defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) + +- if constexpr (nstages <= 1) { ++ // With prefetch the barrier at the start of the next iteration (or after the loop) protects tile_V. ++ if constexpr (nstages <= 1 && !(kv_prefetch && !last_iter)) { + __syncthreads(); // Only needed if tile_K == tile_V. + } + } +@@ -1376,6 +1514,15 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( + + int kb0 = kb0_start; + ++ // Quantized prompt tiles: start the raw K and mask tiles of the first iteration, see flash_attn_ext_f16_iter. ++ if constexpr (ggml_cuda_fattn_mma_kv_prefetch(type_KV, DKQ, DV, ncols1, ncols2)) { ++ char * raw_K = (char *) (tile_mask + ncols1*(nbatch_fa + 8)); ++ flash_attn_ext_f16_load_mask ++ (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, nbatch_fa, jt*ncols1, ne01); ++ flash_attn_ext_f16_load_raw_q ++ ((const char *) K_h2 + int64_t(kb0)*nbatch_fa*stride_K, raw_K, stride_K); ++ } ++ + // Preload mask and K data for first iteration when using cp_async with multiple stages: + if constexpr (nstages > 1) { + static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline"); +@@ -2051,7 +2198,11 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml + const size_t nbytes_shared_mask = ncols1 * (nbatch_fa/2 + 4) * sizeof(half2); + const size_t nbytes_shared_combine = nwarps*cols_per_warp * (nbatch_combine + 4) * sizeof(half2); + +- const size_t nbytes_shared_KV = nstages <= 1 ? nbytes_shared_KV_1stage : nbytes_shared_KV_2stage; ++ // Raw staging tiles of the quantized prompt-tile pipeline, behind the mask tile. ++ const size_t nbytes_shared_raw = ggml_cuda_fattn_mma_kv_prefetch(type_KV, DKQ, DV, ncols1, ncols2, cc) ? ++ nbatch_fa * (ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ) + ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV)) : 0; ++ ++ const size_t nbytes_shared_KV = (nstages <= 1 ? nbytes_shared_KV_1stage : nbytes_shared_KV_2stage) + nbytes_shared_raw; + + const size_t nbytes_shared_total = std::max(nbytes_shared_combine, Q_in_reg ? + std::max(nbytes_shared_Q, nbytes_shared_KV + nbytes_shared_mask) : +@@ -2140,9 +2291,11 @@ static inline bool ggml_cuda_fattn_mma_kv_native_supported(const ggml_tensor * d + if (!(Q->ne[0] == 128 || Q->ne[0] == 256) || V->ne[0] != Q->ne[0]) { + return false; + } +- // Rows must be 32-bit addressable for the block-pair loads. ++ // Rows must be 32-bit addressable for the block-pair loads, and whole 16-byte chunks for the cp.async copies of ++ // the prompt tiles at head size 256 (see ggml_cuda_fattn_mma_kv_prefetch). ++ const size_t align = Q->ne[0] == 256 ? 16 : 4; + for (const ggml_tensor * t : {K, V}) { +- if (t->nb[1] % 4 != 0 || t->nb[2] % 4 != 0 || t->nb[3] % 4 != 0) { ++ if (t->nb[1] % align != 0 || t->nb[2] % align != 0 || t->nb[3] % align != 0 || (uintptr_t) t->data % align != 0) { + return false; + } + } +diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu +index 62e8e9b8..39655868 100644 +--- a/ggml/src/ggml-cuda/fattn.cu ++++ b/ggml/src/ggml-cuda/fattn.cu +@@ -42,6 +42,21 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + } + ++// Whether launch_fattn gives every output tile of ncols1 x ncols2 columns to one block, i.e. does not use stream-k ++// (mirrors its rule, with the occupancy the MMA config targets). ++static bool ggml_cuda_fattn_mma_whole_tiles(const ggml_tensor * dst, const int DKQ, const int DV, const int ncols1, const int ncols2, const int cc) { ++ if (cc >= GGML_CUDA_CC_ADA_LOVELACE || amd_wmma_available(cc)) { ++ return false; ++ } ++ const ggml_tensor * Q = dst->src[0]; ++ const ggml_tensor * K = dst->src[1]; ++ const int64_t gqa_ratio = Q->ne[2] / K->ne[2]; ++ const int64_t ntiles = (Q->ne[1] + ncols1 - 1)/ncols1 * ((gqa_ratio + ncols2 - 1)/ncols2) * K->ne[2] * Q->ne[3]; ++ const int64_t max_blocks = int64_t(ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2, cc)) * ggml_cuda_info().devices[ggml_cuda_get_device()].nsm; ++ const int64_t nwaves = (ntiles + max_blocks - 1) / max_blocks; ++ return 100 * ntiles / (max_blocks*nwaves) >= 75; ++} ++ + template + static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; +@@ -97,6 +112,24 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con + } + } + ++ // Prompt batches over a quantized cache: tile the Q heads of a KV head without padding columns (24 over 4 KV heads ++ // as 3 pairs, not 6 of 8), so no tensor-core work is wasted and every K/V tile is streamed fewer times. Each ++ // column's arithmetic does not depend on ncols2, so the result is the same as long as neither tiling splits a ++ // column's KV range across blocks (stream-k): small batches and decode/verify keep the 8-wide tile. ++ if constexpr (ggml_cuda_fattn_mma_kv_native(type_KV)) { ++ if (use_gqa_opt && Q->ne[1] > 16 && gqa_ratio > 4 && gqa_ratio % 8 != 0 && gqa_ratio % 2 == 0) { ++ if (gqa_ratio % 4 == 0) { ++ if (ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 16, 4, cc) && ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 8, 8, cc)) { ++ ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); ++ return; ++ } ++ } else if (ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 32, 2, cc) && ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 8, 8, cc)) { ++ ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); ++ return; ++ } ++ } ++ } ++ + if (use_gqa_opt && gqa_ratio > 4) { + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); + return; +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index c93cacf5..57ead864 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -10171,6 +10171,15 @@ static std::vector> make_test_cases_eval() { + } + } + ++ // prompt-shaped cases with the same shape: the quantized cache is tiled two query heads per column group ++ for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, }) { ++ for (int nb : { 9, 33, 64, 300, 1024, }) { ++ for (int kv : { 512, 4096, }) { ++ test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV)); ++ } ++ } ++ } ++ + // prefill-shaped cases with long KV (nb >= 32, kv >= 1024): covers the + // XMX/GEMM-accelerated SYCL FA path which only activates for these shapes. + for (int kv : { 1024, 2048, }) { +@@ -10746,6 +10755,24 @@ static std::vector> make_test_cases_perf() { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + ++ // prefill at depth with the Qwen 3.x full-attention shape (24 query heads over 4 KV heads) ++ for (int kv : { 17408, 33792, 66560, }) { ++ for (int nb : { 512, 1024, }) { ++ for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_Q8_0, }) { ++ test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV)); ++ } ++ } ++ } ++ ++ // prompt-shaped cases with the same shape: the quantized cache is tiled two query heads per column group ++ for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, }) { ++ for (int nb : { 9, 33, 64, 300, 1024, }) { ++ for (int kv : { 512, 4096, }) { ++ test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV)); ++ } ++ } ++ } ++ + for (int kv : { 4096, 8192, 16384, }) { + for (int hs : { 64, 128, }) { + for (int nr : { 1, 4, }) { diff --git a/pkgs-set/pkgs/llama-cpp-bonsai-mmq-staged.patch b/pkgs-set/pkgs/llama-cpp-bonsai-mmq-staged.patch new file mode 100644 index 0000000..2d343ba --- /dev/null +++ b/pkgs-set/pkgs/llama-cpp-bonsai-mmq-staged.patch @@ -0,0 +1,508 @@ +diff --git a/ggml/src/ggml-cuda/mmq-ptq1_0.cuh b/ggml/src/ggml-cuda/mmq-ptq1_0.cuh +new file mode 100644 +index 00000000..9100c33d +--- /dev/null ++++ b/ggml/src/ggml-cuda/mmq-ptq1_0.cuh +@@ -0,0 +1,334 @@ ++#pragma once ++ ++// PTQ1_0 MMQ main loop for Ampere and newer (cp.async + int8 mma), used when every warp gets a multiple of 8 columns. ++// The k loop runs one 128-value block per stage over two shared-memory stages, with one barrier per stage. During ++// stage s the warps compute from buffer s%2, decode the raw x block of stage s+1 into the other buffer between their ++// mma tiles, and copy y for stage s+1 and raw x for stage s+2 with cp.async. B fragments come from ldmatrix. ++// Every output keeps the order of ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma with MMQ_Q8_1_DS_LAYOUT_D1 (an int32 sum per ++// block, then sum = fma(C*dA, dB, sum) in k order) and the tile shape stays I x J, so the stream-k split is unchanged: ++// results are bit-identical to the generic loop, which GGML_CUDA_NO_PTQ1_0_STAGED selects. ++ ++#include "mmq.cuh" ++ ++#if !defined(GGML_USE_HIP) ++ ++namespace ggml_cuda_mmq_ptq1_0 { ++ ++using namespace ggml_cuda_mma; ++ ++static constexpr int x_stride = MMQ_TILE_NE_K + 4; // ints per x row: one decoded block plus padding against bank conflicts ++static constexpr int y_stride = MMQ_TILE_Y_K; // ints per y column: four scale slots and one block of int8 values ++static constexpr int x_words_per_row = sizeof(block_ptq1_0) / sizeof(int); ++ ++template struct layout { ++ static constexpr int nthreads = ggml_cuda_mmq_get_nthreads(GGML_TYPE_PTQ1_0, J, fallback); ++ static constexpr int nwarps = nthreads / 32; ++ static constexpr int I = ggml_cuda_mmq_get_I(GGML_TYPE_PTQ1_0, J, fallback); ++ static constexpr int warps_i = I / 32; // each warp owns 32 rows: two 16-row mma tiles ++ static constexpr int warps_j = nwarps / warps_i; ++ static constexpr int cols = J / warps_j; // columns per warp ++ static constexpr int ntiles = cols / 8; // 8-column mma tiles per warp ++ static constexpr int ntasks = I * 8 / nthreads; // x words decoded per thread and stage, 8 lanes per row ++ static constexpr int x_ints = I * x_stride + I; // decoded values, then one float scale per row ++ static constexpr int raw_ints = I * x_words_per_row; // one stage of x as stored in VRAM ++ static constexpr int ncopies = (raw_ints + nthreads - 1) / nthreads; ++ static constexpr int y_ints = J * y_stride; ++ ++ static_assert(ggml_cuda_mmq_ptq1_0_staged(J, fallback), "unsupported PTQ1_0 MMQ layout"); ++}; ++ ++static __device__ __forceinline__ void ldmatrix_x4(int (&r)[4], const int * p) { ++#ifdef TURING_MMA_AVAILABLE ++ asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];" ++ : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) ++ : "r"(ggml_cuda_cvta_generic_to_shared((void *) p))); ++#else ++ GGML_UNUSED_VARS(r, p); ++ NO_DEVICE_CODE; ++#endif // TURING_MMA_AVAILABLE ++} ++ ++static __device__ __forceinline__ void cp_async_ca_4(const unsigned int dst, const void * src) { ++#ifdef CP_ASYNC_AVAILABLE ++ asm volatile("cp.async.ca.shared.global [%0], [%1], 4;" : : "r"(dst), "l"(src)); ++#else ++ GGML_UNUSED_VARS(dst, src); ++ NO_DEVICE_CODE; ++#endif // CP_ASYNC_AVAILABLE ++} ++ ++// Decodes word n of this thread's share of the raw x stage (block_ptq1_0 words, x_words_per_row per row) into xs. ++// Same unpack and layout as ggml_cuda_mmq_ptq1_0_x_stage::store; lane 6 also holds the scale, lane 7 has nothing to do. ++template ++static __device__ __forceinline__ void decode_word(const int * __restrict__ raw, int * __restrict__ xs, const int n) { ++ using L = layout; ++ const int lane = threadIdx.x % 8; ++ const int i = (n * L::nthreads + threadIdx.y * 32 + threadIdx.x) / 8; ++ const uint32_t w = raw[i * x_words_per_row + min(lane, x_words_per_row - 1)]; ++ int * row = xs + i * x_stride; ++ ++ // Branch-free so that decoding does not split the mma code into basic blocks: lanes 6 and 7 send the stores they ++ // do not need to the padding at the end of the row, which ldmatrix never reads. ++ uint32_t v_lo = __byte_perm(w, 0, 0x4140); ++ uint32_t v_hi = __byte_perm(w, 0, 0x4342); ++ const bool full_lane = lane < 6; ++ v_hi = full_lane ? v_hi : v_lo; ++ const int dst_base = lane < 4 ? lane : lane < 6 ? 16 + lane : MMQ_TILE_NE_K; ++ const int dst_stride = lane < 4 ? 4 : lane < 6 ? 2 : 0; ++ int q[5]; ++#pragma unroll ++ for (int t = 0; t < 5; ++t) { ++ const uint32_t w_lo = v_lo * 3; ++ const uint32_t w_hi = v_hi * 3; ++ v_lo = w_lo & 0x00FF00FF; ++ v_hi = w_hi & 0x00FF00FF; ++ // Each byte is a trit in {0, 1, 2}: adding 0x7F never carries, and the xor gives trit - 1 as int8. ++ q[t] = (__byte_perm(w_lo, w_hi, 0x7531) + 0x7F7F7F7F) ^ 0x80808080; ++ row[dst_base + t * dst_stride] = q[t]; ++ } ++ // q[t] = {qh0.t, qh1.t, qh0.t, qh1.t} on lane 6; the layout wants {qh0.t, qh1.t, qh0.t+1, qh1.t+1}. ++ const bool qh_lane = lane == 6; ++ row[qh_lane ? 30 : MMQ_TILE_NE_K + 1] = __byte_perm(q[0], q[1], 0x5410); ++ row[qh_lane ? 31 : MMQ_TILE_NE_K + 2] = __byte_perm(q[2], q[3], 0x5410); ++ float * d_dst = qh_lane ? (float *) (xs + L::I * x_stride) + i : (float *) (row + MMQ_TILE_NE_K + 3); ++ *d_dst = __half2float(__ushort_as_half((unsigned short) (w >> 16))); ++} ++ ++template ++static __device__ __forceinline__ void stage(const int * __restrict__ xs, const int * __restrict__ ys, float * __restrict__ sum, ++ const int * __restrict__ raw_next, int * __restrict__ xs_next) { ++ using L = layout; ++ typedef tile<16, 8, int> tile_A; ++ typedef tile< 8, 8, int> tile_B; ++ typedef tile<16, 8, int> tile_C; ++ ++ const int wi = threadIdx.y / L::warps_j; ++ const int wj = threadIdx.y % L::warps_j; ++ const int i0 = wi * 32; ++ ++ const float * xd = (const float *) (xs + L::I * x_stride); ++ const float * yd = (const float *) ys; ++ const int * yq = ys + 4; ++ ++ tile_A A[2][4]; ++ float dA[2][2]; ++#pragma unroll ++ for (int m = 0; m < 2; ++m) { ++#pragma unroll ++ for (int k = 0; k < 4; ++k) { ++ load_ldmatrix(A[m][k], xs + (i0 + m * 16) * x_stride + k * QI8_0, x_stride); ++ } ++#pragma unroll ++ for (int h = 0; h < 2; ++h) { ++ dA[m][h] = xd[i0 + m * 16 + tile_C::get_i(2 * h)]; ++ } ++ } ++ ++ // Tiles are processed in groups with the k loop outside, so that each IMMA has 2*group independent accumulators ++ // to alternate with instead of waiting on its own previous result. ++ constexpr int group = L::ntiles % 4 == 0 ? 4 : L::ntiles % 3 == 0 ? 3 : L::ntiles % 2 == 0 ? 2 : 1; ++ static_assert(L::ntiles % group == 0, "bad tile group"); ++ ++#pragma unroll ++ for (int t0 = 0; t0 < L::ntiles; t0 += group) { ++ // b[g][h][q] holds ints 4*q + threadIdx.x%4 (+16 for h = 1) of column j0 + threadIdx.x/4 of tile t0 + g: the ++ // B fragments of k-steps 2*h and 2*h + 1. ++ int b[group][2][4]; ++#pragma unroll ++ for (int g = 0; g < group; ++g) { ++ const int j0 = wj * L::cols + (t0 + g) * 8; ++#pragma unroll ++ for (int h = 0; h < 2; ++h) { ++ ldmatrix_x4(b[g][h], yq + (j0 + threadIdx.x % 8) * y_stride + (threadIdx.x / 8) * 4 + h * 16); ++ } ++ } ++ ++ tile_C C[group][2]; ++#pragma unroll ++ for (int k = 0; k < 4; ++k) { ++#pragma unroll ++ for (int g = 0; g < group; ++g) { ++ tile_B B; ++ B.x[0] = b[g][k / 2][2 * (k % 2) + 0]; ++ B.x[1] = b[g][k / 2][2 * (k % 2) + 1]; ++#pragma unroll ++ for (int m = 0; m < 2; ++m) { ++ mma(C[g][m], A[m][k], B); ++ } ++ } ++ } ++ ++#pragma unroll ++ for (int g = 0; g < group; ++g) { ++ const int t = t0 + g; ++ const int j0 = wj * L::cols + t * 8; ++ float dB[2]; ++#pragma unroll ++ for (int l = 0; l < 2; ++l) { ++ dB[l] = yd[(j0 + tile_C::get_j(l)) * y_stride]; ++ } ++#pragma unroll ++ for (int m = 0; m < 2; ++m) { ++#pragma unroll ++ for (int l = 0; l < tile_C::ne; ++l) { ++ // Written as the fma the generic vec_dot contracts to, so that the two stay bit-identical. ++ float & s = sum[(t * 2 + m) * tile_C::ne + l]; ++ s = fmaf(C[g][m].x[l] * dA[m][l / 2], dB[l % 2], s); ++ } ++ } ++ } ++ ++ if constexpr (decode) { ++#pragma unroll ++ for (int n = 0; n < L::ntasks; ++n) { ++ if ((n * (L::ntiles / group)) / L::ntasks == t0 / group) { ++ decode_word(raw_next, xs_next, n); ++ } ++ } ++ } ++ } ++} ++ ++template ++static __device__ __forceinline__ void write_back(const float * __restrict__ sum, const int * __restrict__ ids_dst, ++ float * __restrict__ dst, const int stride, const int i_max, const int j_max) { ++ using L = layout; ++ typedef tile<16, 8, int> tile_C; ++ ++ const int i0 = (threadIdx.y / L::warps_j) * 32; ++ const int wj = threadIdx.y % L::warps_j; ++ ++#pragma unroll ++ for (int t = 0; t < L::ntiles; ++t) { ++#pragma unroll ++ for (int m = 0; m < 2; ++m) { ++#pragma unroll ++ for (int l = 0; l < tile_C::ne; ++l) { ++ const int j = wj * L::cols + t * 8 + tile_C::get_j(l); ++ if (j > j_max) { ++ continue; ++ } ++ const int i = i0 + m * 16 + tile_C::get_i(l); ++ if (fallback && i > i_max) { ++ continue; ++ } ++ dst[ids_dst[j] * stride + i] = sum[(t * 2 + m) * tile_C::ne + l]; ++ } ++ } ++ } ++} ++ ++} // namespace ggml_cuda_mmq_ptq1_0 ++ ++template ++static __device__ __forceinline__ void mul_mat_q_process_tile_ptq1_0( ++ const char * __restrict__ x, const int offset_x, const int * __restrict__ y, ++ const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup, ++ const int stride_row_x, const int ncols_y, const int stride_col_dst, ++ const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) { ++#if defined(CP_ASYNC_AVAILABLE) && defined(TURING_MMA_AVAILABLE) ++ using namespace ggml_cuda_mmq_ptq1_0; ++ using L = layout; ++ ++ extern __shared__ int data_mul_mat_q[]; ++ int * ys[2]; ++ int * xs[2]; ++ int * raw[2]; ++ ys[0] = data_mul_mat_q + J; ++ ys[1] = ys[0] + L::y_ints; ++ xs[0] = ys[1] + L::y_ints; ++ xs[1] = xs[0] + L::x_ints; ++ raw[0] = xs[1] + L::x_ints; ++ raw[1] = raw[0] + L::raw_ints; ++ ++ constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int); ++ const int tid = threadIdx.y * 32 + threadIdx.x; ++ ++ // Source offsets in ints of this thread's raw x words, relative to the block at kb = 0. ++ int x_src[L::ncopies]; ++#pragma unroll ++ for (int c = 0; c < L::ncopies; ++c) { ++ const int l = min(c * L::nthreads + tid, L::raw_ints - 1); ++ const int i = fallback ? min(l / x_words_per_row, tile_x_max_i) : l / x_words_per_row; ++ x_src[c] = (offset_x + i * stride_row_x) * x_words_per_row + l % x_words_per_row; ++ } ++ const int * x_words = (const int *) x; ++ ++ // One cp.async group per stage: the y block for the next stage and the raw x block for the one after. ++ const auto copy_y = [&](int * dst_y, const int kb) { ++ const char * src = reinterpret_cast(y + ncols_y * kb * sz); ++ char * dst_bytes = reinterpret_cast(dst_y); ++#pragma unroll ++ for (int byte0 = 16 * tid; byte0 < L::y_ints * int(sizeof(int)); byte0 += 16 * L::nthreads) { ++ cp_async_cg_16<256>(ggml_cuda_cvta_generic_to_shared(dst_bytes + byte0), src + byte0); ++ } ++ }; ++ const auto copy_x = [&](int * dst_raw, const int kb) { ++#pragma unroll ++ for (int c = 0; c < L::ncopies; ++c) { ++ if (L::raw_ints % L::nthreads == 0 || c * L::nthreads + tid < L::raw_ints) { ++ cp_async_ca_4(ggml_cuda_cvta_generic_to_shared(dst_raw + c * L::nthreads + tid), ++ x_words + x_src[c] + kb * x_words_per_row); ++ } ++ } ++ }; ++ ++ float sum[L::ntiles * 2 * tile<16, 8, int>::ne] = {0.0f}; ++ ++ const int nstages = kb0_stop - kb0_start; ++ ++ copy_y(ys[0], kb0_start); ++ copy_x(raw[0], kb0_start); ++ if (nstages > 1) { ++ copy_x(raw[1], kb0_start + 1); ++ } ++ cp_async_commit_group(); ++ cp_async_wait_all(); ++ __syncthreads(); ++#pragma unroll ++ for (int n = 0; n < L::ntasks; ++n) { ++ decode_word(raw[0], xs[0], n); ++ } ++ ++ // Stage s computes from buffer s%2 and decodes raw x for stage s+1 into the other one. It copies y for stage s+1 ++ // and raw x for stage s+2; the barrier at the top of the next stage makes both visible. ++ const auto run_stage = [&](const int s, const int p) { ++ __syncthreads(); ++ const bool next = s + 1 < nstages; ++ if (next) { ++ copy_y(ys[1 - p], kb0_start + s + 1); ++ } ++ if (s + 2 < nstages) { ++ copy_x(raw[p], kb0_start + s + 2); ++ } ++ cp_async_commit_group(); ++ if (next) { ++ stage(xs[p], ys[p], sum, raw[1 - p], xs[1 - p]); ++ } else { ++ stage(xs[p], ys[p], sum, raw[1 - p], xs[1 - p]); ++ } ++ cp_async_wait_all(); ++ }; ++ ++ int s = 0; ++ for (; s + 1 < nstages; s += 2) { ++ run_stage(s + 0, 0); ++ run_stage(s + 1, 1); ++ } ++ if (s < nstages) { ++ run_stage(s, 0); ++ } ++ ++ if (fixup) { ++ write_back(sum, ids_dst, tmp_fixup + blockIdx.x * (J * L::I), L::I, L::I, J); ++ } else { ++ write_back(sum, ids_dst, dst, stride_col_dst, tile_x_max_i, tile_y_max_j); ++ } ++#else ++ GGML_UNUSED_VARS(x, offset_x, y, ids_dst, dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, ++ tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); ++ NO_DEVICE_CODE; ++#endif // defined(CP_ASYNC_AVAILABLE) && defined(TURING_MMA_AVAILABLE) ++} ++ ++#endif // !defined(GGML_USE_HIP) +diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh +index 99b05d86..d6fe11f1 100644 +--- a/ggml/src/ggml-cuda/mmq.cuh ++++ b/ggml/src/ggml-cuda/mmq.cuh +@@ -919,7 +919,23 @@ static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_b + + // --------------------------------------------------------------------------------------------- + +-template ++#if !defined(GGML_USE_HIP) ++// PTQ1_0 on Ampere and newer runs the staged main loop in mmq-ptq1_0.cuh when every warp gets a multiple of 8 columns. ++static constexpr __device__ bool ggml_cuda_mmq_ptq1_0_staged(const int J, const bool fallback) { ++ const int nwarps = ggml_cuda_mmq_get_nthreads(GGML_TYPE_PTQ1_0, J, fallback) / 32; ++ const int warps_i = ggml_cuda_mmq_get_I(GGML_TYPE_PTQ1_0, J, fallback) / 32; ++ return warps_i > 0 && nwarps % warps_i == 0 && nwarps / warps_i > 0 && J % (8 * (nwarps / warps_i)) == 0; ++} ++ ++template ++static __device__ __forceinline__ void mul_mat_q_process_tile_ptq1_0( ++ const char * __restrict__ x, const int offset_x, const int * __restrict__ y, ++ const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup, ++ const int stride_row_x, const int ncols_y, const int stride_col_dst, ++ const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop); ++#endif // !defined(GGML_USE_HIP) ++ ++template + static __device__ __forceinline__ void mul_mat_q_process_tile( + const char * __restrict__ x, const int offset_x, const int * __restrict__ y, + const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup, +@@ -964,7 +980,11 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( + constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int); + + #if !defined(GGML_USE_HIP) && defined(CP_ASYNC_AVAILABLE) && defined(TURING_MMA_AVAILABLE) +- if constexpr (type == GGML_TYPE_PTQ1_0) { ++ if constexpr (staged && type == GGML_TYPE_PTQ1_0 && ggml_cuda_mmq_ptq1_0_staged(J, fallback)) { ++ mul_mat_q_process_tile_ptq1_0(x, offset_x, y, ids_dst, dst, tmp_fixup, stride_row_x, ncols_y, ++ stride_col_dst, tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); ++ return; ++ } else if constexpr (type == GGML_TYPE_PTQ1_0) { + // Software pipeline over the two y buffers: the y half for the next vec_dot is copied with cp.async while the + // current one is computed, and the next x tile is fetched into registers before the first vec_dot and only + // decoded into shared memory after the second, so no global load latency is exposed between k-iterations. +@@ -1118,7 +1138,8 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( + + // The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598 + +-template ++// staged = false keeps PTQ1_0 on the older main loop (GGML_CUDA_NO_PTQ1_0_STAGED); other types ignore it. ++template + __launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_get_occupancy(type, J, fallback)) + static __global__ void mul_mat_q( + const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst, +@@ -1221,7 +1242,7 @@ static __global__ void mul_mat_q( + const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; + + constexpr bool fixup = false; +- mul_mat_q_process_tile ++ mul_mat_q_process_tile + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, + stride_row_x, ncols_y, stride_col_dst, + tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z); +@@ -1315,7 +1336,7 @@ static __global__ void mul_mat_q( + const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; + + constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. +- mul_mat_q_process_tile ++ mul_mat_q_process_tile + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, + stride_row_x, ncols_y, stride_col_dst, + tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); +@@ -1399,7 +1420,7 @@ static __global__ void mul_mat_q( + const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; + + constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. +- mul_mat_q_process_tile ++ mul_mat_q_process_tile + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, + stride_row_x, ncols_y, stride_col_dst, + tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); +@@ -1561,11 +1582,13 @@ static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const i + (config.type == GGML_TYPE_Q1_0 || config.type == GGML_TYPE_Q2_0 || config.type == GGML_TYPE_PQ2_0)) || + (config.type == GGML_TYPE_PTQ1_0 && cp_async_available(cc) && turing_mma_available(cc)); + const int y_buffers = async_buffer_y ? 2 : 1; +- return nbs_ids + nbs_x + y_buffers*nbs_y_padded; ++ // The staged PTQ1_0 loop (mmq-ptq1_0.cuh) also stages two raw x blocks per row with cp.async. ++ const size_t nbs_x_raw = config.type == GGML_TYPE_PTQ1_0 && async_buffer_y ? 2*config.I*sizeof(block_ptq1_0) : 0; ++ return nbs_ids + nbs_x + y_buffers*nbs_y_padded + nbs_x_raw; + } + +-template +-static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { ++template ++static void launch_mul_mat_q_impl(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { + const int id = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[id].cc; + const int nsm = ggml_cuda_info().devices[id].nsm; +@@ -1578,8 +1601,8 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a + + const dim3 block_dims(warp_size, nwarps, 1); + +- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); +- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); ++ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); ++ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + const int nty = (args.nrows_x + config.I - 1) / config.I; + const int ntx = (args.ncols_max + config.J - 1) / config.J; +@@ -1599,7 +1622,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a + const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio); + + if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) { +- mul_mat_q<<>> ++ mul_mat_q<<>> + (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale, + blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, +@@ -1628,7 +1651,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a + const dim3 block_nums_fixup(block_nums_stream_k.x, config.I/warp_size, 1); + const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z); + +- mul_mat_q<<>> ++ mul_mat_q<<>> + (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale, + blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, +@@ -1646,6 +1669,18 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a + ntx_fd); + } + ++template ++static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { ++ if constexpr (type == GGML_TYPE_PTQ1_0) { ++ static const bool staged = getenv("GGML_CUDA_NO_PTQ1_0_STAGED") == nullptr; ++ if (!staged) { ++ launch_mul_mat_q_impl(ctx, args, stream); ++ return; ++ } ++ } ++ launch_mul_mat_q_impl(ctx, args, stream); ++} ++ + template + void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { + const int id = ggml_cuda_get_device(); +diff --git a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +index 073e4154..3958a045 100755 +--- a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py ++++ b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +@@ -60,6 +60,7 @@ DECL_MMQ_CASE({type}); + SOURCE_MMQ_CUDA_ONLY = """// This file has been autogenerated by generate_cu_files.py, do not edit manually. + + #include "../mmq.cuh" ++#include "../mmq-ptq1_0.cuh" + + #if !defined(GGML_USE_HIP) + DECL_MMQ_CASE({type}); +diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-ptq1_0.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-ptq1_0.cu +index 692e9529..2011f3ac 100644 +--- a/ggml/src/ggml-cuda/template-instances/mmq-instance-ptq1_0.cu ++++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-ptq1_0.cu +@@ -1,6 +1,7 @@ + // This file has been autogenerated by generate_cu_files.py, do not edit manually. + + #include "../mmq.cuh" ++#include "../mmq-ptq1_0.cuh" + + #if !defined(GGML_USE_HIP) + DECL_MMQ_CASE(GGML_TYPE_PTQ1_0); diff --git a/pkgs-set/pkgs/llama-cpp-bonsai.nix b/pkgs-set/pkgs/llama-cpp-bonsai.nix index cfca038..beb19ab 100644 --- a/pkgs-set/pkgs/llama-cpp-bonsai.nix +++ b/pkgs-set/pkgs/llama-cpp-bonsai.nix @@ -75,6 +75,14 @@ cudaPackages.backendStdenv.mkDerivation (finalAttrs: { # MMQ reads, so the matmul skips its own quantize pass (and the f32 copy when nothing else reads it). # Bit-identical; +3% prefill. ./llama-cpp-bonsai-fwht-q8.patch + # PTQ1_0 MMQ main loop staged one 128-value block at a time: one barrier per block, x and y copied with + # cp.async, trits decoded between the MMAs, B fragments via ldmatrix. Bit-identical; +10% prefill. + # GGML_CUDA_NO_PTQ1_0_STAGED=1 restores the old loop. + ./llama-cpp-bonsai-mmq-staged.patch + # Prompt batches over a q8_0 cache tile the 6 query heads of a KV head in pairs instead of padding them + # to 8, stage raw q8_0 rows with cp.async and dequantize exactly in half2. Bit-identical; prefill +8% at + # 16K, +14% at 32K, +21% at 64K. Decode and verify keep the 8-wide tile. + ./llama-cpp-bonsai-fattn-q8-prefill.patch # GET /cache (slots and both prompt-cache tiers) and POST /cache/clear?scope=slots|ram|disk|all, run on the # server loop like the slot actions; a slot that is generating is left alone. ./llama-cpp-bonsai-cache-api.patch -- 2.54.0