diff --git a/CMakeLists.txt b/CMakeLists.txt index 7ba7ee6349..2e74715ca7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -408,6 +408,8 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) src/kernels/cuda/qsa_select.cu src/kernels/cuda/verify_kernels.cu src/kernels/cuda/iq_kernels.cu + src/kernels/cuda/qns_kernels.cu + src/kernels/cuda/qns_gather.cu src/kernels/ngram.cpp src/ngram/ple_reader.cpp src/platform/direct_file.cpp) @@ -1541,7 +1543,8 @@ if(STRATA_NATIVE_EXPERTS) src/kernels/cpu/iq_avx2.cpp src/kernels/cpu/kq_avx2.cpp src/kernels/cpu/kq_avx1.cpp src/kernels/cpu/q2_avx1.cpp src/kernels/cpu/iq_avx1.cpp src/kernels/cpu/nvfp4_avx512.cpp src/kernels/cpu/nvfp4_avx2.cpp - src/kernels/cpu/tq_cpu.cpp src/kernels/cpu/tq_avx2.cpp src/kernels/cpu/tq_avx512.cpp) + src/kernels/cpu/tq_cpu.cpp src/kernels/cpu/tq_avx2.cpp src/kernels/cpu/tq_avx512.cpp + src/kernels/cpu/qns_avx2.cpp) target_include_directories(strata_kernels_cpu PRIVATE third_party/ggml) if(MSVC) set_source_files_properties(src/kernels/cpu/iq_avx512.cpp src/kernels/cpu/nvfp4_avx512.cpp @@ -1551,6 +1554,7 @@ if(STRATA_NATIVE_EXPERTS) set_source_files_properties(src/kernels/cpu/nvfp4_avx2.cpp PROPERTIES COMPILE_OPTIONS "/arch:AVX2") set_source_files_properties(src/kernels/cpu/tq_avx2.cpp PROPERTIES COMPILE_OPTIONS "/arch:AVX2") set_source_files_properties(src/kernels/cpu/tq_avx512.cpp PROPERTIES COMPILE_OPTIONS "/arch:AVX512") + set_source_files_properties(src/kernels/cpu/qns_avx2.cpp PROPERTIES COMPILE_OPTIONS "/arch:AVX2") # the router lookahead's AVX1 dot (older CPUs; called behind cpu_avx1_ok): no FMA, no 256-bit integer ops set_source_files_properties(src/kernels/cpu/kq_avx1.cpp PROPERTIES COMPILE_OPTIONS "/arch:AVX") # the older CPUs' Q2_0 rows and activation quantizer (q2_avx1.cpp, issue #1699): 128-bit integer lanes, @@ -1568,6 +1572,8 @@ if(STRATA_NATIVE_EXPERTS) set_source_files_properties(src/kernels/cpu/tq_avx2.cpp PROPERTIES COMPILE_OPTIONS "-mavx2;-mfma;-mf16c") set_source_files_properties(src/kernels/cpu/tq_avx512.cpp PROPERTIES COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512vl;-mavx512vnni;-mavx512vbmi;-mavx2;-mfma;-mf16c") + # Qwen Next S trellis experts (qns.hpp): AVX2 only, the floor of the target desktops and Zen 2 servers + set_source_files_properties(src/kernels/cpu/qns_avx2.cpp PROPERTIES COMPILE_OPTIONS "-mavx2;-mfma;-mf16c") set_source_files_properties(src/kernels/cpu/kq_avx1.cpp PROPERTIES COMPILE_OPTIONS "-mavx") # the older CPUs' Q2_0 rows and activation quantizer (q2_avx1.cpp, issue #1699): SSE4.2 + AVX with # AVX2/FMA/F16C explicitly off, so a -march=native build host cannot leak AVX2 code into the file that diff --git a/docs/QNS.md b/docs/QNS.md new file mode 100644 index 0000000000..c3fdc2a74e --- /dev/null +++ b/docs/QNS.md @@ -0,0 +1,57 @@ +# Mirai Qwen3.8-Flash-Next-S (QNS trellis) + +Validated on Linux (Ubuntu 22.04, CUDA 12.8) with RTX 3090s at 300 W, an EPYC 7642 (AVX2) and 125 GB RAM on +2026-10-20. + +Mirai's Flash-Next-S is a trellis build of Qwen3.8-Flash-Next: the routed experts take 1 to 5 bits each (2.4 on +average), the backbone and head 5 bits. Mirai packed it for Strata as two files: `Qwen3.8-Flash-Next-S-qns.gguf` +(the experts and a Q8_0 backbone) and `Qwen3.8-Flash-Next-S-qns-dense.gguf` (the backbone and head as trellis, which +takes less VRAM). There is also a 32 GB build, S-small. The weights are on Hugging Face from Mirai +(`trymirai/Qwen3.8-Flash-Next-S-strata-experimental`). + +Every expert is its own blob of its own size, so the RAM arena is compact and each VRAM slot is sized to its expert. +The format, the CUDA and AVX2 kernels and the packer support come from Mirai's `strata-qns.patch` for v0.1.43. + +## Pack + +```sh +.venv/bin/python tools/iq_pack.py --gguf /models/Qwen3.8-Flash-Next-S-qns.gguf --out packs/qns +``` + +It takes seconds and writes `native_experts.txt` v6 with `qns_offsets.bin` and `qns_rotation.bin`. A pack made by the +v0.1.43 patch (v5) is refused with a note to repack: v5 means Hadamard-folded experts in this engine. + +## Run + +The n-gram table is the PLE shard of any ISTA-DASLab size (shard 2, 28.8 GB); the MTP layer and the expert profile +are the original model's. + +```sh +build/strata --serve --pack packs/qns \ + --native /models/Qwen3.8-Flash-Next-S-qns.gguf --native-dense-gguf /models/Qwen3.8-Flash-Next-S-qns.gguf \ + --qns-dense-gguf /models/Qwen3.8-Flash-Next-S-qns-dense.gguf \ + --ple-gguf /models/IQ3_S/Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-00002-of-00002.gguf \ + --expert-profile data/expert-profile.bin --expert-cache auto --prefill auto --spec 4 --mtp --kv int8 +``` + +`--layer-split auto` with `--gpu 0,1` runs it on two cards. + +## Measured + +Greedy, thinking off, 32K context, MTP on, the server's own tok/s per request (six short requests and a 6,210-token +document); one run each, so the ranges are indicative, not a benchmark: + +| Model | Cards | Experts in RAM | Decode | 6,210-token prompt | +|---|---|---:|---:|---:| +| S | 1 | 33.8 GiB | 35-81 tok/s | 4.1 s | +| S | 2 (layer split) | 33.8 GiB | 60-125 tok/s | 4.2 s | +| S-small | 1 | 26.5 GiB | 47-106 tok/s | 3.9 s | +| S-small | 2 (layer split) | 26.5 GiB | 83-166 tok/s | 4.0 s | + +On S-small the greedy answers of this engine and of v0.1.43 with Mirai's patch are the same on five of six probes +(the sixth parts on a near-tie after 10-55 words), on one card and on two. + +## Limits + +- NVIDIA (CUDA) and the AVX2 CPU pool. Not built or run on AMD (HIP), Intel (SYCL) or Mac. +- Setup and the Terminal do not offer these models yet; start them with a config like the command above. diff --git a/include/strata/artifact/gguf_reader.hpp b/include/strata/artifact/gguf_reader.hpp index 5b67f54ba7..46ccb0d06e 100644 --- a/include/strata/artifact/gguf_reader.hpp +++ b/include/strata/artifact/gguf_reader.hpp @@ -230,6 +230,14 @@ inline bool block_geometry(uint32_t t, int& elems, int& bytes) { elems = 32; bytes = 17; return true; + case 26: // I32: a qns backbone's input permutation (.qns_perm) + elems = 1; + bytes = 4; + return true; + case 27: // I64: a qns model's per-layer expert index (blk.N.ffn_exps_qns_idx) + elems = 1; + bytes = 8; + return true; case 42: elems = 64; bytes = 18; diff --git a/include/strata/core/expert_cache.hpp b/include/strata/core/expert_cache.hpp index 9685aa4d5c..c755be443b 100644 --- a/include/strata/core/expert_cache.hpp +++ b/include/strata/core/expert_cache.hpp @@ -232,6 +232,8 @@ class ExpertCache { if (s < 0 || s >= slots_ || base_ == nullptr) return nullptr; return slot_room(s) >= blob + need ? base_ + slot_offset(s) + blob : nullptr; } + /// Bytes slot `s` holds. + uint64_t slot_bytes(int64_t s) const { return slot_offset(s + 1) - slot_offset(s); } private: #if defined(STRATA_USE_HIP) diff --git a/include/strata/core/mtp.hpp b/include/strata/core/mtp.hpp index 7216f015af..5a3faefda2 100644 --- a/include/strata/core/mtp.hpp +++ b/include/strata/core/mtp.hpp @@ -261,6 +261,8 @@ class MtpDrafter { float *h_prob_ = nullptr, *m_prob_ = nullptr; // each draft's probability under the draft layer // the draft head: the main head's rows for a token subset (rt/draft_vocab.bin), or the whole head uint8_t* dhead_ = nullptr; + float* dscales_ = nullptr; ///< a trellis head (kQnsDenseType): the draft subset's row scales + uint8_t* qns_act_ = nullptr; ///< and the drafter's rotated, quantized head input int32_t* dvocab_ = nullptr; int64_t n_dvocab_ = 0; std::string rt_dir_; diff --git a/include/strata/core/native_dense.hpp b/include/strata/core/native_dense.hpp index 5e3907b416..0812ec13a7 100644 --- a/include/strata/core/native_dense.hpp +++ b/include/strata/core/native_dense.hpp @@ -31,6 +31,10 @@ class NativeDense { /// Layer split: load only blocks [lb, le) (every other `blk.N.` projection belongs to another GPU's stage; the /// PLE tensors are loaded everywhere). Process-wide, read by the next `load`; (-1, -1) = all layers. static void set_layer_range(int lb, int le); + /// --qns-dense-gguf: the GGUF that holds Qwen Next S backbone matrices as trellis (.qns_words / .qns_scales / + /// .qns_perm); `load` serves each matrix it has from there instead of the model's Q8_0 blocks. Empty: off. + static void set_qns_dense(const std::string& path); + static const std::string& qns_dense(); /// Layer split: the device bytes `load` would allocate for the matrices of layers [lb, le) - the same walk and /// filters, each matrix rounded up to the granule cudaMalloc backs it with - reading the GGUF headers and `table` /// only, so it allocates nothing and needs no device (#1238). `table` must hold every tensor the shards carry. diff --git a/include/strata/core/native_head.hpp b/include/strata/core/native_head.hpp index 501f6ec276..e44d6eb8cb 100644 --- a/include/strata/core/native_head.hpp +++ b/include/strata/core/native_head.hpp @@ -31,10 +31,15 @@ class NativeHead { int device() const { return device_; } /// Bytes of one vocabulary row. size_t row_bytes() const { return n_out_ > 0 ? (size_t) (bytes_ / (uint64_t) n_out_) : 0; } + /// type() == kQnsDenseType (the head as trellis from --qns-dense-gguf): its row scales and T bits per step. + const float* qns_scales() const { return qns_scales_; } + int qns_bits() const { return qns_bits_; } private: void* weights_ = nullptr; void* scratch_ = nullptr; + float* qns_scales_ = nullptr; + int qns_bits_ = 0; uint64_t bytes_ = 0; int n_in_ = 0, n_out_ = 0; int type_ = -1; diff --git a/include/strata/core/verify.hpp b/include/strata/core/verify.hpp index e6ff842b50..916862f4a6 100644 --- a/include/strata/core/verify.hpp +++ b/include/strata/core/verify.hpp @@ -612,6 +612,9 @@ class Verifier { // unless some layer's gate/up is folded const float* had_sx_ = nullptr; // APR: device sign vectors for x (n_embd) and h (n_ff); nullptr = identity const float* had_sh_ = nullptr; + uint8_t* qns_xq_ = nullptr; // a qns pack's rotated two-plane activations (qns_cuda.hpp) + uint8_t* qns_xd_ = nullptr; // a trellis backbone's 2560-wide layer input, quantized (qns_dense_quant) + uint8_t* qns_od_ = nullptr; // and its 6144-wide output-projection input float* hit_xs_ = nullptr; void* hit_scratch_ = nullptr; float *head_mixed_ = nullptr, *head_inj_ = nullptr, *head_logits_ = nullptr; diff --git a/include/strata/core/weights.hpp b/include/strata/core/weights.hpp index fe37231206..084fb97b00 100644 --- a/include/strata/core/weights.hpp +++ b/include/strata/core/weights.hpp @@ -45,6 +45,10 @@ enum class WeightKind : int { F16InF32 = 3, ///< an f16 value promoted to f32: a real f32->f16 conversion, not a truncation }; +/// A Qwen Next S backbone matrix served as its trellis (--qns-dense-gguf; not a ggml type id): `native_data` holds the +/// row tapes as u32 words in 32-row groups, `qns_scales` the row scales (qns_cuda.hpp). +inline constexpr int kQnsDenseType = 251; + struct WeightRef { const void* data = nullptr; ///< device pointer inside the arena uint64_t bytes = 0; @@ -90,6 +94,11 @@ struct WeightRef { const void* native_data = nullptr; void* native_q8_1 = nullptr; int native_type = -1; + /// native_type == kQnsDenseType: the row scales, T bits per step, and for an input whose order differs from the + /// checkpoint's columns (GDN out_proj) the gather: checkpoint column c is the engine's input index qns_perm[c]. + const float* qns_scales = nullptr; + const int32_t* qns_perm = nullptr; + int qns_bits = 0; /// S23 experiment (STRATA_HC_Q8=1): a hyper-connection projection's Q8_0 bytes as the GGUF stores them, owned by /// NativeDense, for the verify window's read (every other path keeps the pack's BF16 `data`) const void* hc_q8 = nullptr; diff --git a/include/strata/kernels/cpu/expert_layout.hpp b/include/strata/kernels/cpu/expert_layout.hpp index b78528c779..a2c7fe1cc6 100644 --- a/include/strata/kernels/cpu/expert_layout.hpp +++ b/include/strata/kernels/cpu/expert_layout.hpp @@ -15,6 +15,10 @@ namespace strata::kernels::cpu { +/// Qwen Next S trellis experts (not a ggml type id): one self-describing qns blob per expert (header with its own bit +/// widths, trellis words in 32-row groups, row scales). native_experts.txt v6 gives such a layer `kQnsType kQnsType`. +inline constexpr int kQnsType = 250; + struct ExpertLayout { bool native = false; int64_t n_layers = 0, n_expert = NE; @@ -31,14 +35,28 @@ struct ExpertLayout { int version = 0; ///< native_experts.txt's header version (0 = none given) uint64_t max_blob = BLOB; uint64_t total = 0; ///< experts.bin size + /// qns packs (qns_offsets.bin): each expert's blob has its own size, so the arena is compact and a blob's offset + /// is this table's, n_layers x (n_expert + 1) with each layer's end last. `bytes` stays the layer's largest blob: + /// the VRAM slot and the copy size. A copy of `bytes` may run into the next blob; `total` ends with max_blob of + /// slack so it never runs past the arena. + std::vector qns_off; + /// qns packs (qns_rotation.bin): the experts' input rotations, signs[2560], Q[5][5] (2560 = 512 x 5), signs[640], + /// Q[5][5] (640 = 128 x 5). + std::vector qns_rotation; + bool qns() const { return !qns_off.empty(); } uint64_t blob_bytes(int64_t layer) const { return native ? bytes[(size_t) layer] : (uint64_t) BLOB; } uint64_t layer_offset(int64_t layer) const { return native ? offset[(size_t) layer] : (uint64_t) layer * (uint64_t) n_expert * (uint64_t) BLOB; } uint64_t blob_offset(int64_t layer, int64_t expert) const { + if (qns()) return qns_off[(size_t) (layer * (n_expert + 1) + expert)]; return layer_offset(layer) + (uint64_t) expert * blob_bytes(layer); } + /// One expert's own blob: a qns blob's size (its VRAM slot is sized to it), else the layer's. + uint64_t expert_bytes(int64_t layer, int64_t expert) const { + return qns() ? blob_offset(layer, expert + 1) - blob_offset(layer, expert) : blob_bytes(layer); + } }; /// Plan v0.3 P6: whether this CPU (and its OS) runs the AVX-512 kernels (F, BW, VL, VNNI, VBMI). Probed in a @@ -98,6 +116,8 @@ bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n /// written only when some layer's roles are in different shards (every other pack stays v3, byte for byte). /// v5 = v4 plus `hadamard.txt` beside it: Hadamard-folded experts (hadamard.hpp), written only for such models, /// so an engine that cannot rotate the activations refuses the pack instead of running it wrong. -inline constexpr int kExpertLayoutVersion = 5; +/// v6 = qns layers (`kQnsType kQnsType offset max_blob payload_off 0 0`) with qns_offsets.bin and qns_rotation.bin +/// (strata-qns.patch for 0.1.43 wrote these as v5; such a pack is refused with a note to repack). +inline constexpr int kExpertLayoutVersion = 6; } // namespace strata::kernels::cpu diff --git a/include/strata/kernels/cpu/native_expert.hpp b/include/strata/kernels/cpu/native_expert.hpp index fbac350522..791179dfb9 100644 --- a/include/strata/kernels/cpu/native_expert.hpp +++ b/include/strata/kernels/cpu/native_expert.hpp @@ -13,10 +13,11 @@ namespace strata::kernels::cpu { -/// Bytes of the largest quantized activation any native layer uses (2560 values as Q8_K: 10 x 292). -inline constexpr size_t kNativeActBytes = 4096; -/// Bytes of the largest quantized down activation (640 values as Q8_0: 20 x 34, or Q8_K 3 x 292). -inline constexpr size_t kNativeHBytes = 1024; +/// Bytes of the largest quantized activation any native layer uses (2560 values as Q8_K: 10 x 292; a qns layer's +/// rotated two-plane int8 activation is larger, qns.hpp). +inline constexpr size_t kNativeActBytes = 16384; +/// Bytes of the largest quantized down activation (640 values as Q8_0: 20 x 34, or Q8_K 3 x 292; qns: qns.hpp). +inline constexpr size_t kNativeHBytes = 4096; /// One layer's native expert geometry. struct NativeFmt { diff --git a/include/strata/kernels/cpu/qns.hpp b/include/strata/kernels/cpu/qns.hpp new file mode 100644 index 0000000000..400b29da37 --- /dev/null +++ b/include/strata/kernels/cpu/qns.hpp @@ -0,0 +1,38 @@ +// include/strata/kernels/cpu/qns.hpp - Qwen Next S routed experts (trellis "qns" blobs, qns_blob.py) on AVX2. +// +// A blob row is a trellis path in a rotated basis: the input is rotated (x_rot = (H (x) Q)(signs * x)) and quantized +// per token to x_rot ~ step * (q0 + q1 / 254) (two int8 planes), the kernel computes per row the exact integer +// S = sum (level + 54) * (254 q0 + q1) = 254 C + F (C, F the two planes' sums of the GPU kernels), and finishes +// y = scale * ((step / RMS) * (S - 54 * sum (254 q0 + q1)) / 254 + d0 * E + d1 * O), d_j = -MEAN_j / RMS +// with E, O = step * sum (q0 + q1 / 254) over the even / odd columns. Same contract as the vLLM plugin's kernels.cu. +#pragma once + +#include +#include + +namespace strata::kernels::cpu::qns { + +/// One token's quantized activation (4-byte aligned): a 16-byte header {float step; int32 sum of 254 q0 + q1; +/// float E; float O}, then int16 254 q0 + q1 of the rotated columns, each 4 columns 4g .. 4g + 3 stored as 4g, 4g + 2, +/// 4g + 1, 4g + 3. +inline constexpr size_t kActBytes = 16 + 2 * 2560; ///< expert input x (2560) +inline constexpr size_t kHBytes = 16 + 2 * 640; ///< down input h = silu(gate) * up (640) + +/// The rotation factors (rotation.signs.N / rotation.q.N of the package, q row-major 5 x 5). Once, at load. +void set_rotation(const float* signs2560, const float* q2560, const float* signs640, const float* q640); +/// x (2560 floats) -> kActBytes. +void quant_act(const float* x, void* dst); +/// h (640 floats) -> kHBytes. +void quant_h(const float* h, void* dst); + +/// ff[t][r] = silu(gate_r . x_t) * (up_r . x_t) for rows r in [r0, r1) of [0, 640), nt (1..8) tokens. Rows go in +/// blocks of 8: a range that cuts a block computes the whole block (split ranges at multiples of 8). +void gu_rows(const uint8_t* blob, const void* const* act, int nt, float* const* ff, int r0, int r1); +/// out[t][r] = down_r . h_t for rows r in [r0, r1) of [0, 2560), nt (1..8) tokens, blocks of 8 as gu_rows. +void down_rows(const uint8_t* blob, const void* const* hq, int nt, float* const* out, int r0, int r1); + +/// The kernels' core, exposed for the exactness test: s[t][i] = S of row 8 b + i of the gate_up matrix (rows 0-639 +/// gate, 640-1279 up; down = false) or of the down matrix, for act[0 .. nt). +void block_sums(const uint8_t* blob, bool down, int b, const void* const* act, int nt, int64_t (*s)[8]); + +} // namespace strata::kernels::cpu::qns diff --git a/include/strata/kernels/qns_cuda.hpp b/include/strata/kernels/qns_cuda.hpp new file mode 100644 index 0000000000..3e644b5f05 --- /dev/null +++ b/include/strata/kernels/qns_cuda.hpp @@ -0,0 +1,73 @@ +// qns_cuda.hpp - Qwen Next S ("qns") routed experts on CUDA: one expert = one qns blob (qns_blob.py), every weight row a +// 2-bit-class trellis path in a rotated basis (FORMAT.md). Built with nvcc -O3 -arch=sm_86 (no NVRTC). +// +// Activations: a token's 2560-wide input is rotated (x_rot = R (signs * x), R = H_512 (x) Q_5) and quantized to two int8 +// planes, x_rot ~ step (q0 + q1 / 254), in an opaque record of qns_act_bytes() bytes. Every expert row then finishes the +// exact int32 sums C = sum (level + 54) q0 and F = sum (level + 54) q1, so an entry's output does not depend on which +// entries share its group or call. +// +// Every entry point is CUDA-graph safe: no allocation, no host synchronization, launch shapes fixed by the arguments. +#pragma once + +#include + +#include +#include + +namespace strata::kernels { + +/// The rotation factors (rotation.signs.N / rotation.q.N of the package; host float32 pointers: signs [2560], Q [5][5], +/// signs [640], Q [5][5]). Copied to device symbols of the current device: call once per device, before any other call. +void qns_cuda_set_rotation(const float* signs2560, const float* q2560, const float* signs640, const float* q640); + +/// Bytes of one token's quantized 2560 activation (a multiple of 16). +size_t qns_act_bytes(); + +/// x [n_tok][2560] float32 (device) -> act [n_tok] records of qns_act_bytes() (the same as qns_dense_quant(.., 2560, null, ..)). +void qns_quant_act(const float* x, int64_t n_tok, void* act, cudaStream_t s); + +/// Bytes of scratch qns_expert_grouped needs for `cap_entries` entries (gate/up outputs and the down activations). +size_t qns_grouped_scratch_bytes(int64_t cap_entries); + +/// Grouped experts, native_expert_grouped's contract: group g's blob at device address grp_ptr[g] (256-byte aligned); +/// its entries [grp_start[g], grp_start[g+1]) read token ent_tok[e]'s record of `act` and write row ent_dst[e] of `out` +/// (2560 floats: the unweighted expert output). *n_groups (device) groups; entries index the scratch, so they must be +/// < cap_entries. `grid_groups` (1 .. cap_groups; 0 = cap_groups) groups run side by side, a block row each striding +/// over the rest; the results do not depend on it. +void qns_expert_grouped(const unsigned long long* grp_ptr, const int32_t* grp_start, const int32_t* n_groups, + const int32_t* ent_dst, const int32_t* ent_tok, int64_t cap_groups, int64_t cap_entries, + const void* act, void* scratch, float* out, cudaStream_t s, int64_t grid_groups); + +/// One expert's gate/up in the original basis as fp16 [1280][2560], interleaved like iq_dequant_gu_f16: row 2r = gate +/// row r, row 2r+1 = up row r. +void qns_dequant_gu_f16(const uint8_t* blob, uint16_t* dst, cudaStream_t s); + +/// One expert's down projection in the original basis as fp16 [2560][640]. +void qns_dequant_d_f16(const uint8_t* blob, uint16_t* dst, cudaStream_t s); + +// ---- Dense backbone matrices (GDN in_proj_qkv / in_proj_z / out_proj, attention q / k / v / o): the same trellis rows and +// integer contract, stored without a blob: words [rows/32][W][32] u32 (W = cols/64*T + 1, the expert layout) and f32 +// scales [rows]. cols is 2560 (= 512 x 5) or 6144 (= 2048 x 3); rows a multiple of 32; T in 3..10. A 32-row-aligned +// slice of a matrix is a matrix (words + r0/32 * W * 32, scales + r0). +// perm (GDN out_proj only, else null): x_ckpt[c] = x[perm[c]], the engine's column order of checkpoint column c. + +/// The 6144 rotation's factors (host float32: signs [6144], Q [3][3] row-major). Once per device, before any capture. +void qns_cuda_set_rotation_6144(const float* signs6144, const float* q6144); + +/// Bytes of one token's quantized activation for a dense matrix of `cols` columns (a multiple of 16). +size_t qns_dense_act_bytes(int cols); + +/// x [n_tok][cols] float32 (device, engine column order) -> act [n_tok] records of qns_dense_act_bytes(cols): gathered by +/// perm (when not null), rotated and quantized. +void qns_dense_quant(const float* x, int64_t n_tok, int cols, const int32_t* perm, void* act, cudaStream_t s); + +/// y[t * rows + r] = row r . x_t for t < n_tok (1..8): each weight chunk is decoded once for all tokens. act is only read. +void qns_dense_gemv(const uint32_t* words, const float* scales, int rows, int cols, int T, const void* act, int n_tok, + float* y, cudaStream_t s); + +/// dst [rows][cols] fp16 in the original basis and the engine's column order: dst[r][perm[c]] = w_ckpt[r][c] (dst[r][c] +/// without perm). +void qns_dense_dequant_f16(const uint32_t* words, const float* scales, int rows, int cols, int T, const int32_t* perm, + uint16_t* dst, cudaStream_t s); + +} // namespace strata::kernels diff --git a/include/strata/kernels/qns_gather.hpp b/include/strata/kernels/qns_gather.hpp new file mode 100644 index 0000000000..61b66d2d79 --- /dev/null +++ b/include/strata/kernels/qns_gather.hpp @@ -0,0 +1,16 @@ +// include/strata/kernels/qns_gather.hpp - a row subset of a qns trellis matrix (u32 words in 32-row groups, row +// scales beside them): the MTP drafter's draft head, gathered from the trellis head. +#pragma once + +#include + +#include + +namespace strata::kernels { + +/// Row i of dst (words and scales, the same 32-row-group layout, W words a row) = row rows[i] of src, for i < n; +/// rows n .. n_pad - 1 (n_pad a multiple of 32) are zero words with scale 0. +void qns_gather_rows(const uint32_t* words, const float* scales, int W, const int32_t* rows, int n, int n_pad, + uint32_t* dst_words, float* dst_scales, cudaStream_t s); + +} // namespace strata::kernels diff --git a/include/strata/kernels/verify_kernels.hpp b/include/strata/kernels/verify_kernels.hpp index 45bc70e1a6..92671c7138 100644 --- a/include/strata/kernels/verify_kernels.hpp +++ b/include/strata/kernels/verify_kernels.hpp @@ -133,6 +133,9 @@ void fetch_blobs(const unsigned long long* src, const int32_t* n, uint8_t* dst, /// table[pairs[2k]] = pairs[2k + 1] for k < n, `pairs` in mapped host memory: a residency table's changed entries, /// stream-ordered before the window that reads it, without the copy engine (where the tier's swaps may queue). void res_patch(int32_t* table, const int32_t* pairs, int n, void* stream); +/// `fetch_blobs` for a qns pack: the blobs differ in size, so each copies only its own bytes (its header's word 9), +/// into the same `slot_bytes` stride. +void fetch_qns_blobs(const unsigned long long* src, const int32_t* n, uint8_t* dst, int64_t slot_bytes, int cap, void* stream); /// ptr[k] = base + k * blob_bytes for k < *n (the staged copies `fetch_blobs` made). void rebase_ptrs(unsigned long long* ptr, const int32_t* n, uint8_t* base, int64_t blob_bytes, void* stream); diff --git a/src/core/expert_source.cpp b/src/core/expert_source.cpp index 5cd037da24..3080140c4f 100644 --- a/src/core/expert_source.cpp +++ b/src/core/expert_source.cpp @@ -5027,6 +5027,21 @@ bool check_experts_gguf(const std::string& gguf, const strata::kernels::cpu::Exp std::map> files; for (int64_t l = 0; l < lay.n_layers; ++l) { const auto& fm = lay.fmt[(size_t) l]; + if (fm.gu_type == strata::kernels::cpu::kQnsType) { // the layer's blobs as one I8 payload tensor + const std::string path = expert_gguf_file(gguf, lay, l, 0); + auto& f = files[path]; + if (!f) f = std::make_unique(path); + const std::string name = "blk." + std::to_string(l) + ".ffn_exps_qns"; + const strata::TensorInfo* t = f->find(name); + const uint64_t span = lay.blob_offset(l, lay.n_expert) - lay.blob_offset(l, 0); + if (t == nullptr || t->type != 24 /* I8 */ || strata::tensor_payload_bytes(*t) != span || + f->data_start() + t->offset != lay.gguf_off[(size_t) (3 * l)]) { + err = "the pack's native_experts.txt does not match the model: " + name + " in " + path + + " is missing, not I8, not " + std::to_string(span) + " B or elsewhere - repack with tools/iq_pack.py"; + return false; + } + continue; + } const uint64_t blob = lay.bytes[(size_t) l]; const uint64_t per[3] = {fm.up_off, fm.up_off, blob - fm.down_off}; for (int r = 0; r < 3; ++r) { @@ -5095,6 +5110,11 @@ LoadStats load_experts_gguf(const std::string& gguf, uint8_t* dst, const strata: std::atomic next{0}; std::atomic bad{false}; #if defined(_WIN32) + if (unbuffered && lay.qns()) { // this reader scatters gate/up/down roles, which a qns pack does not have + st.ok = false; + st.error = "qns packs are Linux-only for now"; + return st; + } if (unbuffered) { uint64_t max_chunk = 0; for (int64_t l = 0; l < lay.n_layers; ++l) { @@ -5199,6 +5219,19 @@ LoadStats load_experts_gguf(const std::string& gguf, uint8_t* dst, const strata: if (l >= lay.n_layers || bad) break; wait_ready(l); const auto& fm = lay.fmt[(size_t) l]; + if (fm.gu_type == strata::kernels::cpu::kQnsType) { // the payload is the arena's layer, byte for byte + const std::string name = expert_gguf_file(gguf, lay, l, 0); + if (name != open_name) { + if (file.f != nullptr) std::fclose(file.f); + file.f = std::fopen(name.c_str(), "rb"); + if (file.f == nullptr) { bad = true; return; } + open_name = name; + } + const uint64_t start = lay.blob_offset(l, 0), span = lay.blob_offset(l, lay.n_expert) - start; + if (STRATA_FSEEK64(file.f, lay.gguf_off[(size_t) (3 * l)]) != 0 || + std::fread(dst + start, 1, (size_t) span, file.f) != (size_t) span) { bad = true; return; } + continue; + } const uint64_t blob = lay.bytes[(size_t) l]; const uint64_t per[3] = {fm.up_off, fm.up_off, blob - fm.down_off}; const uint64_t at[3] = {0, fm.up_off, fm.down_off}; diff --git a/src/core/mtp.cpp b/src/core/mtp.cpp index 58627029de..a1d0e53935 100644 --- a/src/core/mtp.cpp +++ b/src/core/mtp.cpp @@ -8,6 +8,9 @@ #include "strata/core/progress.hpp" #include "strata/core/native_head.hpp" +#include "strata/core/weights.hpp" +#include "strata/kernels/qns_gather.hpp" +#include "strata/kernels/qns_cuda.hpp" #include "strata/core/peer_experts.hpp" #include "strata/platform/memory.hpp" #include "strata/kernels/bf16_gemv.hpp" @@ -160,6 +163,8 @@ MtpDrafter::~MtpDrafter() { if (arena_) cudaFree(arena_); if (head_logits_) cudaFree(head_logits_); if (owns_draft_head_ && dhead_) { strata::kernels::native_q6_k_unpack(dhead_); cudaFree(dhead_); } + if (owns_draft_head_ && dscales_) cudaFree(dscales_); + if (qns_act_) cudaFree(qns_act_); if (owns_draft_head_ && dvocab_) cudaFree(dvocab_); if (ev_chain_) cudaEventDestroy(ev_chain_); for (cudaEvent_t e : ev_step_) if (e) cudaEventDestroy(e); @@ -684,6 +689,7 @@ bool MtpDrafter::bind(const WeightTable& wt, const NativeHead* head, const float return false; } dhead_ = shared->dhead_; + dscales_ = shared->dscales_; dhead_type_ = shared->dhead_type_; // the subset's type (the main head's, or Q4_0): -1 failed every --batch-mtp step dvocab_ = shared->dvocab_; n_dvocab_ = shared->n_dvocab_; @@ -695,8 +701,11 @@ bool MtpDrafter::bind(const WeightTable& wt, const NativeHead* head, const float if (read_file(vocab_file(), raw) && raw.size() >= 4 && raw.size() % 4 == 0) { n_dvocab_ = (int64_t) (raw.size() / 4); const int64_t row_bytes = (int64_t) head->row_bytes(); // a vocabulary row of the native head + const bool qns = head->type() == kQnsDenseType; // trellis rows: whole 32-row groups + const int64_t n_rows = qns ? (n_dvocab_ + 31) / 32 * 32 : n_dvocab_; if (cudaMalloc((void**) &dvocab_, raw.size()) != cudaSuccess || - cudaMalloc((void**) &dhead_, (size_t) (n_dvocab_ * row_bytes)) != cudaSuccess) { + cudaMalloc((void**) &dhead_, (size_t) (n_rows * row_bytes)) != cudaSuccess || + (qns && cudaMalloc((void**) &dscales_, (size_t) n_rows * sizeof(float)) != cudaSuccess)) { err = "mtp: the draft head does not fit"; draft_head_hint(n_dvocab_, row_bytes); return false; @@ -707,28 +716,40 @@ bool MtpDrafter::bind(const WeightTable& wt, const NativeHead* head, const float const int head_dev = head->device(); if (head_dev >= 0 && head_dev != device_) { // the main head is on another card (the drafter on the first stage of a layer split, 0.1.44): the rows are - // gathered where the head is and the subset copied over once + // gathered where the head is and the subset copied over once (a trellis head: its row scales too) void* tmp = nullptr; + float* tsc = nullptr; int32_t* dv = nullptr; bool ok_x = true; { const OnDevice on_head(head_dev); - ok_x = cudaMalloc((void**) &tmp, (size_t) (n_dvocab_ * row_bytes)) == cudaSuccess && + ok_x = cudaMalloc((void**) &tmp, (size_t) (n_rows * row_bytes)) == cudaSuccess && + (!qns || cudaMalloc((void**) &tsc, (size_t) n_rows * sizeof(float)) == cudaSuccess) && cudaMalloc((void**) &dv, raw.size()) == cudaSuccess && cudaMemcpy(dv, raw.data(), raw.size(), cudaMemcpyHostToDevice) == cudaSuccess; if (ok_x) { - strata::kernels::gather_rows((const uint8_t*) head->weights(), row_bytes, dv, n_dvocab_, - (uint8_t*) tmp, nullptr); + if (qns) + strata::kernels::qns_gather_rows((const uint32_t*) head->weights(), head->qns_scales(), + (int) (row_bytes / 4), dv, (int) n_dvocab_, (int) n_rows, + (uint32_t*) tmp, tsc, nullptr); + else + strata::kernels::gather_rows((const uint8_t*) head->weights(), row_bytes, dv, n_dvocab_, + (uint8_t*) tmp, nullptr); ok_x = cudaDeviceSynchronize() == cudaSuccess; } } - ok_x = ok_x && cudaMemcpy(dhead_, tmp, (size_t) (n_dvocab_ * row_bytes), cudaMemcpyDeviceToDevice) == cudaSuccess; + ok_x = ok_x && cudaMemcpy(dhead_, tmp, (size_t) (n_rows * row_bytes), cudaMemcpyDeviceToDevice) == cudaSuccess && + (!qns || cudaMemcpy(dscales_, tsc, (size_t) n_rows * sizeof(float), cudaMemcpyDeviceToDevice) == cudaSuccess); { const OnDevice on_head(head_dev); if (tmp) cudaFree(tmp); + if (tsc) cudaFree(tsc); if (dv) cudaFree(dv); } if (!ok_x) { err = "mtp: the draft head's rows could not be copied from the head's card"; return false; } + } else if (qns) { + strata::kernels::qns_gather_rows((const uint32_t*) head->weights(), head->qns_scales(), (int) (row_bytes / 4), + dvocab_, (int) n_dvocab_, (int) n_rows, (uint32_t*) dhead_, dscales_, nullptr); } else { strata::kernels::gather_rows((const uint8_t*) head->weights(), row_bytes, dvocab_, n_dvocab_, dhead_, nullptr); } @@ -737,12 +758,18 @@ bool MtpDrafter::bind(const WeightTable& wt, const NativeHead* head, const float std::fprintf(stderr, "strata mtp: draft head over %lld tokens (%.1f MiB)\n", (long long) n_dvocab_, (double) (n_dvocab_ * row_bytes) / 1048576.0); dhead_type_ = head->type(); + if (q4_head_ && qns) { err = "mtp: --mtp-q4 head does not apply to a trellis head"; return false; } if (q4_head_ && !make_q4_head(err)) return false; if (dhead_type_ == 14 && strata::kernels::native_q6_k_packed_enabled()) // STRATA_Q6_PACKED=1 strata::kernels::native_q6_k_pack(dhead_, (int) (row_bytes / 210) * 256, (int) n_dvocab_, "MTP draft head"); } } + if (head->type() == kQnsDenseType && qns_act_ == nullptr && + cudaMalloc((void**) &qns_act_, (size_t) max_t_ * strata::kernels::qns_dense_act_bytes((int) g_->n_embd)) != cudaSuccess) { + err = "mtp: the trellis head's activation does not fit"; + return false; + } if (coupled_draft_env() && cparams_ == nullptr && !setup_coupled(err)) return false; return true; } @@ -1190,11 +1217,18 @@ bool MtpDrafter::record_rest(int step_row, cudaStream_t cs, std::string& err) { sample_ + t * N, dummy_inj_, cs); } g_mp.stamp(13, cs); - native_quantize_q8_1(sample_, xq_, (int) N, T, cs); const bool sub = dhead_ != nullptr; const int64_t nv = sub ? n_dvocab_ : n_vocab_; - native_mmvq(sub ? dhead_type_ : head_->type(), sub ? dhead_ : head_->weights(), xq_, head_logits_, (int) N, - (int) nv, T, cs); + if (head_->type() == kQnsDenseType) { // the head as trellis (the draft subset regrouped at bind) + strata::kernels::qns_dense_quant(sample_, T, (int) N, nullptr, qns_act_, cs); + strata::kernels::qns_dense_gemv((const uint32_t*) (sub ? (const void*) dhead_ : head_->weights()), + sub ? dscales_ : head_->qns_scales(), (int) ((nv + 31) / 32 * 32), (int) N, + head_->qns_bits(), qns_act_, T, head_logits_, cs); + } else { + native_quantize_q8_1(sample_, xq_, (int) N, T, cs); + native_mmvq(sub ? dhead_type_ : head_->type(), sub ? dhead_ : head_->weights(), xq_, head_logits_, (int) N, + (int) nv, T, cs); + } g_mp.stamp(14, cs); if (coupled_rec_) { // coupled draft sampling: the target's chain and Philox draw for the row that will verify this draft diff --git a/src/core/native_dense.cpp b/src/core/native_dense.cpp index aba2882b83..0f2b767d82 100644 --- a/src/core/native_dense.cpp +++ b/src/core/native_dense.cpp @@ -1,4 +1,5 @@ #include "strata/core/native_dense.hpp" +#include "strata/kernels/qns_cuda.hpp" #include #include "strata/core/weights.hpp" #include "strata/artifact/gguf_reader.hpp" @@ -158,7 +159,19 @@ struct Pending { uint64_t bytes; DevicePtr data; DevicePtr packed; // STRATA_Q8_PACKED=1 copy, or null + DevicePtr qns_scales, qns_perm; // kQnsDenseType }; +std::string g_qns_dense; +DevicePtr upload(const strata::GgufFile& gguf, const strata::TensorInfo& t, std::string& err) { + const uint64_t bytes = strata::tensor_payload_bytes(t); + if (bytes == 0) { err = "native dense: " + t.name + " has no readable payload"; return {}; } + void* allocation = nullptr; + cudaError_t status = cudaMalloc(&allocation, bytes); + DevicePtr data(allocation); + if (status == cudaSuccess) status = cudaMemcpy(data.get(), gguf.tensor_data(t), bytes, cudaMemcpyHostToDevice); + if (status != cudaSuccess) { err = "native dense upload " + t.name + ": " + cudaGetErrorString(status); return {}; } + return data; +} } bool NativeDense::served_names(const std::vector& shards, bool include_ple_key, @@ -179,6 +192,8 @@ bool NativeDense::served_names(const std::vector& shards, bool incl } void NativeDense::set_layer_range(int lb, int le) { g_layer_lb = lb; g_layer_le = le; } +void NativeDense::set_qns_dense(const std::string& path) { g_qns_dense = path; } +const std::string& NativeDense::qns_dense() { return g_qns_dense; } // The byte walk `load` does for the layers [lb, le), without the allocations: the same filters in the same order, // reading the GGUF headers and the canonical table only, so it needs no device. A later stage's projections are not // on the card when the split search prices that stage, and `load` runs one cudaMalloc per matrix (#1238). The two @@ -261,6 +276,17 @@ bool NativeDense::load(const std::vector& shards, WeightTable& tabl std::vector pending; std::set seen; int max_in = 0; + std::unique_ptr qns; // --qns-dense-gguf: backbone matrices served as trellis + int qns_bits = 0; + if (!g_qns_dense.empty()) { + qns = std::make_unique(g_qns_dense); + const strata::MetaValue* bits = qns->get("qns.dense_bits"); + const strata::TensorInfo* signs = qns->find("qns.rot_signs_6144"); + const strata::TensorInfo* q = qns->find("qns.rot_q_6144"); + if (!bits || !signs || !q) { err = "native dense: " + g_qns_dense + " is not a qns backbone GGUF"; return false; } + qns_bits = (int) bits->u; + strata::kernels::qns_cuda_set_rotation_6144((const float*) qns->tensor_data(*signs), (const float*) qns->tensor_data(*q)); + } uint64_t total = 0, hc_q8_bytes = 0, allocated = 0; uint64_t split_count = 0, split_tensors = 0; std::set split_numbers; @@ -347,6 +373,22 @@ bool NativeDense::load(const std::vector& shards, WeightTable& tabl tensor.shape[0] != (uint64_t) ref.ne0 || tensor.shape[1] != (uint64_t) ref.ne1) { err = "native dense: incompatible matrix " + tensor.name; return false; } + if (const strata::TensorInfo* words = qns ? qns->find(tensor.name + ".qns_words") : nullptr) { + const strata::TensorInfo* scales = qns->find(tensor.name + ".qns_scales"); + const strata::TensorInfo* perm = qns->find(tensor.name + ".qns_perm"); + if (!scales || scales->shape[0] != (uint64_t) ref.ne1 || (perm && perm->shape[0] != (uint64_t) ref.ne0)) { + err = "native dense: " + tensor.name + " in " + g_qns_dense + " does not match the model"; return false; + } + DevicePtr w = upload(*qns, *words, err), sc = upload(*qns, *scales, err); + DevicePtr pm = perm ? upload(*qns, *perm, err) : DevicePtr(); + if (!w || !sc || (perm && !pm)) return false; + const uint64_t qbytes = strata::tensor_payload_bytes(*words) + strata::tensor_payload_bytes(*scales) + + (perm ? strata::tensor_payload_bytes(*perm) : 0); + total += qbytes; + allocated += alloc_bytes(qbytes); + pending.push_back(Pending{&ref, kQnsDenseType, qbytes, std::move(w), {}, std::move(sc), std::move(pm)}); + continue; + } const auto bytes = strata::kernels::native_mmvq_weight_bytes( tensor.type, (int) ref.ne0, (int) ref.ne1); void* allocation = nullptr; @@ -379,7 +421,7 @@ bool NativeDense::load(const std::vector& shards, WeightTable& tabl return false; } } - pending.push_back(Pending{&ref, (int) tensor.type, bytes, std::move(data), std::move(packed)}); + pending.push_back(Pending{&ref, (int) tensor.type, bytes, std::move(data), std::move(packed), {}, {}}); } if (hc_q8_env()) for (const auto& tensor : gguf.tensors()) { @@ -477,6 +519,13 @@ bool NativeDense::load(const std::vector& shards, WeightTable& tabl item.ref->native_data = item.data.get(); item.ref->native_type = item.type; item.ref->native_q8_1 = scratch.get(); + if (item.type == kQnsDenseType) { + item.ref->qns_scales = (const float*) item.qns_scales.get(); + item.ref->qns_perm = (const int32_t*) item.qns_perm.get(); + item.ref->qns_bits = qns_bits; + weights_.push_back(item.qns_scales.release()); + if (item.qns_perm) weights_.push_back(item.qns_perm.release()); + } if (item.packed) { strata::kernels::native_q8_0_packed_register(item.data.get(), item.packed.get(), (int) item.ref->ne0, (int) item.ref->ne1); diff --git a/src/core/native_head.cpp b/src/core/native_head.cpp index a1fb19441e..36af6aedfe 100644 --- a/src/core/native_head.cpp +++ b/src/core/native_head.cpp @@ -1,4 +1,7 @@ #include "strata/core/native_head.hpp" +#include "strata/core/native_dense.hpp" +#include "strata/core/weights.hpp" +#include "strata/kernels/qns_cuda.hpp" #include "strata/artifact/gguf_reader.hpp" #include "strata/kernels/iq_kernels.hpp" #include "strata/kernels/native_mmvq.hpp" @@ -20,6 +23,7 @@ void NativeHead::unload() { if (weights_) strata::kernels::native_q6_k_unpack(weights_); if (weights_) strata::kernels::native_q8_0_sm60_release(weights_); if (scratch_) cudaFree(scratch_); + if (qns_scales_) cudaFree(qns_scales_); if (weights_) cudaFree(weights_); scratch_ = nullptr; weights_ = nullptr; @@ -27,6 +31,34 @@ void NativeHead::unload() { bool NativeHead::load(const std::vector& shards, int64_t n_in, int64_t n_out, std::string& err) { if (loaded()) { err = "native head is already loaded"; return false; } + if (!NativeDense::qns_dense().empty()) { // the head as trellis, when the qns backbone GGUF holds it + try { + const strata::GgufFile qns(NativeDense::qns_dense()); + const strata::TensorInfo* words = qns.find("output.weight.qns_words"); + const strata::TensorInfo* scales = qns.find("output.weight.qns_scales"); + const strata::MetaValue* bits = qns.get("qns.dense_bits"); + if (words && scales && bits) { + if (scales->shape[0] != (uint64_t) n_out || n_out % 32) { err = "native head: the qns head does not match"; return false; } + const uint64_t wb = strata::tensor_payload_bytes(*words), sb = strata::tensor_payload_bytes(*scales); + if (cudaMalloc(&weights_, wb) != cudaSuccess || cudaMalloc((void**) &qns_scales_, sb) != cudaSuccess || + cudaMalloc(&scratch_, strata::kernels::qns_dense_act_bytes((int) n_in)) != cudaSuccess || + cudaMemcpy(weights_, qns.tensor_data(*words), wb, cudaMemcpyHostToDevice) != cudaSuccess || + cudaMemcpy(qns_scales_, qns.tensor_data(*scales), sb, cudaMemcpyHostToDevice) != cudaSuccess) { + err = "native head: the qns head upload failed"; + return false; + } + bytes_ = wb; // row_bytes(): a row's words (the drafter's gather); the scales are beside them + n_in_ = (int) n_in; + n_out_ = (int) n_out; + type_ = kQnsDenseType; + qns_bits_ = (int) bits->u; + return true; + } + } catch (const std::exception& error) { + err = std::string("native head: ") + error.what(); + return false; + } + } if (n_in <= 0 || n_out <= 0 || n_in > INT_MAX || n_out > INT_MAX || n_in % 256) { err = "native head requires positive int32 dimensions and whole 256-value rows"; return false; @@ -89,7 +121,11 @@ bool NativeHead::run(const float* mixed, float* logits, void* stream, std::strin return false; } try { - if (type_ == 13) { + if (type_ == kQnsDenseType) { + strata::kernels::qns_dense_quant(mixed, 1, n_in_, nullptr, scratch_, (cudaStream_t) stream); + strata::kernels::qns_dense_gemv((const uint32_t*) weights_, qns_scales_, n_out_, n_in_, qns_bits_, scratch_, 1, + logits, (cudaStream_t) stream); + } else if (type_ == 13) { strata::kernels::native_q5_k_f32(weights_, mixed, scratch_, logits, n_in_, n_out_, 1, stream); } else { strata::kernels::native_quantize_q8_1(mixed, scratch_, n_in_, 1, stream); diff --git a/src/core/verify.cpp b/src/core/verify.cpp index e8f747f403..38acc14a94 100644 --- a/src/core/verify.cpp +++ b/src/core/verify.cpp @@ -23,6 +23,7 @@ #include "strata/core/on_device.hpp" #include "strata/core/peer_experts.hpp" #include "strata/kernels/iq_kernels.hpp" +#include "strata/kernels/qns_cuda.hpp" #include "strata/kernels/cpu/expert_layout.hpp" #include "strata/kernels/cpu/hadamard.hpp" #include "strata/kernels/bf16_gemv.hpp" @@ -813,9 +814,16 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& hit_xq_ = b.take(T * (N / 32) * 34); hit_xs_ = b.take(T * (N / 32)); nat_xq_ = b.take(T * (N / 32) * 36); nat_xr_ = had_gu_any ? b.take(T * N) : nullptr; - hit_scratch_ = b.take(std::max( + if (strata::kernels::cpu::expert_layout().qns()) { + qns_xq_ = b.take(T * strata::kernels::qns_act_bytes()); + // a trellis backbone (--qns-dense-gguf): the 2560-wide layer input and the 6144-wide output-projection input + qns_xd_ = b.take(T * strata::kernels::qns_dense_act_bytes(2560)); + qns_od_ = b.take(T * strata::kernels::qns_dense_act_bytes(6144)); + } + hit_scratch_ = b.take(std::max({ strata::kernels::moe_hit_grouped_scratch_bytes((int64_t) (T * K), g.n_embd, g.n_ff), - strata::kernels::native_expert_scratch_bytes((int64_t) (T * K), g.n_ff))); + strata::kernels::native_expert_scratch_bytes((int64_t) (T * K), g.n_ff), + strata::kernels::cpu::expert_layout().qns() ? strata::kernels::qns_grouped_scratch_bytes((int64_t) (T * K)) : 0})); head_mixed_ = b.take(T * N); head_inj_ = b.take(HC); sh_bf16_ = b.take(T * N); sh_gate_ = b.take(T * (uint64_t) g.n_ff); sh_up_ = b.take(T * (uint64_t) g.n_ff); sh_g_ = b.take(T + 4); @@ -1311,6 +1319,11 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { il_ready = false; }; auto mm = [&](const WeightRef* w, float* out, int n_in, int n_out) { + if (w->native_type == kQnsDenseType) { // trellis rows read the rotated input quantized beside xq_ + strata::kernels::qns_dense_gemv((const uint32_t*) w->native_data, w->qns_scales, n_out, n_in, w->qns_bits, + n_in == (int) N ? qns_xd_ : qns_od_, n, out, cs); + return; + } if (g_mmvq_il() && strata::kernels::native_mmvq_il_supported(w->native_type, n, n_out)) { if (!il_ready && n >= 2) { // one column reads the plain q8_1 blocks strata::kernels::native_q8_1_interleave(xq_, xil_, n_in, n, cs); @@ -1348,19 +1361,23 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { if (cudaEventRecord(mjoin_, side_) != cudaSuccess) { err = "verify: the GDN gates' branch could not join"; return false; } } if (!q8_attn) quant_in(xm, (int) N); + if (wqkv->native_type == kQnsDenseType) strata::kernels::qns_dense_quant(xm, n, (int) N, nullptr, qns_xd_, cs); if (side_proj) { // the gate projection (compute-bound for i-quants) beside qkv and the conv // the branch reads F4's interleaved copy too (written on cs before the fork; bitwise native_mmvq's // output), and qkv is recorded before the branch's launch, so its grid is dispatched first and the // conv / recurrence that wait for it start sooner (z is read only after the recurrence) - const bool g_il = g_mmvq_il() && strata::kernels::native_mmvq_il_supported(wg->native_type, n, (int) ZV); + const bool g_qns = wg->native_type == kQnsDenseType; // trellis rows read qns_xd_, written above + const bool g_il = !g_qns && g_mmvq_il() && strata::kernels::native_mmvq_il_supported(wg->native_type, n, (int) ZV); if (g_il && !il_ready && n >= 2) { strata::kernels::native_q8_1_interleave(xq_, xil_, (int) N, n, cs); il_ready = true; } if (!branch(pfork_, cs, side2_, "fork")) return false; mm(wqkv, qkv + (size_t) tb * C, (int) N, (int) C); - if (g_il) strata::kernels::native_mmvq_il(wg->native_type, wg->native_data, xq_, xil_, z_ + (size_t) tb * ZV, - (int) N, (int) ZV, n, side2_); + if (g_qns) strata::kernels::qns_dense_gemv((const uint32_t*) wg->native_data, wg->qns_scales, (int) ZV, (int) N, + wg->qns_bits, qns_xd_, n, z_ + (size_t) tb * ZV, side2_); + else if (g_il) strata::kernels::native_mmvq_il(wg->native_type, wg->native_data, xq_, xil_, z_ + (size_t) tb * ZV, + (int) N, (int) ZV, n, side2_); else native_mmvq(wg->native_type, wg->native_data, xq_, z_ + (size_t) tb * ZV, (int) N, (int) ZV, n, side2_); if (cudaEventRecord(pjoin_, side2_) != cudaSuccess) { err = "verify: the gate projection's branch could not join"; return false; } } @@ -1370,8 +1387,12 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { (const float*) wsa->data, gate + (size_t) tb * HV, beta + (size_t) tb * HV, (int) N, (int) HV, n, df_side_[0]); fork(1); - native_mmvq(wg->native_type, wg->native_data, xq_, z_ + (size_t) tb * ZV, (int) N, (int) ZV, n, - df_side_[1]); + if (wg->native_type == kQnsDenseType) + strata::kernels::qns_dense_gemv((const uint32_t*) wg->native_data, wg->qns_scales, (int) ZV, (int) N, + wg->qns_bits, qns_xd_, n, z_ + (size_t) tb * ZV, df_side_[1]); + else + native_mmvq(wg->native_type, wg->native_data, xq_, z_ + (size_t) tb * ZV, (int) N, (int) ZV, n, + df_side_[1]); } const bool fuse_ab = fuse_convab_on() && !br && !batch_rec_ && !prof_on_ && !sidem; // the side mixer computes the gates on its branch if (!side_proj) mm(wqkv, qkv + (size_t) tb * C, (int) N, (int) C); // with side_proj: above, first @@ -1456,6 +1477,8 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { // destination (#1139: the out-projection read stale bytes), so that case quantizes here as without it if (!g_qfuse() || batch_rec_) quant_in(y_ + (size_t) tb * ZV, (int) ZV); else il_ready = false; + if (wout->native_type == kQnsDenseType) // the engine's GDN output is a column permutation of the checkpoint's + strata::kernels::qns_dense_quant(y_ + (size_t) tb * ZV, n, (int) ZV, wout->qns_perm, qns_od_, cs); mm(wout, bo_ + tb * N, (int) ZV, (int) N); } else { // ======================= QSA ======================= @@ -1503,13 +1526,18 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { }(); const bool qb = dec_batch && n > 1 && native_qsa_enabled() && native_rope_enabled() && (!st.kv_q4 || qb_q4); if (!q8_attn) quant_in(xm, (int) N); + if (wq->native_type == kQnsDenseType) strata::kernels::qns_dense_quant(xm, n, (int) N, nullptr, qns_xd_, cs); // STRATA_DF_BRANCH: the query side (q and its norm/RoPE/rotation on side 0, the indexer query on side 1) // beside the K/V side; the indexer query joins before the block scores, the query before attention const bool qbr = br && qb; if (qbr) { fork(0); - native_mmvq(wq->native_type, wq->native_data, xq_, qfull_ + tb * NH * 2 * HD, (int) N, - (int) (NH * 2 * HD), n, df_side_[0]); + if (wq->native_type == kQnsDenseType) + strata::kernels::qns_dense_gemv((const uint32_t*) wq->native_data, wq->qns_scales, (int) (NH * 2 * HD), + (int) N, wq->qns_bits, qns_xd_, n, qfull_ + tb * NH * 2 * HD, df_side_[0]); + else + native_mmvq(wq->native_type, wq->native_data, xq_, qfull_ + tb * NH * 2 * HD, (int) N, + (int) (NH * 2 * HD), n, df_side_[0]); copy_rows_strided(qcur_ + tb * NH * HD, qfull_ + tb * NH * 2 * HD, (int64_t) n * NH, HD, 2 * HD, df_side_[0]); norm_rope_on(qcur_ + tb * NH * HD, wqn, (int) (n * NH), (int) HD, pos_ + tb * NH, df_side_[0]); @@ -1584,10 +1612,13 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { // the attention's queries: the projection (mm on cs, which may read F4's interleaved copy; on a branch // F4's kernel too, its copy written on cs before the fork - bitwise the same), the q/gate // split with norm and RoPE, the rotation - const bool q_il = side_proj && g_mmvq_il() && + const bool q_il = side_proj && g_mmvq_il() && wq->native_type != kQnsDenseType && strata::kernels::native_mmvq_il_supported(wq->native_type, n, (int) (NH * 2 * HD)); auto q_chain = [&](cudaStream_t ns) -> bool { if (ns == cs) mm(wq, qfull_ + tb * NH * 2 * HD, (int) N, (int) (NH * 2 * HD)); + else if (wq->native_type == kQnsDenseType) // trellis rows read qns_xd_ (written on cs before the fork) + strata::kernels::qns_dense_gemv((const uint32_t*) wq->native_data, wq->qns_scales, (int) (NH * 2 * HD), + (int) N, wq->qns_bits, qns_xd_, n, qfull_ + tb * NH * 2 * HD, ns); else if (q_il) strata::kernels::native_mmvq_il(wq->native_type, wq->native_data, xq_, xil_, qfull_ + tb * NH * 2 * HD, (int) N, (int) (NH * 2 * HD), n, ns); else native_mmvq(wq->native_type, wq->native_data, xq_, qfull_ + tb * NH * 2 * HD, (int) N, @@ -1825,6 +1856,8 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { } stamp(l, 14, grp); quant_in(attn32_ + tb * NH * HD, (int) (NH * HD)); + if (wo->native_type == kQnsDenseType) + strata::kernels::qns_dense_quant(attn32_ + tb * NH * HD, n, (int) (NH * HD), nullptr, qns_od_, cs); mm(wo, bo_ + tb * N, (int) (NH * HD), (int) N); } } catch (const std::exception& e) { @@ -1971,6 +2004,9 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { } else if (!qdedup && !q8_ffn) { quantize_q8_1_rows(xm, n, N, nat_xq_ + (size_t) tb * (N / 32) * 36, cs); } + // qns experts read their own rotated two-plane activation (the q8_1 image above still feeds the shared expert) + if (strata::kernels::cpu::expert_layout().qns()) + strata::kernels::qns_quant_act(xm, n, qns_xq_ + (size_t) tb * strata::kernels::qns_act_bytes(), cs); } else quantize_q8_0_scaled(xm, hit_xq_ + (size_t) tb * (N / 32) * 34, hit_xs_ + (size_t) tb * (N / 32), (int64_t) n * N, cs); stamp(l, 18, grp); @@ -2003,7 +2039,11 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { // plan v0.3 P6: the VRAM groups now; the PCIe groups once the copy engine has landed them in staging. // `gy`: the native launch's groups side by side (0: cap, one block row per possible group). auto grouped = [&](const unsigned long long* gp, const int32_t* gs, const int32_t* gn, int64_t gy, float* dst_buf) { - if (lay.native) { + if (lay.qns()) { + strata::kernels::qns_expert_grouped(gp, gs, gn, p_dst, p_tok, cap, cap, + qns_xq_ + (size_t) tb * strata::kernels::qns_act_bytes(), hit_scratch_, + dst_buf, cs, gy); + } else if (lay.native) { // the layer's GGUF formats (i-quant gate/up, Q2_0 / IQ4_NL down) const auto& f = lay.fmt[(size_t) l]; NativeExpertLayout L = native_expert_layout(f.gu_type, f.d_type, f.n_embd, f.n_ff); @@ -2045,7 +2085,8 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { if (sink_.pcie_mode == 2) { // stage it with a copy kernel, then point at staging const int64_t per = G == 2 ? kStagingBlobs / 2 : kStagingBlobs; uint8_t* stage = staging_ + (size_t) (grp * per) * lay.max_blob; - fetch_blobs(p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), (int) per, cs); + if (lay.qns()) fetch_qns_blobs(p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), (int) per, cs); + else fetch_blobs(p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), (int) per, cs); rebase_ptrs((unsigned long long*) p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), cs); } stamp(l, 21, grp); @@ -2210,7 +2251,11 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { } } } - if (head_ != nullptr && head_->loaded()) { + if (head_ != nullptr && head_->loaded() && head_->type() == kQnsDenseType) { // the head as trellis + strata::kernels::qns_dense_quant(head_mixed_, T, (int) N, nullptr, qns_xd_, cs); + strata::kernels::qns_dense_gemv((const uint32_t*) head_->weights(), head_->qns_scales(), (int) n_vocab_, (int) N, + head_->qns_bits(), qns_xd_, T, head_logits_, cs); + } else if (head_ != nullptr && head_->loaded()) { try { native_quantize_q8_1(head_mixed_, xq_, (int) N, T, cs); if (g_mmvq_il() && strata::kernels::native_mmvq_il_supported(head_->type(), (int) T, (int) n_vocab_)) { diff --git a/src/kernels/cpu/expert_layout.cpp b/src/kernels/cpu/expert_layout.cpp index b6924950da..57ae862e40 100644 --- a/src/kernels/cpu/expert_layout.cpp +++ b/src/kernels/cpu/expert_layout.cpp @@ -1,6 +1,7 @@ // src/kernels/cpu/expert_layout.cpp - plan v0.3 P6: the per-layer expert table. See the header. #include "strata/kernels/cpu/expert_layout.hpp" #include "strata/kernels/cpu/hadamard.hpp" +#include "strata/kernels/cpu/qns.hpp" #include #include @@ -438,6 +439,7 @@ bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n } NativeFmt f; if (!native_fmt((int) gt, (int) dt, H, FF, f, err)) return false; + if (gt == kQnsType) f.bytes = blob; // a qns layer's blobs differ in size: this is the largest if (f.bytes != blob) { err = "native_experts.txt: layer " + std::to_string(l) + " blob is " + std::to_string(blob) + " B but its formats make " + std::to_string(f.bytes); @@ -480,25 +482,52 @@ bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n L.bytes[(size_t) l] = blob; if (blob > L.max_blob) L.max_blob = blob; } + const bool qns = L.fmt[0].gu_type == kQnsType; // a qns pack has qns experts in every layer + // v6 marks a qns pack (v5 is the Hadamard-folded experts below); strata-qns.patch for 0.1.43 wrote it as v5 + if (qns != (L.version == 6)) { + err = qns ? "native_experts.txt v" + std::to_string(L.version) + " has qns layers, which this engine reads " + "from v6 packs only (a pack written for strata-qns.patch on 0.1.43?): repack with tools/iq_pack.py" + : std::string("native_experts.txt v6 declares a qns pack, but its layers are not qns"); + return false; + } + if (qns) { + auto read_all = [&](const std::string& name, void* dst, size_t bytes) { + std::ifstream f(pack_dir + "/" + name, std::ios::binary); + f.read((char*) dst, (std::streamsize) bytes); + if (f.gcount() == (std::streamsize) bytes && f.peek() == EOF) return true; + err = "a qns pack needs " + name + " of " + std::to_string(bytes) + " bytes (tools/iq_pack.py writes it)"; + return false; + }; + L.qns_off.resize((size_t) (n_layers * (L.n_expert + 1))); + L.qns_rotation.resize(2560 + 25 + 640 + 25); + if (!read_all("qns_offsets.bin", L.qns_off.data(), L.qns_off.size() * 8) || + !read_all("qns_rotation.bin", L.qns_rotation.data(), L.qns_rotation.size() * 4)) + return false; + } uint64_t at = 0; for (int64_t l = 0; l < n_layers; ++l) { - if (L.offset[(size_t) l] != at) { + if (L.offset[(size_t) l] != at || (L.fmt[(size_t) l].gu_type == kQnsType) != qns || + (qns && L.blob_offset(l, 0) != at)) { err = "native_experts.txt: layer " + std::to_string(l) + " is missing or not contiguous"; return false; } - at += L.bytes[(size_t) l] * (uint64_t) L.n_expert; + at = qns ? L.blob_offset(l, L.n_expert) : at + L.bytes[(size_t) l] * (uint64_t) L.n_expert; + } + L.total = qns ? at + L.max_blob : at; // qns: slack for a layer-max copy of the last blob + if (qns) { + const float* r = L.qns_rotation.data(); + qns::set_rotation(r, r + 2560, r + 2585, r + 3225); } - L.total = at; // v5: Agention's Hadamard-folded experts (prism.hadamard, hadamard.hpp). The pack must say which layers are // rotated and with what signs; a v5 pack without hadamard.txt (or the reverse) is refused, never run unrotated. { std::ifstream probe(pack_dir + "/hadamard.txt"); const bool have = (bool) probe; - if (L.version >= 5 && !have) { + if (L.version == 5 && !have) { err = "native_experts.txt v5 declares Hadamard-folded experts, but " + pack_dir + "/hadamard.txt is missing"; return false; } - if (have && L.version < 5) { + if (have && L.version != 5) { err = pack_dir + "/hadamard.txt exists but native_experts.txt is v" + std::to_string(L.version) + ": the pack is inconsistent (repack with tools/iq_pack.py)"; return false; diff --git a/src/kernels/cpu/native_expert.cpp b/src/kernels/cpu/native_expert.cpp index 57cb6dfa93..3fd3c7fa14 100644 --- a/src/kernels/cpu/native_expert.cpp +++ b/src/kernels/cpu/native_expert.cpp @@ -12,6 +12,7 @@ #include "strata/kernels/cpu/expert_layout.hpp" #include "strata/kernels/cpu/hadamard.hpp" #include "strata/kernels/cpu/tq_cpu.hpp" +#include "strata/kernels/cpu/qns.hpp" #include "ggml.h" #include "ggml-cpu.h" @@ -60,6 +61,20 @@ bool iq256_supported(int type) noexcept { bool kq256_supported(int type) noexcept { return type == 12 || type == 7 || type == 8; } bool native_fmt(int gu_type, int d_type, int64_t n_embd, int64_t n_ff, NativeFmt& f, std::string& err) { + if (gu_type == kQnsType || d_type == kQnsType) { // Qwen Next S trellis experts (qns.hpp); blob size per layer + static_assert(qns::kActBytes <= kNativeActBytes && qns::kHBytes <= kNativeHBytes, "qns activations fit the pool"); + if (gu_type != d_type || n_embd != 2560 || n_ff != 640) { + err = "native experts: a qns layer is qns gate/up and down at 2560 x 640"; + return false; + } + f = NativeFmt{}; + f.gu_type = f.d_type = f.gu_act = f.d_act = kQnsType; + f.n_embd = n_embd; + f.n_ff = n_ff; + f.act_bytes = qns::kActBytes; + f.h_bytes = qns::kHBytes; + return true; + } init_once(); // The trellis types (Gyro, ggml types 144-146) are not in a ggml built from upstream (and an upstream ggml indexes its // type table by the id, so asking it about 144 reads past the table): their geometry and activation format (Q8_0) @@ -130,6 +145,7 @@ bool q8k_avx1(int type) { } // namespace void native_quant_act(const NativeFmt& f, const float* x, void* dst) { + if (f.gu_type == kQnsType) { qns::quant_act(x, dst); return; } if (f.had_gu) { // Hadamard-folded gate/up: quantize H (s * x), never x itself (hadamard.hpp) thread_local std::vector tmp; tmp.assign(x, x + f.n_embd); @@ -142,6 +158,7 @@ void native_quant_act(const NativeFmt& f, const float* x, void* dst) { } void native_quant_h(const NativeFmt& f, const float* h, void* dst) { + if (f.d_type == kQnsType) { qns::quant_h(h, dst); return; } if (f.had_d) { thread_local std::vector tmp; tmp.assign(h, h + f.n_ff); @@ -154,6 +171,7 @@ void native_quant_h(const NativeFmt& f, const float* h, void* dst) { } int native_gu_mt_min(int gu_type) { + if (gu_type == kQnsType) return 1; // one kernel for every group size: a token's rows never depend on the drafting // #152: from how many tokens the multi-token kernels run (ggml's vec_dot below that). The default 2 is the // measured-fastest rule, but a token's expert rows then round differently alone than in a group, so greedy output // can depend on how many drafts a verify window held. STRATA_IQ_MT_MIN=1 (opt-in, 0.1.30) uses the multi-token @@ -174,6 +192,7 @@ int native_gu_mt_min(int gu_type) { void native_gu_rows(const NativeFmt& f, const uint8_t* blob, const void* const* act, int nt, float* const* ff, int r0, int r1) { + if (f.gu_type == kQnsType) { qns::gu_rows(blob, act, nt, ff, r0, r1); return; } // the multi-token kernels decode the weights once for all tokens: 2.0-2.4x ggml-cpu at three tokens, no faster // at one (all are bound by the codebook lookups, ~5 GB/s per core), measured by native_expert_parity. AVX-512 // first, then the AVX-2 one (Zen 2/3, Intel 12th-14th gen). STRATA_NO_IQ512 drops an AVX-512 CPU to the @@ -286,6 +305,7 @@ void native_gu_rows(const NativeFmt& f, const uint8_t* blob, const void* const* void native_down_rows(const NativeFmt& f, const uint8_t* blob, const void* const* hq, int nt, float* const* out, int r0, int r1) { + if (f.d_type == kQnsType) { qns::down_rows(blob, hq, nt, out, r0, r1); return; } // IQ4_NL down rows: the AVX-2 multi-token kernel decodes the nibbles and absolutises the weights once per // block instead of once per token; ggml-cpu's dot is single-token. STRATA_NO_IQ4NL falls back to it. static const bool iq4nl_mt = std::getenv("STRATA_NO_IQ4NL") == nullptr; diff --git a/src/kernels/cpu/pool.cpp b/src/kernels/cpu/pool.cpp index fd6952eb57..b0e741ffed 100644 --- a/src/kernels/cpu/pool.cpp +++ b/src/kernels/cpu/pool.cpp @@ -826,7 +826,8 @@ void ExpertPool::drain(int id, ExpertScratch& scratch, uint32_t epoch) { // plan v0.3 P6: native layers, 5 = gate/up rows, 6 = down rows (the layer's own n_ff / n_embd: Flash-Next's // FF / H, or another family's smaller expert) const int per = mode_ == 5 ? (int) nfmt_->n_ff : (int) nfmt_->n_embd; - const int64_t g0 = mrows_ * (int64_t) i / mtasks_, g1 = mrows_ * (int64_t) (i + 1) / mtasks_; + int64_t g0 = mrows_ * (int64_t) i / mtasks_, g1 = mrows_ * (int64_t) (i + 1) / mtasks_; + if (nfmt_->gu_type == kQnsType) { g0 &= ~7ll; g1 &= ~7ll; } // whole 8-row blocks (qns.hpp) for (int64_t r = g0; r < g1;) { const int e = (int) (r / per), r0 = (int) (r % per); const int r1 = (int) std::min(per, r0 + (g1 - r)); diff --git a/src/kernels/cpu/qns_avx2.cpp b/src/kernels/cpu/qns_avx2.cpp new file mode 100644 index 0000000000..90afdd8b56 --- /dev/null +++ b/src/kernels/cpu/qns_avx2.cpp @@ -0,0 +1,350 @@ +// src/kernels/cpu/qns_avx2.cpp - qns.hpp on AVX2 + FMA (Haswell / Zen 2 and later; no AVX-512). +// +// Rows are stored in groups of 32 as [group][word][lane] u32 words (qns_blob.py), so word k of 8 consecutive rows is +// one 32-byte load and a vector lane is a row. A chunk = 32 steps = 64 columns = T words, and with the next word every +// state of the chunk is a 16-bit window at a compile-time bit offset. Per pair of steps: two fmix32 (vpmulld), both +// hashes' bytes 0, 1 folded into one dword, level + 54 of the 4 bytes by two vpshufb nibble lookups, split into the +// even and odd columns as int16 and multiplied by the token's int16 254 q0 + q1 of the same columns (vpmaddwd). +// Zen 2 (EPYC 7H12) per 8 rows x 2 steps (32 weights): ~14 cycles for one token, +2.3 per more token sharing the expert +// (a decoded chunk is reused). The floor is the one integer-multiply pipe: 4 vpmulld (2 cycles each) + 2 vpmaddwd. +#include "strata/kernels/cpu/qns.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace strata::kernels::cpu::qns { +namespace { + +constexpr float kMean0 = 1.5769348f, kMean1 = 1.4910126f, kRms = 19.183672f; + +struct ActHeader { + float step; + int32_t sum; // sum of 254 q0 + q1 over all columns + float even, odd; // E, O +}; +static_assert(sizeof(ActHeader) == 16); + +// One input width N = 5 n: the signs, and the 5 x 5 factor Q as 9 diagonals over 8 consecutive columns from a column +// r (mod 5): diag[r][d + 4][l] = Q[o][o + d] with o = (r + l) % 5 (0 outside Q), so (I (x) Q) y at columns j .. j + 7 is +// the sum over d of diag[j % 5][d + 4] * y[j + d .. j + d + 7]. +template +struct Rotation { + alignas(32) float signs[N]; + alignas(32) float diag[5][9][8]; +}; +Rotation<2560> g_rot2560; +Rotation<640> g_rot640; + +template +void set(Rotation& rot, const float* signs, const float* q) { + std::memcpy(rot.signs, signs, sizeof(rot.signs)); + for (int r = 0; r < 5; ++r) + for (int d = -4; d <= 4; ++d) + for (int l = 0; l < 8; ++l) { + const int o = (r + l) % 5, c = o + d; + rot.diag[r][d + 4][l] = c >= 0 && c < 5 ? q[o * 5 + c] : 0.0f; + } +} + +// Walsh-Hadamard butterflies of stride 5, 10, 20 for 8 columns from j: which lanes are the upper half of their pair. +// The pattern depends on j % 40, i.e. on (j / 8) % 5. +constexpr std::array kUpper = [] { + std::array t{}; + for (int s = 0; s < 3; ++s) + for (int k = 0; k < 5; ++k) + for (int l = 0; l < 8; ++l) t[(s * 5 + k) * 8 + l] = (8 * k + l) % (10 << s) >= (5 << s) ? -1 : 0; + return t; +}(); + +// x_rot = (H_{N / 5} (x) Q)(signs * x) / sqrt(N / 5), quantized into dst (layout in qns.hpp). Every float operation is +// the plain scalar recipe's, in its order (Q's sum over c = 0 .. 4 from 0, butterflies a + b and a - b, |v * norm|, +// v * norm * inv), so the bytes do not depend on the vectorization. That needs mul and add kept apart: GCC fuses them +// into FMAs depending on its tuning, Clang and MSVC do not across intrinsics. +#if defined(__GNUC__) && !defined(__clang__) +#pragma GCC push_options +#pragma GCC optimize("fp-contract=off") +#endif +template +void quantize(const float* x, const Rotation& rot, void* dst) { + constexpr int pad = 24; // the diagonals read columns j - 4 .. j + 11, the small butterflies j - 20 .. j + 27 + alignas(32) float buf_a[pad + N + pad], buf_b[pad + N + pad]; + std::memset(buf_a, 0, pad * sizeof(float)); // the diagonals multiply the pads by 0: they must be finite + std::memset(buf_a + pad + N, 0, pad * sizeof(float)); + std::memset(buf_b, 0, pad * sizeof(float)); + std::memset(buf_b + pad + N, 0, pad * sizeof(float)); + float* a = buf_a + pad; + float* b = buf_b + pad; + for (int j = 0; j < N; j += 8) _mm256_store_ps(a + j, _mm256_mul_ps(_mm256_load_ps(rot.signs + j), _mm256_loadu_ps(x + j))); + for (int j = 0; j < N; j += 32) { // b = (I (x) Q) a, four sums at once (each is a chain of 9 dependent adds) + const auto term = [&](int k, int d) { + return _mm256_mul_ps(_mm256_load_ps(rot.diag[(j + 8 * k) % 5][d + 4]), _mm256_loadu_ps(a + j + 8 * k + d)); + }; + __m256 acc0 = _mm256_setzero_ps(), acc1 = acc0, acc2 = acc0, acc3 = acc0; + for (int d = -4; d <= 4; ++d) { + acc0 = _mm256_add_ps(acc0, term(0, d)); + acc1 = _mm256_add_ps(acc1, term(1, d)); + acc2 = _mm256_add_ps(acc2, term(2, d)); + acc3 = _mm256_add_ps(acc3, term(3, d)); + } + _mm256_store_ps(b + j, acc0); + _mm256_store_ps(b + j + 8, acc1); + _mm256_store_ps(b + j + 16, acc2); + _mm256_store_ps(b + j + 24, acc3); + } + // Walsh-Hadamard over i of column 5 i + o: strides 5, 10, 20 from one buffer into the other (b -> a -> b -> a), as + // a + b (lower lanes) and a - b computed from the column one stride below (upper lanes); 40 and up in place in a + for (int s = 0; s < 3; ++s) { + const int len = 5 << s; + const float* from = s == 1 ? a : b; + float* to = s == 1 ? b : a; + for (int j = 0; j < N; j += 8) { + const __m256 here = _mm256_load_ps(from + j); + const __m256 sum = _mm256_add_ps(here, _mm256_loadu_ps(from + j + len)); + const __m256 diff = _mm256_sub_ps(_mm256_loadu_ps(from + j - len), here); + const __m256i upper = _mm256_loadu_si256(reinterpret_cast(kUpper.data() + (s * 5 + (j / 8) % 5) * 8)); + _mm256_store_ps(to + j, _mm256_blendv_ps(sum, diff, _mm256_castsi256_ps(upper))); + } + } + static_assert(N == 640 || N == 2560, "strides 40 .. N / 2 go two per pass: an even count"); + for (int len = 40; len < N; len *= 4) // strides len and 2 len per pass (the same sums as one stride at a time) + for (int i = 0; i < N; i += 4 * len) + for (int j = i; j < i + len; j += 8) { + float* p = a + j; + const __m256 x0 = _mm256_load_ps(p), x1 = _mm256_load_ps(p + len); + const __m256 x2 = _mm256_load_ps(p + 2 * len), x3 = _mm256_load_ps(p + 3 * len); + const __m256 y0 = _mm256_add_ps(x0, x1), y1 = _mm256_sub_ps(x0, x1); + const __m256 y2 = _mm256_add_ps(x2, x3), y3 = _mm256_sub_ps(x2, x3); + _mm256_store_ps(p, _mm256_add_ps(y0, y2)); + _mm256_store_ps(p + len, _mm256_add_ps(y1, y3)); + _mm256_store_ps(p + 2 * len, _mm256_sub_ps(y0, y2)); + _mm256_store_ps(p + 3 * len, _mm256_sub_ps(y1, y3)); + } + const __m256 norm = _mm256_set1_ps(1.0f / std::sqrt(float(N / 5))); + const __m256 magnitude = _mm256_castsi256_ps(_mm256_set1_epi32(0x7FFFFFFF)); + __m256 max8 = _mm256_setzero_ps(); + for (int j = 0; j < N; j += 8) max8 = _mm256_max_ps(max8, _mm256_and_ps(_mm256_mul_ps(_mm256_load_ps(a + j), norm), magnitude)); + alignas(32) float lanes[8]; + _mm256_store_ps(lanes, max8); + const float max = *std::max_element(lanes, lanes + 8); + const float step = max > 0.0f ? max / 127.0f : 1.0f; + const __m256 inv = _mm256_set1_ps(1.0f / step), k254 = _mm256_set1_ps(254.0f); + const auto plane = [&](__m256 v) { // clamp(nearbyint(v), -127, 127) + const __m256 r = _mm256_round_ps(v, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + return _mm256_min_ps(_mm256_max_ps(r, _mm256_set1_ps(-127.0f)), _mm256_set1_ps(127.0f)); + }; + const auto q16 = [&](int j) { // 254 q0 + q1 of columns j .. j + 7 (exact in float) + const __m256 scaled = _mm256_mul_ps(_mm256_mul_ps(_mm256_load_ps(a + j), norm), inv); + const __m256 q0 = plane(scaled), q1 = plane(_mm256_mul_ps(_mm256_sub_ps(scaled, q0), k254)); + return _mm256_cvtps_epi32(_mm256_add_ps(_mm256_mul_ps(q0, k254), q1)); + }; + ActHeader* h = static_cast(dst); + __m256i sums = _mm256_setzero_si256(); // lane l: columns = l (mod 8), so even lanes hold the even columns + for (int j = 0; j < N; j += 16) { + const __m256i lo = q16(j), hi = q16(j + 8); + sums = _mm256_add_epi32(sums, _mm256_add_epi32(lo, hi)); + // per 4 columns the order 0, 2, 1, 3 (the kernel's even and odd level words), then int16 + const __m256i packed = _mm256_packs_epi32(_mm256_shuffle_epi32(lo, _MM_SHUFFLE(3, 1, 2, 0)), + _mm256_shuffle_epi32(hi, _MM_SHUFFLE(3, 1, 2, 0))); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(reinterpret_cast(h + 1) + j), + _mm256_permute4x64_epi64(packed, _MM_SHUFFLE(3, 1, 2, 0))); + } + alignas(32) int32_t s[8]; + _mm256_store_si256(reinterpret_cast<__m256i*>(s), sums); + const int even = s[0] + s[2] + s[4] + s[6], odd = s[1] + s[3] + s[5] + s[7]; + *h = {step, even + odd, step * float(even) / 254.0f, step * float(odd) / 254.0f}; +} +#if defined(__GNUC__) && !defined(__clang__) +#pragma GCC pop_options +#endif + +// The 16-bit state at bit B of the chunk whose words k = 0 .. T of the 8 lanes are at w + 32 k. +template +inline __m256i state(const uint32_t* w) { + constexpr int k = B / 32, s = B % 32; + const __m256i mask = _mm256_set1_epi32(0xFFFF); + const __m256i lo = _mm256_loadu_si256(reinterpret_cast(w + 32 * k)); + if constexpr (s <= 16) { + return _mm256_and_si256(_mm256_srli_epi32(lo, 16 - s), mask); + } else { + const __m256i hi = _mm256_loadu_si256(reinterpret_cast(w + 32 * (k + 1))); + return _mm256_and_si256(_mm256_or_si256(_mm256_slli_epi32(lo, s - 16), _mm256_srli_epi32(hi, 48 - s)), mask); + } +} + +// level + 54 of byte b = kNibbles[b & 15] + kNibbles[32 + (b >> 4)]: 8 * (sum of b's 2-bit fields) + ((3 * b) & 15). +constexpr std::array kNibbles = [] { + std::array t{}; + for (int i = 0; i < 32; ++i) { + const int n = i % 16, fields = (n & 3) + (n >> 2); + t[i] = int8_t(8 * fields + ((3 * n) & 15)); + t[32 + i] = int8_t(8 * fields); + } + return t; +}(); + +// x0, x1 = fmix32 of steps 2g, 2g + 1 before its last x ^= x >> 16 -> level + 54 of columns 4g .. 4g + 3 as int16 +// pairs: out[0] = columns {4g, 4g + 2}, out[1] = {4g + 1, 4g + 3} (the activation's pair order). +inline void pair_levels(__m256i x0, __m256i x1, __m256i* out) { + // the hashes' low halves lo16 ^ hi16 of both in one dword: {x0 lo ^ x0 hi, x1 hi ^ x1 lo} + const __m256i swap16 = _mm256_setr_epi8(2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13, + 2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13); + const __m256i h = _mm256_xor_si256(_mm256_blend_epi16(x0, x1, 0xAA), + _mm256_shuffle_epi8(_mm256_blend_epi16(x1, x0, 0xAA), swap16)); + const __m256i nibble = _mm256_set1_epi8(0x0F); + const __m256i low = _mm256_loadu_si256(reinterpret_cast(kNibbles.data())); + const __m256i high = _mm256_loadu_si256(reinterpret_cast(kNibbles.data() + 32)); + const __m256i l = _mm256_add_epi8(_mm256_shuffle_epi8(low, _mm256_and_si256(h, nibble)), + _mm256_shuffle_epi8(high, _mm256_and_si256(_mm256_srli_epi32(h, 4), nibble))); + out[0] = _mm256_and_si256(l, _mm256_set1_epi32(0x00FF00FF)); + out[1] = _mm256_srli_epi16(l, 8); +} + +// Steps 4G .. 4G + 3 of a chunk: use(lv, G) with lv[0 .. 3] = pair_levels of pairs 2G, 2G + 1. The four hashes go stage +// by stage: one step alone is a ~25-cycle dependency chain and the compiler keeps source order (10% faster). +template +inline void quad_levels(const uint32_t* w, Use& use) { + __m256i x0 = state<4 * G * T>(w), x1 = state<(4 * G + 1) * T>(w), x2 = state<(4 * G + 2) * T>(w), + x3 = state<(4 * G + 3) * T>(w); + const auto each = [&](auto f) { x0 = f(x0), x1 = f(x1), x2 = f(x2), x3 = f(x3); }; + each([](__m256i x) { return _mm256_add_epi32(_mm256_mullo_epi32(x, _mm256_set1_epi32(int(0xCFCCB83Fu))), _mm256_set1_epi32(0x584B4AA3)); }); + each([](__m256i x) { return _mm256_xor_si256(x, _mm256_srli_epi32(x, 16)); }); + each([](__m256i x) { return _mm256_mullo_epi32(x, _mm256_set1_epi32(int(0x85EBCA6Bu))); }); + __m256i lv[4]; + pair_levels(x0, x1, lv); + pair_levels(x2, x3, lv + 2); + use(lv, G); +} + +template +inline void decode_chunk(const uint32_t* w, Use& use, std::integer_sequence) { + (quad_levels(w, use), ...); +} + +// acc, acc2 += levels lv[0 .. 3] . the token's activation pairs a[0 .. 3]. |acc + acc2| <= 64 * 109 * 32385 < 2^31 +// over a chunk. +inline void dot4(const __m256i* lv, const int32_t* a, __m256i& acc, __m256i& acc2) { + acc = _mm256_add_epi32(acc, _mm256_madd_epi16(lv[0], _mm256_set1_epi32(a[0]))); + acc2 = _mm256_add_epi32(acc2, _mm256_madd_epi16(lv[1], _mm256_set1_epi32(a[1]))); + acc = _mm256_add_epi32(acc, _mm256_madd_epi16(lv[2], _mm256_set1_epi32(a[2]))); + acc2 = _mm256_add_epi32(acc2, _mm256_madd_epi16(lv[3], _mm256_set1_epi32(a[3]))); +} + +// s[t][i] for the 8 rows whose lanes start at w (word k of chunk c at w + 32 (c T + k)). +template +void block_sums_t(const uint32_t* w, int chunks, const void* const* act, int nt, int64_t (*s)[8]) { + __m256i lo[8], hi[8]; // per token, int64 sums of lanes 0-3 and 4-7 + for (int t = 0; t < nt; ++t) lo[t] = hi[t] = _mm256_setzero_si256(); + const auto pairs = [&](int t, int c) { return reinterpret_cast(static_cast(act[t]) + 1) + 32 * c; }; + const auto add = [&](int t, __m256i acc) { + lo[t] = _mm256_add_epi64(lo[t], _mm256_cvtepi32_epi64(_mm256_castsi256_si128(acc))); + hi[t] = _mm256_add_epi64(hi[t], _mm256_cvtepi32_epi64(_mm256_extracti128_si256(acc, 1))); + }; + for (int c = 0; c < chunks; ++c) { + __m256i acc = _mm256_setzero_si256(), acc2 = _mm256_setzero_si256(); + if (nt == 1) { // the levels go straight into the dot product (10% faster than through lv) + auto dot = [&](const __m256i* l, int g) { dot4(l, pairs(0, c) + 4 * g, acc, acc2); }; + decode_chunk(w + 32 * T * c, dot, std::make_integer_sequence()); + add(0, _mm256_add_epi32(acc, acc2)); + continue; + } + __m256i lv[32]; + auto keep = [&](const __m256i* l, int g) { std::memcpy(lv + 4 * g, l, 4 * sizeof(__m256i)); }; + decode_chunk(w + 32 * T * c, keep, std::make_integer_sequence()); + for (int t = 0; t < nt; ++t) { + acc = acc2 = _mm256_setzero_si256(); + for (int g = 0; g < 8; ++g) dot4(lv + 4 * g, pairs(t, c) + 4 * g, acc, acc2); + add(t, _mm256_add_epi32(acc, acc2)); + } + } + for (int t = 0; t < nt; ++t) { + _mm256_storeu_si256(reinterpret_cast<__m256i*>(s[t]), lo[t]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(s[t] + 4), hi[t]); + } +} + +struct Matrix { + const uint32_t* words; + const float* scales; + int bits, chunks; +}; + +// Header u32[16]: magic, T_gu, T_d, W_gu, W_d, byte offsets of gu words, d words, gu scales, d scales (qns_blob.py). +Matrix matrix(const uint8_t* blob, bool down) { + const uint32_t* h = reinterpret_cast(blob); + assert(h[0] == 0x31534E51u); + const int bits = int(h[down ? 2 : 1]), chunks = down ? 10 : 40; + assert(int(h[down ? 4 : 3]) == chunks * bits + 1); + return {reinterpret_cast(blob + h[down ? 6 : 5]), + reinterpret_cast(blob + h[down ? 8 : 7]), bits, chunks}; +} + +// y = scale * ((step / RMS) * ((C - 54 sum q0) + (F - 54 sum q1) / 254) + d0 E + d1 O), with S = 254 C + F. +float finish(int64_t s, const void* act, float scale) { + const ActHeader& a = *static_cast(act); + const float dot = float(double(s - 54 * int64_t(a.sum)) / 254.0); + return scale * (a.step / kRms * dot - kMean0 / kRms * a.even - kMean1 / kRms * a.odd); +} + +} // namespace + +void set_rotation(const float* signs2560, const float* q2560, const float* signs640, const float* q640) { + set(g_rot2560, signs2560, q2560); + set(g_rot640, signs640, q640); +} + +void quant_act(const float* x, void* dst) { quantize<2560>(x, g_rot2560, dst); } + +void quant_h(const float* h, void* dst) { quantize<640>(h, g_rot640, dst); } + +void block_sums(const uint8_t* blob, bool down, int b, const void* const* act, int nt, int64_t (*s)[8]) { + assert(1 <= nt && nt <= 8); + const Matrix m = matrix(blob, down); + const uint32_t* w = m.words + size_t(b / 4) * (m.chunks * m.bits + 1) * 32 + 8 * (b % 4); + switch (m.bits) { + case 3: return block_sums_t<3>(w, m.chunks, act, nt, s); + case 4: return block_sums_t<4>(w, m.chunks, act, nt, s); + case 5: return block_sums_t<5>(w, m.chunks, act, nt, s); + case 6: return block_sums_t<6>(w, m.chunks, act, nt, s); + case 7: return block_sums_t<7>(w, m.chunks, act, nt, s); + case 8: return block_sums_t<8>(w, m.chunks, act, nt, s); + case 9: return block_sums_t<9>(w, m.chunks, act, nt, s); + case 10: return block_sums_t<10>(w, m.chunks, act, nt, s); + } + std::abort(); // T outside 3..10: not a qns blob +} + +void gu_rows(const uint8_t* blob, const void* const* act, int nt, float* const* ff, int r0, int r1) { + assert(0 <= r0 && r0 < r1 && r1 <= 640); + const float* scales = matrix(blob, false).scales; + for (int b = r0 / 8; 8 * b < r1; ++b) { + int64_t gate[8][8], up[8][8]; + block_sums(blob, false, b, act, nt, gate); + block_sums(blob, false, 80 + b, act, nt, up); + for (int t = 0; t < nt; ++t) + for (int r = std::max(r0, 8 * b); r < std::min(r1, 8 * b + 8); ++r) { + const float g = finish(gate[t][r % 8], act[t], scales[r]); + ff[t][r] = g / (1.0f + std::exp(-g)) * finish(up[t][r % 8], act[t], scales[640 + r]); + } + } +} + +void down_rows(const uint8_t* blob, const void* const* hq, int nt, float* const* out, int r0, int r1) { + assert(0 <= r0 && r0 < r1 && r1 <= 2560); + const float* scales = matrix(blob, true).scales; + for (int b = r0 / 8; 8 * b < r1; ++b) { + int64_t s[8][8]; + block_sums(blob, true, b, hq, nt, s); + for (int t = 0; t < nt; ++t) + for (int r = std::max(r0, 8 * b); r < std::min(r1, 8 * b + 8); ++r) + out[t][r] = finish(s[t][r % 8], hq[t], scales[r]); + } +} + +} // namespace strata::kernels::cpu::qns diff --git a/src/kernels/cuda/qns_gather.cu b/src/kernels/cuda/qns_gather.cu new file mode 100644 index 0000000000..a4430f8988 --- /dev/null +++ b/src/kernels/cuda/qns_gather.cu @@ -0,0 +1,23 @@ +// src/kernels/cuda/qns_gather.cu - see include/strata/kernels/qns_gather.hpp. +#include "strata/kernels/qns_gather.hpp" + +namespace strata::kernels { +namespace { + +__global__ void qns_gather_rows_kernel(const uint32_t* __restrict__ words, const float* __restrict__ scales, int W, + const int32_t* __restrict__ rows, int n, uint32_t* __restrict__ dst, + float* __restrict__ dst_scales) { + const int i = blockIdx.x, r = i < n ? rows[i] : -1; + for (int k = threadIdx.x; k < W; k += blockDim.x) + dst[((size_t) (i / 32) * W + k) * 32 + i % 32] = r < 0 ? 0u : words[((size_t) (r / 32) * W + k) * 32 + r % 32]; + if (threadIdx.x == 0) dst_scales[i] = r < 0 ? 0.0f : scales[r]; +} + +} // namespace + +void qns_gather_rows(const uint32_t* words, const float* scales, int W, const int32_t* rows, int n, int n_pad, + uint32_t* dst_words, float* dst_scales, cudaStream_t s) { + qns_gather_rows_kernel<<>>(words, scales, W, rows, n, dst_words, dst_scales); +} + +} // namespace strata::kernels diff --git a/src/kernels/cuda/qns_kernels.cu b/src/kernels/cuda/qns_kernels.cu new file mode 100644 index 0000000000..52d718aefa --- /dev/null +++ b/src/kernels/cuda/qns_kernels.cu @@ -0,0 +1,603 @@ +// qns_kernels.cu - Qwen Next S routed experts on CUDA (qns_cuda.hpp). A port of the vLLM plugin's kernels +// (qwen_next_s/kernels.cu): the same word layout, hash, level map and exact int32 dot products. +// +// A blob row of `cols` rotated weights is a 16-bit trellis path of cols / 2 states (FORMAT.md): state t is the 16-bit +// window at bit t T of the row's MSB-first tape and emits columns 2t, 2t + 1 as scale * (level_j - MEAN_j) / RMS, +// level_j = level(byte j of fmix(state)). The blob holds each row's tape as big-endian u32 words, 32 rows interleaved +// per word ([group][word][lane], W = cols / 64 * T + 1 words a row): a warp with lane = row reads one word of its 32 +// rows as one 128-byte load. A 64-column chunk is exactly T words, so chunk k starts at word k T and, with the next +// word, every state of it is a funnel shift at a compile-time offset. +// +// Activations are rotated, x_rot = (H (x) Q)(signs * x), and quantized per token: x_rot ~ step (q0 + q1 / 254). A row's +// output finishes the exact sums C = sum (level + 54) q0, F = sum (level + 54) q1 (finish()). + +#include "strata/kernels/qns_cuda.hpp" + +#include + +#include +#include + +namespace strata::kernels { +namespace { + +typedef unsigned int u32; + +constexpr u32 kMagic = 0x31534E51u; // 'QNS1' +constexpr float kMean0 = 1.5769348f, kMean1 = 1.4910126f, kRms = 19.183672f; + +// The odd factors: Q_5 of the 2560 (0) and 640 (1) rotations and Q_3 of the 6144 one, [o][c]; every lane reads the +// same entry (a constant-cache broadcast). The signs: each lane reads its own columns. +__constant__ float c_q5[2][25], c_q3[9]; +__device__ float g_signs2560[2560], g_signs640[640], g_signs6144[6144]; + +// Rotation widths: 2560 = 512 x 5 (the hidden state), 640 = 128 x 5 (the expert down input), 6144 = 2048 x 3. +template +constexpr int kOdd = COLS == 6144 ? 3 : 5; +template +constexpr float kNorm = COLS == 2560 ? 0.044194173824159216f : COLS == 640 ? 0.088388347648318441f : 0.022097086912079608f; // 1 / sqrt(COLS / kOdd) +template +__device__ __forceinline__ float q_entry(int i) { + if constexpr (COLS == 6144) return c_q3[i]; + else return c_q5[COLS == 2560 ? 0 : 1][i]; +} +template +__device__ __forceinline__ const float* signs_of() { + return COLS == 2560 ? g_signs2560 : COLS == 640 ? g_signs640 : g_signs6144; +} + +// One token's rotated, quantized activation: q[g] = the plane-0 and the plane-1 bytes of columns 4g .. 4g + 3. +struct Stats { + float step; + int sum_q0, sum_q1; + float even, odd; // step (sum q0 + sum q1 / 254) over the even / the odd columns + int pad[3]; +}; +template +struct Act { + uint2 q[COLS / 4]; + Stats st; +}; +static_assert(sizeof(Act<2560>) % 16 == 0 && sizeof(Act<640>) % 16 == 0 && sizeof(Act<6144>) % 16 == 0, "records stay 16-byte aligned"); + +void check(const char* what) { + const cudaError_t e = cudaGetLastError(); + if (e == cudaSuccess) return; + std::fprintf(stderr, "%s: %s\n", what, cudaGetErrorString(e)); + std::abort(); +} + +// Bytes 0, 1 of fmix32(s0 * 0xCFCCB83F + 0x584B4AA3) and of the same for s1, as one word. Only the low half of the last +// step x ^ (x >> 16) is used, and it is lo16(x) ^ hi16(x): two byte_perms and one xor for both states. +__device__ __forceinline__ u32 fmix_hash_pair(u32 s0, u32 s1) { + u32 x0 = s0 * 0xCFCCB83Fu + 0x584B4AA3u, x1 = s1 * 0xCFCCB83Fu + 0x584B4AA3u; + x0 ^= x0 >> 16; + x1 ^= x1 >> 16; + x0 *= 0x85EBCA6Bu; + x1 *= 0x85EBCA6Bu; + return __byte_perm(x0, x1, 0x5410) ^ __byte_perm(x0, x1, 0x7632); +} + +// 2 (level(b) + 54) for each byte b of h, in [0, 218]: 16 (sum of b's four 2-bit fields) + 2 ((3 (b & 15)) & 15). +// Doubled because then each term takes one mask (2 fewer ALU-pipe ops; the decode is ALU-bound). Per nibble, +// h - 3 ((h >> 2) & 3) is the sum of its two fields; times 17 adds the low nibble's sum into the high one (no carries). +__device__ __forceinline__ u32 levels_x2(u32 h) { + const u32 pairs = h - 3u * ((h >> 2) & 0x33333333u); + return ((pairs * 17u) & 0xF0F0F0F0u) + (((h & 0x0F0F0F0Fu) * 6u) & 0x1E1E1E1Eu); +} + +// c += a (16 x 32 u8, row-major) * b (32 x 8 s8, column-major) on the int8 tensor cores; exact int32 sums. +__device__ __forceinline__ void mma_u8s8(int (&c)[4], u32 a0, u32 a1, u32 a2, u32 a3, u32 b0, u32 b1) { + asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.u8.s8.s32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+r"(c[0]), "+r"(c[1]), "+r"(c[2]), "+r"(c[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// One decoded chunk (lane = row: 16 words of 4 doubled levels, columns 4g .. 4g + 3) times 8 B columns on the tensor +// cores. The warp's tile (row stride 20 words: conflict-free 16-byte writes and fragment reads) turns the rows into +// mma m16n8k32 A fragments (16 rows x 32 columns); xq is this lane's B source, the plane words of B column gid (b0: its +// columns 4 tig .. 4 tig + 3, b1: 16 more). acc: [row tile: rows 0-15, 16-31][C fragment]. +__device__ __forceinline__ void mma_chunk(u32 (&tile)[32][20], const u32 (&levels)[16], const u32* __restrict__ xq, int chunk, + int (&acc)[2][4]) { + const int lane = threadIdx.x & 31, gid = lane >> 2, tig = lane & 3; +#pragma unroll + for (int k = 0; k < 4; ++k) + reinterpret_cast(tile[lane])[k] = make_uint4(levels[4 * k], levels[4 * k + 1], levels[4 * k + 2], levels[4 * k + 3]); + __syncwarp(); +#pragma unroll + for (int kt = 0; kt < 2; ++kt) { // columns 32 kt .. 32 kt + 31 of the chunk + const int word = chunk * 16 + 8 * kt + tig; + const u32 b0 = __ldg(xq + 2 * word), b1 = __ldg(xq + 2 * (word + 4)); +#pragma unroll + for (int rt = 0; rt < 2; ++rt) + mma_u8s8(acc[rt], tile[16 * rt + gid][8 * kt + tig], tile[16 * rt + gid + 8][8 * kt + tig], + tile[16 * rt + gid][8 * kt + 4 + tig], tile[16 * rt + gid + 8][8 * kt + 4 + tig], b0, b1); + } + __syncwarp(); // the next chunk rewrites the tile +} + + +// The 32 states of one chunk from its T words and the next one, as 16 words of four 2 (level + 54) (columns 4g .. 4g + 3). +template +__device__ __forceinline__ void decode_chunk(const u32 (&w)[T + 1], u32 (&levels)[16]) { +#pragma unroll + for (int g = 0; g < 16; ++g) { + const int b0 = 2 * g * T, b1 = (2 * g + 1) * T; + const u32 s0 = __funnelshift_l(w[b0 / 32 + 1], w[b0 / 32], b0 % 32) >> 16; + const u32 s1 = __funnelshift_l(w[b1 / 32 + 1], w[b1 / 32], b1 % 32) >> 16; + levels[g] = levels_x2(fmix_hash_pair(s0, s1)); + } +} + +// The row's output from its exact sums. Written with explicit fmaf so that a host mirror can reproduce it bit for bit. +__device__ __forceinline__ float finish(int coarse, int fine, const Stats& st, float scale) { + constexpr float c = 1.0f / kRms, d0 = -kMean0 / kRms, d1 = -kMean1 / kRms; + coarse -= 54 * st.sum_q0; + fine -= 54 * st.sum_q1; + const float dot = fmaf(static_cast(fine), 1.0f / 254.0f, static_cast(coarse)); + return scale * fmaf(d1, st.odd, fmaf(d0, st.even, c * st.step * dot)); +} + +// Unnormalized rotation of one vector of COLS = 32 WARPS R Hadamard indices x ODD spread over WARPS warps (the whole CTA +// when WARPS > 1): thread t of them holds indices t R .. t R + R - 1 with their ODD odd-factor components, i.e. the +// contiguous columns ODD R t .. ODD R (t + 1) - 1. v <- (H (x) Q) v, or with TRANSPOSE (H (x) Q^T) v; H the +-1 +// Sylvester Hadamard (the caller applies kNorm). Strides below R run inside the thread, R .. 16 R across lanes, the +// larger ones through `exchange` (shared, COLS floats; null with one warp), which the caller may have used before. +template +__device__ __forceinline__ void rotate(float (&v)[COLS / 32 / WARPS], float* __restrict__ exchange) { + constexpr int ODD = kOdd, N = COLS / 32 / WARPS, R = N / ODD, THREADS = 32 * WARPS; + static_assert(R * ODD == N, "whole Hadamard indices a thread"); + const int lane = threadIdx.x & 31; +#pragma unroll + for (int r = 0; r < R; ++r) { + float in[ODD]; +#pragma unroll + for (int c = 0; c < ODD; ++c) in[c] = v[ODD * r + c]; +#pragma unroll + for (int o = 0; o < ODD; ++o) { + float acc = 0.0f; +#pragma unroll + for (int c = 0; c < ODD; ++c) acc += q_entry(TRANSPOSE ? ODD * c + o : ODD * o + c) * in[c]; + v[ODD * r + o] = acc; + } + } +#pragma unroll + for (int stride = 1; stride < R; stride <<= 1) +#pragma unroll + for (int r = 0; r < R; ++r) + if ((r & stride) == 0) +#pragma unroll + for (int o = 0; o < ODD; ++o) { + const float a = v[ODD * r + o], b = v[ODD * (r + stride) + o]; + v[ODD * r + o] = a + b; + v[ODD * (r + stride) + o] = a - b; + } +#pragma unroll + for (int m = 1; m < 32; m <<= 1) { + const float sign = (lane & m) ? -1.0f : 1.0f; // a + b in the lower lane, a - b in the upper one +#pragma unroll + for (int j = 0; j < N; ++j) v[j] = fmaf(v[j], sign, __shfl_xor_sync(0xffffffffu, v[j], m)); + } + if constexpr (WARPS > 1) { + const int warp = threadIdx.x >> 5; + __syncthreads(); // the caller is done with the buffer +#pragma unroll + for (int j = 0; j < N; ++j) exchange[j * THREADS + threadIdx.x] = v[j]; + __syncthreads(); +#pragma unroll + for (int j = 0; j < N; ++j) { + float acc = 0.0f; +#pragma unroll + for (int w = 0; w < WARPS; ++w) { + const float other = exchange[j * THREADS + 32 * w + lane]; + acc += (__popc(warp & w) & 1) ? -other : other; + } + v[j] = acc; + } + } +} + +// Quantizes x_rot of one vector spread over the CTA (thread t: columns N t .. N t + N - 1) into act: +// step = max |x_rot| / 127 (1 if 0), q0 = clamp(rn(x / step)), q1 = clamp(rn((x / step - q0) 254)), x / step taken as +// x * (1 / step) like the plugin. The two planes are gathered into whole words of 4 columns in `buffer` (shared, COLS +// floats; written only after every thread is past its first barrier here). +template +__device__ __forceinline__ void quantize(const float (&v)[N], Act* __restrict__ act, float* __restrict__ buffer) { + constexpr int THREADS = COLS / N, WARPS = THREADS / 32; + __shared__ float maxima[WARPS]; + __shared__ int partial[WARPS][4]; + signed char* planes[2] = {reinterpret_cast(buffer), reinterpret_cast(buffer) + COLS}; + const int t = threadIdx.x, lane = t & 31, warp = t >> 5; + float maximum = 0.0f; +#pragma unroll + for (int j = 0; j < N; ++j) maximum = fmaxf(maximum, fabsf(v[j])); +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) maximum = fmaxf(maximum, __shfl_xor_sync(0xffffffffu, maximum, offset)); + if (lane == 0) maxima[warp] = maximum; + __syncthreads(); +#pragma unroll + for (int w = 0; w < WARPS; ++w) maximum = fmaxf(maximum, maxima[w]); + const float step = maximum > 0.0f ? maximum / 127.0f : 1.0f, inverse = 1.0f / step; + int sums[4] = {}; // q0, q1 over all columns, q0, q1 over the odd ones +#pragma unroll + for (int j = 0; j < N; ++j) { + const int column = N * t + j, odd = column & 1; + const float scaled = v[j] * inverse; + const int q0 = min(127, max(-127, __float2int_rn(scaled))); + const int q1 = min(127, max(-127, __float2int_rn((scaled - q0) * 254.0f))); + planes[0][column] = static_cast(q0); + planes[1][column] = static_cast(q1); + sums[0] += q0; + sums[1] += q1; + sums[2] += odd * q0; + sums[3] += odd * q1; + } +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) sums[i] += __shfl_xor_sync(0xffffffffu, sums[i], offset); + if (lane == 0) +#pragma unroll + for (int i = 0; i < 4; ++i) partial[warp][i] = sums[i]; + __syncthreads(); + const u32* words0 = reinterpret_cast(planes[0]); + const u32* words1 = reinterpret_cast(planes[1]); // COLS is a multiple of 16: both stay word-aligned + for (int g = t; g < COLS / 4; g += THREADS) act->q[g] = make_uint2(words0[g], words1[g]); + if (t != 0) return; +#pragma unroll + for (int w = 1; w < WARPS; ++w) +#pragma unroll + for (int i = 0; i < 4; ++i) sums[i] += partial[w][i]; + act->st.step = step; + act->st.sum_q0 = sums[0]; + act->st.sum_q1 = sums[1]; + act->st.even = step * fmaf(static_cast(sums[1] - sums[3]), 1.0f / 254.0f, static_cast(sums[0] - sums[2])); + act->st.odd = step * fmaf(static_cast(sums[3]), 1.0f / 254.0f, static_cast(sums[2])); +} + +// quant_act: grid (n_tok), block 32 WARPS: CTA = token of x [n_tok][COLS] -> act[token]. The checkpoint-ordered, +// signed input signs[c] x[perm[c]] (x[c] without perm) is staged in shared memory with whole-warp loads, then thread t +// takes columns N t .. N t + N - 1. +template +__global__ void __launch_bounds__(32 * WARPS) quant_act(const float* __restrict__ x, const int32_t* __restrict__ perm, + Act* __restrict__ act) { + constexpr int N = COLS / 32 / WARPS, THREADS = 32 * WARPS; + __shared__ __align__(16) float buffer[COLS]; + const float* row = x + static_cast(blockIdx.x) * COLS; + for (int c = threadIdx.x; c < COLS; c += THREADS) buffer[c] = row[perm ? perm[c] : c] * signs_of()[c]; + __syncthreads(); + float v[N]; +#pragma unroll + for (int k = 0; k < N / 4; ++k) { + const float4 a = reinterpret_cast(buffer + N * threadIdx.x)[k]; + v[4 * k] = a.x; + v[4 * k + 1] = a.y; + v[4 * k + 2] = a.z; + v[4 * k + 3] = a.w; + } + rotate(v, buffer); +#pragma unroll + for (int j = 0; j < N; ++j) v[j] *= kNorm; + quantize(v, act + blockIdx.x, buffer); +} + +// act_640: grid (cap_entries), block 128: CTA = entry e of the call ([grp_start[0], grp_start[*n_groups])): +// h = silu(gate) * up of row e of gu [.][1280] (fp32), rotated and quantized into hq[e] (thread t: columns 5 t .. 5 t + 4). +__global__ void __launch_bounds__(128) act_640(const int32_t* __restrict__ grp_start, const int32_t* __restrict__ n_groups, + const float* __restrict__ gu, Act<640>* __restrict__ hq) { + __shared__ float buffer[640]; + const int e = blockIdx.x; + if (e < grp_start[0] || e >= grp_start[*n_groups]) return; + const float* gate = gu + static_cast(e) * 1280 + 5 * threadIdx.x; + float v[5]; +#pragma unroll + for (int j = 0; j < 5; ++j) v[j] = __fdividef(gate[j], 1.0f + __expf(-gate[j])) * gate[640 + j] * g_signs640[5 * threadIdx.x + j]; + rotate<640, 4, false>(v, buffer); +#pragma unroll + for (int j = 0; j < 5; ++j) v[j] *= kNorm<640>; + quantize<640, 5>(v, hq + e, buffer); +} + +// Block x of a matrix (32 rows, lane = row) against the entries [e0, e1) of its group, 4 at a time: the warps stride +// over the 64-column chunks, each decoded once per pass into the warp's tile and multiplied on the tensor cores with the +// pass's entries (as dense_rows below, one n tile: B column n = entry n / 2, plane n % 2). +// Gate/up (COLS 2560): entry e reads token ent_map[e]'s Act<2560>, writes row e of out [.][1280]. +// Down (COLS 640): entry e reads Act<640> e, writes row ent_map[e] of out [.][2560]. +template +__device__ __forceinline__ void gemv_block(const u32* __restrict__ words, float rowscale, int e0, int e1, + const int32_t* __restrict__ ent_map, const Act* __restrict__ acts, + float* __restrict__ out, u32 (&tile)[WARPS][32][20]) { + constexpr bool DOWN = COLS == 640; + constexpr int ROWS = DOWN ? 2560 : 1280; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, gid = lane >> 2, tig = lane & 3, row = blockIdx.x * 32 + lane; + const u32* lane_words = words + blockIdx.x * (COLS / 64 * T + 1) * 32 + lane; + for (int first = e0; first < e1; first += 4) { + const int n = min(4, e1 - first), mine = first + min(gid / 2, n - 1); // an absent entry reads the last one + const u32* xq = reinterpret_cast(acts[DOWN ? mine : ent_map[mine]].q) + gid % 2; + int acc[2][4] = {}; // [row tile][C fragment] + for (int chunk = warp; chunk < COLS / 64; chunk += WARPS) { + u32 w[T + 1]; +#pragma unroll + for (int i = 0; i <= T; ++i) w[i] = __ldg(lane_words + (chunk * T + i) * 32); + u32 levels[16]; + decode_chunk(w, levels); + mma_chunk(tile[warp], levels, xq, chunk, acc); + } + __syncthreads(); // every warp is done with its tile + int(*partial)[4][WARPS][32] = reinterpret_cast(&tile[0][0][0]); // [plane][entry][warp][row] +#pragma unroll + for (int rt = 0; rt < 2; ++rt) +#pragma unroll + for (int half = 0; half < 2; ++half) { // C: rows gid, gid + 8; columns 2 tig (plane 0), 2 tig + 1 (plane 1) + partial[0][tig][warp][16 * rt + gid + 8 * half] = acc[rt][2 * half]; + partial[1][tig][warp][16 * rt + gid + 8 * half] = acc[rt][2 * half + 1]; + } + __syncthreads(); + if (warp == 0) + for (int b = 0; b < n; ++b) { + int c = 0, f = 0; +#pragma unroll + for (int w = 0; w < WARPS; ++w) { + c += partial[0][b][w][lane]; + f += partial[1][b][w][lane]; + } + const int e = first + b; + const size_t dst = DOWN ? static_cast(ent_map[e]) * ROWS : static_cast(e) * ROWS; + out[dst + row] = finish(c >> 1, f >> 1, acts[DOWN ? e : ent_map[e]].st, rowscale); // doubled levels, halved exactly + } + __syncthreads(); // the tile memory is reused by the next pass + } +} + +#define QNS_SWITCH_T(t, ...) \ + switch (t) { \ + case 3: { constexpr int T = 3; __VA_ARGS__; break; } \ + case 4: { constexpr int T = 4; __VA_ARGS__; break; } \ + case 5: { constexpr int T = 5; __VA_ARGS__; break; } \ + case 6: { constexpr int T = 6; __VA_ARGS__; break; } \ + case 7: { constexpr int T = 7; __VA_ARGS__; break; } \ + case 8: { constexpr int T = 8; __VA_ARGS__; break; } \ + case 9: { constexpr int T = 9; __VA_ARGS__; break; } \ + case 10: { constexpr int T = 10; __VA_ARGS__; break; } \ + default: __trap(); \ + } + +// expert_gemv: grid (rows / 32, grid groups), block 32 WARPS: block (x, y) takes 32-row block x of the gate/up +// (COLS 2560) or the down (COLS 640) matrix of groups y, y + gridDim.y, ... < *n_groups (as native_gu_kernel). +// Header words: 1 T_gu, 2 T_d, 5 / 6 the gate/up / down words' byte offset, 7 / 8 their scales' byte offset. +template +__global__ void __launch_bounds__(32 * WARPS, COLS == 2560 ? 3 : 4) // the registers to decode without spills + expert_gemv(const unsigned long long* __restrict__ grp_ptr, const int32_t* __restrict__ grp_start, + const int32_t* __restrict__ n_groups, const int32_t* __restrict__ ent_map, + const Act* __restrict__ acts, float* __restrict__ out) { + constexpr int M = COLS == 640; // 0 gate/up, 1 down + __shared__ __align__(16) u32 tile[WARPS][32][20]; + const int ng = *n_groups; + for (int g = blockIdx.y; g < ng; g += gridDim.y) { + const u32* blob = reinterpret_cast(grp_ptr[g]); + if (blob[0] != kMagic) __trap(); + const u32* words = blob + blob[5 + M] / 4; + const float rowscale = reinterpret_cast(blob + blob[7 + M] / 4)[blockIdx.x * 32 + (threadIdx.x & 31)]; + QNS_SWITCH_T(blob[1 + M], gemv_block(words, rowscale, grp_start[g], grp_start[g + 1], ent_map, + acts, out, tile)) + } +} + +// dense_gemv: grid (rows / 32), block 32 WARPS: block x takes rows 32 x .. 32 x + 31 of a dense matrix [rows][COLS] for +// n_tok <= 8 tokens; the warps stride over the 64-column chunks, each decoded once (lane = row) and multiplied with 4 +// tokens a mma_chunk (B column n of n tile nt: token 4 nt + n / 2, plane n % 2): 1 to 4 tokens cost the same, 5 to 8 +// twice that. After the chunks the tile memory holds the warps' sums for the cross-warp reduction. y[t rows + r] = row r . x_t. +template +__device__ __forceinline__ void dense_rows(const u32* __restrict__ words, const float* __restrict__ scales, int rows, + const Act* __restrict__ acts, int n_tok, float* __restrict__ y, + u32 (&tile)[WARPS][32][20]) { + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, gid = lane >> 2, tig = lane & 3; + const u32* lane_words = words + blockIdx.x * (COLS / 64 * T + 1) * 32 + lane; + const int ntiles = (n_tok + 3) / 4; + const u32* xq[2]; // this lane's B source: a token's plane words (an absent token reads the last one) +#pragma unroll + for (int nt = 0; nt < 2; ++nt) xq[nt] = reinterpret_cast(acts[min(4 * nt + gid / 2, n_tok - 1)].q) + gid % 2; + int acc[2][2][4] = {}; // [n tile][row tile: rows 0-15, 16-31][C fragment] + for (int chunk = warp; chunk < COLS / 64; chunk += WARPS) { + u32 w[T + 1]; +#pragma unroll + for (int i = 0; i <= T; ++i) w[i] = __ldg(lane_words + (chunk * T + i) * 32); + u32 levels[16]; + decode_chunk(w, levels); +#pragma unroll + for (int nt = 0; nt < 2; ++nt) + if (nt < ntiles) mma_chunk(tile[warp], levels, xq[nt], chunk, acc[nt]); // uniform + } + __syncthreads(); // every warp is done with its tile + int(*partial)[8][WARPS][32] = reinterpret_cast(&tile[0][0][0]); // [plane][token][warp][row] +#pragma unroll + for (int nt = 0; nt < 2; ++nt) { + if (nt >= ntiles) break; +#pragma unroll + for (int rt = 0; rt < 2; ++rt) +#pragma unroll + for (int half = 0; half < 2; ++half) { // C: rows gid, gid + 8; columns 2 tig (plane 0), 2 tig + 1 (plane 1) + partial[0][4 * nt + tig][warp][16 * rt + gid + 8 * half] = acc[nt][rt][2 * half]; + partial[1][4 * nt + tig][warp][16 * rt + gid + 8 * half] = acc[nt][rt][2 * half + 1]; + } + } + __syncthreads(); + if (warp != 0) return; + const int row = blockIdx.x * 32 + lane; + for (int b = 0; b < n_tok; ++b) { + int c = 0, f = 0; +#pragma unroll + for (int w = 0; w < WARPS; ++w) { + c += partial[0][b][w][lane]; + f += partial[1][b][w][lane]; + } + y[static_cast(b) * rows + row] = finish(c >> 1, f >> 1, acts[b].st, scales[row]); // doubled levels, halved exactly + } +} + +template +__global__ void __launch_bounds__(32 * WARPS) dense_gemv(const u32* __restrict__ words, const float* __restrict__ scales, + int rows, int T, const Act* __restrict__ acts, int n_tok, + float* __restrict__ y) { + static_assert(sizeof(int[2][8][WARPS][32]) <= sizeof(u32[WARPS][32][20]), "the sums fit in the tiles"); + __shared__ __align__(16) u32 tile[WARPS][32][20]; + QNS_SWITCH_T(T, dense_rows(words, scales, rows, acts, n_tok, y, tile)) +} + +// Thread t of a row's CTA decodes states S t .. S t + S - 1 (S = COLS / 64 / WARPS), i.e. the rotated columns rotate() +// expects in it, from its own window of the row's words (clamped to the row), realigned so that state i sits at bit +// i T. The fp16 row is gathered in shared memory (at its engine column, perm[c], when there is a permutation), so that +// every warp store is one contiguous 128 bytes. +template +__device__ __forceinline__ void dequant_row(const u32* __restrict__ words, int row, float scale, const int32_t* __restrict__ perm, + __half* __restrict__ dst, __half (&stage)[COLS], float* __restrict__ exchange) { + constexpr int THREADS = 32 * WARPS, S = COLS / 64 / WARPS, W = COLS / 64 * T + 1; + constexpr int A = ((S - 1) * T + 16) / 32 + 2; // realigned words: state i needs i T / 32 + 1 + const int t = threadIdx.x, bit = t * S * T; + const u32* row_words = words + (row / 32) * W * 32 + row % 32; + u32 raw[A + 1], a[A]; +#pragma unroll + for (int i = 0; i <= A; ++i) raw[i] = __ldg(row_words + 32 * min(bit / 32 + i, W - 1)); // past the row: unused bits +#pragma unroll + for (int i = 0; i < A; ++i) a[i] = __funnelshift_l(raw[i + 1], raw[i], bit % 32); + const float c = scale / kRms * kNorm; // the rotation's normalization folded in + float v[2 * S]; +#pragma unroll + for (int i = 0; i < S; i += 2) { // states i, i + 1 (an odd S decodes its last state twice) + const int b0 = i * T, b1 = min(i + 1, S - 1) * T; + const u32 s0 = __funnelshift_l(a[b0 / 32 + 1], a[b0 / 32], b0 % 32) >> 16; + const u32 s1 = __funnelshift_l(a[b1 / 32 + 1], a[b1 / 32], b1 % 32) >> 16; + const u32 levels = levels_x2(fmix_hash_pair(s0, s1)); +#pragma unroll + for (int k = 0; k < 4; ++k) + if (2 * i + k < 2 * S) + v[2 * i + k] = (static_cast(static_cast((levels >> (8 * k + 1)) & 127u) - 54) - (k & 1 ? kMean1 : kMean0)) * c; + } + rotate(v, exchange); +#pragma unroll + for (int j = 0; j < 2 * S; ++j) { + const int column = 2 * S * t + j; + stage[perm ? perm[column] : column] = __float2half_rn(v[j] * signs_of()[column]); + } + __syncthreads(); + for (int k = t; k < COLS / 2; k += THREADS) reinterpret_cast(dst)[k] = reinterpret_cast(stage)[k]; +} + +// expert_dequant: grid (rows), block 32 WARPS: CTA = one row of a blob's gate/up (COLS 2560, interleaved output rows) or +// down (COLS 640) matrix in the original basis, w = (w_rot @ R) * signs. +template +__global__ void __launch_bounds__(32 * WARPS) expert_dequant(const u32* __restrict__ blob, __half* __restrict__ dst) { + constexpr int M = COLS == 640; + __shared__ __align__(16) __half stage[COLS]; + __shared__ float exchange[COLS]; + if (blob[0] != kMagic) __trap(); + const u32* words = blob + blob[5 + M] / 4; + const float* scales = reinterpret_cast(blob + blob[7 + M] / 4); + const int row = blockIdx.x, out_row = M ? row : (row < 640 ? 2 * row : 2 * (row - 640) + 1); + QNS_SWITCH_T(blob[1 + M], dequant_row(words, row, scales[row], nullptr, dst + static_cast(out_row) * COLS, + stage, exchange)) +} + +// dense_dequant: grid (rows), block 32 WARPS: CTA = row r of a dense matrix -> dst[r][perm[c]] (dst[r][c] without perm). +template +__global__ void __launch_bounds__(32 * WARPS) dense_dequant(const u32* __restrict__ words, const float* __restrict__ scales, int T, + const int32_t* __restrict__ perm, __half* __restrict__ dst) { + __shared__ __align__(16) __half stage[COLS]; + __shared__ float exchange[COLS]; + const int row = blockIdx.x; + QNS_SWITCH_T(T, dequant_row(words, row, scales[row], perm, dst + static_cast(row) * COLS, stage, exchange)) +} + +} // namespace + +void qns_cuda_set_rotation(const float* signs2560, const float* q2560, const float* signs640, const float* q640) { + cudaMemcpyToSymbol(c_q5, q2560, 25 * sizeof(float)); + cudaMemcpyToSymbol(c_q5, q640, 25 * sizeof(float), 25 * sizeof(float)); + cudaMemcpyToSymbol(g_signs2560, signs2560, sizeof(g_signs2560)); + cudaMemcpyToSymbol(g_signs640, signs640, sizeof(g_signs640)); + check("qns_cuda_set_rotation"); +} + +void qns_cuda_set_rotation_6144(const float* signs6144, const float* q6144) { + cudaMemcpyToSymbol(c_q3, q6144, 9 * sizeof(float)); + cudaMemcpyToSymbol(g_signs6144, signs6144, sizeof(g_signs6144)); + check("qns_cuda_set_rotation_6144"); +} + +size_t qns_act_bytes() { return sizeof(Act<2560>); } + +void qns_quant_act(const float* x, int64_t n_tok, void* act, cudaStream_t s) { + if (n_tok <= 0) return; + quant_act<2560, 4><<(n_tok), 128, 0, s>>>(x, nullptr, static_cast*>(act)); + check("qns_quant_act"); +} + +size_t qns_grouped_scratch_bytes(int64_t cap_entries) { + return static_cast(cap_entries) * (1280 * sizeof(float) + sizeof(Act<640>)); +} + +void qns_expert_grouped(const unsigned long long* grp_ptr, const int32_t* grp_start, const int32_t* n_groups, + const int32_t* ent_dst, const int32_t* ent_tok, int64_t cap_groups, int64_t cap_entries, + const void* act, void* scratch, float* out, cudaStream_t s, int64_t grid_groups) { + if (cap_groups <= 0 || cap_entries <= 0) return; + const unsigned gy = static_cast(grid_groups <= 0 || grid_groups > cap_groups ? cap_groups : grid_groups); + float* gu = static_cast(scratch); // [cap_entries][1280], then the down activations + Act<640>* hq = reinterpret_cast*>(gu + cap_entries * 1280); + expert_gemv<2560, 8><<>>(grp_ptr, grp_start, n_groups, ent_tok, + static_cast*>(act), gu); + act_640<<(cap_entries), 128, 0, s>>>(grp_start, n_groups, gu, hq); + expert_gemv<640, 5><<>>(grp_ptr, grp_start, n_groups, ent_dst, hq, out); + check("qns_expert_grouped"); +} + +void qns_dequant_gu_f16(const uint8_t* blob, uint16_t* dst, cudaStream_t s) { + expert_dequant<2560, 4><<<1280, 128, 0, s>>>(reinterpret_cast(blob), reinterpret_cast<__half*>(dst)); + check("qns_dequant_gu_f16"); +} + +void qns_dequant_d_f16(const uint8_t* blob, uint16_t* dst, cudaStream_t s) { + expert_dequant<640, 2><<<2560, 64, 0, s>>>(reinterpret_cast(blob), reinterpret_cast<__half*>(dst)); + check("qns_dequant_d_f16"); +} + +static void dense_shape(const char* what, int rows, int cols, int T) { + if ((cols == 2560 || cols == 6144) && rows > 0 && rows % 32 == 0 && T >= 3 && T <= 10) return; + std::fprintf(stderr, "%s: rows %d, cols %d, T %d\n", what, rows, cols, T); + std::abort(); +} + +size_t qns_dense_act_bytes(int cols) { + dense_shape("qns_dense_act_bytes", 32, cols, 10); + return cols == 2560 ? sizeof(Act<2560>) : sizeof(Act<6144>); +} + +void qns_dense_quant(const float* x, int64_t n_tok, int cols, const int32_t* perm, void* act, cudaStream_t s) { + dense_shape("qns_dense_quant", 32, cols, 10); + if (n_tok <= 0) return; + if (cols == 2560) quant_act<2560, 4><<(n_tok), 128, 0, s>>>(x, perm, static_cast*>(act)); + else quant_act<6144, 8><<(n_tok), 256, 0, s>>>(x, perm, static_cast*>(act)); + check("qns_dense_quant"); +} + +void qns_dense_gemv(const uint32_t* words, const float* scales, int rows, int cols, int T, const void* act, int n_tok, float* y, + cudaStream_t s) { + dense_shape("qns_dense_gemv", rows, cols, T); + if (n_tok < 1 || n_tok > 8) { std::fprintf(stderr, "qns_dense_gemv: n_tok %d (1..8)\n", n_tok); std::abort(); } + const bool tall = rows > 2560; // enough row blocks to fill the GPU with 8 warps a block; else 16 warps split the columns + if (cols == 2560 && tall) dense_gemv<2560, 8><<>>(words, scales, rows, T, static_cast*>(act), n_tok, y); + else if (cols == 2560) dense_gemv<2560, 16><<>>(words, scales, rows, T, static_cast*>(act), n_tok, y); + else if (tall) dense_gemv<6144, 8><<>>(words, scales, rows, T, static_cast*>(act), n_tok, y); + else dense_gemv<6144, 16><<>>(words, scales, rows, T, static_cast*>(act), n_tok, y); + check("qns_dense_gemv"); +} + +void qns_dense_dequant_f16(const uint32_t* words, const float* scales, int rows, int cols, int T, const int32_t* perm, + uint16_t* dst, cudaStream_t s) { + dense_shape("qns_dense_dequant_f16", rows, cols, T); + if (cols == 2560) dense_dequant<2560, 4><<>>(words, scales, T, perm, reinterpret_cast<__half*>(dst)); + else dense_dequant<6144, 8><<>>(words, scales, T, perm, reinterpret_cast<__half*>(dst)); + check("qns_dense_dequant_f16"); +} + +} // namespace strata::kernels diff --git a/src/kernels/cuda/verify_kernels.cu b/src/kernels/cuda/verify_kernels.cu index a451d9bd56..0d66bdef2d 100644 --- a/src/kernels/cuda/verify_kernels.cu +++ b/src/kernels/cuda/verify_kernels.cu @@ -602,6 +602,24 @@ __global__ void res_patch_kernel(int32_t* __restrict__ table, const volatile int for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < n; k += gridDim.x * blockDim.x) table[pairs[2 * k]] = pairs[2 * k + 1]; } +// qns blobs: each copies its own bytes (header word 9, qns_blob.py), cut into kPieces ranges so a few blobs still keep +// the whole grid's loads in flight over PCIe. +__global__ void fetch_qns_blobs_kernel(const unsigned long long* __restrict__ src, const int32_t* __restrict__ n, + uint4* __restrict__ dst, long long per) { + constexpr int kPieces = 16; + __shared__ long long words; + const int blobs = *n; + for (int task = blockIdx.x; task < blobs * kPieces; task += gridDim.x) { + const int k = task % blobs, piece = task / blobs; + const uint4* s = (const uint4*) src[k]; + if (threadIdx.x == 0) words = (((const uint32_t*) s)[9] + 15) / 16; + __syncthreads(); + const long long lo = words * piece / kPieces, hi = words * (piece + 1) / kPieces; + for (long long i = lo + threadIdx.x; i < hi; i += blockDim.x) dst[k * per + i] = s[i]; + __syncthreads(); + } +} + __global__ void rebase_ptrs_kernel(unsigned long long* ptr, const int32_t* n, unsigned long long base, long long bytes) { const int k = threadIdx.x; if (k < *n) ptr[k] = base + (unsigned long long) k * (unsigned long long) bytes; @@ -703,6 +721,13 @@ void res_patch(int32_t* table, const int32_t* pairs, int n, void* stream) { check("res_patch"); } +void fetch_qns_blobs(const unsigned long long* src, const int32_t* n, uint8_t* dst, int64_t slot_bytes, int cap, void* stream) { + if (cap <= 0) return; + if (slot_bytes % 16 != 0) { std::fprintf(stderr, "fetch_qns_blobs: slot size must be a multiple of 16\n"); std::exit(1); } + fetch_qns_blobs_kernel<<<48 * 8, 256, 0, (cudaStream_t) stream>>>(src, n, (uint4*) dst, (long long) (slot_bytes / 16)); + check("fetch_qns_blobs"); +} + void rebase_ptrs(unsigned long long* ptr, const int32_t* n, uint8_t* base, int64_t blob_bytes, void* stream) { rebase_ptrs_kernel<<<1, 128, 0, (cudaStream_t) stream>>>(ptr, n, (unsigned long long) base, (long long) blob_bytes); check("rebase_ptrs"); diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index f00f1cc134..b1d9d36770 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -23,6 +23,7 @@ #include "strata/kernels/ple.hpp" #include "strata/kernels/native_ple_postops.hpp" #include "strata/kernels/iq_kernels.hpp" +#include "strata/kernels/qns_cuda.hpp" #include "strata/kernels/cpu/expert_layout.hpp" #include "strata/kernels/cpu/hadamard.hpp" #include "strata/kernels/qsa.hpp" @@ -2418,6 +2419,19 @@ ColdTrace g_cold; bool native_proj(Gemm& gm, const core::WeightRef* w, const uint16_t* X, float* Y, int64_t T, const std::string& name, std::string& err, int64_t ldy = 0, int64_t ldx = 0) { if (!w->native_data) { err = "prefill: " + name + " has no native GGUF blocks (run with --native)"; return false; } + if (w->native_type == core::kQnsDenseType) { // trellis rows, dequantized to fp16 in slices of whole 32-row groups + const int64_t N = w->ne1, K = w->ne0, rows = gm.scratch_elems() / K / 32 * 32; + if (rows <= 0 || (ldx != 0 && ldx != K)) { err = "prefill: " + name + ": a qns matrix needs 32 scratch rows and an unpadded X"; return false; } + const int64_t group_words = (K / 64 * w->qns_bits + 1) * 32; + for (int64_t r0 = 0; r0 < N; r0 += rows) { + const int64_t n = std::min(rows, N - r0); + strata::kernels::qns_dense_dequant_f16((const uint32_t*) w->native_data + r0 / 32 * group_words, w->qns_scales + r0, + (int) n, (int) K, w->qns_bits, w->qns_perm, gm.scratch(), + (cudaStream_t) gm.stream()); + gm.f16(X, gm.scratch(), Y + r0, T, n, K, ldy > 0 ? ldy : N, 0.0f); + } + return true; + } const bool ct = g_cold.begin(); gm.native(X, w->native_type, w->native_data, Y, T, w->ne1, w->ne0, ldy, 0.0f, ldx); if (ct) g_cold.end("native", name, T, w->ne1, w->ne0); @@ -2939,7 +2953,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri const StreamEntry& en = seq[issued]; const int sl = (int) (issued % (size_t) m.ring); const auto th = Clock::now(); - const size_t bytes = (size_t) lay0.blob_bytes(en.l); + const size_t bytes = (size_t) lay0.expert_bytes(en.l, en.e); // a qns blob's own size if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[m.used_of[sl]], 0); if (en.job < 0) { cudaMemcpyAsync(m.stage_dev[sl], en.blob, bytes, cudaMemcpyHostToDevice, m.copy); @@ -2987,7 +3001,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri const StreamEntry& en = seq[idx]; const int sl = (int) (idx % (size_t) m.ring); const auto th = Clock::now(); - const size_t bytes = (size_t) lay0.blob_bytes(en.l); + const size_t bytes = (size_t) lay0.expert_bytes(en.l, en.e); if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[m.used_of[sl]], 0); if (en.job < 0) { cudaMemcpyAsync(m.stage_dev[sl], en.blob, bytes, cudaMemcpyHostToDevice, m.copy); @@ -4386,13 +4400,13 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri if (pinned) { // DMA straight from the page-locked arena: the copy stream only waits for the slot if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[m.used_of[sl]], 0); - cudaMemcpyAsync(m.stage_dev[sl], b, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, m.copy); + cudaMemcpyAsync(m.stage_dev[sl], b, (size_t) lay.expert_bytes(l, e), cudaMemcpyHostToDevice, m.copy); ++stats_.experts_dma; } else { // copied to a pinned buffer by the stager (waits only if it is behind), then DMA const uint8_t* hb = m.stager->wait(job_of[j]); if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[m.used_of[sl]], 0); - cudaMemcpyAsync(m.stage_dev[sl], hb, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, m.copy); + cudaMemcpyAsync(m.stage_dev[sl], hb, (size_t) lay.expert_bytes(l, e), cudaMemcpyHostToDevice, m.copy); m.stager->issued_one(job_of[j], m.copy); } cudaEventRecord(m.copied[sl], m.copy); @@ -4579,7 +4593,10 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri return true; } const int q = (int) (j % DQ); - if (lay.native) { + if (lay.qns()) { // trellis decode, un-rotated to the original basis (qns_cuda.hpp) + strata::kernels::qns_dequant_gu_f16(blob_dev, m.dq_gu[q], m.cs); + strata::kernels::qns_dequant_d_f16(blob_dev, m.dq_d[q], m.cs); + } else if (lay.native) { // plan v0.3 P6: a native pack's layer, dequantized by llama.cpp's own formulas const auto& f = lay.fmt[(size_t) l]; strata::kernels::iq_dequant_gu_f16(f.gu_type, blob_dev, blob_dev + f.up_off, f.n_ff, f.n_embd, diff --git a/src/program/generate.cpp b/src/program/generate.cpp index 01b99f06cf..fe13c72059 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -48,6 +48,7 @@ #include "strata/kernels/cpu/hadamard.hpp" #include "strata/kernels/ngram.hpp" #include "strata/kernels/iq_kernels.hpp" +#include "strata/kernels/qns_cuda.hpp" #include "strata/kernels/s2_expert_grouped.hpp" #include "strata/kernels/sampler.hpp" #include "strata/kernels/verify_kernels.hpp" @@ -142,9 +143,36 @@ namespace strata::metal_engine { void set_model_shard(const char* shard1); } / #include #include #include +#include #include namespace { + +// The adaptive tier's swaps in one layer of a qns pack. Every VRAM slot holds the blob it was sized for (qns blobs +// differ per expert), so each incoming expert (`cand`, most used first) takes the least-used resident expert +// (`vict`) whose slot holds it, while the gain is at least `min_gain`. `res`: the layer's row of the residency table. +// Returns (gain, incoming, outgoing). +std::vector> qns_swaps(int64_t layer, const std::vector>& cand, + std::vector> vict, float min_gain, + const int32_t* res, const strata::core::ExpertCache& cache) { + const auto& lay = strata::kernels::cpu::expert_layout(); + std::sort(vict.begin(), vict.end(), [](auto& a, auto& b) { return a.first < b.first; }); + std::vector taken(vict.size(), 0); + std::vector> out; + size_t first_free = 0; + for (const auto& [usage, in] : cand) { + while (first_free < vict.size() && taken[first_free]) ++first_free; + if (first_free == vict.size() || usage < vict[first_free].first + min_gain) break; // no victim gains enough + const uint64_t need = lay.expert_bytes(layer, in); + for (size_t j = first_free; j < vict.size() && usage >= vict[j].first + min_gain; ++j) { + if (taken[j] || cache.slot_bytes(res[vict[j].second]) < need) continue; + taken[j] = 1; + out.emplace_back(usage - vict[j].first, in, vict[j].second); + break; + } + } + return out; +} // Windows' WDDM driver model: native Windows, or WSL2 (its GPU goes through /dev/dxg to the Windows driver). There, // pinning a large arena into two CUDA contexts leaves WDDM refusing every later allocation (the 5080 + 3090 rig); // a Linux driver has no such limit (#253: the 8 GiB cap cost a 4090 + 3060 split 3x of its prompt speed). @@ -461,6 +489,7 @@ struct Options { bool cpu_oracle_q8_0 = false; // pinned x86 activation scales/codes at both expert stages std::string native_head_gguf; // native output.weight experiment; same model shard as the pack std::vector native_dense_gguf; // repeat for native GDN/QSA projection shards + std::string qns_dense_gguf; // a qns backbone GGUF: those matrices served as trellis (NativeDense) /// Every shard of --native's model (strata::gguf_split_paths: the metadata shard first; a missing shard is an /// error) and of --native-head-gguf's (the same list unless that names another model). std::vector native_shards, native_head_shards; @@ -844,6 +873,7 @@ void usage() { " --embd-gguf PATH the token embedding from this GGUF instead of --native's (tools/embd_bf16_pack.py:\n" " BF16 as the checkpoint ships it; mapped host memory, no VRAM)\n" " --native-dense-gguf PATH native GDN/QSA/shared projections; repeat for each source model shard\n" + " --qns-dense-gguf PATH a qns model's backbone as trellis (the GDN/QSA projections it holds; less VRAM)\n" " --expert-cache-cpu-order experimental GPU expert reduction matching CPU order\n" " --max-new N tokens to generate (default 16)\n" " --max-context N KV/state capacity (default 4096)\n" @@ -2386,6 +2416,7 @@ int main(int argc, char** argv) { else if (a == "--native-head-gguf") o.native_head_gguf = next("--native-head-gguf"); else if (a == "--embd-gguf") o.embd_gguf = next("--embd-gguf"); else if (a == "--native-dense-gguf") o.native_dense_gguf.push_back(next("--native-dense-gguf")); + else if (a == "--qns-dense-gguf") o.qns_dense_gguf = next("--qns-dense-gguf"); else if (a == "--no-capture") o.no_capture = true; else if (a == "--no-pool") o.no_pool = true; else if (a == "--sync-every-layer") o.sync_every_layer = true; @@ -3285,7 +3316,14 @@ int main(int argc, char** argv) { // anything is allocated: an unsupported down type used to exit from inside the first verify window, and // an unsupported dequant type left the prompt path's fp16 buffer unwritten. const auto& lay = strata::kernels::cpu::expert_layout(); - for (int64_t l = 0; lay.native && l < (int64_t) lay.fmt.size(); ++l) { + if (lay.qns()) { // Qwen Next S trellis experts: the GPU kernels take the rotations once + const float* r = lay.qns_rotation.data(); + strata::kernels::qns_cuda_set_rotation(r, r + 2560, r + 2585, r + 3225); + std::fprintf(stderr, "strata generate: qns pack: trellis experts, compact arena %.2f GiB, each VRAM slot " + "sized to its expert's blob (up to %.2f MB)\n", + (double) lay.total / 1073741824.0, (double) lay.max_blob / 1e6); + } + for (int64_t l = 0; lay.native && !lay.qns() && l < (int64_t) lay.fmt.size(); ++l) { const auto& f = lay.fmt[(size_t) l]; if (!strata::kernels::native_expert_supported(f.gu_type, f.d_type, f.n_embd, f.n_ff)) { std::fprintf(stderr, "strata generate: layer %lld's experts are %s/%s (ggml types %d/%d), which this " @@ -3728,6 +3766,7 @@ int main(int argc, char** argv) { strata::core::NativeDense native_dense; if (!o.native_dense_gguf.empty()) { + strata::core::NativeDense::set_qns_dense(o.qns_dense_gguf); if (!native_dense.load(o.native_dense_gguf, wt, err, o.native_ple_key)) { std::fprintf(stderr, "strata generate: native dense projections: %s\n", err.c_str()); return 1; @@ -4122,6 +4161,11 @@ int main(int argc, char** argv) { std::vector> stages; // One stage's dense weights on its device (the caller has made it current): `trim` = only layers [lb, le). auto load_stage_weights = [&](GpuStage& st, int64_t lb, int64_t le, bool trim) -> bool { + if (strata::kernels::cpu::expert_layout().qns()) { // the expert rotations live in __constant__ memory, one copy per card + const strata::core::OnDevice on(st.dev); + const float* r = strata::kernels::cpu::expert_layout().qns_rotation.data(); + strata::kernels::qns_cuda_set_rotation(r, r + 2560, r + 2585, r + 3225); + } std::set skip_s = skip_base; uint64_t pool_s = pool_bytes; if (trim) { @@ -5965,11 +6009,11 @@ int main(int argc, char** argv) { const size_t keep_free = ((size_t) o.vram_reserve_mib << 20) + (size_t) pipe_first; size_t free_room = free_b > keep_free ? free_b - keep_free : 0; const uint64_t cap = std::min(budget, (uint64_t) free_room); - for (const auto& pr : profile) { - const uint64_t b = (uint64_t) strata::core::ExpertCache::slot_stride((int64_t) lay.blob_bytes(pr.first), slot_guard); + for (const auto& pr : profile) { // qns: each slot holds its pair's own blob (they differ per expert) + const uint64_t b = (uint64_t) strata::core::ExpertCache::slot_stride((int64_t) lay.expert_bytes(pr.first, pr.second), slot_guard); if (used + b > cap) break; used += b; - sized_slots.push_back((int64_t) lay.blob_bytes(pr.first)); + sized_slots.push_back((int64_t) lay.expert_bytes(pr.first, pr.second)); } // #831: the sized-slots loop walks the profile's pairs, so the profile's length is the ceiling; an explicit // --expert-cache above it used to be cut with no word. A warning only (what is allocated does not change). @@ -6178,12 +6222,12 @@ int main(int argc, char** argv) { } const uint8_t* b = srcp->blob_stable(profile[(size_t) i].first, profile[(size_t) i].second); if (b == nullptr || !xcache.fill_slot_blocking(slot, b, err, - (int64_t) strata::kernels::cpu::expert_layout().blob_bytes(profile[(size_t) i].first))) { + (int64_t) strata::kernels::cpu::expert_layout().expert_bytes(profile[(size_t) i].first, profile[(size_t) i].second))) { std::fprintf(stderr, "strata generate: the profile fill failed at pair %lld: %s\n", (long long) i, err.c_str()); return 1; } - fill_bytes += strata::kernels::cpu::expert_layout().blob_bytes(profile[(size_t) i].first); + fill_bytes += strata::kernels::cpu::expert_layout().expert_bytes(profile[(size_t) i].first, profile[(size_t) i].second); ++prefilled; } const double fill_s = std::chrono::duration(std::chrono::steady_clock::now() - fill_t0).count(); @@ -6192,7 +6236,7 @@ int main(int argc, char** argv) { // cache's own `verify_slot` is the check and it costs one 1.38 MB D2H at startup. if (prefilled > 0 && !xcache.verify_slot(xcache.slot_of(profile[0].first, profile[0].second), srcp->blob_stable(profile[0].first, profile[0].second), err, - (int64_t) strata::kernels::cpu::expert_layout().blob_bytes(profile[0].first))) { + (int64_t) strata::kernels::cpu::expert_layout().expert_bytes(profile[0].first, profile[0].second))) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } @@ -7732,10 +7776,10 @@ int main(int argc, char** argv) { for (; pi < profile.size(); ++pi) { const int32_t l = profile[pi].first, e = profile[pi].second; const size_t i = (size_t) l * (size_t) g.n_expert + (size_t) e; - if (host_res[i] >= 0 || lay.blob_bytes(l) > room) continue; + if (host_res[i] >= 0 || lay.expert_bytes(l, e) > room) continue; const uint8_t* b = srcp->blob_stable(l, e); if (b == nullptr) continue; - cudaMemcpyAsync(xcache.device_slot((int32_t) s), b, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, + cudaMemcpyAsync(xcache.device_slot((int32_t) s), b, (size_t) lay.expert_bytes(l, e), cudaMemcpyHostToDevice, nullptr); filled.emplace_back(i, (int32_t) s); ++pi; @@ -9693,6 +9737,8 @@ int main(int argc, char** argv) { pcie_auto.on = true; pcie_auto.link_gbps = o.pcie_probe_gbps; pcie_auto.blob_bytes = g.n_layers > 0 ? sum / (double) g.n_layers : 0.0; + if (pl.qns()) // a PCIe fetch moves a qns blob's own bytes (fetch_qns_blobs): the mean blob, not the layer's largest + pcie_auto.blob_bytes = (double) pl.blob_offset(g.n_layers - 1, pl.n_expert) / (double) (g.n_layers * pl.n_expert); } // PcieAuto after a decode window (the serial loop, and 0.1.44 the pipelined one at every verdict): the share, refined from // the CPU pool's speed. A layer split: every later stage takes the share its own probed link gives against the same pool @@ -9977,6 +10023,11 @@ int main(int argc, char** argv) { const size_t nc = std::min(cand.size(), vict.size()); std::partial_sort(vict.begin(), vict.begin() + (ptrdiff_t) nc, vict.end(), [](auto& a, auto& b) { return a.first < b.first; }); + if (strata::kernels::cpu::expert_layout().qns()) { + for (const auto& [gain, in, out] : qns_swaps(l, cand, vict, o.adapt_min_gain, r, xcache)) + swaps.push_back({gain, (int32_t) l, in, out}); + continue; + } for (size_t i = 0; i < nc; ++i) { if (cand[i].first < vict[i].first + o.adapt_min_gain) break; swaps.push_back({cand[i].first - vict[i].first, (int32_t) l, cand[i].second, vict[i].second}); @@ -10048,7 +10099,7 @@ int main(int argc, char** argv) { if (dma_mode != 0 && slot >= 0 && b != nullptr) { cp_dst.push_back(xcache.device_slot(slot)); cp_src.push_back(b); - cp_bytes.push_back((size_t) strata::kernels::cpu::expert_layout().blob_bytes(s.layer)); + cp_bytes.push_back((size_t) strata::kernels::cpu::expert_layout().expert_bytes(s.layer, s.in)); main_live = true; host_res[out] = strata::core::kNotResident; srcp->prefetch(s.layer, s.out); @@ -10057,7 +10108,7 @@ int main(int argc, char** argv) { } if (slot < 0 || b == nullptr || adapt_copy_h2d(gs ? gs->cache.device_slot(slot) : xcache.device_slot(slot), b, - (size_t) strata::kernels::cpu::expert_layout().blob_bytes(s.layer), + (size_t) strata::kernels::cpu::expert_layout().expert_bytes(s.layer, s.in), gs ? gs->adapt_stream : adapt_stream) != cudaSuccess) { std::fprintf(stderr, "strata serve: adaptive swap copy failed (layer %d, slot %d, pinned %d): %s\n", (int) s.layer, (int) slot, pin_live.empty() ? 0 : 1, cudaGetErrorString(cudaGetLastError())); @@ -10241,6 +10292,11 @@ int main(int argc, char** argv) { const size_t nc = std::min(cand.size(), vict.size()); std::partial_sort(vict.begin(), vict.begin() + (ptrdiff_t) nc, vict.end(), [](auto& a, auto& b) { return a.first < b.first; }); + if (strata::kernels::cpu::expert_layout().qns()) { + for (const auto& [gain, in, out] : qns_swaps(l, cand, vict, o.adapt_min_gain, r, xcache)) + all.push_back({gain, (int32_t) l, in, out}); + continue; + } for (size_t i = 0; i < nc; ++i) { if (cand[i].first < vict[i].first + o.adapt_min_gain) break; all.push_back({cand[i].first - vict[i].first, (int32_t) l, cand[i].second, vict[i].second}); @@ -10926,7 +10982,7 @@ int main(int argc, char** argv) { float best = -1.0f; for (int64_t ex = 0; ex < g.n_expert; ++ex) { const size_t j = (size_t) (layer * g.n_expert + ex); - if (host_res[j] >= 0) continue; + if (host_res[j] >= 0 || strata::kernels::cpu::expert_layout().expert_bytes(layer, ex) > xcache.slot_bytes(slot)) continue; const float u = drive.d.usage.empty() ? 0.0f : drive.d.usage[j]; if (u > best) { best = u; pick = (int64_t) j; } } @@ -10934,7 +10990,7 @@ int main(int argc, char** argv) { if (pick < 0) continue; const uint8_t* b = srcp->blob_stable(layer, pick % g.n_expert); if (b == nullptr || !xcache.fill_slot_queued(slot, b, ferr, - (int64_t) strata::kernels::cpu::expert_layout().blob_bytes(layer))) { + (int64_t) strata::kernels::cpu::expert_layout().expert_bytes(layer, pick % g.n_expert))) { e = "VRAM: refilling the cache failed: " + ferr; return false; } @@ -12910,7 +12966,7 @@ int main(int argc, char** argv) { continue; } const uint8_t* b = srcp->blob_stable(i / g.n_expert, i % g.n_expert); - const int64_t nb = (int64_t) strata::kernels::cpu::expert_layout().blob_bytes(i / g.n_expert); + const int64_t nb = (int64_t) strata::kernels::cpu::expert_layout().expert_bytes(i / g.n_expert, i % g.n_expert); if (b == nullptr || !(refill_blocking() ? p.cache->fill_slot_blocking(slot, b, e, nb) : p.cache->fill_slot_queued(slot, b, e, nb))) return false; @@ -15442,7 +15498,7 @@ int main(int argc, char** argv) { const Clock::time_point tr = Clock::now(); for (const auto& [i, slot] : lent) { // D-4: queued, one wait (STRATA_REFILL_BLOCKING=1: each) const uint8_t* b = srcp->blob_stable(i / g.n_expert, i % g.n_expert); - const int64_t nb = (int64_t) strata::kernels::cpu::expert_layout().blob_bytes(i / g.n_expert); + const int64_t nb = (int64_t) strata::kernels::cpu::expert_layout().expert_bytes(i / g.n_expert, i % g.n_expert); if (b == nullptr || !(refill_blocking() ? xcache.fill_slot_blocking(slot, b, err, nb) : xcache.fill_slot_queued(slot, b, err, nb))) { std::fprintf(stderr, "strata generate: refilling a lent slot failed: %s\n", err.c_str()); @@ -15907,6 +15963,11 @@ int main(int argc, char** argv) { const size_t nc = std::min(cand.size(), vict.size()); std::partial_sort(vict.begin(), vict.begin() + (ptrdiff_t) nc, vict.end(), [](auto& a, auto& b) { return a.first < b.first; }); + if (strata::kernels::cpu::expert_layout().qns()) { + for (const auto& [gain, in, out] : qns_swaps(l, cand, vict, 1.5f, r, xcache)) + swaps.push_back({gain, (int32_t) l, in, out}); + continue; + } for (size_t i = 0; i < nc; ++i) { if (cand[i].first < vict[i].first + 1.5f) break; swaps.push_back({cand[i].first - vict[i].first, (int32_t) l, cand[i].second, vict[i].second}); @@ -15937,7 +15998,7 @@ int main(int argc, char** argv) { const size_t in = (size_t) s.layer * g.n_expert + s.in, out = (size_t) s.layer * g.n_expert + s.out; const int32_t slot = host_res[out]; const uint8_t* b = srcp->blob_stable(s.layer, s.in); - const size_t blob_n = (size_t) strata::kernels::cpu::expert_layout().blob_bytes(s.layer); + const size_t blob_n = (size_t) strata::kernels::cpu::expert_layout().expert_bytes(s.layer, s.in); // asynchronous: the copies run while the MTP drafts; the next window waits for them if (slot < 0 || b == nullptr) { std::fprintf(stderr, "strata generate: an adaptive refill failed\n"); diff --git a/tools/gguf_reader.py b/tools/gguf_reader.py index 4da9d2d248..24fb2243cc 100644 --- a/tools/gguf_reader.py +++ b/tools/gguf_reader.py @@ -49,6 +49,7 @@ "IQ1_M": (256, 56), "Q2_0": (64, 18), "MXFP4": (32, 17), "NVFP4": (64, 36), # fp16 scale + trellis stream (32 steps x 8 / 6 / 7 bits); PTQ1_0: 26 B of trits + fp16 scale "TQ2_T": (128, 34), "TQK6": (128, 26), "TQK7": (128, 30), "PTQ1_0": (128, 28), + "I8": (1, 1), "I64": (1, 8), } diff --git a/tools/iq_pack.py b/tools/iq_pack.py index b48af1aad3..feee15cd1b 100644 --- a/tools/iq_pack.py +++ b/tools/iq_pack.py @@ -221,7 +221,7 @@ def tensor_bytes(mm, g, t) -> np.ndarray: def is_expert(name: str) -> bool: - return name.startswith("blk.") and name.endswith(("_exps.weight",)) + return name.startswith("blk.") and name.endswith(("_exps.weight", "_exps_qns", "_exps_qns_idx")) def convert(name: str, type_name: str, raw: np.ndarray, compat_bf16: bool): @@ -277,7 +277,7 @@ def index_standalone(src, out, model: Model, compat_bf16: bool = False) -> int: reads them (FORM; converted when stored otherwise, see above), quantized ones served natively from the GGUF.""" todo, problems = [], [] for name, (g, t, mm, p) in model.where.items(): - if is_expert(t.name) or t.name in NOT_IN_PACK: + if is_expert(t.name) or t.name in NOT_IN_PACK or t.name.startswith("qns."): # qns.*: qns_rotation.bin continue if len(t.shape) > 2: print("tensor %s has %d dimensions; the index holds two" % (t.name, len(t.shape))) @@ -433,6 +433,8 @@ def expert_layout(model: Model, src: pathlib.Path): shards do not start their data section at the same byte, so one role's data_start must not be used for another's (per-role data_start as in #255, gopinath87607).""" T = {n: w[1] for n, w in model.where.items()} + if "blk.0.ffn_exps_qns" in T: + return qns_layout(model, src) exps = [n for n in T if n.startswith("blk.") and n.endswith("_exps.weight")] if not exps: return "the model has no expert tensors (blk.N.ffn_{gate,up,down}_exps.weight)" @@ -572,6 +574,36 @@ def hadamard_spec(meta: dict, layout, n_embd: int, n_ff: int): if r: lines.append("layer %d %d %d %d" % (l, r & 1, (r >> 1) & 1, (r >> 2) & 1)) return "\n".join(lines) + "\n" +QNS_TYPE = 250 # the engine's kQnsType: Qwen Next S trellis experts, one self-describing blob per expert + + +def qns_layout(model: Model, src: pathlib.Path): + """A qns model (Qwen Next S, strata/qns_convert.py): each layer's 512 blobs are one I8 payload (blk.N.ffn_exps_qns) + with an I64 [512, 4] index (offset, bytes, T_gu, T_d). The arena keeps the payloads back to back, so it is + compact; native_experts.txt v6 gives each layer its offset, its largest blob (the VRAM slot) and the payload's + absolute offset in its shard, and qns_offsets.bin every blob's arena offset (written by main).""" + T = {n: w[1] for n, w in model.where.items()} + n_layers = 1 + max(int(n.split(".")[1]) for n in T if n.endswith("_exps_qns")) + n_expert = int(T["blk.0.ffn_gate_inp.weight"].shape[1]) + layout, lines, offsets, offset = [], [], [], 0 + for l in range(n_layers): + g, t, mm, p = model.where["blk.%d.ffn_exps_qns" % l] + idx = model.bytes("blk.%d.ffn_exps_qns_idx" % l).view(np.int64).reshape(n_expert, 4) + assert idx[0, 0] == 0 and (idx[1:, 0] == idx[:-1, 0] + idx[:-1, 1]).all() and idx[-1, 0] + idx[-1, 1] == t.elements + blob = int(idx[:, 1].max()) + offsets.append(np.append(idx[:, 0], t.elements) + offset) + layout.append((l, QNS_TYPE, QNS_TYPE, offset, blob, [t])) + lines.append("%d %d %d %d %d %d 0 0%s" % (l, QNS_TYPE, QNS_TYPE, offset, blob, g.data_start + t.offset, + "" if p == src else " " + p.name)) + offset += t.elements + head = ("# strata native experts v6: layer gu_type d_type offset blob_bytes payload_off 0 0 [shard] (n_expert %d, " + "total %d; qns layers: blob_bytes is the layer's largest blob, every blob's offset in qns_offsets.bin; " + "absolute offsets in %s, or in the named shard beside it)\n" % (n_expert, offset, src.name)) + rotation = np.concatenate([model.bytes(n).view(np.float32) for n in + ("qns.rot_signs_2560", "qns.rot_q_2560", "qns.rot_signs_640", "qns.rot_q_640")]) + layout_files = {"qns_offsets.bin": np.concatenate(offsets).astype(np.int64).tobytes(), + "qns_rotation.bin": rotation.astype(np.float32).tobytes()} + return layout, head + "".join(line + "\n" for line in lines), n_expert, offset, layout_files def experts_source(model: Model, text: str, total: int) -> dict: @@ -622,13 +654,17 @@ def main() -> int: if isinstance(got, str): print(got) return 1 - layout, text, n_expert, offset = got + layout, text, n_expert, offset, *qns_files = got + if qns_files and a.experts_bin: + print("a qns pack reads its experts from the GGUF payloads; --experts-bin is not supported for it") + return 1 if any(t.type_name == "NVFP4" for *_, ts in layout for t in ts): a.experts_bin = True # the GGUF holds the weights without the scale tails: the engine must read experts.bin # ---- prism.hadamard: rotated experts need hadamard.txt beside native_experts.txt, which becomes v5 so an engine - # that cannot rotate refuses the pack (instead of running the folded experts on unrotated activations) + # that cannot rotate refuses the pack (instead of running the folded experts on unrotated activations). A qns + # pack (v6) has its own rotation in qns_rotation.bin and no per-role tensors to fold. t0 = layout[0][5] - had = hadamard_spec(g.metadata, layout, int(t0[0].shape[0]), int(t0[2].shape[0])) + had = None if qns_files else hadamard_spec(g.metadata, layout, int(t0[0].shape[0]), int(t0[2].shape[0])) if isinstance(had, str) and not had.startswith("# strata hadamard"): print("prism.hadamard: " + had) return 1 @@ -716,6 +752,8 @@ def layer_blobs(l, blob, ts): htmp.replace(out / "hadamard.txt") else: (out / "hadamard.txt").unlink(missing_ok=True) # a previous Hadamard model's, packed into this folder + for name, data in (qns_files[0].items() if qns_files else ()): # before native_experts.txt, the marker + (out / name).write_bytes(data) tmp = out / "native_experts.txt.tmp" with open(tmp, "w", encoding="utf-8", newline="\n") as fo: fo.write(text)