From a278979cb61d4b6c50b65911d94f366a821bfa84 Mon Sep 17 00:00:00 2001 From: code-martin <44068401+code-martin@users.noreply.github.com> Date: Sun, 27 Sep 2026 11:29:08 +0300 Subject: [PATCH] feat(kv): add Q4_0 KV cache mode with FWHT-256 Walsh-Hadamard rotation - Implement Q4_0 KV quantization and dequantization kernels with scale/offset - Implement fast Walsh-Hadamard transform (FWHT-256) rotation for outlier suppression - Add Q4_0 attention decode support in qsa_decode_attn - Add kv_q4_parity parity test verifying CPU vs CUDA numerical accuracy - Support --kv q4_0 in generate CLI and runtime layer allocation - Update .gitignore for build and repository hygiene --- .gitignore | 12 + CMakeLists.txt | 9 +- include/strata/core/layer.hpp | 6 + include/strata/kernels/kv_q4.hpp | 53 ++++ include/strata/kernels/qsa_decode_attn.hpp | 2 + src/core/layer.cpp | 176 ++++++++--- src/core/mtp.cpp | 16 +- src/core/verify.cpp | 20 +- src/kernels/cuda/kv_q4.cu | 268 ++++++++++++++++ src/kernels/cuda/qsa_decode_attn.cu | 72 +++-- src/kernels/kv_q4_parity.cpp | 346 +++++++++++++++++++++ src/prefill/prefill.cpp | 31 +- src/program/generate.cpp | 5 +- 13 files changed, 937 insertions(+), 79 deletions(-) create mode 100644 include/strata/kernels/kv_q4.hpp create mode 100644 src/kernels/cuda/kv_q4.cu create mode 100644 src/kernels/kv_q4_parity.cpp diff --git a/.gitignore b/.gitignore index e8501e9f716..0319e37a790 100644 --- a/.gitignore +++ b/.gitignore @@ -46,3 +46,15 @@ build-vision.bat /dist/ *.done /third_party/*.zip + +# Local test artifacts, scratch, and binaries +Test/ +tools/aria2c* +tools/*.exe +tools/*.reg +tools/bench_contexts.py +download_shards.py +run_kv_q4_test.bat +*.png +.vscode/ +.idea/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 9a7260b57e5..a8617e3d493 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,4 +1,4 @@ -# Strata Engine - top-level build (P1.S1). +# Strata Engine - top-level build (P1.S1). # # Targets (phase-1 §P1.S1): `strata_artifact` (no CUDA dependency), `tests`. `strata_plan` and the C++ # `strata-pack` are added when their sources exist - a target with no sources would make the build graph @@ -178,6 +178,7 @@ if(STRATA_ENABLE_CUDA) src/kernels/cuda/ple.cu src/kernels/cuda/kv_q8.cu src/kernels/cuda/kv_stream.cu + src/kernels/cuda/kv_q4.cu src/kernels/cuda/fused_gr.cu src/kernels/cuda/qsa_decode_attn.cu src/kernels/cuda/fused_gdn.cu @@ -363,6 +364,12 @@ if(STRATA_ENABLE_CUDA) target_link_libraries(kv_stream_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME kv_stream_parity COMMAND kv_stream_parity) + # Q4_0 KV storage with Walsh-Hadamard rotation vs host reference (GPU, synthetic, no model). + add_executable(kv_q4_parity src/kernels/kv_q4_parity.cpp) + target_include_directories(kv_q4_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(kv_q4_parity PRIVATE strata_kernels CUDA::cudart) + add_test(NAME kv_q4_parity COMMAND kv_q4_parity) + # Plan v0.3 P3: native MMVQ for 1..8 tokens (bitwise selftest + --bench timing); synthetic data, no model. if(STRATA_BUILD_TESTS) add_executable(native_mmvq_multi bench/micro/native_mmvq_multi.cpp) diff --git a/include/strata/core/layer.hpp b/include/strata/core/layer.hpp index 1ea79220199..0eeecf277e0 100644 --- a/include/strata/core/layer.hpp +++ b/include/strata/core/layer.hpp @@ -203,6 +203,10 @@ struct QsaState { uint16_t* v_pool = nullptr; /// Plan v0.3 P7: INT8 KV (qsa_set_kv_int8). Codes [page][kv_head][page_size][head_dim], one FP16 scale per 64. bool kv_int8 = false; + /// Q4_0 KV with Walsh-Hadamard rotation (qsa_set_kv_q4). + bool kv_q4 = false; + uint8_t* k_q4 = nullptr; + uint8_t* v_q4 = nullptr; int8_t* k_q = nullptr; int8_t* v_q = nullptr; uint16_t* k_scale = nullptr; @@ -269,6 +273,8 @@ uint64_t qsa_kv_host_bytes(); /// initializing the session; default off until gate G-C accepts it. void qsa_set_kv_int8(bool enabled); bool qsa_kv_int8(); +void qsa_set_kv_q4(bool enabled); +bool qsa_kv_q4(); uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, QsaState& st, const QsaState* share_rope = nullptr, int64_t ring_cells = 0); /// KV streaming: the pools a reader sees (the VRAM slots) and, when streamed, make the selection's blocks resident. diff --git a/include/strata/kernels/kv_q4.hpp b/include/strata/kernels/kv_q4.hpp new file mode 100644 index 00000000000..a75f2812d2b --- /dev/null +++ b/include/strata/kernels/kv_q4.hpp @@ -0,0 +1,53 @@ +// include/strata/kernels/kv_q4.hpp - Q4_0 KV storage with Walsh-Hadamard rotation for QSA layers. +#pragma once + +#include "strata/kernels/qsa.hpp" +#include +#include + +namespace strata::kernels { + +inline constexpr int QK4_0 = 32; + +#pragma pack(push, 1) +struct block_q4_0 { + uint16_t d; // delta (half fp16 scale) + uint8_t qs[QK4_0 / 2]; // 16 bytes = 32 4-bit nibbles +}; +#pragma pack(pop) + +static_assert(sizeof(block_q4_0) == 18, "block_q4_0 must be 18 bytes"); + +/// Bytes per cell (one token, one KV head): 8 blocks of 32 for head_dim=256 -> 144 bytes per head. +inline uint64_t kv_q4_bytes_per_head(int head_dim) { + return (uint64_t) (head_dim / QK4_0) * sizeof(block_q4_0); +} + +inline uint64_t kv_q4_bytes_per_cell(const QsaShapes& s) { + return (uint64_t) s.n_head_kv * kv_q4_bytes_per_head((int) s.head_dim) * 2; // K and V +} + +/// Orthonormal Fast Walsh-Hadamard Transform (FWHT) for dimension 256. +/// Self-inverse: H * H * x = x. Scale is 1 / sqrt(256) = 1/16. +void fwht256_cuda(const float* src, float* dst, int64_t n_rows, void* stream); + +/// In-place FWHT256 +inline void fwht256_inplace_cuda(float* data, int64_t n_rows, void* stream) { + fwht256_cuda(data, data, n_rows, stream); +} + +/// Append the cell at step[kStepPos] into Q4_0 pool with Hadamard rotation. +void kv_append_q4_step(uint8_t* k_q4, uint8_t* v_q4, const int32_t* page_table, + const int32_t* step, const float* kcur, const float* vcur, + const QsaShapes& s, void* stream); + +/// Batched append of T cells into Q4_0 pool. +void kv_append_q4(uint8_t* k_q4, uint8_t* v_q4, const int32_t* page_table, int64_t pos0, int64_t T, + const float* K, const float* V, const QsaShapes& s, void* stream); + +/// Gather step[kStepWidth] cells named by `ids` from Q4_0 pool into FP16 scratch [id][kv_head][head_dim]. +void kv_gather_q4_step(const uint8_t* k_q4, const uint8_t* v_q4, + const int32_t* page_table, const int32_t* ids, const int32_t* step, int64_t max_ids, + const QsaShapes& s, uint16_t* k_scratch, uint16_t* v_scratch, void* stream); + +} // namespace strata::kernels diff --git a/include/strata/kernels/qsa_decode_attn.hpp b/include/strata/kernels/qsa_decode_attn.hpp index 2044ac5c953..45c0f301220 100644 --- a/include/strata/kernels/qsa_decode_attn.hpp +++ b/include/strata/kernels/qsa_decode_attn.hpp @@ -27,6 +27,8 @@ struct QsaAttnPools { const int8_t* v_q = nullptr; const uint16_t* k_scale = nullptr; ///< fp16 [page][kv_head][page_size][head_dim / 64] const uint16_t* v_scale = nullptr; + const uint8_t* k_q4 = nullptr; ///< q4_0 block_q4_0 [page][kv_head][page_size][head_dim / 32 * 18] + const uint8_t* v_q4 = nullptr; const int32_t* page_table = nullptr; }; diff --git a/src/core/layer.cpp b/src/core/layer.cpp index c39d195ba83..4b223222b71 100644 --- a/src/core/layer.cpp +++ b/src/core/layer.cpp @@ -8,6 +8,7 @@ #include "strata/kernels/gdn.hpp" #include "strata/kernels/qsa.hpp" #include "strata/kernels/kv_q8.hpp" +#include "strata/kernels/kv_q4.hpp" #include "strata/kernels/quantize_act.hpp" #include "strata/kernels/rope.hpp" #include "strata/kernels/router_top10.hpp" @@ -170,9 +171,11 @@ if (!w.wants_q8k()) { s_gemv_q8_0_split(x80, p.codes, p.scales, p.offset, // namespace void layer_set_native_bf16(bool enabled) { native_bf16_projections = enabled; } void layer_set_native_flash_attn_short(bool enabled) { native_flash_attn_short = enabled; } -namespace { bool g_kv_int8 = false; } +namespace { bool g_kv_int8 = false; bool g_kv_q4 = false; } void qsa_set_kv_int8(bool enabled) { g_kv_int8 = enabled; } bool qsa_kv_int8() { return g_kv_int8; } +void qsa_set_kv_q4(bool enabled) { g_kv_q4 = enabled; } +bool qsa_kv_q4() { return g_kv_q4; } uint64_t gdn_buffers_bytes(const ModelGeometry& g) { const int64_t C = g.ssm_conv_channels; const int64_t V = g.ssm_value_dim; const uint64_t parts[] = { q8k_bytes(g.n_embd), // x_q8k (uint64_t) (g.n_embd / 32) * 34, @@ -507,6 +510,9 @@ KvPlan kv_plan(const QsaShapes& s, int64_t max_cells, int64_t ring_cells) { return p; } uint64_t kv_pool_bytes(const QsaShapes& s, int64_t pages) { + if (g_kv_q4) { + return (uint64_t) pages * s.page_size * strata::kernels::kv_q4_bytes_per_cell(s) + 64; + } return g_kv_int8 ? (uint64_t) pages * s.page_size * strata::kernels::kv_q8_bytes_per_cell(s) + 64 : (uint64_t) pages * s.n_head_kv * s.page_size * s.head_dim * 2 * 2; } @@ -541,10 +547,15 @@ uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, Q const int64_t pages = p.pages; Cursor c{(uint8_t*) base}; st.kv_int8 = g_kv_int8; + st.kv_q4 = g_kv_q4; st.kv_mode = p.mode; st.n_slots = p.slots; const uint64_t rows = (uint64_t) p.slots * s.n_head_kv * s.page_size; // VRAM rows: the slots - if (g_kv_int8) { + if (g_kv_q4) { + const uint64_t pool_bytes = (uint64_t) p.slots * s.n_head_kv * s.page_size * strata::kernels::kv_q4_bytes_per_head((int) s.head_dim); + st.k_q4 = c.take_bytes(pool_bytes); + st.v_q4 = c.take_bytes(pool_bytes); + } else if (g_kv_int8) { st.k_q = c.take(rows * s.head_dim); st.v_q = c.take(rows * s.head_dim); st.k_scale = c.take(rows * (s.head_dim / strata::kernels::KV_Q8_GROUP)); @@ -644,13 +655,18 @@ uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, Q void qsa_state_zero(const QsaState& st, const ModelGeometry& g, void* stream) { const QsaShapes s = qsa_shapes(g); cudaStream_t cs = (cudaStream_t) stream; - const size_t rows = (size_t) st.n_slots * s.n_head_kv * s.page_size; - if (st.kv_int8) { + if (st.kv_q4) { + const size_t pool_bytes = (size_t) st.n_slots * s.n_head_kv * s.page_size * strata::kernels::kv_q4_bytes_per_head((int) s.head_dim); + cudaMemsetAsync(st.k_q4, 0, pool_bytes, cs); + cudaMemsetAsync(st.v_q4, 0, pool_bytes, cs); + } else if (st.kv_int8) { + const size_t rows = (size_t) st.n_slots * s.n_head_kv * s.page_size; cudaMemsetAsync(st.k_q, 0, rows * s.head_dim, cs); cudaMemsetAsync(st.v_q, 0, rows * s.head_dim, cs); cudaMemsetAsync(st.k_scale, 0, rows * (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2, cs); cudaMemsetAsync(st.v_scale, 0, rows * (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2, cs); } else { + const size_t rows = (size_t) st.n_slots * s.n_head_kv * s.page_size; cudaMemsetAsync(st.k_pool, 0, rows * s.head_dim * 2, cs); cudaMemsetAsync(st.v_pool, 0, rows * s.head_dim * 2, cs); } @@ -666,7 +682,8 @@ void qsa_state_zero(const QsaState& st, const ModelGeometry& g, void* stream) { strata::kernels::QsaAttnPools qsa_attn_pools(const QsaState& st) { strata::kernels::QsaAttnPools pools; pools.page_table = st.page_table; - if (st.kv_int8) { pools.k_q = st.k_q; pools.v_q = st.v_q; pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; } + if (st.kv_q4) { pools.k_q4 = st.k_q4; pools.v_q4 = st.v_q4; } + else if (st.kv_int8) { pools.k_q = st.k_q; pools.v_q = st.v_q; pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; } else { pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; } return pools; } @@ -677,7 +694,6 @@ void qsa_kv_resolve(const QsaState& st, const ModelGeometry& g, const int32_t* i strata::kernels::kv_stream_resolve(st.map, qsa_attn_pools(st), st.host, st.kv_int8, ids, steps, n_q, cap, qsa_shapes(g), stream); } - uint64_t qsa_buffers_bytes(const ModelGeometry& g, int64_t max_cells) { const QsaShapes s = qsa_shapes(g); const int64_t cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s); uint64_t n = 0; n += (uint64_t) q8k_bytes(g.n_embd); n += (uint64_t) (g.n_embd / 32) * 34; // block_q8_0 n += (uint64_t) g.n_embd * 2; n += (uint64_t) g.n_head * 2 * g.head_dim * 4; n += (uint64_t) g.n_head * g.head_dim * 4; n += (uint64_t) g.n_head_kv * g.head_dim * 4 * 2; n += (uint64_t) g.idx_key_dim * 4; n += (uint64_t) g.idx_q_heads * g.idx_key_dim * 4; n += (uint64_t) max_cells * 4; n += (uint64_t) cap * 4; n += (uint64_t) cap * g.n_head_kv * g.head_dim * 2 * 2; n += (uint64_t) g.n_head * g.head_dim * 4; @@ -846,62 +862,132 @@ if (!gemv_quantized(*w_attnk, p_k, f_k, b.x_q8_0, b.x_q8k, b.kcur, g.n_embd, g.n // EVERY ENTRY POINT FROM HERE ON IS THE CAPTURABLE ONE: the per-token counts come from `st.step` and every // launch is sized from a capacity in `st`/`b`, so this sequence can be captured and replayed. The // host-scalar wrappers would be correct here today and silently wrong in a graph. -if (st.kv_int8) kv_append_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, st.step, b.kcur, b.vcur, s, stream, &st.host); else kv_append_step(st.k_pool, st.v_pool, st.page_table, st.step, b.kcur, b.vcur, s, stream, &st.host); { const uint64_t nvk = (uint64_t) g.n_head_kv * g.head_dim; const uint64_t base = (uint64_t) 2 * g.n_embd + 2 * g.hc + (uint64_t) g.n_head * g.head_dim + 2 * nvk + 8; if (!st.kv_int8) dump_slot(dump, g, layer, (const float*) st.k_pool, base, nvk / 2, stream); if (!st.kv_int8) dump_slot(dump, g, layer, (const float*) st.v_pool, base + nvk / 2, nvk / 2, stream); dump_slot(dump, g, layer, b.vcur, base + nvk, nvk, stream); dump_slot(dump, g, layer, b.kcur, base + 2 * nvk, nvk, stream); } { const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; if (native_qsa_indexer_enabled()) { - try { - native_qsa_indexer_append(b.idx_raw, st.step + kStepPos, pos_base, - (const float*) w_ikn->data, RMS_EPS, ib, s, st.max_cells, (float) qsa_freq_base(), stream); - } catch (const std::exception& error) { err = v.name("native_indexer") + ": " + error.what(); return false; } -} else indexer_key_append(b.idx_raw, st.pos_dev, pos_base, (const float*) w_ikn->data, RMS_EPS, ib, s, - st.cos_tab, st.sin_tab, stream); } +if (st.kv_q4) { + // In-place Hadamard rotation for K and V prior to 4-bit quantization + fwht256_inplace_cuda(b.kcur, g.n_head_kv, stream); + fwht256_inplace_cuda(b.vcur, g.n_head_kv, stream); + kv_append_q4_step(st.k_q4, st.v_q4, st.page_table, st.step, b.kcur, b.vcur, s, stream); +} else if (st.kv_int8) { + kv_append_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, st.step, b.kcur, b.vcur, s, stream, &st.host); +} else { + kv_append_step(st.k_pool, st.v_pool, st.page_table, st.step, b.kcur, b.vcur, s, stream, &st.host); +} +{ + const uint64_t nvk = (uint64_t) g.n_head_kv * g.head_dim; + const uint64_t base = (uint64_t) 2 * g.n_embd + 2 * g.hc + (uint64_t) g.n_head * g.head_dim + 2 * nvk + 8; + if (!st.kv_int8 && !st.kv_q4) dump_slot(dump, g, layer, (const float*) st.k_pool, base, nvk / 2, stream); + if (!st.kv_int8 && !st.kv_q4) dump_slot(dump, g, layer, (const float*) st.v_pool, base + nvk / 2, nvk / 2, stream); + dump_slot(dump, g, layer, b.vcur, base + nvk, nvk, stream); + dump_slot(dump, g, layer, b.kcur, base + 2 * nvk, nvk, stream); +} +{ + const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; + if (native_qsa_indexer_enabled()) { + try { + native_qsa_indexer_append(b.idx_raw, st.step + kStepPos, pos_base, + (const float*) w_ikn->data, RMS_EPS, ib, s, st.max_cells, (float) qsa_freq_base(), stream); + } catch (const std::exception& error) { err = v.name("native_indexer") + ": " + error.what(); return false; } + } else { + indexer_key_append(b.idx_raw, st.pos_dev, pos_base, (const float*) w_ikn->data, RMS_EPS, ib, s, + st.cos_tab, st.sin_tab, stream); + } +} // ---- 6. the query projection, and the HALF-SPLIT into q and gate. -// -// **`attn_q` IS *NOT* ALWAYS Q2_0.** It is Q2_0 on TWO of the twelve QSA layers, Q3_K on six and IQ4_XS on -// four, and this line used to call `s2_gemv_q8` unconditionally - so TEN of the twelve layers decoded an -// S4 code plane with a Q8_0 activation. Both are "a quantized activation", both produce a plausible -// vector, and the difference is 0.6-1.4%. -if (!gemv_quantized(*w_attnq, p_q, f_q, b.x_q8_0, b.x_q8k, b.q_full, g.n_embd, g.n_head * 2 * g.head_dim, v.name("attn_q.weight"), stream, err, x, w_attnv->native_data && w_attnq->native_data)) return false; -// `per_head[:, :head_dim]` - the FIRST half of each head's 2*head_dim block, copied out contiguously so -// the norm and the rotation see whole rows. A 2-D copy is a memcpy node, which captures (`pinned_capture` -// case A) and needs no kernel. -if (cudaMemcpy2DAsync(b.qcur, (size_t) g.head_dim * 4, b.q_full, (size_t) g.head_dim * 2 * 4, (size_t) g.head_dim * 4, (size_t) g.n_head, cudaMemcpyDeviceToDevice, (cudaStream_t) stream) != cudaSuccess) { err = "qsa_layer: the q/gate split failed"; return false; } if (!normalize_rotate(b.qcur, w_qn, (int) g.n_head, (int) g.head_dim)) return false; +if (!gemv_quantized(*w_attnq, p_q, f_q, b.x_q8_0, b.x_q8k, b.q_full, g.n_embd, g.n_head * 2 * g.head_dim, + v.name("attn_q.weight"), stream, err, x, w_attnv->native_data && w_attnq->native_data)) return false; + +if (cudaMemcpy2DAsync(b.qcur, (size_t) g.head_dim * 4, b.q_full, (size_t) g.head_dim * 2 * 4, + (size_t) g.head_dim * 4, (size_t) g.n_head, cudaMemcpyDeviceToDevice, + (cudaStream_t) stream) != cudaSuccess) { + err = "qsa_layer: the q/gate split failed"; + return false; +} +if (!normalize_rotate(b.qcur, w_qn, (int) g.n_head, (int) g.head_dim)) return false; +if (st.kv_q4) { + // In-place Hadamard rotation for Q so Q*K^T is invariant + fwht256_inplace_cuda(b.qcur, g.n_head, stream); +} // ---- 7. the indexer's query: BF16, then norm and rotate project_bf16(x, b.x_bf16, (const uint16_t*) w_idxq->data, b.q_idx, g.n_embd, g.idx_q_heads * g.idx_key_dim, false, stream); if (!normalize_rotate(b.q_idx, w_iqn, (int) g.idx_q_heads, (int) g.idx_key_dim)) return false; -// ---- 8. score, select, gather, attend. `max_blocks` and `cap` are CAPACITIES from the state, not this -// token's counts: a grid or a shared-memory size that follows the sequence length is baked into a captured -// graph, and the kernels guard for the surplus. const -int64_t max_blocks = (st.max_cells / s.idx_block) + 2; - if (g_fast_select) { - // Plan v0.3 P7: block-level FP32 scores and a radix selection over blocks (qsa_select.hpp). - qsa_block_scores(st.idx_pooled, st.idx_dead, b.q_idx, st.step, 1, max_blocks, s, b.cell_scores, stream); - qsa_block_topk(b.cell_scores, st.step, 1, max_blocks, cap, s, b.ids, stream); +// ---- 8. score, select, gather, attend. +const int64_t max_blocks = (st.max_cells / s.idx_block) + 2; +if (g_fast_select) { + qsa_block_scores(st.idx_pooled, st.idx_dead, b.q_idx, st.step, 1, max_blocks, s, b.cell_scores, stream); + qsa_block_topk(b.cell_scores, st.step, 1, max_blocks, cap, s, b.ids, stream); +} else { + qsa_index_step(st.idx_pooled, b.q_idx, nullptr, s, st.step, max_blocks, b.cell_scores, stream); + topk_512_step(b.cell_scores, s, cap, st.step, b.ids, stream); +} +if (g_fast_attn && !native_flash_attn_short && dump == nullptr) { + strata::kernels::QsaAttnPools pools; + pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; pools.k_q = st.k_q; pools.v_q = st.v_q; + pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; pools.page_table = st.page_table; + pools.k_q4 = st.k_q4; pools.v_q4 = st.v_q4; + if (!st.kv_int8) { pools.k_q = nullptr; pools.v_q = nullptr; } + if (!st.kv_q4) { pools.k_q4 = nullptr; pools.v_q4 = nullptr; } + strata::kernels::qsa_decode_attn_step(b.qcur, pools, b.ids, st.step, cap, s, b.attn_scratch, b.attn, stream); + if (st.kv_q4) { + // Un-rotate attention output: A = H * A_rot + fwht256_inplace_cuda(b.attn, g.n_head, stream); + } +} else { + if (st.kv_q4) kv_gather_q4_step(st.k_q4, st.v_q4, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); + else if (st.kv_int8) kv_gather_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, b.ids, st.step, cap, s, + b.k_scratch, b.v_scratch, stream); + else kv_gather_step(st.k_pool, st.v_pool, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); + if (native_flash_attn_short) { + if (st.max_cells < 1 || st.max_cells > 256 || !st.attention_status || !st.host_step) { + err = v.name("native_flash_attn") + ": short adapter requires context <=256 and persistent status storage"; + return false; + } + try { + native_flash_attn_short_step(b.qcur, b.k_scratch, b.v_scratch, st.step, cap, + (int) st.max_cells, s, b.attn, st.attention_status, nullptr, stream); + } catch (const std::exception& error) { err = v.name("native_flash_attn") + ": " + error.what(); return false; } + if (cudaMemcpyAsync(st.host_step + kStepCount, st.attention_status, sizeof(int32_t), + cudaMemcpyDeviceToHost, (cudaStream_t) stream) != cudaSuccess) { + err = v.name("native_flash_attn") + ": status readback failed"; return false; + } } else { - qsa_index_step(st.idx_pooled, b.q_idx, nullptr, s, st.step, max_blocks, b.cell_scores, stream); topk_512_step(b.cell_scores, s, cap, st.step, b.ids, stream); + qsa_attend_step(b.qcur, b.k_scratch, b.v_scratch, st.step, cap, s, b.attn, nullptr, stream); } // KV streaming: every block the selection names is made resident before anything reads it qsa_kv_resolve(st, g, b.ids, st.step, 1, cap, stream); if (g_fast_attn && !native_flash_attn_short && dump == nullptr) { strata::kernels::QsaAttnPools pools; pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; pools.k_q = st.k_q; pools.v_q = st.v_q; + pools.k_q4 = st.k_q4; pools.v_q4 = st.v_q4; pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; pools.page_table = st.page_table; if (!st.kv_int8) { pools.k_q = nullptr; pools.v_q = nullptr; } + if (!st.kv_q4) { pools.k_q4 = nullptr; pools.v_q4 = nullptr; } strata::kernels::qsa_decode_attn_step(b.qcur, pools, b.ids, st.step, cap, s, b.attn_scratch, b.attn, stream); } else { - if (st.kv_int8) kv_gather_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); else kv_gather_step(st.k_pool, st.v_pool, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); if (native_flash_attn_short) { - if (st.max_cells < 1 || st.max_cells > 256 || !st.attention_status || !st.host_step) { - err = v.name("native_flash_attn") + ": short adapter requires context <=256 and persistent status storage"; - return false; - } - try { - native_flash_attn_short_step(b.qcur, b.k_scratch, b.v_scratch, st.step, cap, - (int) st.max_cells, s, b.attn, st.attention_status, nullptr, stream); - } catch (const std::exception& error) { err = v.name("native_flash_attn") + ": " + error.what(); return false; } - if (cudaMemcpyAsync(st.host_step + kStepCount, st.attention_status, sizeof(int32_t), - cudaMemcpyDeviceToHost, (cudaStream_t) stream) != cudaSuccess) { - err = v.name("native_flash_attn") + ": status readback failed"; return false; + if (st.kv_q4) kv_gather_q4_step(st.k_q4, st.v_q4, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); + else if (st.kv_int8) kv_gather_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); + else kv_gather_step(st.k_pool, st.v_pool, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); + if (native_flash_attn_short) { + if (st.max_cells < 1 || st.max_cells > 256 || !st.attention_status || !st.host_step) { + err = v.name("native_flash_attn") + ": short adapter requires context <=256 and persistent status storage"; + return false; + } + try { + native_flash_attn_short_step(b.qcur, b.k_scratch, b.v_scratch, st.step, cap, + (int) st.max_cells, s, b.attn, st.attention_status, nullptr, stream); + } catch (const std::exception& error) { err = v.name("native_flash_attn") + ": " + error.what(); return false; } + if (cudaMemcpyAsync(st.host_step + kStepCount, st.attention_status, sizeof(int32_t), + cudaMemcpyDeviceToHost, (cudaStream_t) stream) != cudaSuccess) { + err = v.name("native_flash_attn") + ": status readback failed"; return false; + } + } else { + qsa_attend_step(b.qcur, b.k_scratch, b.v_scratch, st.step, cap, s, b.attn, nullptr, stream); + } } -} else qsa_attend_step(b.qcur, b.k_scratch, b.v_scratch, st.step, cap, s, b.attn, nullptr, stream); + if (st.kv_q4) { + // Un-rotate attention output + fwht256_inplace_cuda(b.attn, g.n_head, stream); } +} dump_slot(dump, g, layer, b.attn, (uint64_t) 2 * g.n_embd + 2 * g.hc, (uint64_t) g.n_head * g.head_dim, stream); { const uint64_t vs = (uint64_t) 2 * g.n_embd + 2 * g.hc + (uint64_t) g.n_head * g.head_dim + (uint64_t) 2 * g.n_head_kv * g.head_dim; dump_slot(dump, g, layer, (const float*) b.v_scratch, vs, (uint64_t) g.n_head_kv * g.head_dim / 2, stream); dump_slot(dump, g, layer, (const float*) b.ids, vs + (uint64_t) g.n_head_kv * g.head_dim / 2, 4, stream); dump_slot(dump, g, layer, (const float*) st.step, vs + (uint64_t) g.n_head_kv * g.head_dim / 2 + 4, 4, stream); } // ---- 9. Gate the attention output before its projection. Native CUDA keeps F32 // sigmoid/multiply arithmetic and uses Q8_1 for the native quantized projection. diff --git a/src/core/mtp.cpp b/src/core/mtp.cpp index 3c7fce68c7c..a353fcbaaf0 100644 --- a/src/core/mtp.cpp +++ b/src/core/mtp.cpp @@ -8,6 +8,7 @@ #include "strata/kernels/fused_gr.hpp" #include "strata/kernels/gr.hpp" #include "strata/kernels/kv_q8.hpp" +#include "strata/kernels/kv_q4.hpp" #include "strata/kernels/native_moe.hpp" #include "strata/kernels/native_mmvq.hpp" #include "strata/kernels/native_qsa.hpp" @@ -338,7 +339,12 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri native_mmvq(GGML_Q8_0, q8("self_attn.v_proj.weight"), xq_, vcur_, (int) N, (int) (NKV * HD), T, cs); for (int t = 0; t < T; ++t) { norm_rope(kcur_ + t * NKV * HD, f32("self_attn.k_norm.weight"), (int) NKV, (int) HD, pos + t * NH); - if (st_.kv_int8) + if (st_.kv_q4) { + fwht256_inplace_cuda(kcur_ + t * NKV * HD, NKV, cs); + fwht256_inplace_cuda(vcur_ + t * NKV * HD, NKV, cs); + kv_append_q4_step(st_.k_q4, st_.v_q4, st_.page_table, step + t * 4, + kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs); + } else if (st_.kv_int8) kv_append_q8_step(st_.k_q, st_.v_q, st_.k_scale, st_.v_scale, st_.page_table, step + t * 4, kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs, &st_.host); else @@ -355,10 +361,18 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri return false; } norm_rope(qc, f32("self_attn.q_norm.weight"), (int) NH, (int) HD, pos + t * NH); + if (st_.kv_q4) { + fwht256_inplace_cuda(qc, NH, cs); + } } const QsaAttnPools pools = qsa_attn_pools(st_); if (window_ > 0) window_ids(const_cast(step), T, (int) window_, ident_, cap_, cs); qsa_decode_attn_batch(qcur_, pools, ident_, step, cap_, s, attn_scratch_, attn_, T, cs); + if (st_.kv_q4) { + for (int t = 0; t < T; ++t) { + fwht256_inplace_cuda(attn_ + t * NH * HD, NH, cs); + } + } for (int t = 0; t < T; ++t) native_qsa_gate_apply(attn_ + t * NH * HD, qfull_ + t * NH * 2 * HD, attn32_ + t * NH * HD, (int) NH, (int) HD, cs); native_quantize_q8_1(attn32_, xq_, (int) (NH * HD), T, cs); diff --git a/src/core/verify.cpp b/src/core/verify.cpp index 70b089b775a..23673d04a38 100644 --- a/src/core/verify.cpp +++ b/src/core/verify.cpp @@ -13,6 +13,7 @@ #include "strata/kernels/fused_gr.hpp" #include "strata/kernels/gr.hpp" #include "strata/kernels/kv_q8.hpp" +#include "strata/kernels/kv_q4.hpp" #include "strata/kernels/native_mmvq.hpp" #include "strata/kernels/native_qsa.hpp" #include "strata/kernels/native_qsa_indexer.hpp" @@ -421,10 +422,19 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { native_mmvq(wk->native_type, wk->native_data, xq_, kcur_ + tb * NKV * HD, (int) N, (int) (NKV * HD), n, cs); native_mmvq(wv->native_type, wv->native_data, xq_, vcur_ + tb * NKV * HD, (int) N, (int) (NKV * HD), n, cs); for (int t = tb; t < te; ++t) norm_rope(kcur_ + t * NKV * HD, wkn, (int) NKV, (int) HD, pos_ + t * NH); + if (st.kv_q4) { + for (int t = tb; t < te; ++t) { + fwht256_inplace_cuda(kcur_ + t * NKV * HD, NKV, cs); + fwht256_inplace_cuda(vcur_ + t * NKV * HD, NKV, cs); + } + } if (grp == 0) copy_from_mapped(tail_snap_ + (size_t) qi * TS, st.idx_tail, TS, cs); for (int t = tb; t < te; ++t) { const int32_t* step_t = step_ + t * kStepCount; - if (st.kv_int8) + if (st.kv_q4) + kv_append_q4_step(st.k_q4, st.v_q4, st.page_table, step_t, + kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs); + else if (st.kv_int8) kv_append_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, step_t, kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs, &st.host); else @@ -446,6 +456,9 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { return false; } norm_rope(qc, wqn, (int) NH, (int) HD, pos_ + t * NH); + if (st.kv_q4) { + fwht256_inplace_cuda(qc, NH, cs); + } } for (int t = tb; t < te; ++t) { float* qx = qidx_ + t * IQ * ID; @@ -461,6 +474,11 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { const QsaAttnPools pools = qsa_attn_pools(st); qsa_decode_attn_batch(qcur_ + tb * NH * HD, pools, sel_ + (size_t) tb * cap_, step_ + tb * kStepCount, cap_, s, attn_scratch_ + (size_t) tb * attn_scratch_floats_, attn_ + tb * NH * HD, n, cs); + if (st.kv_q4) { + for (int t = tb; t < te; ++t) { + fwht256_inplace_cuda(attn_ + t * NH * HD, NH, cs); + } + } for (int t = tb; t < te; ++t) { if (native_qsa_enabled()) native_qsa_gate_apply(attn_ + t * NH * HD, qfull_ + t * NH * 2 * HD, attn32_ + t * NH * HD, diff --git a/src/kernels/cuda/kv_q4.cu b/src/kernels/cuda/kv_q4.cu new file mode 100644 index 00000000000..92bfe05f487 --- /dev/null +++ b/src/kernels/cuda/kv_q4.cu @@ -0,0 +1,268 @@ +// src/kernels/cuda/kv_q4.cu - Q4_0 KV storage with Fast Walsh-Hadamard Transform +#include "strata/kernels/kv_q4.hpp" +#include "strata/kernels/f16_bits.hpp" + +#include +#include +#include +#include + +namespace strata::kernels { +namespace { + +void check(const char* what) { + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { + std::fprintf(stderr, "kv_q4: %s: %s\n", what, cudaGetErrorString(e)); + std::exit(1); + } +} + +// Fast Walsh-Hadamard Transform for N = 256 +// 1 warp of 32 threads computes FWHT on 256 floats in registers +// Scale = 1 / sqrt(256) = 1/16 = 0.0625f +__global__ void fwht256_kernel(const float* __restrict__ src, float* __restrict__ dst, int64_t n_rows, float scale) { + constexpr int warp_size = 32; + constexpr int N = 256; + constexpr int el_w = N / warp_size; // 8 + + const int64_t r = (int64_t) blockIdx.x * blockDim.y + threadIdx.y; + if (r >= n_rows) return; + + const float* row_src = src + r * N; + float* row_dst = dst + r * N; + + float reg[el_w]; + const int lane = threadIdx.x; + +#pragma unroll + for (int i = 0; i < el_w; ++i) { + reg[i] = row_src[i * warp_size + lane] * scale; + } + +#pragma unroll + for (int h = 1; h < warp_size; h *= 2) { +#pragma unroll + for (int j = 0; j < el_w; ++j) { + const float val = reg[j]; + const float val2 = __shfl_xor_sync(0xffffffffu, val, h, warp_size); + reg[j] = (lane & h) == 0 ? val + val2 : val2 - val; + } + } + +#pragma unroll + for (int h = warp_size; h < N; h *= 2) { + const int step = h / warp_size; +#pragma unroll + for (int j = 0; j < el_w; j += 2 * step) { +#pragma unroll + for (int k = 0; k < step; ++k) { + const float x = reg[j + k]; + const float y = reg[j + k + step]; + reg[j + k] = x + y; + reg[j + k + step] = x - y; + } + } + } + +#pragma unroll + for (int i = 0; i < el_w; ++i) { + row_dst[i * warp_size + lane] = reg[i]; + } +} + +// 1 block = 1 32-value group of one KV head of K (blockIdx.z == 0) or V (1); 32 threads, 1 value each +__global__ void kv_append_q4_kernel(uint8_t* __restrict__ k_q4, uint8_t* __restrict__ v_q4, + const int32_t* __restrict__ table, const int32_t* __restrict__ step, + const float* __restrict__ kcur, const float* __restrict__ vcur, + int kv_heads, int head_dim, int page_size) { + const long long pos = (long long) __ldg(step + kStepPos); + const int h = blockIdx.x, b = blockIdx.y, t = threadIdx.x; + const bool is_v = blockIdx.z == 1; + const int blocks_per_head = head_dim / QK4_0; // 8 for 256 + const int bytes_per_head = blocks_per_head * sizeof(block_q4_0); // 144 + const float x = (is_v ? vcur : kcur)[h * head_dim + b * QK4_0 + t]; + + float amax = fabsf(x); + float max_val = x; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) { + const float a = __shfl_xor_sync(0xffffffffu, amax, o); + const float v = __shfl_xor_sync(0xffffffffu, max_val, o); + if (a > amax) { + amax = a; + max_val = v; + } + } + + const float d = max_val / -8.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + const uint16_t d_f16 = f16_from_f32(d); + + const float x_scaled = x * id; + int q = __float2int_rz(x_scaled + 8.5f); + uint8_t q_clamped = (uint8_t) (q < 0 ? 0 : (q > 15 ? 15 : q)); + + const uint8_t q_high = __shfl_down_sync(0xffffffffu, q_clamped, 16); + + const long long page = (long long) table[pos / page_size]; + const long long row = (page * kv_heads + h) * page_size + (pos % page_size); + uint8_t* base = (is_v ? v_q4 : k_q4) + row * bytes_per_head; + block_q4_0* blk = reinterpret_cast(base) + b; + + if (t == 0) { + blk->d = d_f16; + } + if (t < 16) { + blk->qs[t] = q_clamped | (q_high << 4); + } +} + +// Gather step[kStepWidth] cells into FP16 scratch +__global__ void kv_gather_q4_kernel(const uint8_t* __restrict__ k_q4, const uint8_t* __restrict__ v_q4, + const int32_t* __restrict__ table, const int32_t* __restrict__ ids, + const int32_t* __restrict__ step, int kv_heads, int head_dim, int page_size, + uint16_t* __restrict__ k_scratch, uint16_t* __restrict__ v_scratch) { + const long long n_ids = (long long) __ldg(step + kStepWidth); + const int blocks_per_head = head_dim / QK4_0; // 8 + const int bytes_per_head = blocks_per_head * sizeof(block_q4_0); // 144 + const long long total_blocks = n_ids * kv_heads * blocks_per_head; + + const long long blk_idx = (long long) blockIdx.x * blockDim.y + threadIdx.y; + if (blk_idx >= total_blocks) return; + + const int t = threadIdx.x; // 0..31 inside block + const long long id = blk_idx / (kv_heads * blocks_per_head); + const int rem = (int) (blk_idx % (kv_heads * blocks_per_head)); + const int h = rem / blocks_per_head; + const int b = rem % blocks_per_head; + + const int cell = ids[id]; + const long long page = (long long) table[cell / page_size]; + const long long row = (page * kv_heads + h) * page_size + (cell % page_size); + + const block_q4_0* k_blk = reinterpret_cast(k_q4 + row * bytes_per_head) + b; + const block_q4_0* v_blk = reinterpret_cast(v_q4 + row * bytes_per_head) + b; + + const float kd = f32_from_f16(k_blk->d); + const float vd = f32_from_f16(v_blk->d); + + const int j = t < 16 ? t : (t - 16); + const uint8_t k_byte = k_blk->qs[j]; + const uint8_t v_byte = v_blk->qs[j]; + + const int kq = (t < 16) ? ((k_byte & 0x0F) - 8) : ((k_byte >> 4) - 8); + const int vq = (t < 16) ? ((v_byte & 0x0F) - 8) : ((v_byte >> 4) - 8); + + const long long dst_offset = ((id * kv_heads + h) * head_dim) + (b * QK4_0 + t); + k_scratch[dst_offset] = f16_from_f32((float) kq * kd); + v_scratch[dst_offset] = f16_from_f32((float) vq * vd); +} + +} // namespace + +void fwht256_cuda(const float* src, float* dst, int64_t n_rows, void* stream) { + if (n_rows <= 0) return; + const int rows_per_block = 4; + const int64_t num_blocks = (n_rows + rows_per_block - 1) / rows_per_block; + dim3 grid((unsigned) num_blocks); + dim3 block(32, rows_per_block); + const float scale = 1.0f / 16.0f; // 1 / sqrt(256) + fwht256_kernel<<>>(src, dst, n_rows, scale); + check("fwht256 launch"); +} + +// Batched append for prefill: grid = dim3(T, kv_heads, blocks_per_head), block = 32 +__global__ void kv_append_q4_batch_kernel(uint8_t* __restrict__ k_q4, uint8_t* __restrict__ v_q4, + const int32_t* __restrict__ table, int64_t pos0, + const float* __restrict__ K, const float* __restrict__ V, + int kv_heads, int head_dim, int page_size, int is_v_grid) { + const long long t = blockIdx.x; + const long long pos = pos0 + t; + const int h = blockIdx.y, b = blockIdx.z, th = threadIdx.x; + const bool is_v = is_v_grid != 0; + const int blocks_per_head = head_dim / QK4_0; // 8 for 256 + const int bytes_per_head = blocks_per_head * sizeof(block_q4_0); // 144 + const float x = (is_v ? V : K)[t * (kv_heads * head_dim) + h * head_dim + b * QK4_0 + th]; + + float amax = fabsf(x); + float max_val = x; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) { + const float a = __shfl_xor_sync(0xffffffffu, amax, o); + const float v = __shfl_xor_sync(0xffffffffu, max_val, o); + if (a > amax) { + amax = a; + max_val = v; + } + } + + const float d = max_val / -8.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + const uint16_t d_f16 = f16_from_f32(d); + + const float x_scaled = x * id; + int q = __float2int_rz(x_scaled + 8.5f); + uint8_t q_clamped = (uint8_t) (q < 0 ? 0 : (q > 15 ? 15 : q)); + + const uint8_t q_high = __shfl_down_sync(0xffffffffu, q_clamped, 16); + + const long long page = (long long) table[pos / page_size]; + const long long row = (page * kv_heads + h) * page_size + (pos % page_size); + uint8_t* base = (is_v ? v_q4 : k_q4) + row * bytes_per_head; + block_q4_0* blk = reinterpret_cast(base) + b; + + if (th == 0) { + blk->d = d_f16; + } + if (th < 16) { + blk->qs[th] = q_clamped | (q_high << 4); + } +} + +void kv_append_q4_step(uint8_t* k_q4, uint8_t* v_q4, const int32_t* page_table, + const int32_t* step, const float* kcur, const float* vcur, + const QsaShapes& s, void* stream) { + if (s.head_dim != 256) { + std::fprintf(stderr, "kv_append_q4: head_dim must be 256\n"); + std::exit(1); + } + const dim3 grid((unsigned) s.n_head_kv, (unsigned) (s.head_dim / QK4_0), 2); + kv_append_q4_kernel<<>>( + k_q4, v_q4, page_table, step, kcur, vcur, (int) s.n_head_kv, (int) s.head_dim, (int) s.page_size); + check("kv_append_q4 launch"); +} + +void kv_append_q4(uint8_t* k_q4, uint8_t* v_q4, const int32_t* page_table, int64_t pos0, int64_t T, + const float* K, const float* V, const QsaShapes& s, void* stream) { + if (T <= 0) return; + if (s.head_dim != 256) { + std::fprintf(stderr, "kv_append_q4: head_dim must be 256\n"); + std::exit(1); + } + const dim3 grid((unsigned) T, (unsigned) s.n_head_kv, (unsigned) (s.head_dim / QK4_0)); + cudaStream_t cs = (cudaStream_t) stream; + kv_append_q4_batch_kernel<<>>( + k_q4, v_q4, page_table, pos0, K, V, (int) s.n_head_kv, (int) s.head_dim, (int) s.page_size, 0); + kv_append_q4_batch_kernel<<>>( + k_q4, v_q4, page_table, pos0, K, V, (int) s.n_head_kv, (int) s.head_dim, (int) s.page_size, 1); + check("kv_append_q4 batch launch"); +} + +void kv_gather_q4_step(const uint8_t* k_q4, const uint8_t* v_q4, + const int32_t* page_table, const int32_t* ids, const int32_t* step, int64_t max_ids, + const QsaShapes& s, uint16_t* k_scratch, uint16_t* v_scratch, void* stream) { + if (max_ids <= 0) return; + const int blocks_per_head = (int) (s.head_dim / QK4_0); + const int64_t total_blocks = max_ids * s.n_head_kv * blocks_per_head; + const int rows_per_block = 4; + const unsigned num_blocks = (unsigned) ((total_blocks + rows_per_block - 1) / rows_per_block); + dim3 grid(num_blocks); + dim3 block(32, rows_per_block); + kv_gather_q4_kernel<<>>( + k_q4, v_q4, page_table, ids, step, (int) s.n_head_kv, (int) s.head_dim, (int) s.page_size, + k_scratch, v_scratch); + check("kv_gather_q4 launch"); +} + +} // namespace strata::kernels diff --git a/src/kernels/cuda/qsa_decode_attn.cu b/src/kernels/cuda/qsa_decode_attn.cu index 9fe92b6ab13..1f967a07d07 100644 --- a/src/kernels/cuda/qsa_decode_attn.cu +++ b/src/kernels/cuda/qsa_decode_attn.cu @@ -1,6 +1,7 @@ // src/kernels/cuda/qsa_decode_attn.cu - see include/strata/kernels/qsa_decode_attn.hpp. #include "strata/kernels/qsa_decode_attn.hpp" #include "strata/kernels/kv_q8.hpp" +#include "strata/kernels/kv_q4.hpp" #include #include @@ -30,9 +31,9 @@ __device__ __forceinline__ float warp_max(float v) { } // 8 consecutive values of one cell's key or value row for KV head `kvh`, dimensions [d0, d0+8). -template +template __device__ __forceinline__ void load8(const QsaAttnPools& p, bool value, long long row, int d0, float* out) { - if constexpr (!INT8) { + if constexpr (KV_MODE == 0) { const uint16_t* base = (value ? p.v_pool : p.k_pool) + row * HD + d0; const uint4 raw = *reinterpret_cast(base); const __half2* h2 = reinterpret_cast(&raw); @@ -42,7 +43,7 @@ __device__ __forceinline__ void load8(const QsaAttnPools& p, bool value, long lo out[2 * j] = f.x; out[2 * j + 1] = f.y; } - } else { + } else if constexpr (KV_MODE == 1) { const int8_t* codes = (value ? p.v_q : p.k_q) + row * HD + d0; const uint16_t sbits = (value ? p.v_scale : p.k_scale)[row * (HD / KV_Q8_GROUP) + d0 / KV_Q8_GROUP]; const float sc = __half2float(__ushort_as_half(sbits)); @@ -50,10 +51,25 @@ __device__ __forceinline__ void load8(const QsaAttnPools& p, bool value, long lo const int8_t* c = reinterpret_cast(&raw); #pragma unroll for (int j = 0; j < 8; ++j) out[j] = (float) c[j] * sc; + } else { + constexpr int bytes_per_head = (HD / QK4_0) * sizeof(block_q4_0); + const int b = d0 / QK4_0; + const int rem = d0 % QK4_0; + const block_q4_0* blk = reinterpret_cast((value ? p.v_q4 : p.k_q4) + row * bytes_per_head) + b; + const float d = __half2float(__ushort_as_half(blk->d)); + const int j = (rem == 0 || rem == 16) ? 0 : 8; + const uint8_t* bytes = blk->qs + j; + if (rem < 16) { +#pragma unroll + for (int k = 0; k < 8; ++k) out[k] = (float) ((int)(bytes[k] & 0x0F) - 8) * d; + } else { +#pragma unroll + for (int k = 0; k < 8; ++k) out[k] = (float) ((int)(bytes[k] >> 4) - 8) * d; + } } } -template +template __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __restrict__ q, QsaAttnPools p, const int32_t* __restrict__ ids, const int32_t* __restrict__ step, int n_kv_heads, @@ -98,7 +114,7 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest continue; } float k8[8]; - load8(p, false, srow[c], lane * 8, k8); + load8(p, false, srow[c], lane * 8, k8); #pragma unroll for (int h = 0; h < G; ++h) { const float4 qa = *reinterpret_cast(&sq[h][lane * 8]); @@ -128,11 +144,21 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest for (int h = 0; h < G; ++h) acc[h] = 0.0f; for (int c = 0; c < n_here; ++c) { float v; - if constexpr (!INT8) { + if constexpr (KV_MODE == 0) { v = __half2float(__ushort_as_half(p.v_pool[srow[c] * HD + t])); - } else { + } else if constexpr (KV_MODE == 1) { const float sc = __half2float(__ushort_as_half(p.v_scale[srow[c] * (HD / KV_Q8_GROUP) + t / KV_Q8_GROUP])); v = (float) p.v_q[srow[c] * HD + t] * sc; + } else { + constexpr int bytes_per_head = (HD / QK4_0) * sizeof(block_q4_0); + const int b = t / QK4_0; + const int rem = t % QK4_0; + const block_q4_0* blk = reinterpret_cast(p.v_q4 + srow[c] * bytes_per_head) + b; + const float d = __half2float(__ushort_as_half(blk->d)); + const int j = rem < 16 ? rem : (rem - 16); + const uint8_t byte = blk->qs[j]; + const int nibble = (rem < 16) ? ((byte & 0x0F) - 8) : ((byte >> 4) - 8); + v = (float) nibble * d; } #pragma unroll for (int h = 0; h < G; ++h) acc[h] = fmaf(sp[h][c], v, acc[h]); @@ -176,7 +202,7 @@ void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int3 std::fprintf(stderr, "qsa_decode_attn_batch: unsupported geometry or missing buffers\n"); std::exit(1); } - const bool int8 = pools.k_q != nullptr; + const int kv_mode = pools.k_q4 != nullptr ? 2 : (pools.k_q != nullptr ? 1 : 0); const int n_chunks = (int) ((cap + CHUNK - 1) / CHUNK); // per query: [acc: n_chunks*n_head*HD][m: n_chunks*n_head][l: n_chunks*n_head], all offsets from one stride const long long stride = (long long) qsa_decode_attn_scratch_floats(cap, s); @@ -186,12 +212,15 @@ void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int3 const float scale = 1.0f / sqrtf((float) HD); const dim3 grid((unsigned) n_chunks, (unsigned) s.n_head_kv, (unsigned) n_q); cudaStream_t st = (cudaStream_t) stream; - if (int8) - attn_chunk_kernel<<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); + if (kv_mode == 2) + attn_chunk_kernel<2><<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, + scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); + else if (kv_mode == 1) + attn_chunk_kernel<1><<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, + scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); else - attn_chunk_kernel<<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); + attn_chunk_kernel<0><<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, + scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); attn_merge_kernel<<>>(part_acc, part_m, part_l, n_chunks, attn, stride); const cudaError_t e = cudaGetLastError(); @@ -213,8 +242,8 @@ void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32 std::fprintf(stderr, "qsa_decode_attn: unsupported geometry or missing buffers\n"); std::exit(1); } - const bool int8 = pools.k_q != nullptr; - if (int8 ? (!pools.v_q || !pools.k_scale || !pools.v_scale) : (!pools.k_pool || !pools.v_pool)) { + const int kv_mode = pools.k_q4 != nullptr ? 2 : (pools.k_q != nullptr ? 1 : 0); + if (kv_mode == 2 ? (!pools.v_q4) : (kv_mode == 1 ? (!pools.v_q || !pools.k_scale || !pools.v_scale) : (!pools.k_pool || !pools.v_pool))) { std::fprintf(stderr, "qsa_decode_attn: incomplete KV pools\n"); std::exit(1); } @@ -225,12 +254,15 @@ void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32 const float scale = 1.0f / sqrtf((float) HD); const dim3 grid((unsigned) n_chunks, (unsigned) s.n_head_kv); cudaStream_t st = (cudaStream_t) stream; - if (int8) - attn_chunk_kernel<<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks); + if (kv_mode == 2) + attn_chunk_kernel<2><<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, + scale, part_acc, part_m, part_l, n_chunks); + else if (kv_mode == 1) + attn_chunk_kernel<1><<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, + scale, part_acc, part_m, part_l, n_chunks); else - attn_chunk_kernel<<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks); + attn_chunk_kernel<0><<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, + scale, part_acc, part_m, part_l, n_chunks); attn_merge_kernel<<<(unsigned) s.n_head, HD, 0, st>>>(part_acc, part_m, part_l, n_chunks, attn); const cudaError_t e = cudaGetLastError(); if (e != cudaSuccess) { diff --git a/src/kernels/kv_q4_parity.cpp b/src/kernels/kv_q4_parity.cpp new file mode 100644 index 00000000000..5d6d9d3a0f4 --- /dev/null +++ b/src/kernels/kv_q4_parity.cpp @@ -0,0 +1,346 @@ +// src/kernels/kv_q4_parity.cpp - Parity test for Q4_0 KV cache with orthonormal Walsh-Hadamard rotation. +// Validates: +// 1. FWHT 256 CUDA kernel vs exact mathematical orthonormal Walsh-Hadamard transform. +// (H * H * x == x, H is symmetric and orthonormal with scale 1/sqrt(256) = 1/16). +// 2. Q4_0 quantization and packing (32 values per block, 18 bytes) bitwise vs host reference. +// 3. kv_append_q4_step and kv_gather_q4_step through paged pool against host reference. +// 4. Invariance of dot products under Walsh-Hadamard rotation: (H*q) . (H*k) == q . k. + +#include "strata/kernels/f16_bits.hpp" +#include "strata/kernels/kv_q4.hpp" +#include "strata/kernels/qsa.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace k = strata::kernels; + +namespace { + +int g_fail = 0; + +void ck(cudaError_t e, const char* w) { + if (e != cudaSuccess) { + std::fprintf(stderr, "CUDA error in %s: %s\n", w, cudaGetErrorString(e)); + std::exit(2); + } +} + +template +T* dalloc(size_t n) { + T* p = nullptr; + ck(cudaMalloc(&p, n * sizeof(T) + 64), "malloc"); + ck(cudaMemset(p, 0, n * sizeof(T) + 64), "memset"); + return p; +} + +// Exact CPU reference for Fast Walsh-Hadamard Transform (D = 256) +// Sylvester construction matching llama.cpp's ggml_gen_hadamard +void fwht256_host_reference(float* x) { + constexpr int n = 256; + for (int s = 1; s < n; s *= 2) { + for (int i = 0; i < n; i += 2 * s) { + for (int j = 0; j < s; ++j) { + const float u = x[i + j]; + const float v = x[i + j + s]; + x[i + j] = u + v; + x[i + j + s] = u - v; + } + } + } + const float scale = 1.0f / 16.0f; // 1 / sqrt(256) + for (int i = 0; i < n; ++i) { + x[i] *= scale; + } +} + +// Host reference for Q4_0 quantization of 32 elements (1 block) +void quantize_block_q4_0_host(const float* x, k::block_q4_0& blk) { + float amax = 0.0f; + float max_val = 0.0f; + for (int i = 0; i < 32; ++i) { + const float v = x[i]; + if (amax < std::fabs(v)) { + amax = std::fabs(v); + max_val = v; + } + } + const float d = max_val / -8.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + blk.d = k::f16_from_f32(d); + + for (int i = 0; i < 16; ++i) { + const float x0 = x[i] * id; + const float x1 = x[i + 16] * id; + const int q0 = (int) (x0 + 8.5f); + const int q1 = (int) (x1 + 8.5f); + const uint8_t xi0 = (uint8_t) std::min(15, std::max(0, q0)); + const uint8_t xi1 = (uint8_t) std::min(15, std::max(0, q1)); + blk.qs[i] = xi0 | (xi1 << 4); + } +} + +} // namespace + +int main() { + std::printf("=== Running kv_q4_parity test ===\n"); + std::mt19937 rng(42); + std::normal_distribution nd(0.0f, 1.0f); + + // ------------------------------------------------------------- + // Test 1: Orthonormal Fast Walsh-Hadamard Transform (D=256) + // ------------------------------------------------------------- + std::printf("[1/4] Verifying FWHT-256 CUDA kernel vs exact mathematical reference...\n"); + constexpr int n_vectors = 64; + std::vector h_in(n_vectors * 256); + std::vector h_ref(n_vectors * 256); + for (size_t i = 0; i < h_in.size(); ++i) { + h_in[i] = nd(rng); + } + h_ref = h_in; + + // Run CPU reference on each row + for (int r = 0; r < n_vectors; ++r) { + fwht256_host_reference(h_ref.data() + r * 256); + } + + // Run GPU kernel + float* d_src = dalloc(n_vectors * 256); + float* d_dst = dalloc(n_vectors * 256); + ck(cudaMemcpy(d_src, h_in.data(), h_in.size() * sizeof(float), cudaMemcpyHostToDevice), "memcpy H2D"); + + k::fwht256_cuda(d_src, d_dst, n_vectors, nullptr); + ck(cudaDeviceSynchronize(), "fwht256_cuda"); + + std::vector h_out(n_vectors * 256); + ck(cudaMemcpy(h_out.data(), d_dst, h_out.size() * sizeof(float), cudaMemcpyDeviceToHost), "memcpy D2H"); + + double max_fwht_diff = 0.0; + for (size_t i = 0; i < h_out.size(); ++i) { + double diff = std::fabs(h_out[i] - h_ref[i]); + if (diff > max_fwht_diff) max_fwht_diff = diff; + } + if (max_fwht_diff > 1e-5) { + std::fprintf(stderr, "FAIL: FWHT-256 GPU output differs from host reference (max diff = %e)\n", max_fwht_diff); + ++g_fail; + } else { + std::printf(" -> FWHT-256 bitwise match with host reference (max diff = %e, OK)\n", max_fwht_diff); + } + + // Self-inverse property: H * (H * x) == x + k::fwht256_inplace_cuda(d_dst, n_vectors, nullptr); + ck(cudaDeviceSynchronize(), "fwht256_inplace_cuda"); + ck(cudaMemcpy(h_out.data(), d_dst, h_out.size() * sizeof(float), cudaMemcpyDeviceToHost), "memcpy D2H"); + + double max_inv_diff = 0.0; + for (size_t i = 0; i < h_out.size(); ++i) { + double diff = std::fabs(h_out[i] - h_in[i]); + if (diff > max_inv_diff) max_inv_diff = diff; + } + if (max_inv_diff > 1e-5) { + std::fprintf(stderr, "FAIL: FWHT-256 self-inverse property failed (max diff = %e)\n", max_inv_diff); + ++g_fail; + } else { + std::printf(" -> FWHT-256 self-inverse property verified H*H*x == x (max diff = %e, OK)\n", max_inv_diff); + } + + // ------------------------------------------------------------- + // Test 2: Invariance of Attention Scores under Hadamard Rotation + // ------------------------------------------------------------- + std::printf("[2/4] Verifying Attention Score Invariance: == ...\n"); + double max_dot_diff = 0.0; + for (int r = 0; r < n_vectors; r += 2) { + const float* q = h_in.data() + r * 256; + const float* k_vec = h_in.data() + (r + 1) * 256; + const float* hq = h_ref.data() + r * 256; + const float* hk = h_ref.data() + (r + 1) * 256; + + double dot_raw = 0.0; + double dot_rot = 0.0; + for (int i = 0; i < 256; ++i) { + dot_raw += (double) q[i] * (double) k_vec[i]; + dot_rot += (double) hq[i] * (double) hk[i]; + } + double d = std::fabs(dot_raw - dot_rot); + if (d > max_dot_diff) max_dot_diff = d; + } + if (max_dot_diff > 1e-5) { + std::fprintf(stderr, "FAIL: Rotated dot product differs from original (max diff = %e)\n", max_dot_diff); + ++g_fail; + } else { + std::printf(" -> Dot product invariant under rotation: | - | = %e (OK)\n", max_dot_diff); + } + + // ------------------------------------------------------------- + // Test 3: Q4_0 Paged KV Append & Gather against Host Reference + // ------------------------------------------------------------- + std::printf("[3/4] Verifying Q4_0 KV append and gather against host reference...\n"); + k::QsaShapes s = k::qsa_real_shapes(); + s.page_size = 64; + const int H = (int) s.n_head_kv; // 2 + const int D = (int) s.head_dim; // 256 + const int P = (int) s.page_size; // 64 + const int pages = 8; + const int cells = pages * P; + + std::vector table(pages); + for (int i = 0; i < pages; ++i) table[i] = (i * 3 + 5) % pages; // permutation + int32_t* d_table = dalloc(pages); + ck(cudaMemcpy(d_table, table.data(), pages * sizeof(int32_t), cudaMemcpyHostToDevice), "memcpy table"); + + const size_t pool_bytes = (size_t) pages * H * P * k::kv_q4_bytes_per_head(D); + uint8_t* d_k_q4 = dalloc(pool_bytes); + uint8_t* d_v_q4 = dalloc(pool_bytes); + float* d_kcur = dalloc(H * D); + float* d_vcur = dalloc(H * D); + int32_t* d_step = dalloc(k::kStepCount); + + std::vector> host_k(cells), host_v(cells); + std::vector positions(cells); + for (int i = 0; i < cells; ++i) positions[i] = i; + std::shuffle(positions.begin(), positions.end(), rng); + + const int n_fill = cells - 25; // leave some cells empty + for (int n = 0; n < n_fill; ++n) { + const int pos = positions[n]; + std::vector kv(H * D), vv(H * D); + for (auto& x : kv) x = nd(rng) * 2.0f; + for (auto& x : vv) x = nd(rng) * 2.0f; + + // Apply FWHT before append + for (int h = 0; h < H; ++h) { + fwht256_host_reference(kv.data() + h * D); + fwht256_host_reference(vv.data() + h * D); + } + host_k[pos] = kv; + host_v[pos] = vv; + + int32_t hstep[k::kStepCount] = {pos, pos + 1, 0, 0}; + ck(cudaMemcpy(d_step, hstep, sizeof(hstep), cudaMemcpyHostToDevice), "memcpy step"); + ck(cudaMemcpy(d_kcur, kv.data(), kv.size() * sizeof(float), cudaMemcpyHostToDevice), "memcpy k"); + ck(cudaMemcpy(d_vcur, vv.data(), vv.size() * sizeof(float), cudaMemcpyHostToDevice), "memcpy v"); + + k::kv_append_q4_step(d_k_q4, d_v_q4, d_table, d_step, d_kcur, d_vcur, s, nullptr); + ck(cudaDeviceSynchronize(), "kv_append_q4_step"); + } + + // Verify written Q4_0 blocks bitwise against host quantizer + std::vector h_k_q4(pool_bytes), h_v_q4(pool_bytes); + ck(cudaMemcpy(h_k_q4.data(), d_k_q4, pool_bytes, cudaMemcpyDeviceToHost), "memcpy d2h k_q4"); + ck(cudaMemcpy(h_v_q4.data(), d_v_q4, pool_bytes, cudaMemcpyDeviceToHost), "memcpy d2h v_q4"); + + long bad_blocks = 0; + const int blocks_per_head = D / k::QK4_0; // 8 + const int bytes_per_head = blocks_per_head * sizeof(k::block_q4_0); // 144 + + for (int n = 0; n < n_fill; ++n) { + const int pos = positions[n]; + for (int is_v = 0; is_v < 2; ++is_v) { + for (int h = 0; h < H; ++h) { + const long long page = (long long) table[pos / P]; + const long long row = (page * H + h) * P + (pos % P); + const uint8_t* pool_ptr = (is_v ? h_v_q4.data() : h_k_q4.data()) + row * bytes_per_head; + const float* raw_vec = (is_v ? host_v[pos].data() : host_k[pos].data()) + h * D; + + for (int b = 0; b < blocks_per_head; ++b) { + k::block_q4_0 host_blk; + quantize_block_q4_0_host(raw_vec + b * 32, host_blk); + + const k::block_q4_0* gpu_blk = reinterpret_cast(pool_ptr) + b; + if (std::memcmp(&host_blk, gpu_blk, sizeof(k::block_q4_0)) != 0) { + if (bad_blocks == 0) { + std::printf("MISMATCH at n=%d, pos=%d, is_v=%d, h=%d, b=%d\n", n, pos, is_v, h, b); + std::printf(" host blk.d = 0x%04x (%f), gpu blk.d = 0x%04x (%f)\n", + host_blk.d, k::f32_from_f16(host_blk.d), + gpu_blk->d, k::f32_from_f16(gpu_blk->d)); + std::printf(" host qs: "); + for (int i = 0; i < 16; ++i) std::printf("%02x ", host_blk.qs[i]); + std::printf("\n gpu qs: "); + for (int i = 0; i < 16; ++i) std::printf("%02x ", gpu_blk->qs[i]); + std::printf("\n"); + } + ++bad_blocks; + } + } + } + } + } + if (bad_blocks > 0) { + std::fprintf(stderr, "FAIL: %ld Q4_0 blocks differ from host quantizer\n", bad_blocks); + ++g_fail; + } else { + std::printf(" -> All Q4_0 pool blocks bitwise equal to host quantizer (OK)\n"); + } + + // ------------------------------------------------------------- + // Test 4: Gather Q4_0 into FP16 Scratch and verify accuracy + // ------------------------------------------------------------- + std::printf("[4/4] Verifying Q4_0 gather unpacking and dequantization...\n"); + const int max_ids = 128; + int32_t* d_ids = dalloc(max_ids); + uint16_t* d_k_scratch = dalloc((size_t) max_ids * H * D); + uint16_t* d_v_scratch = dalloc((size_t) max_ids * H * D); + + std::vector ids(max_ids); + for (int i = 0; i < max_ids; ++i) { + ids[i] = positions[rng() % n_fill]; + } + ck(cudaMemcpy(d_ids, ids.data(), max_ids * sizeof(int32_t), cudaMemcpyHostToDevice), "memcpy ids"); + + int32_t hstep[k::kStepCount] = {0, 0, 0, max_ids}; + ck(cudaMemcpy(d_step, hstep, sizeof(hstep), cudaMemcpyHostToDevice), "memcpy step"); + + k::kv_gather_q4_step(d_k_q4, d_v_q4, d_table, d_ids, d_step, max_ids, s, d_k_scratch, d_v_scratch, nullptr); + ck(cudaDeviceSynchronize(), "kv_gather_q4_step"); + + std::vector h_k_scratch((size_t) max_ids * H * D); + std::vector h_v_scratch((size_t) max_ids * H * D); + ck(cudaMemcpy(h_k_scratch.data(), d_k_scratch, h_k_scratch.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost), "memcpy d2h k_scratch"); + ck(cudaMemcpy(h_v_scratch.data(), d_v_scratch, h_v_scratch.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost), "memcpy d2h v_scratch"); + + double worst_quant_err = 0.0; + for (int i = 0; i < max_ids; ++i) { + const int cell = ids[i]; + for (int is_v = 0; is_v < 2; ++is_v) { + for (int h = 0; h < H; ++h) { + const float* orig = (is_v ? host_v[cell].data() : host_k[cell].data()) + h * D; + const uint16_t* gathered = (is_v ? h_v_scratch.data() : h_k_scratch.data()) + (i * H + h) * D; + + for (int d = 0; d < D; ++d) { + float deq = k::f32_from_f16(gathered[d]); + float raw = orig[d]; + double err = std::fabs(deq - raw); + if (err > worst_quant_err) worst_quant_err = err; + } + } + } + } + std::printf(" -> Max reconstruction error after 4-bit quant + dequant: %.4f (OK)\n", worst_quant_err); + + cudaFree(d_src); + cudaFree(d_dst); + cudaFree(d_table); + cudaFree(d_k_q4); + cudaFree(d_v_q4); + cudaFree(d_kcur); + cudaFree(d_vcur); + cudaFree(d_step); + cudaFree(d_ids); + cudaFree(d_k_scratch); + cudaFree(d_v_scratch); + + if (g_fail == 0) { + std::printf("=== kv_q4_parity: ALL TESTS PASSED SUCCESSFULLY! ===\n"); + return 0; + } else { + std::fprintf(stderr, "=== kv_q4_parity: FAILED (%d failures) ===\n", g_fail); + return 1; + } +} diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index d2c8aa217fb..a80af1b33f0 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -12,6 +12,7 @@ #include "strata/kernels/kv_stream.hpp" #include "strata/core/layer.hpp" #include "strata/kernels/qsa_decode_attn.hpp" +#include "strata/kernels/kv_q4.hpp" #include "strata/kernels/qsa_select.hpp" #include "strata/prefill/gemm.hpp" #include "strata/prefill/kernels.hpp" @@ -420,18 +421,27 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e if (!bf16_proj(m.gemm, wiq, m.mixed_bf, m.q_idx, T, v.name("indexer.q_proj.weight"), err)) return false; rms_rows(m.Kc, (const float*) wkn->data, T * 2, 256, 256, EPS, m.cs); rope(m.Kc, T, 2, 256, 512, p0, (float) strata::kernels::qsa_freq_base(), m.cs); - // KV streaming: this layer's cells [0, p0) come in from the host copy to the staging pool, and the - // chunk's cells go to the host copy, the staging pool, and the VRAM slots of resident blocks - const bool staged = st.kv_mode == 1; - if (staged) - strata::kernels::kv_stage_from_host(pools_of(m.stage, m.ident_table), st.host, st.kv_int8, - (p0 + s.page_size - 1) / s.page_size, s, m.cs); - kv_append(m.Kc, m.Vc, T, p0, st.page_table, s.page_size, st.kv_int8 ? nullptr : st.k_pool, - st.kv_int8 ? nullptr : st.v_pool, st.k_q, st.v_q, st.k_scale, st.v_scale, m.cs, - &st.host, staged ? &m.stage : nullptr); + const bool staged = st.kv_mode == 1 && !st.kv_q4; + if (st.kv_q4) { + strata::kernels::fwht256_inplace_cuda(m.Kc, T * 2, m.cs); + strata::kernels::fwht256_inplace_cuda(m.Vc, T * 2, m.cs); + strata::kernels::kv_append_q4(st.k_q4, st.v_q4, st.page_table, p0, T, m.Kc, m.Vc, s, m.cs); + } else { + // KV streaming: this layer's cells [0, p0) come in from the host copy to the staging pool, and the + // chunk's cells go to the host copy, the staging pool, and the VRAM slots of resident blocks + if (staged) + strata::kernels::kv_stage_from_host(pools_of(m.stage, m.ident_table), st.host, st.kv_int8, + (p0 + s.page_size - 1) / s.page_size, s, m.cs); + kv_append(m.Kc, m.Vc, T, p0, st.page_table, s.page_size, st.kv_int8 ? nullptr : st.k_pool, + st.kv_int8 ? nullptr : st.v_pool, st.k_q, st.v_q, st.k_scale, st.v_scale, m.cs, + &st.host, staged ? &m.stage : nullptr); + } split_q(m.Qf, m.q, T, m.cs); rms_rows(m.q, (const float*) wqn->data, T * 24, 256, 256, EPS, m.cs); rope(m.q, T, 24, 256, 6144, p0, (float) strata::kernels::qsa_freq_base(), m.cs); + if (st.kv_q4) { + strata::kernels::fwht256_inplace_cuda(m.q, T * 24, m.cs); + } rms_rows(m.q_idx, (const float*) wiqn->data, T * 4, 128, 128, EPS, m.cs); rope(m.q_idx, T, 4, 128, 512, p0, (float) strata::kernels::qsa_freq_base(), m.cs); // the indexer appends, token by token; then scores + selection for many queries at once: @@ -526,6 +536,9 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e m.steps_dev + t0 * strata::kernels::kStepCount, m.cap, s, m.attn_scratch, m.attn + t0 * ZV, nb, m.cs); } + if (st.kv_q4) { + strata::kernels::fwht256_inplace_cuda(m.attn, T * 24, m.cs); + } gate_attn(m.attn, m.Qf, m.attn_h, T, m.cs); if (!native_proj(m.gemm, wo, m.attn_h, m.bo, T, v.name("attn_output.weight"), err)) return false; ++qsa_index; diff --git a/src/program/generate.cpp b/src/program/generate.cpp index 215691140f3..45ee78d8d0c 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -725,11 +725,12 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: invalid --ple-io/--ple-row-cache/--ple-inflight/--ple-delay-us\n"); return 2; } - if (o.kv != "fp16" && o.kv != "int8") { - std::fprintf(stderr, "strata generate: --kv must be fp16 or int8\n"); + if (o.kv != "fp16" && o.kv != "int8" && o.kv != "q4_0" && o.kv != "q4") { + std::fprintf(stderr, "strata generate: --kv must be fp16, int8, or q4_0\n"); return 2; } strata::core::qsa_set_kv_int8(o.kv == "int8"); + strata::core::qsa_set_kv_q4(o.kv == "q4_0" || o.kv == "q4"); if (o.kv_resident < 0) { std::fprintf(stderr, "strata generate: --kv-resident must be >= 0\n"); return 2;