diff --git a/common/chat.cpp b/common/chat.cpp index 24e58ab06400..a585b0c6e8c7 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -1670,7 +1670,15 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat auto reasoning = p.eps(); if (extract_reasoning) { - reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END); + // LFM2.5's template does not prefill an opening into the generation prompt, + // so the model must emit it itself each turn. It sometimes skips the opening tag and + // dives straight into reasoning, closing only with . Without forced-open + // handling that raw reasoning bleeds into content. Make the opening tag + // optional so a leading reasoning block terminated by is captured whether or + // not was emitted. A pure-content turn has no , so until() fails to + // find the terminator, the sequence fails, and the outer optional collapses -> the + // whole output falls through to content exactly as before. + reasoning = p.optional(p.optional(p.literal(THINK_START)) + p.reasoning(p.until(THINK_END)) + p.literal(THINK_END)); } if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { diff --git a/docs/paged-attn-debug.md b/docs/paged-attn-debug.md index 8032e96c81b2..bfa134f63b89 100644 --- a/docs/paged-attn-debug.md +++ b/docs/paged-attn-debug.md @@ -170,3 +170,90 @@ exercised). * Reverting the multi-seq commits — they aren't the cause. * Rewriting from scratch — the existing kernel is mostly right; it's one or two layout / shape assumption bugs. + +--- + +## 2026-06-15 session — non-WMMA prefill, gfx803 perf, checkpoint+paged, tool bleed + +Investigation on **mad-lab-2026** (GTX 1070 / sm_61 / CUDA `build-army`; +RX 480 / gfx803 / HIP `build-rocm-gfx803` in docker `rx480-army`). + +### A. Paged PREFILL was on the slow scalar kernel for ALL non-WMMA GPUs (fixed) + +`nsys` on a 2501-tok paged prefill (1070): `mt_paged_attention_kernel` += **84.7% of GPU time, 623 ms/launch**, grid `n_heads×num_seqs` (NO query +parallelism). Root cause: the fast query-tiled prefill kernel +(`mt_pagedattn_tile.cu`, grid `n_heads×num_seqs×n_q_tiles`) is gated at +`mt_pagedattn.cu` behind `tile_gate_on = wmma_ok && tile_env_on`, and +`wmma_ok = amd_wmma_available(cc)` is **RDNA3/4-only**. So Pascal (1070) +and GCN/gfx803 (480) both fall through to the scalar O(n_q·n_kv) kernel. +Decode was unaffected (`launch_paged_attn_decode` is not WMMA-gated). + +**Fix (uncommitted, working tree):** +1. Added a portable FMA `#else` body to `mt_paged_attention_tile_kernel` + (was `NO_DEVICE_CODE`): QK^T as smem dot-products, scores→`smem_s`, + scores·V as smem GEMM. Same lane layout as the WMMA path; shfl + reductions pass explicit `WARP_SIZE` (correct on gfx803 64-lane waves). + Footprint ~13 KiB @ HS128 (fits Pascal's 48 KiB). +2. Dispatch: `tile_gate_on = tile_env_on` (drop `wmma_ok`); force + `mw_on=false` on non-WMMA (multi-warp tile still WMMA-only). +3. **Also fixed a latent grid-sizing bug** (same class as the decode + `93a1e6a11` fix): the tile launches sized `n_q_tiles` from + `avg_q_len = total_q_tokens/num_seqs`, which floors with idle + `--parallel` slots (1 active + N idle) → high Q-tiles never launch → + uncomputed rows → token-salad. Now use `total_q_tokens`. **This also + fixes the WMMA path and is very likely the core of MAD-288.** + +**Verified:** 1070 (CUDA) paged prefill **160 → 328 t/s** (non-paged ref +395), output byte-matches the scalar oracle (`GGML_PAGED_TILE=0`) for +Qwen3.5-9B (HS128) and LFM2.5-8B-A1B (HS64 **padded to 128** in the paged +cache → uses the tile path; no separate HS64 work needed — clarifies +MAD-298). 480 (gfx803): output correct, no crash. Revert: `GGML_PAGED_TILE=0`. + +### B. gfx803 (RX 480) ROCm is ~10× slower than Vulkan — and it's NOT the attn kernel + +Same card, Qwen3.5-9B, `-p512 -n128`: + +| Path | Prefill | Decode | +|---|---|---| +| gfx803 ROCm/HIP, paged tile | 16.6 t/s | — | +| gfx803 ROCm/HIP, paged scalar | 13.8 t/s | — | +| **480 via Vulkan0 (RADV, f16 KV)** | **162 t/s** | 20.4 t/s | +| 1070 via Vulkan | 181 t/s | — | + +Tile ≈ scalar on gfx803 (both ~16) → the bottleneck is **gfx803 ROCm/HIP +itself**, not the attention kernel and not the Polaris silicon. Suspects: +turbo4 per-element dequant in `stage_k/v_tile` (`ops::k_load`), paged-block +gather, or `gated_delta_net` hybrid layers under HIP. **Decision (Kurtis): +keep the custom paged/turbo4/tiered features on gfx803, find the ROCm perf +bug — do NOT port to Vulkan** (Vulkan lacks all custom kernels; only viable +for standard dense models that fit VRAM). NEXT: profile the gfx803 paged +prefill (rocprof not installed in `rx480-army` — needs a profiling path). + +### C. Checkpoint reuse vs paged cache — block-alignment crash (NOT fixed) + +`--ctx-checkpoints>0` would let hybrid/recurrent models reuse the prompt +prefix (the restore path at `server-context.cpp` ~3397 already exists, and +with paged the checkpoints are tiny — **0.282 MiB**, blocks referenced not +copied). But restore sets `n_past` to the checkpoint pos (e.g. 1364), and +`llama_kv_cache_paged::seq_rm` requires **block-aligned (×16)** ranges → +unaligned trim leaves stale partial blocks → `GGML_ASSERT compute_slot_mapping +failed` abort. This is why LFM2.5 runs `--ctx-checkpoints 0` and full- +reprocesses every turn (~2.5 min for ~3700 tok). **Two fix attempts failed:** +round-at-restore is wrong (recurrent state can't be partially rolled back — +must reuse exactly `pos_max+1`); block-aligning the checkpoint batch-break at +creation crashed turn-1 prefill. **Reverted.** Needs a deeper design (paged +`seq_rm`/`compute_slot_mapping` tolerant of the boundary, or a hybrid+paged- +aware checkpoint path). Upstream unmerged references: PRs #20955, #21099, #20428. + +### D. Tool-schema bleed (LFM2.5) + +"hello!" → model regurgitates tool-schema JSON fragments then hallucinates. +The LFM2.5 jinja template renders the tool list as a **plain-text system +message** `List of tools: [{json}]` (faithful to the official template; +`common/chat.cpp:2226` detects this variant, vs LFM2's native +`<|tool_list_start|>` tokens at :622). Tool *calls* are correct +(grammar-constrained `<|tool_call_start|>`). With `thinking=1` the model +also reasons verbosely about the tools. Fix TBD (limit/disable reasoning on +tool turns, or revisit tool presentation). Single/multi-turn tool calls +otherwise correct on 1070. diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index ae6523a9f270..d1c5b3dbe5de 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -439,7 +439,8 @@ extern "C" { GGML_TYPE_F8_E4M3 = 49, // MAD-223 Phase G: fp8 e4m3 1-byte storage, used as sidecar dtype for ml8 centroids GGML_TYPE_ML8_4_SOA = 50, // MAD-223 Phase G.7 / MAD-244: ml8-4 with stored-as-repacked SOA layout (b_packed bytes followed by b_scale fp32 per expert row). Same numerics as ML8_4; eliminates the runtime MoE repack cache. GGML_TYPE_ML8_FP8 = 51, // MAD Task 9: ml8 fp8 weight quant, 32-element blocks, fp16 per-block scale + 32 e4m3 bytes (34 bytes/block). On-disk id=51; matches gguf-py GGML_QUANT_SIZES[ML8_FP8]=(32,34). - GGML_TYPE_COUNT = 52, + GGML_TYPE_TURBO4_64 = 52, // MAD-301C Lever B: native head_dim-64 turbo4 KV cache, 64-element block (34 bytes), no 64->128 pad. Runtime KV-cache-only type (never serialized to GGUF). + GGML_TYPE_COUNT = 53, }; // precision diff --git a/ggml/src/ggml-common.h b/ggml/src/ggml-common.h index 5ad2f2a5f209..282e5d852676 100644 --- a/ggml/src/ggml-common.h +++ b/ggml/src/ggml-common.h @@ -329,6 +329,19 @@ static_assert(sizeof(block_turbo4_0) == 2*sizeof(ggml_half) + QK_TURBO4*3/8 + QK static_assert(QK_TURBO4 == 128, "turbo4 kernels assume QK_TURBO4 == 128"); +// MAD-301C Lever B: native head_dim-64 turbo4 KV cache block. +// Identical 4-bit PolarQuant + per-block norm as block_turbo4_0, but a 64-element +// block so a head_dim-64 head (LFM2.5, gpt-oss) is exactly one block with NO +// 64->128 zero-padding. The paged path skips RHT, so dequant returns +// centroid*norm and dot products match the padded-128 path bit-for-bit at ~half +// the storage (34 vs 68 bytes/head). 4-bit centroid layout only. +#define QK_TURBO4_64 64 +typedef struct { + ggml_half norm; // 2 bytes: per-block scale + uint8_t qs[QK_TURBO4_64 / 2]; // 32 bytes: 4-bit PolarQuant indices (nibble packed) +} block_turbo4_64; // 34 bytes total +static_assert(sizeof(block_turbo4_64) == 2 + QK_TURBO4_64/2, "wrong turbo4_64 block size"); + // TurboQuant 2-bit: 2-bit PolarQuant indices only (no QJL) // Per block: norm(fp16) + 2-bit indices (8 bytes) = 10 bytes per 32 values // = 2.5 bits/value → 6.4× compression vs fp16 diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 789fc4fdb671..6d4f58fd22db 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2930,7 +2930,12 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * } } - if (ggml_cuda_should_use_mmq(src0->type, cc, ne12, /*n_experts=*/ne02)) { + // routing_active forces MMQ: it is the only mul_mat_id path that honors the + // weight-pager expert pointers here (MMVQ handled above for small batch). On + // gfx80x ggml_cuda_should_use_mmq() now returns false for large batches (no + // hardware dp4a -> route to dequant+hipBLAS), so without this OR the routing + // case would fall through to the GGML_ABORT below. + if (routing_active || ggml_cuda_should_use_mmq(src0->type, cc, ne12, /*n_experts=*/ne02)) { ggml_cuda_mul_mat_q(ctx, src0, src1, ids, dst); return; } @@ -5940,7 +5945,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g || op->src[1]->type == GGML_TYPE_Q8_0 || op->src[1]->type == GGML_TYPE_TURBO4_0 || op->src[1]->type == GGML_TYPE_TURBO3_0 - || op->src[1]->type == GGML_TYPE_TURBO4_FP8_BS256) // MAD-214 + || op->src[1]->type == GGML_TYPE_TURBO4_FP8_BS256 // MAD-214 + || op->src[1]->type == GGML_TYPE_TURBO4_64) // MAD-301C Lever B && op->src[6] // k_cur (fused scatter) && op->src[6]->type == GGML_TYPE_F16 && op->src[7] // v_cur (fused scatter) diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 8ef5e6a68588..789f7d5295f8 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -477,5 +477,13 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t return true; } + // gfx80x (Tonga/Fiji/Polaris, GGML_CUDA_CC_GCN4 and earlier) have no hardware dp4a + // (v_dot4): ggml_cuda_dp4a falls back to fully-scalar byte emulation, which makes MMQ + // ~8x slower than dequantization + hipBLAS for prefill-sized batches (rocBLAS ships + // gfx803 fp16/fp32 Tensile kernels). Restrict MMQ to small (decode) batches only. + if (GGML_CUDA_CC_IS_GCN(cc) && cc <= GGML_CUDA_CC_GCN4) { + return ne11 < MMQ_DP4A_MAX_BATCH_SIZE; + } + return (!GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE; } diff --git a/ggml/src/ggml-cuda/mt_pagedattn.cu b/ggml/src/ggml-cuda/mt_pagedattn.cu index 9257824a2213..fbf33325d414 100644 --- a/ggml/src/ggml-cuda/mt_pagedattn.cu +++ b/ggml/src/ggml-cuda/mt_pagedattn.cu @@ -454,6 +454,113 @@ __global__ void mt_scatter_kv_turbo4_0_kernel( } } +// Turbo4_64 scatter (MAD-301C Lever B): native head_dim-64 cooperative quantize. +// Identical pipeline to mt_scatter_kv_turbo4_0_kernel but a 64-element block +// (64 threads / 2 warps), writing block_turbo4_64 (no rnorm field). No RHT. +template +__launch_bounds__(QK_TURBO4_64) +__global__ void mt_scatter_kv_turbo4_64_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_TURBO4_64; // 64 + constexpr int N_QBLOCKS_PER_TOKEN = HEAD_SIZE / Q_BLOCK; + constexpr int N_WARPS = Q_BLOCK / WARP_SIZE; // 2 + static_assert(HEAD_SIZE % Q_BLOCK == 0, "HEAD_SIZE must be divisible by QK_TURBO4_64"); + + const int j = threadIdx.x; // 0..63 + 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; + + 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_turbo4_64 * blk = (block_turbo4_64 *) dst_buf + block_ib; + + __shared__ float x[Q_BLOCK]; + x[j] = __half2float(src[src_off]); + __syncthreads(); + + __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; + + x[j] *= inv_norm; + __syncthreads(); + + // No RHT (see mt_scatter_kv_turbo4_0_kernel step-4 comment): centroid-quant + // normalized K directly; dequant returns centroid*norm ~= K. + const float rv = x[j]; + const uint8_t idx = turbo_nearest_centroid_4bit(rv); + + const int lane = j % WARP_SIZE; + const uint8_t my_nibble = idx & 0xF; + const uint8_t partner_nibble = __shfl_sync(0xffffffffu, my_nibble, lane ^ 1, WARP_SIZE); + if ((j & 1) == 0) { + blk->qs[j / 2] = my_nibble | (partner_nibble << 4); + } + + { + const float c = TURBO_CENTROIDS_4BIT[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; + + if (j == 0) { + blk->norm = __float2half(corrected_norm); + } +} + // Turbo3_0 scatter: per-128-element-block cooperative quantize. // // Same threading and pipeline as turbo4 above, but packs 3-bit indices into @@ -641,6 +748,18 @@ static void launch_scatter_kv( mt_scatter_kv_turbo4_0_kernel <<>>( k_cache, v_cache, k_cur, v_cur, slot_mapping, n_kv_heads); + } else if constexpr (CACHE_TYPE == GGML_TYPE_TURBO4_64) { + // MAD-301C Lever B: native head_dim-64 turbo4. Same grid topology as + // TURBO4_0 but 64 threads/block (one 64-element block per token/head). + constexpr int Q_BLOCK = QK_TURBO4_64; + 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_turbo4_64_kernel + <<>>( + 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. @@ -1097,15 +1216,24 @@ void ggml_cuda_op_paged_attn_mt(ggml_backend_cuda_context & ctx, ggml_tensor * 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_TURBO3_0)) { + // MAD-180: WMMA tile FA gate. HS=128/256 keep F16/TURBO4_0/TURBO3_0. + // MAD-301C Lever B adds HS=64 paired ONLY with GGML_TYPE_TURBO4_64 + // (native head_dim-64 turbo4) — HS=64 for other cache types stays on + // the scalar path, so no <64,*> tile instances are needed for them. + if constexpr ((BS == 16) + && ((((HS == 128) || (HS == 256)) + && (CT == GGML_TYPE_F16 || CT == GGML_TYPE_TURBO4_0 || CT == GGML_TYPE_TURBO3_0)) + || (HS == 64 && CT == GGML_TYPE_TURBO4_64))) { 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; const bool wmma_ok = amd_wmma_available(cc); const bool tile_env_on = get_paged_tile_mode() != 0; - const bool tile_gate_on = wmma_ok && tile_env_on; + // Tile kernel now has a portable FMA #else path (non-WMMA: Pascal + // sm_61, GCN gfx803) in addition to the RDNA3/4 WMMA path, so it is + // no longer gated on wmma_ok. GGML_PAGED_TILE=0 reverts to the slow + // scalar prefill kernel. (Multi-warp variant stays WMMA-only below.) + const bool tile_gate_on = tile_env_on; if (tile_gate_on && avg_q_len >= 16) { if (probe_on) { int n = probe_tile.fetch_add(1, std::memory_order_relaxed); @@ -1123,6 +1251,11 @@ void ggml_cuda_op_paged_attn_mt(ggml_backend_cuda_context & ctx, ggml_tensor * d (const int32_t *) slot_mapping->data, (const int32_t *) q_lens->data, num_seqs, (int) k_cur->ne[2], n_kv_heads, stream); + // Multi-warp tile kernel now has a portable FMA #else path + // (MAD-301C), so non-WMMA hardware (gfx803, Pascal) uses it too: + // K/V staged once per block + shared across warps (HBM-reuse win) + // and N_WARPS*32 threads fill gfx803's 64-lane wave (occupancy win). + // GGML_PAGED_TILE_MULTIWARP=0 reverts to the single-warp tile kernel. const bool mw_on = get_paged_tile_multiwarp_mode() != 0; if (mw_on) { launch_paged_attn_tile_mw( @@ -1134,7 +1267,7 @@ void ggml_cuda_op_paged_attn_mt(ggml_backend_cuda_context & ctx, ggml_tensor * d (const int32_t *) context_lens->data, (const int32_t *) q_lens->data, num_seqs, n_heads, n_kv_heads, max_bps, - avg_q_len, + total_q_tokens, // see note below: size grid.z by max, not avg, q_len scale, stream); } else { launch_paged_attn_tile( @@ -1146,7 +1279,12 @@ void ggml_cuda_op_paged_attn_mt(ggml_backend_cuda_context & ctx, ggml_tensor * d (const int32_t *) context_lens->data, (const int32_t *) q_lens->data, num_seqs, n_heads, n_kv_heads, max_bps, - avg_q_len, // approximation; tile kernel skips q_tile_start >= q_len + total_q_tokens, // grid.z (n_q_tiles) must cover the LARGEST per-seq + // q_len. avg_q_len = total/num_seqs floors with idle + // parallel slots (1 active + N idle) and under-sizes + // the grid -> high Q-tiles never launch -> uncomputed + // rows -> garbage. total_q_tokens is a safe upper bound + // (kernel still skips tiles with q_tile_start >= q_len). scale, stream); } return; @@ -1165,6 +1303,9 @@ void ggml_cuda_op_paged_attn_mt(ggml_backend_cuda_context & ctx, ggml_tensor * d // (rare) heterogeneous case we use it as max_q_len so // the partials buffer's inner stride fits the largest // seq's queries. + // MAD-301C Lever B: GGML_TYPE_TURBO4_64 now has a flash-decode + // path (decode_coop_stage_turbo4_64), so it uses the same gate as + // the other turbo types — no exclusion. const bool decode_env_on = get_paged_decode_mode() != 0; // **Decode regression root cause (2026-05-18 rocprof hunt)**: // the original gate `avg_q_len >= 1` used integer division @@ -1293,6 +1434,16 @@ 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{}); break; + case GGML_TYPE_TURBO4_64: + // MAD-301C Lever B: native head_dim-64 turbo4 — only valid at HS=64. + // Guard so run_typed (and its kernel instances) is only + // instantiated for HS=64, not 128/256. + if constexpr (HS == 64) { + run_typed(std::integral_constant{}); + } else { + GGML_ABORT("mt_paged_attn: GGML_TYPE_TURBO4_64 only supports head_size=64 (got %d)", HS); + } + break; case GGML_TYPE_TURBO3_0: run_typed(std::integral_constant{}); break; diff --git a/ggml/src/ggml-cuda/mt_pagedattn_decode.cu b/ggml/src/ggml-cuda/mt_pagedattn_decode.cu index 2b876dc65038..278f0f355028 100644 --- a/ggml/src/ggml-cuda/mt_pagedattn_decode.cu +++ b/ggml/src/ggml-cuda/mt_pagedattn_decode.cu @@ -70,7 +70,7 @@ static constexpr int DECODE_K_TILE_N = 16; // 1024 = 64 sub-chunks/block. At 400K ctx → ~400 chunks × n_heads = 12800 // blocks (well above RDNA4 R9700's ~240 concurrent block ceiling, so we // run in waves and amortize scratch). -static constexpr int CHUNK_KV = 1024; +static constexpr int CHUNK_KV = 128; // MAD-301A: was 1024 (tuned for 400K ctx / RDNA4 240-block ceiling); too coarse for gfx803's 36 CUs at mid ctx (ctx 1751 -> only 2 chunks -> few blocks -> CU under-utilization). CHUNK_KV sweep on gfx803: 128 (56/38 t/s) > 256 (52/37) > 512 (45/23) @ctx 1751/10501. TODO: make arch-tunable. // Threads/block. 4 warps (32 lanes each) — fits LDS budget for 2× tile // (smem_k + smem_v at HEAD_SIZE=128 = 4 KiB each = 8 KiB) plus Q+logits. @@ -257,6 +257,70 @@ static __device__ __forceinline__ void decode_coop_stage_turbo4( } } +// TURBO4_64 cooperative dequant for the decode tile (MAD-301C Lever B). +// Native head_dim-64: 64-element block => 32 lanes × 2 elements (1 qs byte/lane, +// 2 nibbles). Matches mt_scatter_kv_turbo4_64_kernel packing. +template +static __device__ __forceinline__ void decode_coop_stage_turbo4_64( + __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_TURBO4_64; // 64 + 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_TURBO4_64=64"); + static_assert(Q_BLOCK == 64, "cooperative dequant expects QK_TURBO4_64=64 (32 lanes × 2 elements)"); + + const block_turbo4_64 * blocks = (const block_turbo4_64 *) 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_turbo4_64 * 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, WARP_SIZE); + + uint8_t packed = 0; + if (blk != nullptr) { + packed = blk->qs[lane_id]; // 1 byte = 2 nibbles (elements 2*lane, 2*lane+1) + } + + const int smem_row_base = row * HEAD_SIZE; + const int smem_col_base = qb_in_token * Q_BLOCK + lane_id * 2; + + #pragma unroll + for (int l = 0; l < 2; ++l) { + const uint8_t idx_nib = (packed >> (l * 4)) & 0xF; + const float val = TURBO_CENTROIDS_4BIT[idx_nib] * norm_f; + smem_dst[smem_row_base + smem_col_base + l] = __float2half(val); + } + } +} + // 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). @@ -346,6 +410,10 @@ static __device__ __forceinline__ void decode_stage_k( decode_coop_stage_turbo3( 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_TURBO4_64) { + decode_coop_stage_turbo4_64( + 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( smem_dst, (const __half *) cache, seq_block_table, tile_start, valid_ctx, @@ -373,6 +441,10 @@ static __device__ __forceinline__ void decode_stage_v( decode_coop_stage_turbo3( 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_TURBO4_64) { + decode_coop_stage_turbo4_64( + 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( smem_dst, (const __half *) cache, seq_block_table, tile_start, valid_ctx, @@ -519,21 +591,12 @@ __global__ void mt_paged_attention_decode_kernel( }; if (chunk_start >= valid_ctx_max) { - // No visible tokens for any query — write neutral partials so the - // reducer's pass over all chunks doesn't see uninitialized memory. - // Loop over (qh, qi) so we cover every head this block owns. - for (int qh = 0; qh < num_queries_per_kv; ++qh) { - const int head_idx = head_base + qh; - const size_t base = partial_chunk_base_for_head(head_idx); - for (int qi = 0; qi < q_len; ++qi) { - const size_t off = base + (size_t) qi * (HEAD_SIZE + 2); - for (int d = tid; d < HEAD_SIZE; d += DECODE_NUM_THREADS) partials[off + d] = 0.0f; - if (tid == 0) { - partials[off + HEAD_SIZE] = -INFINITY; - partials[off + HEAD_SIZE + 1] = 0.0f; - } - } - } + // MAD-301A: this chunk is entirely beyond the seq's real context + // (num_chunks is sized by ALLOCATED ctx, not actual — and with a small + // CHUNK_KV + large --ctx-size there are thousands of these). The reduce + // kernel now bounds its loop to ceil(context_lens/CHUNK_KV), so it never + // reads partials from these chunks — skip the neutral-partial writes and + // just retire the block. Makes decode cost scale with actual depth. return; } const int chunk_end = min(chunk_start + CHUNK_KV, valid_ctx_max); @@ -1058,9 +1121,11 @@ __global__ void mt_paged_attention_decode_reduce_kernel( __half * __restrict__ out, const float * __restrict__ partials, const int32_t * __restrict__ q_lens, + const int32_t * __restrict__ context_lens, + int num_chunks, + int chunk_kv, int n_heads, int n_seqs, - int num_chunks, int max_q_len) { const int head_idx = blockIdx.x; const int seq_idx = blockIdx.y; @@ -1071,6 +1136,16 @@ __global__ void mt_paged_attention_decode_reduce_kernel( for (int s = 0; s < seq_idx; ++s) seq_q_offset += (size_t) q_lens[s]; const int q_len = q_lens[seq_idx]; + // MAD-301A: num_chunks is sized by ALLOCATED ctx (max_blocks_per_seq*block_size). + // With small CHUNK_KV + a large --ctx-size, that's thousands of chunks while a + // decode step only fills context_lens[seq] tokens. Bound the reduction to the + // chunks that actually hold data so cost scales with real depth, not allocated + // capacity. Pass-1 chunk blocks beyond this range early-return without writing + // neutral partials, so they must never be read here. + const int ctx_len = context_lens[seq_idx]; + const int chunks_full = (ctx_len + chunk_kv - 1) / chunk_kv; + const int valid_chunks = chunks_full < num_chunks ? chunks_full : num_chunks; + // Stride from chunk c → chunk c+1 in the partials buffer. const size_t chunk_stride_q = (size_t) max_q_len * (size_t) (HEAD_SIZE + 2); // (head, seq) base. @@ -1082,7 +1157,7 @@ __global__ void mt_paged_attention_decode_reduce_kernel( // Pass 1: global max across chunks for this query position. float global_max = -INFINITY; - for (int c = 0; c < num_chunks; ++c) { + for (int c = 0; c < valid_chunks; ++c) { const float m = partials[partial_seq_base + (size_t) c * chunk_stride_q + qi_stride + HEAD_SIZE]; global_max = max(global_max, m); } @@ -1090,7 +1165,7 @@ __global__ void mt_paged_attention_decode_reduce_kernel( // Pass 2: merge across chunks. float global_sum = 0.0f; float v_d = 0.0f; - for (int c = 0; c < num_chunks; ++c) { + for (int c = 0; c < valid_chunks; ++c) { const size_t cbase = partial_seq_base + (size_t) c * chunk_stride_q + qi_stride; const float c_max = partials[cbase + HEAD_SIZE]; if (c_max == -INFINITY) continue; @@ -1187,7 +1262,8 @@ void launch_paged_attn_decode( dim3 block2(HEAD_SIZE); mt_paged_attention_decode_reduce_kernel <<>>( - out, partials_scratch, q_lens, n_heads, num_seqs, num_chunks, max_q_len); + out, partials_scratch, q_lens, context_lens, num_chunks, CHUNK_KV, + n_heads, num_seqs, max_q_len); } // ── explicit instantiations ──────────────────────────────────────────── @@ -1224,5 +1300,10 @@ 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); +// HEAD_SIZE=64 — MAD-301C Lever B native head_dim-64 turbo4 flash-decode. +template void launch_paged_attn_decode<64, 16, GGML_TYPE_TURBO4_64>( + __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 diff --git a/ggml/src/ggml-cuda/mt_pagedattn_ops.cuh b/ggml/src/ggml-cuda/mt_pagedattn_ops.cuh index 096ca18527ca..915757ecd1d9 100644 --- a/ggml/src/ggml-cuda/mt_pagedattn_ops.cuh +++ b/ggml/src/ggml-cuda/mt_pagedattn_ops.cuh @@ -152,6 +152,41 @@ struct paged_cache_ops { } }; +// Turbo4_64 specialization (MAD-301C Lever B): native head_dim-64 4-bit PolarQuant, +// 64-element blocks. Same layout shape as TURBO4_0 (1 block/token for head_dim 64, +// since N_QBLOCKS_PER_TOKEN = 64/64 = 1) and same no-RHT dequant; only the block +// width (64 vs 128) and struct (block_turbo4_64, 34 bytes) differ. No 64->128 pad. +// kv_store omitted — handled by mt_scatter_kv_turbo4_64_kernel. +template +struct paged_cache_ops { + static constexpr int Q_BLOCK = QK_TURBO4_64; // 64 + 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_TURBO4_64"); + + __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_turbo4_64 * blocks = (const block_turbo4_64 *) 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 turbo4_64_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); + } +}; + // 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). diff --git a/ggml/src/ggml-cuda/mt_pagedattn_tile.cu b/ggml/src/ggml-cuda/mt_pagedattn_tile.cu index 6c41d2744b5a..e269f2246d6a 100644 --- a/ggml/src/ggml-cuda/mt_pagedattn_tile.cu +++ b/ggml/src/ggml-cuda/mt_pagedattn_tile.cu @@ -64,6 +64,13 @@ static constexpr int K_INNER = 16; template struct TileConfig; +template <> +struct TileConfig<64> { + // MAD-301C Lever B: native head_dim-64 turbo4. LDS at 8 tiles: + // 8*16*64*2 + 2*16*64*2 = 16 KiB + 4 KiB = 20 KiB — well under the 64 KiB cap. + static constexpr int Q_TILES_PER_BLOCK = 8; +}; + template <> struct TileConfig<128> { static constexpr int Q_TILES_PER_BLOCK = 8; @@ -571,13 +578,159 @@ __global__ void mt_paged_attention_tile_kernel( } } #else - // Non-WMMA hardware should not be dispatched here; existing scalar - // kernel is the fallback. NO_DEVICE_CODE traps if we somehow are. - GGML_UNUSED(out); GGML_UNUSED(q); GGML_UNUSED(k_cache); GGML_UNUSED(v_cache); - GGML_UNUSED(block_tables); GGML_UNUSED(context_lens); GGML_UNUSED(q_lens); - GGML_UNUSED(max_blocks_per_seq); GGML_UNUSED(n_kv_heads); GGML_UNUSED(n_heads); - GGML_UNUSED(scale); - NO_DEVICE_CODE; + // ── Portable FMA tile flash-attention (non-WMMA: NVIDIA Pascal sm_61, + // AMD GCN gfx803). Mirrors the WMMA path's Q-tile + online-softmax + // structure, but the two matmuls run as explicit shared-memory GEMMs + // instead of matrix-core fragments. Lane ownership is kept identical to + // the WMMA fragment layout (row = tid%16, cols 8*(tid/16)+l), so the + // masking / softmax / writeback logic matches the #if branch line-for-line. + // One 32-lane warp per block; shuffles pass explicit WARP_SIZE so a 32-wide + // logical reduction is correct on gfx803's 64-lane waves too. + static_assert(HEAD_SIZE % K_INNER == 0, "HEAD_SIZE must be multiple of K_INNER=16"); + constexpr int N_INNER = HEAD_SIZE / K_INNER; + + const int head_idx = blockIdx.x; + const int seq_idx = blockIdx.y; + const int q_tile_idx = blockIdx.z; + const int tid = threadIdx.x; + + const int q_len = q_lens[seq_idx]; + const int q_tile_start = q_tile_idx * Q_TILE_M; + if (q_tile_start >= q_len) { + return; + } + const int q_tile_actual = (q_tile_start + Q_TILE_M <= q_len) ? Q_TILE_M : (q_len - q_tile_start); + + const int kv_head_idx = head_idx / (n_heads / n_kv_heads); + const int ctx_len_after_q = context_lens[seq_idx]; + const int * seq_block_table = block_tables + seq_idx * max_blocks_per_seq; + + size_t seq_q_offset = 0; + for (int s = 0; s < seq_idx; ++s) { + seq_q_offset += (size_t) q_lens[s]; + } + + const int q_pos_base = (ctx_len_after_q - q_len) + q_tile_start; + + // smem: Q[16,HS] + K[16,HS] + V[16,HS] half, then scores[16,16] f32. + extern __shared__ unsigned char smem_raw[]; + __half * smem_q = (__half *)(smem_raw); + __half * smem_k = smem_q + Q_TILE_M * HEAD_SIZE; + __half * smem_v = smem_k + K_TILE_N * HEAD_SIZE; + float * smem_s = (float *)(smem_v + K_TILE_N * HEAD_SIZE); // [Q_TILE_M * K_TILE_N] + + const size_t q_global_base = ((seq_q_offset + (size_t) q_tile_start) * (size_t) n_heads + (size_t) head_idx) + * (size_t) HEAD_SIZE; + stage_q_tile(smem_q, q, q_global_base, q_tile_actual, n_heads, tid); + __syncthreads(); + + // Lane ownership: row = tid%16 (Q row), half = tid/16 -> cols 8*half+l. + const int row = tid % 16; + const int half = tid / 16; + + // Online-softmax state + output accumulators (this lane's 8 cols per HS block). + float running_max = -INFINITY; + float running_sum = 0.0f; + float acc[N_INNER][8]; + #pragma unroll + for (int n = 0; n < N_INNER; ++n) { + #pragma unroll + for (int l = 0; l < 8; ++l) acc[n][l] = 0.0f; + } + + const int q_pos_last = q_pos_base + (q_tile_actual - 1); + const int valid_ctx = q_pos_last + 1; + const int q_pos = q_pos_base + row; + const bool row_valid = (row < q_tile_actual); + + for (int k_tile_start = 0; k_tile_start < valid_ctx; k_tile_start += K_TILE_N) { + stage_k_tile( + smem_k, k_cache, seq_block_table, k_tile_start, valid_ctx, + kv_head_idx, n_kv_heads, tid); + __syncthreads(); + + // scores[row, col] = scale * (Q[row] . K[col]) for this lane's 8 cols. + float sc[8]; + #pragma unroll + for (int l = 0; l < 8; ++l) { + const int col = 8 * half + l; + float dot = 0.0f; + #pragma unroll + for (int d = 0; d < HEAD_SIZE; ++d) { + dot += __half2float(smem_q[row * HEAD_SIZE + d]) * __half2float(smem_k[col * HEAD_SIZE + d]); + } + const int k_pos = k_tile_start + col; + const bool visible = row_valid && (k_pos <= q_pos) && (k_pos < valid_ctx); + sc[l] = visible ? (dot * scale) : -INFINITY; + } + + // Per-row max: 8-wide local then pair-lane (tid^16) via shfl_xor. + float local_max = -INFINITY; + #pragma unroll + for (int l = 0; l < 8; ++l) local_max = max(local_max, sc[l]); + const float row_max = max(local_max, __shfl_xor_sync(0xFFFFFFFF, local_max, 16, WARP_SIZE)); + const float new_max = max(running_max, row_max); + + // Rescale running state on new max. + if (running_max > -INFINITY) { + const float rescale = __expf(running_max - new_max); + running_sum *= rescale; + #pragma unroll + for (int n = 0; n < N_INNER; ++n) { + #pragma unroll + for (int l = 0; l < 8; ++l) acc[n][l] *= rescale; + } + } + + // exp(scores - new_max) + per-row sum. + float local_sum = 0.0f; + #pragma unroll + for (int l = 0; l < 8; ++l) { + const float e = (sc[l] == -INFINITY) ? 0.0f : __expf(sc[l] - new_max); + sc[l] = e; + local_sum += e; + smem_s[row * K_TILE_N + (8 * half + l)] = e; // publish full row for PV + } + running_sum += local_sum + __shfl_xor_sync(0xFFFFFFFF, local_sum, 16, WARP_SIZE); + running_max = new_max; + + stage_v_tile( + smem_v, v_cache, seq_block_table, k_tile_start, valid_ctx, + kv_head_idx, n_kv_heads, tid); + __syncthreads(); // smem_s (all 16 cols) + smem_v visible to the warp + + // acc[n][l] += sum_k scores[row, k] * V[k, d], d = n*16 + 8*half + l. + #pragma unroll + for (int n = 0; n < N_INNER; ++n) { + #pragma unroll + for (int l = 0; l < 8; ++l) { + const int d = n * K_INNER + 8 * half + l; + float pv = 0.0f; + #pragma unroll + for (int k = 0; k < K_TILE_N; ++k) { + pv += smem_s[row * K_TILE_N + k] * __half2float(smem_v[k * HEAD_SIZE + d]); + } + acc[n][l] += pv; + } + } + __syncthreads(); // before next iter overwrites smem_k / smem_v / smem_s + } + + // Writeback (same layout as the WMMA branch). + const float inv_sum = 1.0f / (running_sum + 1e-6f); + if (row < q_tile_actual) { + const int q_row_global = q_tile_start + row; + const size_t out_row_base = + ((seq_q_offset + (size_t) q_row_global) * (size_t) n_heads + (size_t) head_idx) * (size_t) HEAD_SIZE; + #pragma unroll + for (int n = 0; n < N_INNER; ++n) { + #pragma unroll + for (int l = 0; l < 8; ++l) { + const int d = n * K_INNER + 8 * half + l; + out[out_row_base + (size_t) d] = __float2half(acc[n][l] * inv_sum); + } + } + } #endif // AMD_WMMA_AVAILABLE } @@ -603,7 +756,10 @@ void launch_paged_attn_tile( dim3 grid(n_heads, num_seqs, n_q_tiles); dim3 block(TILE_NUM_THREADS); - const size_t smem_bytes = (size_t)(Q_TILE_M + 2 * K_TILE_N) * (size_t) HEAD_SIZE * sizeof(__half); + // half Q/K/V tiles + (non-WMMA path only) an f32 scores[Q_TILE_M*K_TILE_N] + // scratch. Allocated unconditionally; the WMMA branch simply ignores it. + const size_t smem_bytes = (size_t)(Q_TILE_M + 2 * K_TILE_N) * (size_t) HEAD_SIZE * sizeof(__half) + + (size_t) Q_TILE_M * (size_t) K_TILE_N * sizeof(float); mt_paged_attention_tile_kernel <<>>( @@ -965,12 +1121,255 @@ __global__ void mt_paged_attention_tile_mw_kernel( } } #else - GGML_UNUSED(out); GGML_UNUSED(q); GGML_UNUSED(k_cache); GGML_UNUSED(v_cache); - GGML_UNUSED(block_tables); GGML_UNUSED(context_lens); GGML_UNUSED(q_lens); - GGML_UNUSED(max_blocks_per_seq); GGML_UNUSED(n_kv_heads); GGML_UNUSED(n_heads); - GGML_UNUSED(scale); - NO_DEVICE_CODE; -#endif + // ── Portable FMA multi-warp tile flash-attention (non-WMMA: NVIDIA + // Pascal sm_61, AMD GCN gfx803). Mirrors the WMMA mw path's block layout + // (Q_TILES warps/block, K/V staged once into smem and shared across warps + // = the HBM-reuse win) but runs the two matmuls as explicit shared-memory + // GEMMs per warp, identical math to the single-warp FMA #else above. On + // gfx803 the N_WARPS*32 threads fill the 64-lane wave (2 logical warps per + // physical wave), recovering the ~2x SIMD utilization the 32-thread + // single-warp tile kernel left on the table. No extra LDS over the WMMA + // layout: the PV step rebuilds each Q row's full 16-wide score vector from + // the pair lane via __shfl_xor(.,16) instead of an smem scratch. + static_assert(HEAD_SIZE % K_INNER == 0, "HEAD_SIZE must be multiple of K_INNER=16"); + constexpr int Q_TILES = TileConfig::Q_TILES_PER_BLOCK; + constexpr int N_WARPS = Q_TILES; + constexpr int N_THREADS = N_WARPS * 32; + constexpr int N_INNER = HEAD_SIZE / K_INNER; + + using ops = paged_cache_ops; + + const int head_idx = blockIdx.x; + const int seq_idx = blockIdx.y; + const int q_tile_grp_idx = blockIdx.z; + const int tid = threadIdx.x; + const int warp_id = tid >> 5; // tid / 32 + const int lane_id = tid & 31; // tid % 32 + + const int q_len = q_lens[seq_idx]; + const int q_tile_base = q_tile_grp_idx * Q_TILES; + + if (q_tile_base * Q_TILE_M >= q_len) { + return; + } + + const int my_q_tile_idx = q_tile_base + warp_id; + const int my_q_tile_start = my_q_tile_idx * Q_TILE_M; + const bool warp_active = (my_q_tile_start < q_len); + const int my_q_tile_actual = warp_active + ? ((my_q_tile_start + Q_TILE_M <= q_len) ? Q_TILE_M : (q_len - my_q_tile_start)) + : 0; + + const int kv_head_idx = head_idx / (n_heads / n_kv_heads); + const int ctx_len_after_q = context_lens[seq_idx]; + const int * seq_block_table = block_tables + seq_idx * max_blocks_per_seq; + + size_t seq_q_offset = 0; + for (int s = 0; s < seq_idx; ++s) { + seq_q_offset += (size_t) q_lens[s]; + } + + const int my_q_pos_base = (ctx_len_after_q - q_len) + my_q_tile_start; + + // Uniform K-loop bound across the block (all warps must hit the same + // cooperative __syncthreads); per-warp causal masking handles the rest. + const int block_last_q_row = min(q_tile_base * Q_TILE_M + Q_TILES * Q_TILE_M, q_len) - 1; + const int block_last_q_pos = (ctx_len_after_q - q_len) + block_last_q_row; + const int block_valid_ctx = block_last_q_pos + 1; + + extern __shared__ unsigned char smem_raw[]; + __half * smem_q = (__half *)(smem_raw); + __half * smem_k = smem_q + Q_TILES * Q_TILE_M * HEAD_SIZE; + __half * smem_v = smem_k + K_TILE_N * HEAD_SIZE; + + // ── cooperative Q load (all threads, all Q_TILES tiles) ── + { + constexpr int TOTAL_Q_ELEMS = Q_TILES * Q_TILE_M * HEAD_SIZE; + for (int idx = tid; idx < TOTAL_Q_ELEMS; idx += N_THREADS) { + const int qt = idx / (Q_TILE_M * HEAD_SIZE); + const int qt_off = idx % (Q_TILE_M * HEAD_SIZE); + const int qrow = qt_off / HEAD_SIZE; + const int qcol = qt_off % HEAD_SIZE; + + const int q_tile_idx_local = q_tile_base + qt; + const int q_row_global = q_tile_idx_local * Q_TILE_M + qrow; + + __half val = __float2half(0.0f); + if (q_row_global < q_len) { + const size_t base = ((seq_q_offset + (size_t) q_row_global) * (size_t) n_heads + + (size_t) head_idx) * (size_t) HEAD_SIZE; + val = q[base + (size_t) qcol]; + } + smem_q[idx] = val; + } + } + __syncthreads(); + + // Lane ownership within a warp: row = lane%16 (Q row), half = lane/16 -> + // this lane owns cols 8*half + l, l in [0,8). + const __half * my_smem_q = smem_q + warp_id * Q_TILE_M * HEAD_SIZE; + const int row = lane_id & 15; + const int half_id = lane_id >> 4; + + float running_max = -INFINITY; + float running_sum = 0.0f; + float acc[N_INNER][8]; + #pragma unroll + for (int n = 0; n < N_INNER; ++n) { + #pragma unroll + for (int l = 0; l < 8; ++l) acc[n][l] = 0.0f; + } + + const int q_pos = my_q_pos_base + row; + const bool row_valid = (row < my_q_tile_actual); + + for (int k_tile_start = 0; k_tile_start < block_valid_ctx; k_tile_start += K_TILE_N) { + // Cooperative K load (shared across all warps in the block). + if constexpr (CACHE_TYPE == GGML_TYPE_TURBO4_0) { + coop_stage_turbo4_tile( + smem_k, k_cache, seq_block_table, k_tile_start, block_valid_ctx, + kv_head_idx, n_kv_heads, warp_id, lane_id); + } else if constexpr (CACHE_TYPE == GGML_TYPE_TURBO3_0) { + coop_stage_turbo3_tile( + smem_k, k_cache, seq_block_table, k_tile_start, block_valid_ctx, + kv_head_idx, n_kv_heads, warp_id, lane_id); + } else { + for (int idx = tid; idx < K_TILE_N * HEAD_SIZE; idx += N_THREADS) { + const int krow = idx / HEAD_SIZE; + const int kcol = idx % HEAD_SIZE; + const int token = k_tile_start + krow; + float val = 0.0f; + if (token < block_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) { + val = ops::k_load(k_cache, physical, kv_head_idx, n_kv_heads, tok_in_block, kcol); + } + } + smem_k[idx] = __float2half(val); + } + } + __syncthreads(); + + // QK: scores[row, col] = scale * (Q[row] . K[col]) for this lane's 8 cols. + float sc[8]; + if (warp_active) { + #pragma unroll + for (int l = 0; l < 8; ++l) { + const int col = 8 * half_id + l; + float dot = 0.0f; + #pragma unroll + for (int d = 0; d < HEAD_SIZE; ++d) { + dot += __half2float(my_smem_q[row * HEAD_SIZE + d]) * __half2float(smem_k[col * HEAD_SIZE + d]); + } + const int k_pos = k_tile_start + col; + const bool visible = row_valid && (k_pos <= q_pos) && (k_pos < block_valid_ctx); + sc[l] = visible ? (dot * scale) : -INFINITY; + } + + // Per-row max: 8-wide local then pair-lane (lane^16) via shfl_xor. + float local_max = -INFINITY; + #pragma unroll + for (int l = 0; l < 8; ++l) local_max = max(local_max, sc[l]); + const float row_max = max(local_max, __shfl_xor_sync(0xFFFFFFFF, local_max, 16, WARP_SIZE)); + const float new_max = max(running_max, row_max); + + // Rescale running state on new max. + if (running_max > -INFINITY) { + const float rescale = __expf(running_max - new_max); + running_sum *= rescale; + #pragma unroll + for (int n = 0; n < N_INNER; ++n) { + #pragma unroll + for (int l = 0; l < 8; ++l) acc[n][l] *= rescale; + } + } + + // exp(scores - new_max) + per-row sum. + float local_sum = 0.0f; + #pragma unroll + for (int l = 0; l < 8; ++l) { + const float e = (sc[l] == -INFINITY) ? 0.0f : __expf(sc[l] - new_max); + sc[l] = e; + local_sum += e; + } + running_sum += local_sum + __shfl_xor_sync(0xFFFFFFFF, local_sum, 16, WARP_SIZE); + running_max = new_max; + } + + // Cooperative V load (same dispatch as K; turbo* V shares K layout). + if constexpr (CACHE_TYPE == GGML_TYPE_TURBO4_0) { + coop_stage_turbo4_tile( + smem_v, v_cache, seq_block_table, k_tile_start, block_valid_ctx, + kv_head_idx, n_kv_heads, warp_id, lane_id); + } else if constexpr (CACHE_TYPE == GGML_TYPE_TURBO3_0) { + coop_stage_turbo3_tile( + smem_v, v_cache, seq_block_table, k_tile_start, block_valid_ctx, + kv_head_idx, n_kv_heads, warp_id, lane_id); + } else { + for (int idx = tid; idx < K_TILE_N * HEAD_SIZE; idx += N_THREADS) { + const int vrow = idx / HEAD_SIZE; + const int vcol = idx % HEAD_SIZE; + const int token = k_tile_start + vrow; + float val = 0.0f; + if (token < block_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) { + val = ops::v_load(v_cache, physical, kv_head_idx, n_kv_heads, tok_in_block, vcol); + } + } + smem_v[idx] = __float2half(val); + } + } + __syncthreads(); // smem_v visible to the warp + + // PV: acc[n][l] += sum_c scores[row, c] * V[c, d], d = n*16 + 8*half + l. + // This lane owns only its 8 exp-scores (cols 8*half+l); fetch the pair + // lane's 8 via shfl_xor(.,16) to rebuild the full 16-wide row. + if (warp_active) { + float s16[K_TILE_N]; + #pragma unroll + for (int l = 0; l < 8; ++l) { + const float partner = __shfl_xor_sync(0xFFFFFFFF, sc[l], 16, WARP_SIZE); + s16[8 * half_id + l] = sc[l]; + s16[8 * (1 - half_id) + l] = partner; + } + #pragma unroll + for (int n = 0; n < N_INNER; ++n) { + #pragma unroll + for (int l = 0; l < 8; ++l) { + const int d = n * K_INNER + 8 * half_id + l; + float pv = 0.0f; + #pragma unroll + for (int c = 0; c < K_TILE_N; ++c) { + pv += s16[c] * __half2float(smem_v[c * HEAD_SIZE + d]); + } + acc[n][l] += pv; + } + } + } + __syncthreads(); // before next iter overwrites smem_k / smem_v + } + + // ── per-warp output writeback ─────────────────────────────────────── + if (warp_active && row < my_q_tile_actual) { + const float inv_sum = 1.0f / (running_sum + 1e-6f); + const int q_row_global = my_q_tile_start + row; + const size_t out_row_base = + ((seq_q_offset + (size_t) q_row_global) * (size_t) n_heads + (size_t) head_idx) * (size_t) HEAD_SIZE; + #pragma unroll + for (int n = 0; n < N_INNER; ++n) { + #pragma unroll + for (int l = 0; l < 8; ++l) { + const int d = n * K_INNER + 8 * half_id + l; + out[out_row_base + (size_t) d] = __float2half(acc[n][l] * inv_sum); + } + } + } +#endif // AMD_WMMA_AVAILABLE } template @@ -1036,6 +1435,11 @@ template void launch_paged_attn_tile<256, 16, GGML_TYPE_TURBO3_0>( __half *, const __half *, const void *, const void *, const int32_t *, const int32_t *, const int32_t *, int, int, int, int, int, float, cudaStream_t); +// HEAD_SIZE=64 — MAD-301C Lever B native head_dim-64 turbo4 (LFM2.5, gpt-oss). +template void launch_paged_attn_tile<64, 16, GGML_TYPE_TURBO4_64>( + __half *, const __half *, const void *, const void *, + const int32_t *, const int32_t *, const int32_t *, + int, int, int, int, int, float, cudaStream_t); // Multi-warp launcher instantiations — same (HEAD_SIZE, BLOCK_SIZE, CACHE_TYPE) // matrix as the single-warp path. TileConfig picks Q_TILES_PER_BLOCK at @@ -1064,5 +1468,10 @@ template void launch_paged_attn_tile_mw<256, 16, GGML_TYPE_TURBO3_0>( __half *, const __half *, const void *, const void *, const int32_t *, const int32_t *, const int32_t *, int, int, int, int, int, float, cudaStream_t); +// HEAD_SIZE=64 — MAD-301C Lever B native head_dim-64 turbo4. +template void launch_paged_attn_tile_mw<64, 16, GGML_TYPE_TURBO4_64>( + __half *, const __half *, const void *, const void *, + const int32_t *, const int32_t *, const int32_t *, + int, int, int, int, int, float, cudaStream_t); } // namespace mt diff --git a/ggml/src/ggml-cuda/turbo-quant.cuh b/ggml/src/ggml-cuda/turbo-quant.cuh index 9edd913ae6a1..cdbc3075116f 100644 --- a/ggml/src/ggml-cuda/turbo-quant.cuh +++ b/ggml/src/ggml-cuda/turbo-quant.cuh @@ -351,6 +351,14 @@ static __device__ __forceinline__ float turbo4_dequant_element( return TURBO_CENTROIDS_4BIT[idx] * norm; } +// MAD-301C Lever B: native head_dim-64 turbo4 dequant. Same nibble layout as +// block_turbo4_0 (j/2 byte, (j%2)*4 shift); only the block struct width differs. +static __device__ __forceinline__ float turbo4_64_dequant_element( + const block_turbo4_64 * __restrict__ x, int j, float norm) { + uint8_t idx = (x->qs[j / 2] >> ((j % 2) * 4)) & 0xF; + return TURBO_CENTROIDS_4BIT[idx] * norm; +} + // ---- Nearest 3-bit centroid index ---- static __device__ __forceinline__ uint8_t turbo_nearest_centroid_3bit(float val) { diff --git a/ggml/src/ggml-quants.h b/ggml/src/ggml-quants.h index db271e517931..adb0440d3f28 100644 --- a/ggml/src/ggml-quants.h +++ b/ggml/src/ggml-quants.h @@ -107,6 +107,11 @@ GGML_API void quantize_row_turbo3_0_ref(const float * GGML_RESTRICT x, block_tur GGML_API void quantize_row_turbo4_0_ref(const float * GGML_RESTRICT x, block_turbo4_0 * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_turbo3_0(const block_turbo3_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_turbo4_0(const block_turbo4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +// MAD-301C Lever B: native head_dim-64 turbo4 (64-element block). CPU refs exist to +// satisfy ggml type-traits; the hot path is CUDA paged (mt_pagedattn). No-RHT, matching +// the paged scatter convention (dequant = centroid*norm). +GGML_API void quantize_row_turbo4_64_ref(const float * GGML_RESTRICT x, block_turbo4_64 * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_turbo4_64(const block_turbo4_64 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API size_t quantize_turbo3_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API size_t quantize_turbo4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API void quantize_row_turbo2_0_ref(const float * GGML_RESTRICT x, block_turbo2_0 * GGML_RESTRICT y, int64_t k); diff --git a/ggml/src/ggml-turbo-quant.c b/ggml/src/ggml-turbo-quant.c index 992c6aa1249f..bba63572c649 100644 --- a/ggml/src/ggml-turbo-quant.c +++ b/ggml/src/ggml-turbo-quant.c @@ -609,6 +609,67 @@ void dequantize_row_turbo4_0(const block_turbo4_0 * GGML_RESTRICT x, float * GGM #endif } +/* MAD-301C Lever B: native head_dim-64 turbo4 (64-element block, no RHT). + * Mirrors the CUDA paged scatter convention (mt_scatter_kv_turbo4_64_kernel): + * L2-norm -> normalize -> 4-bit centroid quant (NO WHT) -> recon-norm correction; + * dequant returns centroid*norm. These CPU refs exist only to satisfy ggml + * type-traits for GGML_TYPE_TURBO4_64; the hot path is CUDA paged attention. */ +void quantize_row_turbo4_64_ref(const float * GGML_RESTRICT x, block_turbo4_64 * GGML_RESTRICT y, int64_t k) { + static const float CENTROIDS_4BIT[16] = { + -0.173926f, -0.117195f, -0.089527f, -0.068756f, + -0.051262f, -0.035597f, -0.020989f, -0.006938f, + 0.006938f, 0.020989f, 0.035597f, 0.051262f, + 0.068756f, 0.089527f, 0.117195f, 0.173926f + }; + assert(k % QK_TURBO4_64 == 0); + const int nb = k / QK_TURBO4_64; + const int d = QK_TURBO4_64; + for (int block = 0; block < nb; block++) { + const float * src = x + block * d; + float norm_sq = 0.0f; + for (int i = 0; i < d; i++) norm_sq += src[i] * src[i]; + const float norm = sqrtf(norm_sq); + float normalized[QK_TURBO4_64]; + if (norm > 1e-10f) { + const float inv = 1.0f / norm; + for (int i = 0; i < d; i++) normalized[i] = src[i] * inv; + } else { + memset(normalized, 0, d * sizeof(float)); + } + uint8_t indices[QK_TURBO4_64]; + for (int i = 0; i < d; i++) indices[i] = (uint8_t)nearest_centroid_4bit(normalized[i]); + float recon_sq = 0.0f; + for (int i = 0; i < d; i++) recon_sq += CENTROIDS_4BIT[indices[i]] * CENTROIDS_4BIT[indices[i]]; + const float recon_norm = sqrtf(recon_sq); + const float corrected_norm = (recon_norm > 1e-10f) ? norm / recon_norm : norm; + y[block].norm = GGML_FP32_TO_FP16(corrected_norm); + memset(y[block].qs, 0, d / 2); + for (int i = 0; i < d; i++) { + y[block].qs[i / 2] |= (uint8_t)((indices[i] & 0xF) << ((i % 2) * 4)); + } + } +} + +void dequantize_row_turbo4_64(const block_turbo4_64 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const float CENTROIDS_4BIT[16] = { + -0.173926f, -0.117195f, -0.089527f, -0.068756f, + -0.051262f, -0.035597f, -0.020989f, -0.006938f, + 0.006938f, 0.020989f, 0.035597f, 0.051262f, + 0.068756f, 0.089527f, 0.117195f, 0.173926f + }; + assert(k % QK_TURBO4_64 == 0); + const int nb = k / QK_TURBO4_64; + const int d = QK_TURBO4_64; + for (int block = 0; block < nb; block++) { + const float norm = GGML_FP16_TO_FP32(x[block].norm); + float * dst = y + block * d; + for (int i = 0; i < d; i++) { + uint8_t idx = (x[block].qs[i / 2] >> ((i % 2) * 4)) & 0xF; + dst[i] = CENTROIDS_4BIT[idx] * norm; + } + } +} + size_t quantize_turbo4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { GGML_UNUSED(imatrix); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 4f44aa726b31..1606630e42ef 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -765,6 +765,14 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { .to_float = (ggml_to_float_t) dequantize_row_turbo4_0, .from_float_ref = (ggml_from_float_t) quantize_row_turbo4_0_ref, }, + [GGML_TYPE_TURBO4_64] = { + .type_name = "turbo4_64", + .blck_size = QK_TURBO4_64, + .type_size = sizeof(block_turbo4_64), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_turbo4_64, + .from_float_ref = (ggml_from_float_t) quantize_row_turbo4_64_ref, + }, [GGML_TYPE_TURBO2_0] = { .type_name = "turbo2", .blck_size = QK_TURBO2, diff --git a/src/llama-kv-cache-paged.cpp b/src/llama-kv-cache-paged.cpp index 450c3075f098..4ce2c0fd994d 100644 --- a/src/llama-kv-cache-paged.cpp +++ b/src/llama-kv-cache-paged.cpp @@ -16,6 +16,7 @@ #endif #include +#include #include #include #include @@ -120,9 +121,26 @@ llama_kv_cache_paged::llama_kv_cache_paged( } } + // MAD-301C Lever B: native head_dim-64 turbo4. Remap TURBO4_0 -> TURBO4_64 + // (64-element block) for head_dim-64 models (LFM2.5, gpt-oss). Each head is + // then exactly one 64-wide block with NO 64->128 zero-pad, ~halving the + // turbo4 KV footprint. Must run before the pad check + sizing below; the + // graph-level pad in build_attn keys off layer.k->type so it auto-skips for + // TURBO4_64 (not in its TURBO{2,3,4}_0 list). + // GGML_PAGED_TURBO4_64=0 forces the legacy padded-128 path (A/B + rollback). + { + const char * t464_env = std::getenv("GGML_PAGED_TURBO4_64"); + const bool t464_on = (t464_env == nullptr || t464_env[0] != '0'); + if (t464_on && head_dim == 64) { + if (type_k == GGML_TYPE_TURBO4_0) { type_k = GGML_TYPE_TURBO4_64; type_k_ = GGML_TYPE_TURBO4_64; } + if (type_v == GGML_TYPE_TURBO4_0) { type_v = GGML_TYPE_TURBO4_64; type_v_ = GGML_TYPE_TURBO4_64; } + } + } + // Turbo cache quantizes 128-element blocks; pad a sub-128 head_dim up to // 128 so each head is exactly one turbo block (matches the graph-level - // padding in llm_graph_context::build_attn for the paged path). + // padding in llm_graph_context::build_attn for the paged path). TURBO4_64 + // is excluded (its block is exactly head_dim 64 — no pad needed). const bool paged_cache_is_turbo = (type_k == GGML_TYPE_TURBO2_0 || type_k == GGML_TYPE_TURBO3_0 || type_k == GGML_TYPE_TURBO4_0); if (paged_cache_is_turbo && head_dim % 128 != 0) { diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index c1be9eb5a99f..b450dc18e919 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -4167,6 +4167,24 @@ static void test_template_output_peg_parsers(bool detailed_debug) { )) .run(); + // Missing opening : the model dives straight into reasoning and only closes + // with . Must still be captured as reasoning, not bled into content + // (regression test for the n_thinking=0 bleed on plain chat). + tst.test("I'm\nthinkingHello, world!\nWhat's up?") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(message_assist_thoughts) + .run(); + + // Missing opening with a trailing tool call after the reasoning block. + tst.test("I need to call a function" + "Let me check the time.<|tool_call_start|>[get_time(city=\"Paris\")]<|tool_call_end|>") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ get_time_tool }) + .expect(message_with_reasoning_content_and_multiple_tool_calls( + "I need to call a function", "Let me check the time.", { { "get_time", "{\"city\":\"Paris\"}" } } + )) + .run(); + // Fake tool call marker inside reasoning is not parsed as a call tst.test("Let me think about <|tool_call_start|>[special_function(arg1=1)]<|tool_call_end|> hmm" "<|tool_call_start|>[special_function(arg1=1)]<|tool_call_end|>")