mt_pagedattn: wire TURBO3_0 K/V cache type through paged framework

Adds paged_cache_ops<TURBO3_0> + cooperative dequant in tile + decode kernels
+ scatter kernel + dispatch case + supports_op whitelist + bench axis. Mirrors
the existing TURBO4_0 plumbing; TURBO3_0 differs only in block payload
(qs[32] 2-bit low + signs[16] 1-bit high vs turbo4's nibble-packed qs[64])
and 3-bit centroid table. RHT is skipped on the paged path with the same
JL-bypass justification as turbo4 (see mt_scatter_kv_turbo4_0_kernel step 4
comment).

Files:
- ggml-cuda.cu: TURBO3_0 in PAGED_ATTN_MT supports_op whitelist (without
  this, the scheduler routes the op to CPU which doesn't implement it
  and aborts in ggml_get_n_tasks).
- mt_pagedattn_ops.cuh: paged_cache_ops<TURBO3_0> specialization.
- mt_pagedattn.cu: mt_scatter_kv_turbo3_0_kernel (128-thread cooperative
  quant, warp-cooperative qs/signs packing via __shfl_xor_sync), launch
  branch, HEAD_SIZE validation, dispatch case.
- mt_pagedattn_tile.cu: coop_stage_turbo3_tile + K/V dispatch + tile/tile_mw
  template instantiations at HS=128 and HS=256.
- mt_pagedattn_decode.cu: decode_coop_stage_turbo3 + K/V dispatch + decode
  template instantiations at HS=128 and HS=256.

Bench (Qwen3.6-35B-A3B-UD-Q4_K_XL @ R9700/gfx1201, n_predict=128):
  paged_turbo4  8K: 1948 prefill / 32.0 decode | 64K: 952 / 30.7
  paged_turbo3  8K: 1924 prefill / 32.5 decode | 64K: 936 / 30.9
TURBO3 == TURBO4 in throughput within run-to-run noise (~1.5%); 24% less
KV memory (3.125 vs 4.125 bpv) for the same effective speed. Quality
validation (PPL+NIAH sweep) pending — added as task #11.

Smoke-tested with --kv-tiered 100,0,0 --kv-tier-paged-blocks
--cache-type-k turbo3 --cache-type-v turbo3: server boots, completes a
53-token prompt at 57.7 decode t/s with coherent output. Locked baseline
JSON committed alongside as the new perf reference at HEAD.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: kmbandy

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -5663,13 +5663,14 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
             // K/V cache may be quantized; mt::ggml_cuda_op_paged_attn_mt
             // dispatch checks for a specialized paged_cache_ops<TYPE> and
             // aborts with a clear message if unsupported. Currently
-            // F16 + Q8_0 + TURBO4_0 are wired (MAD-116).
+            // F16 + Q8_0 + TURBO4_0 + TURBO3_0 are wired (MAD-116).
             return op->type == GGML_TYPE_F16
                 && op->src[0]->type == GGML_TYPE_F16   // q
                 && op->src[1]->type == op->src[2]->type  // k/v cache same type
                 && (op->src[1]->type == GGML_TYPE_F16
                     || op->src[1]->type == GGML_TYPE_Q8_0
-                    || op->src[1]->type == GGML_TYPE_TURBO4_0)
+                    || op->src[1]->type == GGML_TYPE_TURBO4_0
+                    || op->src[1]->type == GGML_TYPE_TURBO3_0)
                 && op->src[6]                          // k_cur (fused scatter)
                 && op->src[6]->type == GGML_TYPE_F16
                 && op->src[7]                          // v_cur (fused scatter)

ggml/src/ggml-cuda/mt_pagedattn.cu

@@ -454,6 +454,153 @@ __global__ void mt_scatter_kv_turbo4_0_kernel(
     }
 }
 
