Skip to content

fix(cuda): turbo KV fused-MMA — use_sparse off-by-one and missing SMEM swizzle #12

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wilky2005 wants to merge 2 commits into
base: perf
Choose a base branch
from
Open

fix(cuda): turbo KV fused-MMA — use_sparse off-by-one and missing SMEM swizzle #12

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8bcacff
fix(cuda): pass use_sparse to launch_fattn in turbo MMA case
wilky2005 Sep 17, 2026
1fdcce6
fix(cuda): apply SMEM swizzle in turbo KV tile loaders
wilky2005 Sep 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 27 additions & 12 deletions ggml/src/ggml-cuda/fattn-mma-f16.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -512,7 +512,7 @@ static __constant__ float TURBO_CENTROIDS_4BIT_FATTN[16] = {
// low nibble = elem 2c, high nibble = elem 2c+1. Hence one byte qs[col_offset+c] yields
// the half2 for tile column c. sizeof(block_turbo4_0)-driven pointer math; never assume
// 66/68 or a qs offset constant.
template<int stride_tile, int nbatch_fa, int nthreads, bool oob_check>
template<int stride_tile, bool swz, int nbatch_fa, int nthreads, bool oob_check>
static __device__ __forceinline__ void flash_attn_ext_turbo4_load_tile(
const char * const __restrict__ KV_raw, half2 * const __restrict__ tile_KV,
const int D2, const int stride_bytes, const int col_offset, const int i_sup) {
Expand All @@ -538,7 +538,12 @@ static __device__ __forceinline__ void flash_attn_ext_turbo4_load_tile(
const uint8_t byte = blk->qs[in_blk];
const half lo = __float2half(TURBO_CENTROIDS_4BIT_FATTN[byte & 0xF] * norm);
const half hi = __float2half(TURBO_CENTROIDS_4BIT_FATTN[byte >> 4] * norm);
tile_KV[row*stride_tile + c] = __halves2half2(lo, hi);
if constexpr (swz) {
const int byte_off = ggml_cuda_fattn_smem_swizzle::bytes_rc<stride_tile>(row, c);
*reinterpret_cast<half2 *>(reinterpret_cast<char *>(tile_KV) + byte_off) = __halves2half2(lo, hi);
} else {
tile_KV[row*stride_tile + c] = __halves2half2(lo, hi);
}
}
}
}
Expand All @@ -550,7 +555,7 @@ static __constant__ float TURBO_CENTROIDS_3BIT_FATTN[8] = {
-0.190207f, -0.118786f, -0.066822f, -0.021663f,
0.021663f, 0.066822f, 0.118786f, 0.190207f
};
template<int stride_tile, int nbatch_fa, int nthreads, bool oob_check>
template<int stride_tile, bool swz, int nbatch_fa, int nthreads, bool oob_check>
static __device__ __forceinline__ void flash_attn_ext_turbo3_load_tile(
const char * const __restrict__ KV_raw, half2 * const __restrict__ tile_KV,
const int D2, const int stride_bytes, const int col_offset, const int i_sup) {
Expand All @@ -577,7 +582,12 @@ static __device__ __forceinline__ void flash_attn_ext_turbo3_load_tile(
const uint8_t idx1 = ((qs_byte >> (shift+2)) & 0x3) | (((sgn_byte >> (j0 % 8 + 1)) & 0x1) << 2);
const half lo = __float2half(TURBO_CENTROIDS_3BIT_FATTN[idx0] * norm);
const half hi = __float2half(TURBO_CENTROIDS_3BIT_FATTN[idx1] * norm);
tile_KV[row*stride_tile + c] = __halves2half2(lo, hi);
if constexpr (swz) {
const int byte_off = ggml_cuda_fattn_smem_swizzle::bytes_rc<stride_tile>(row, c);
*reinterpret_cast<half2 *>(reinterpret_cast<char *>(tile_KV) + byte_off) = __halves2half2(lo, hi);
} else {
tile_KV[row*stride_tile + c] = __halves2half2(lo, hi);
}
}
}
}
Expand All @@ -586,7 +596,7 @@ static __device__ __forceinline__ void flash_attn_ext_turbo3_load_tile(
static __constant__ float TURBO_CENTROIDS_2BIT_FATTN[4] = {
-0.133462f, -0.039994f, 0.039994f, 0.133462f
};
template<int stride_tile, int nbatch_fa, int nthreads, bool oob_check>
template<int stride_tile, bool swz, int nbatch_fa, int nthreads, bool oob_check>
static __device__ __forceinline__ void flash_attn_ext_turbo2_load_tile(
const char * const __restrict__ KV_raw, half2 * const __restrict__ tile_KV,
const int D2, const int stride_bytes, const int col_offset, const int i_sup) {
Expand All @@ -612,7 +622,12 @@ static __device__ __forceinline__ void flash_attn_ext_turbo2_load_tile(
const uint8_t idx1 = (qs_byte >> (shift+2)) & 0x3;
const half lo = __float2half(TURBO_CENTROIDS_2BIT_FATTN[idx0] * norm);
const half hi = __float2half(TURBO_CENTROIDS_2BIT_FATTN[idx1] * norm);
tile_KV[row*stride_tile + c] = __halves2half2(lo, hi);
if constexpr (swz) {
const int byte_off = ggml_cuda_fattn_smem_swizzle::bytes_rc<stride_tile>(row, c);
*reinterpret_cast<half2 *>(reinterpret_cast<char *>(tile_KV) + byte_off) = __halves2half2(lo, hi);
} else {
tile_KV[row*stride_tile + c] = __halves2half2(lo, hi);
}
}
}
}
Expand Down Expand Up @@ -798,13 +813,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
constexpr int nthreads_turbo = nwarps * ggml_cuda_get_physical_warp_size();
const char * K_raw = (const char *) K_h2 + int64_t(k_VKQ_0) * stride_K;
if constexpr (type_K == GGML_TYPE_TURBO4_0) {
flash_attn_ext_turbo4_load_tile<stride_tile_K, nbatch_fa, nthreads_turbo, oob_check>
flash_attn_ext_turbo4_load_tile<stride_tile_K, swz_K, nbatch_fa, nthreads_turbo, oob_check>
(K_raw, tile_K, k0_diff, stride_K, k0_start, k_VKQ_sup);
} else if constexpr (type_K == GGML_TYPE_TURBO3_0) {
flash_attn_ext_turbo3_load_tile<stride_tile_K, nbatch_fa, nthreads_turbo, oob_check>
flash_attn_ext_turbo3_load_tile<stride_tile_K, swz_K, nbatch_fa, nthreads_turbo, oob_check>
(K_raw, tile_K, k0_diff, stride_K, k0_start, k_VKQ_sup);
} else {
flash_attn_ext_turbo2_load_tile<stride_tile_K, nbatch_fa, nthreads_turbo, oob_check>
flash_attn_ext_turbo2_load_tile<stride_tile_K, swz_K, nbatch_fa, nthreads_turbo, oob_check>
(K_raw, tile_K, k0_diff, stride_K, k0_start, k_VKQ_sup);
}
__syncthreads();
Expand Down Expand Up @@ -1171,13 +1186,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
constexpr int nthreads_turbo = nwarps * ggml_cuda_get_physical_warp_size();
const char * V_raw = (const char *) V_h2 + int64_t(k_VKQ_0) * stride_V;
if constexpr (type_V == GGML_TYPE_TURBO4_0) {
flash_attn_ext_turbo4_load_tile<stride_tile_V, nbatch_fa, nthreads_turbo, oob_check>
flash_attn_ext_turbo4_load_tile<stride_tile_V, swz_V, nbatch_fa, nthreads_turbo, oob_check>
(V_raw, tile_V, i0_diff/2, stride_V, i0_start/2, k_VKQ_sup);
} else if constexpr (type_V == GGML_TYPE_TURBO3_0) {
flash_attn_ext_turbo3_load_tile<stride_tile_V, nbatch_fa, nthreads_turbo, oob_check>
flash_attn_ext_turbo3_load_tile<stride_tile_V, swz_V, nbatch_fa, nthreads_turbo, oob_check>
(V_raw, tile_V, i0_diff/2, stride_V, i0_start/2, k_VKQ_sup);
} else {
flash_attn_ext_turbo2_load_tile<stride_tile_V, nbatch_fa, nthreads_turbo, oob_check>
flash_attn_ext_turbo2_load_tile<stride_tile_V, swz_V, nbatch_fa, nthreads_turbo, oob_check>
(V_raw, tile_V, i0_diff/2, stride_V, i0_start/2, k_VKQ_sup);
}
__syncthreads();
Expand Down
2 changes: 1 addition & 1 deletion ggml/src/ggml-cuda/fattn-mma-turbo.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ void ggml_cuda_flash_attn_ext_mma_turbo_case(ggml_backend_cuda_context & ctx, gg
// the kernel receives raw quantized KV + the true byte pitch. stream_k = true.
launch_fattn<DV, ncols1, ncols2>
(ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa,
/*need_f16_K=*/false, /*need_f16_V=*/false, /*stream_k=*/true, warp_size_host);
/*need_f16_K=*/false, /*need_f16_V=*/false, /*stream_k=*/true, /*use_sparse=*/false, warp_size_host);
}


Expand Down