+// Turbo3_0 scatter: per-128-element-block cooperative quantize.
+//
+// Same threading and pipeline as turbo4 above, but packs 3-bit indices into
+// qs[32] (low 2 bits, 4 per byte) + signs[16] (high 1 bit, 8 per byte).
+// 50 bytes payload per 128 elements → 3.125 bpv (vs turbo4's 4.125 bpv).
+//
+// As with turbo4, RHT is intentionally skipped on the paged path — see the
+// step-4 comment in mt_scatter_kv_turbo4_0_kernel for the JL/dot-product
+// reasoning. Turbo3 centroids (TURBO_CENTROIDS_3BIT) are Lloyd-Max for
+// N(0, 1/d) and remain valid without rotation.
+template <int HEAD_SIZE, int BLOCK_SIZE>
+__launch_bounds__(QK_TURBO3)
+__global__ void mt_scatter_kv_turbo3_0_kernel(
+    void           * __restrict__ k_cache,
+    void           * __restrict__ v_cache,
+    const __half   * __restrict__ k_cur,
+    const __half   * __restrict__ v_cur,
+    const int32_t  * __restrict__ slot_mapping,
+    int             n_kv_heads) {
+
+    constexpr int Q_BLOCK             = QK_TURBO3;          // 128
+    constexpr int N_QBLOCKS_PER_TOKEN = HEAD_SIZE / Q_BLOCK;
+    constexpr int N_WARPS             = Q_BLOCK / WARP_SIZE;  // 4
+    static_assert(HEAD_SIZE % Q_BLOCK == 0, "HEAD_SIZE must be divisible by QK_TURBO3");
+    static_assert(Q_BLOCK == 128, "this kernel assumes QK_TURBO3 == 128");
+
+    const int j                = threadIdx.x;  // 0..127, element index within turbo3 block
+    const int global_token_idx = blockIdx.x;
+    const int y_idx            = blockIdx.y;
+    const int kv_select        = blockIdx.z;   // 0 = K, 1 = V
+    const int kv_head_idx      = y_idx / N_QBLOCKS_PER_TOKEN;
+    const int qb_idx           = y_idx % N_QBLOCKS_PER_TOKEN;
+
+    const int slot = slot_mapping[global_token_idx];
+    if (slot < 0) return;  // padding token
+
+    const int paged_block   = slot / BLOCK_SIZE;
+    const int slot_in_block = slot % BLOCK_SIZE;
+
+    const int    d = qb_idx * Q_BLOCK + j;
+    const __half * src = (kv_select == 0) ? k_cur : v_cur;
+    const size_t src_off = (size_t) global_token_idx * n_kv_heads * HEAD_SIZE
+                         + (size_t) kv_head_idx * HEAD_SIZE
+                         + (size_t) d;
+
+    // Output block pointer (matches paged_cache_ops<TURBO3_0>::element_block_index)
+    void * dst_buf = (kv_select == 0) ? k_cache : v_cache;
+    const int64_t block_ib = ((int64_t) paged_block * n_kv_heads + kv_head_idx) * BLOCK_SIZE * N_QBLOCKS_PER_TOKEN
+                           + (int64_t) slot_in_block * N_QBLOCKS_PER_TOKEN
+                           + (int64_t) qb_idx;
+    block_turbo3_0 * blk = (block_turbo3_0 *) dst_buf + block_ib;
+
+    // ---- Step 1: Load into smem (explicit __half2float for HIP safety) ----
+    __shared__ float x[Q_BLOCK];
+    x[j] = __half2float(src[src_off]);
+    __syncthreads();
+
+    // ---- Step 2: Parallel L2 norm ----
+    __shared__ float warp_accum[N_WARPS];
+    {
+        float v_sq = x[j] * x[j];
+        for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
+            v_sq += __shfl_xor_sync(0xffffffffu, v_sq, offset);
+        }
+        if (j % WARP_SIZE == 0) warp_accum[j / WARP_SIZE] = v_sq;
+    }
+    __syncthreads();
+
+    __shared__ float s_norm_sq;
+    if (j == 0) {
+        float total = 0.0f;
+        for (int w = 0; w < N_WARPS; ++w) total += warp_accum[w];
+        s_norm_sq = total;
+    }
+    __syncthreads();
+    const float grp_norm = sqrtf(s_norm_sq);
+    const float inv_norm = (grp_norm > 1e-10f) ? (1.0f / grp_norm) : 0.0f;
+
+    // ---- Step 3: Normalize ----
+    x[j] *= inv_norm;
+    __syncthreads();
+
+    // ---- Step 4: (intentionally NO Randomized Hadamard Transform) ----
+    // See mt_scatter_kv_turbo4_0_kernel step 4 comment for justification.
+
+    // ---- Step 5: Quantize element j to 3-bit centroid index ----
+    const float   rv  = x[j];
+    const uint8_t idx = turbo_nearest_centroid_3bit(rv);  // 0..7
+
+    // ---- Step 6: Pack qs[] (4 elements per byte, low 2 bits each) ----
+    //
+    // 32 threads in a warp produce 8 bytes of qs. Each byte aggregates 4
+    // adjacent threads' 2-bit values via warp shuffles. Each warp writes
+    // to its own non-overlapping chunk of qs[] — no cross-warp conflicts.
+    {
+        const int     lane    = j % WARP_SIZE;
+        const int     warp_id = j / WARP_SIZE;
+        const uint8_t my2     = idx & 0x3;
+        uint8_t byte_val      = my2 << ((lane & 3) * 2);
+        byte_val |= __shfl_xor_sync(0xffffffffu, byte_val, 1);
+        byte_val |= __shfl_xor_sync(0xffffffffu, byte_val, 2);
+        if ((lane & 3) == 0) {
+            blk->qs[warp_id * (WARP_SIZE / 4) + lane / 4] = byte_val;
+        }
+    }
+
+    // ---- Step 7: Pack signs[] (8 elements per byte, high 1 bit each) ----
+    {
+        const int     lane    = j % WARP_SIZE;
+        const int     warp_id = j / WARP_SIZE;
+        const uint8_t my1     = (idx >> 2) & 0x1;
+        uint8_t bits          = my1 << (lane & 7);
+        bits |= __shfl_xor_sync(0xffffffffu, bits, 1);
+        bits |= __shfl_xor_sync(0xffffffffu, bits, 2);
+        bits |= __shfl_xor_sync(0xffffffffu, bits, 4);
+        if ((lane & 7) == 0) {
+            blk->signs[warp_id * (WARP_SIZE / 8) + lane / 8] = bits;
+        }
+    }
+
+    // ---- Step 8: Reconstruction norm (parallel) ----
+    {
+        const float c = TURBO_CENTROIDS_3BIT[idx];
+        float rc = c * c;
+        for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
+            rc += __shfl_xor_sync(0xffffffffu, rc, offset);
+        }
+        if (j % WARP_SIZE == 0) warp_accum[j / WARP_SIZE] = rc;
+    }
+    __syncthreads();
+
+    __shared__ float s_recon_sq;
+    if (j == 0) {
+        float total = 0.0f;
+        for (int w = 0; w < N_WARPS; ++w) total += warp_accum[w];
+        s_recon_sq = total;
+    }
+    __syncthreads();
+    const float recon_norm     = sqrtf(s_recon_sq);
+    const float corrected_norm = (recon_norm > 1e-10f) ? (grp_norm / recon_norm) : grp_norm;
+
+    // ---- Step 9: Lane 0 writes per-block scalar ----
+    if (j == 0) {
+        blk->norm = __float2half(corrected_norm);
+    }
+}
+
 template <ggml_type CACHE_TYPE, typename scalar_t, int HEAD_SIZE, int BLOCK_SIZE>
 static void launch_scatter_kv(
     void           * k_cache,
@@ -494,6 +641,18 @@ static void launch_scatter_kv(
         mt_scatter_kv_turbo4_0_kernel<HEAD_SIZE, BLOCK_SIZE>
             <<<grid, block, 0, stream>>>(
                 k_cache, v_cache, k_cur, v_cur, slot_mapping, n_kv_heads);
+    } else if constexpr (CACHE_TYPE == GGML_TYPE_TURBO3_0) {
+        // Same grid topology as TURBO4_0 — one CUDA block per
+        // (token, (kv_head, qb_idx), K-or-V); 128 threads per block.
+        constexpr int Q_BLOCK             = QK_TURBO3;
+        constexpr int N_QBLOCKS_PER_TOKEN = HEAD_SIZE / Q_BLOCK;
+        dim3 grid(num_tokens_total, n_kv_heads * N_QBLOCKS_PER_TOKEN, 2);
+        dim3 block(Q_BLOCK);
+        GGML_UNUSED(q_lens);
+        GGML_UNUSED(num_seqs);
+        mt_scatter_kv_turbo3_0_kernel<HEAD_SIZE, BLOCK_SIZE>
+            <<<grid, block, 0, stream>>>(
+                k_cache, v_cache, k_cur, v_cur, slot_mapping, n_kv_heads);
     }
     // Fall-through for unsupported types is intentional — dispatch in
     // ggml_cuda_op_paged_attn_mt switches only over types that have a
@@ -926,18 +1085,21 @@ void ggml_cuda_op_paged_attn_mt(ggml_backend_cuda_context & ctx, ggml_tensor * d
         auto run_typed = [&](auto cache_type_const) {
             constexpr ggml_type CT = decltype(cache_type_const)::value;
 
-            // Cache-type / head-size compatibility: TURBO4_0 needs
-            // HEAD_SIZE divisible by QK_TURBO4 (128). For smaller heads
-            // (e.g. HS=64), abort at runtime — we don't even instantiate
-            // the kernel templates for invalid combinations.
+            // Cache-type / head-size compatibility: TURBO4_0 / TURBO3_0 need
+            // HEAD_SIZE divisible by their respective QK (both 128). For
+            // smaller heads (e.g. HS=64), abort at runtime — we don't even
+            // instantiate the kernel templates for invalid combinations.
             if constexpr (CT == GGML_TYPE_TURBO4_0 && HS % QK_TURBO4 != 0) {
                 GGML_ABORT("mt_paged_attn: HEAD_SIZE=%d too small for TURBO4_0 (requires multiple of QK_TURBO4=%d)",
                            HS, (int) QK_TURBO4);
+            } else if constexpr (CT == GGML_TYPE_TURBO3_0 && HS % QK_TURBO3 != 0) {
+                GGML_ABORT("mt_paged_attn: HEAD_SIZE=%d too small for TURBO3_0 (requires multiple of QK_TURBO3=%d)",
+                           HS, (int) QK_TURBO3);
             } else {
 
             // MAD-180: WMMA tile FA gate.
             if constexpr (((HS == 128) || (HS == 256)) && (BS == 16)
-                          && (CT == GGML_TYPE_F16 || CT == GGML_TYPE_TURBO4_0)) {
+                          && (CT == GGML_TYPE_F16 || CT == GGML_TYPE_TURBO4_0 || CT == GGML_TYPE_TURBO3_0)) {
                 const int  cc                = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
                 const int  total_q_tokens    = (int) k_cur->ne[2];
                 const int  avg_q_len         = num_seqs > 0 ? (total_q_tokens / num_seqs) : 0;
@@ -1131,6 +1293,9 @@ void ggml_cuda_op_paged_attn_mt(ggml_backend_cuda_context & ctx, ggml_tensor * d
             case GGML_TYPE_TURBO4_0:
                 run_typed(std::integral_constant<ggml_type, GGML_TYPE_TURBO4_0>{});
                 break;
+            case GGML_TYPE_TURBO3_0:
+                run_typed(std::integral_constant<ggml_type, GGML_TYPE_TURBO3_0>{});
+                break;
             default:
                 GGML_ABORT("mt_paged_attn: unsupported cache type %s — add a paged_cache_ops specialization",
                            ggml_type_name(k_cache->type));

ggml/src/ggml-cuda/mt_pagedattn_decode.cu

@@ -257,6 +257,75 @@ static __device__ __forceinline__ void decode_coop_stage_turbo4(
     }
 }
 
+// TURBO3_0 cooperative dequant for the decode tile. Same threading shape as
+// decode_coop_stage_turbo4 (32 lanes × 4 elements per qblock), but unpacks
+// the 3-bit index as (qs low-2 | signs hi-1).
+template <int HEAD_SIZE, int BLOCK_SIZE>
+static __device__ __forceinline__ void decode_coop_stage_turbo3(
+        __half        * __restrict__ smem_dst,
+        const void    * __restrict__ cache,
+        const int     * __restrict__ seq_block_table,
+        int            tile_start,
+        int            valid_ctx,
+        int            kv_head_idx,
+        int            n_kv_heads,
+        int            warp_id,
+        int            lane_id) {
+    constexpr int Q_BLOCK            = QK_TURBO3;
+    constexpr int QBLOCKS_PER_TOKEN  = HEAD_SIZE / Q_BLOCK;
+    constexpr int N_QBLOCKS_PER_TILE = DECODE_K_TILE_N * QBLOCKS_PER_TOKEN;
+    static_assert(HEAD_SIZE % Q_BLOCK == 0, "HEAD_SIZE must be multiple of QK_TURBO3=128");
+    static_assert(Q_BLOCK == 128, "cooperative dequant expects QK_TURBO3=128 (32 lanes × 4 elements)");
+
+    const block_turbo3_0 * blocks = (const block_turbo3_0 *) cache;
+
+    #pragma unroll
+    for (int qb = warp_id; qb < N_QBLOCKS_PER_TILE; qb += DECODE_NUM_WARPS) {
+        const int row         = qb / QBLOCKS_PER_TOKEN;
+        const int qb_in_token = qb % QBLOCKS_PER_TOKEN;
+        const int token       = tile_start + row;
+
+        const block_turbo3_0 * blk = nullptr;
+        float norm_f = 0.0f;
+
+        if (token < valid_ctx) {
+            const int logical_block = token / BLOCK_SIZE;
+            const int tok_in_block  = token % BLOCK_SIZE;
+            const int physical      = seq_block_table[logical_block];
+            if (physical != kInvalidBlockTableEntry) {
+                const int64_t ib = ((int64_t) physical * n_kv_heads + kv_head_idx) * BLOCK_SIZE * QBLOCKS_PER_TOKEN
+                                 + (int64_t) tok_in_block * QBLOCKS_PER_TOKEN
+                                 + (int64_t) qb_in_token;
+                blk = &blocks[ib];
+                if (lane_id == 0) {
+                    norm_f = __half2float(blk->norm);
+                }
+            }
+        }
+        norm_f = __shfl_sync(0xFFFFFFFF, norm_f, 0);
+
+        uint8_t qs_byte    = 0;
+        uint8_t signs_byte = 0;
+        if (blk != nullptr) {
+            qs_byte    = blk->qs[lane_id];
+            signs_byte = blk->signs[lane_id / 2];
+        }
+        const int signs_shift_base = (lane_id % 2) * 4;
+
+        const int smem_row_base = row * HEAD_SIZE;
+        const int smem_col_base = qb_in_token * Q_BLOCK + lane_id * 4;
+
+        #pragma unroll
+        for (int l = 0; l < 4; ++l) {
+            const uint8_t low2 = (qs_byte    >> (l * 2)) & 0x3;
+            const uint8_t hi1  = (signs_byte >> (signs_shift_base + l)) & 0x1;
+            const uint8_t idx  = low2 | (hi1 << 2);
+            const float   val  = TURBO_CENTROIDS_3BIT[idx] * norm_f;
+            smem_dst[smem_row_base + smem_col_base + l] = __float2half(val);
+        }
+    }
+}
+
 template <int HEAD_SIZE, int BLOCK_SIZE, ggml_type CACHE_TYPE>
 static __device__ __forceinline__ void decode_stage_k(
         __half        * __restrict__ smem_dst,
@@ -273,6 +342,10 @@ static __device__ __forceinline__ void decode_stage_k(
         decode_coop_stage_turbo4<HEAD_SIZE, BLOCK_SIZE>(
             smem_dst, cache, seq_block_table, tile_start, valid_ctx,
             kv_head_idx, n_kv_heads, warp_id, lane_id);
+    } else if constexpr (CACHE_TYPE == GGML_TYPE_TURBO3_0) {
+        decode_coop_stage_turbo3<HEAD_SIZE, BLOCK_SIZE>(
+            smem_dst, cache, seq_block_table, tile_start, valid_ctx,
+            kv_head_idx, n_kv_heads, warp_id, lane_id);
     } else {
         decode_stage_kv_f16<HEAD_SIZE, BLOCK_SIZE>(
             smem_dst, (const __half *) cache, seq_block_table, tile_start, valid_ctx,
@@ -296,6 +369,10 @@ static __device__ __forceinline__ void decode_stage_v(
         decode_coop_stage_turbo4<HEAD_SIZE, BLOCK_SIZE>(
             smem_dst, cache, seq_block_table, tile_start, valid_ctx,
             kv_head_idx, n_kv_heads, warp_id, lane_id);
+    } else if constexpr (CACHE_TYPE == GGML_TYPE_TURBO3_0) {
+        decode_coop_stage_turbo3<HEAD_SIZE, BLOCK_SIZE>(
+            smem_dst, cache, seq_block_table, tile_start, valid_ctx,
+            kv_head_idx, n_kv_heads, warp_id, lane_id);
     } else {
         decode_stage_kv_f16<HEAD_SIZE, BLOCK_SIZE>(
             smem_dst, (const __half *) cache, seq_block_table, tile_start, valid_ctx,
@@ -1065,6 +1142,11 @@ template void launch_paged_attn_decode<128, 16, GGML_TYPE_TURBO4_0>(
     const int32_t *, const int32_t *, const int32_t *,
     float *, int, int, int, int, int, int, float, cudaStream_t);
 
+template void launch_paged_attn_decode<128, 16, GGML_TYPE_TURBO3_0>(
+    __half *, const __half *, const void *, const void *,
+    const int32_t *, const int32_t *, const int32_t *,
+    float *, int, int, int, int, int, int, float, cudaStream_t);
+
 template void launch_paged_attn_decode<256, 16, GGML_TYPE_F16>(
     __half *, const __half *, const void *, const void *,
     const int32_t *, const int32_t *, const int32_t *,
@@ -1075,4 +1157,9 @@ template void launch_paged_attn_decode<256, 16, GGML_TYPE_TURBO4_0>(
     const int32_t *, const int32_t *, const int32_t *,
     float *, int, int, int, int, int, int, float, cudaStream_t);
 
+template void launch_paged_attn_decode<256, 16, GGML_TYPE_TURBO3_0>(
+    __half *, const __half *, const void *, const void *,
+    const int32_t *, const int32_t *, const int32_t *,
+    float *, int, int, int, int, int, int, float, cudaStream_t);
+
 } // namespace mt

ggml/src/ggml-cuda/mt_pagedattn_ops.cuh

@@ -152,4 +152,38 @@ struct paged_cache_ops<GGML_TYPE_TURBO4_0, HEAD_SIZE, BLOCK_SIZE> {
     }
 };
 
+// Turbo3_0 specialization: 3-bit PolarQuant with WHT rotation, 128-element blocks.
+// Same N_QBLOCKS_PER_TOKEN layout as TURBO4_0; differs in block payload size
+// (14 vs 68 bytes) and dequant centroid table (3-bit Lloyd-Max).
+// kv_store omitted — handled by a dedicated cooperative scatter kernel.
+template <int HEAD_SIZE, int BLOCK_SIZE>
+struct paged_cache_ops<GGML_TYPE_TURBO3_0, HEAD_SIZE, BLOCK_SIZE> {
+    static constexpr int Q_BLOCK = QK_TURBO3;  // 128
+    static constexpr int N_QBLOCKS_PER_TOKEN = HEAD_SIZE / Q_BLOCK;
+    static_assert(HEAD_SIZE % Q_BLOCK == 0, "HEAD_SIZE must be divisible by QK_TURBO3");
+
+    __device__ __forceinline__ static int64_t element_block_index(
+            int paged_block, int kv_head, int n_kv_heads, int token_in_block, int d) {
+        return ((int64_t) paged_block * n_kv_heads + kv_head) * BLOCK_SIZE * N_QBLOCKS_PER_TOKEN
+             + (int64_t) token_in_block * N_QBLOCKS_PER_TOKEN
+             + (int64_t) (d / Q_BLOCK);
+    }
+
+    __device__ __forceinline__ static float k_load(
+            const void * buf, int paged_block, int kv_head, int n_kv_heads,
+            int token_in_block, int d) {
+        const block_turbo3_0 * blocks = (const block_turbo3_0 *) buf;
+        const int64_t ib  = element_block_index(paged_block, kv_head, n_kv_heads, token_in_block, d);
+        const int     iqs = d % Q_BLOCK;
+        const float   norm = __half2float(blocks[ib].norm);
+        return turbo3_dequant_element(&blocks[ib], iqs, norm);
+    }
+
+    __device__ __forceinline__ static float v_load(
+            const void * buf, int paged_block, int kv_head, int n_kv_heads,
+            int token_in_block, int d) {
+        return k_load(buf, paged_block, kv_head, n_kv_heads, token_in_block, d);
+    }
+};
+
 } // namespace mt