From a4c41a92ee9e7c85e46bd6fbe1b5f99781bb420c Mon Sep 17 00:00:00 2001 From: ferstar <2854276+ferstar@users.noreply.github.com> Date: Tue, 27 Oct 2026 01:11:53 +0100 Subject: [PATCH 1/2] Build the engine for ARM64 (DGX Spark / GB10) without changing the x86 build (#3622) The engine did not build on aarch64: the CPU expert path assumed the x86 AVX kernels, the CPUID probes and the x86 spin/fence intrinsics. After this change the strata target compiles on aarch64 (CUDA host code, native GGUF packs). On x86 nothing is selected differently. What changed: - CMake: STRATA_CPU_X86 is decided from CMAKE_SYSTEM_PROCESSOR (an unset processor stays the x86 build). On aarch64 src/kernels/cpu/portable.cpp stands in for the AVX kernel sources; native experts run on ggml-cpu's NEON dot products. -ffp-contract=off on aarch64. STRATA_ISA_FLOOR is ignored with a warning there. - include/strata/platform/cpu_relax.hpp: strata_cpu_pause / strata_store_fence map to _mm_pause / _mm_sfence on x86 and to yield / dmb oshst on aarch64. verify.cpp, pool.cpp, session.cpp, mtp.cpp and graph.cpp use them. - expert_layout.cpp: the CPUID/XCR0 probes compile on x86 only and report no AVX on other CPUs. - generate.cpp: --cpu-oracle-q8-0 is rejected on a non-x86 CPU. - setup.py: nvidia-smi's "[N/A]" memory on GB10 (unified memory) takes the pool size from /proc/meminfo on ARM64 only; the CPU probe no longer needs the x86 "flags" line; setup warns that the ARM64 engine is compiled on the PC and which packs work. Changes to the contributor's patch, found while moving it onto the 0.2.7 line: - portable.cpp had no iq512x_rows / iq512x_supported / iq512_gather_on (added to the x86 kernels after the PR was written): the aarch64 link would fail. - The Metal build (Apple silicon) shares these files. It keeps its barrier (dmb ishst), its plain router loop (its bf16_rows_dot_multi is a stub), its Metal start-up block in generate.cpp and its "unknown" CPU name. - The CPU side of the GPU doorbell: strata_load_fence() after the spins that wait for a flag the GPU raised (dmb ishld on ARM, nothing on x86 and on Apple silicon), so the data written before the flag is not read early. - A card whose memory nvidia-smi does not report stays unlisted on x86 (the new fallback is for ARM64 only). - tests/terminal/test_setup_arm64.py. Checked: with the x86 CUDA build (sm_86, gcc 13 -O3, z440) the objects of verify, session, mtp, graph, expert_source, pool, expert_layout and generate are byte-identical before and after (the x86 wrappers are macros for the same intrinsics). Every C++ file of the strata target compiles for aarch64 with clang, and the symbols of the CPU kernel library link up (syntax and symbols only: no ARM machine, no ARM run). The Metal-shared files compile for aarch64-macos with STRATA_GPU_METAL. Not covered: no ARM hardware here. The tests that include x86 headers (iq_avx2_parity and the other AVX parity tools) and the DeepSeek target do not build on aarch64; the strata target does not need them. Memory ordering on ARM is reasoned, not run. The Q2_0 and Gyro packs need the x86 kernels. --- CMakeLists.txt | 32 +++- include/strata/platform/cpu_relax.hpp | 35 +++++ setup.py | 39 ++++- src/core/expert_source.cpp | 5 + src/core/graph.cpp | 4 +- src/core/mtp.cpp | 12 +- src/core/session.cpp | 15 +- src/core/verify.cpp | 38 ++--- src/kernels/cpu/expert_layout.cpp | 44 +++++- src/kernels/cpu/pool.cpp | 10 +- src/kernels/cpu/portable.cpp | 214 ++++++++++++++++++++++++++ src/program/generate.cpp | 9 ++ tests/terminal/test_setup_arm64.py | 68 ++++++++ 13 files changed, 472 insertions(+), 53 deletions(-) create mode 100644 include/strata/platform/cpu_relax.hpp create mode 100644 src/kernels/cpu/portable.cpp create mode 100644 tests/terminal/test_setup_arm64.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 9c6f2879..57c6b411 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -59,6 +59,17 @@ function(strata_utf8_manifest target) endif() endfunction() +# The AVX kernels exist for x86 only. On aarch64 (DGX Spark / GB10) portable.cpp stands in for them and +# native experts run on ggml-cpu's NEON dot products. +if(NOT CMAKE_SYSTEM_PROCESSOR OR CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|AMD64|amd64|x64|X64|i[3-6]86|x86|X86)$") # an unset processor stays the x86 build, as before + set(STRATA_CPU_X86 ON) +else() + set(STRATA_CPU_X86 OFF) +endif() +if(STRATA_ISA_FLOOR AND NOT STRATA_CPU_X86) + message(WARNING "STRATA_ISA_FLOOR is an x86 older-CPU build and does not apply to ${CMAKE_SYSTEM_PROCESSOR}; ignored") + set(STRATA_ISA_FLOOR "" CACHE STRING "" FORCE) +endif() set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_CXX_EXTENSIONS OFF) @@ -1644,8 +1655,13 @@ if(STRATA_BUILD_TESTS OR STRATA_BUILD_CONVERSATION_TESTS) add_test(NAME batch_groups_test COMMAND batch_groups_test) endif() -add_library(strata_kernels_cpu STATIC src/kernels/cpu/expert.cpp src/kernels/cpu/pool.cpp - src/kernels/cpu/expert_layout.cpp src/kernels/cpu/q2_avx2.cpp src/kernels/cpu/hadamard.cpp) +if(STRATA_CPU_X86) + add_library(strata_kernels_cpu STATIC src/kernels/cpu/expert.cpp src/kernels/cpu/pool.cpp + src/kernels/cpu/expert_layout.cpp src/kernels/cpu/q2_avx2.cpp src/kernels/cpu/hadamard.cpp) +else() + add_library(strata_kernels_cpu STATIC src/kernels/cpu/portable.cpp src/kernels/cpu/pool.cpp + src/kernels/cpu/expert_layout.cpp src/kernels/cpu/hadamard.cpp) +endif() target_include_directories(strata_kernels_cpu PUBLIC include) # Plan v0.3 P6: only the kernel files get the AVX-512 / AVX2 flags; the pool and the probe run on any x86-64, # and the engine picks the kernels at run time (cpu_avx512_ok). @@ -1757,11 +1773,13 @@ if(STRATA_NATIVE_EXPERTS) endif() # part of strata_kernels_cpu (the pool calls these and they call its CPU probe: one library, so GNU ld's # one-pass static linking has no cycle to resolve) - target_sources(strata_kernels_cpu PRIVATE src/kernels/cpu/native_expert.cpp src/kernels/cpu/iq_avx512.cpp - 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) + target_sources(strata_kernels_cpu PRIVATE src/kernels/cpu/native_expert.cpp) + if(STRATA_CPU_X86) + target_sources(strata_kernels_cpu PRIVATE src/kernels/cpu/iq_avx512.cpp 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) + endif() 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 diff --git a/include/strata/platform/cpu_relax.hpp b/include/strata/platform/cpu_relax.hpp new file mode 100644 index 00000000..c5dba7d3 --- /dev/null +++ b/include/strata/platform/cpu_relax.hpp @@ -0,0 +1,35 @@ +// include/strata/platform/cpu_relax.hpp - the spin-wait hint and the store fence, per architecture. +// +// x86: `_mm_pause` and `_mm_sfence` (the fence orders the write-combining stores into pinned host memory before +// the doorbell the GPU polls). aarch64: `yield` and `dmb oshst` - the outer-shareable store barrier is what +// makes host stores visible to a coherent device (GB10 / Grace NVLink-C2C) before a later store. Apple silicon (the +// Metal build) keeps `dmb ishst`, the barrier its compat header (metal/include/metal_compat/immintrin.h) has always used. +// Local port of the approach in https://github.com/Niko1221/Strata/pull/1077 onto main e8ca9af. +#pragma once + +#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86) +#include +#define STRATA_X86 1 +// Macros on x86, so the code the compiler sees is exactly the intrinsic that stood there before (the objects of the +// files that use them are identical before and after this header). x86 does not reorder a load before an older load: +// the load fence has nothing to emit. +#define strata_cpu_pause() _mm_pause() +#define strata_store_fence() _mm_sfence() +#define strata_load_fence() ((void) 0) +#elif defined(__aarch64__) +inline void strata_cpu_pause() { __asm__ __volatile__("yield" ::: "memory"); } +#if defined(__APPLE__) +inline void strata_store_fence() { __asm__ __volatile__("dmb ishst" ::: "memory"); } +#define strata_load_fence() ((void) 0) // the Metal build has never had one +#else +inline void strata_store_fence() { __asm__ __volatile__("dmb oshst" ::: "memory"); } +// After a spin saw a flag the GPU raised, the data the GPU wrote before it must be read after the flag: an ARM core +// may otherwise load that data early. `dmb ishld` is the load barrier. +inline void strata_load_fence() { __asm__ __volatile__("dmb ishld" ::: "memory"); } +#endif +#else +#include +inline void strata_cpu_pause() {} +inline void strata_store_fence() { std::atomic_thread_fence(std::memory_order_seq_cst); } +inline void strata_load_fence() { std::atomic_thread_fence(std::memory_order_acquire); } +#endif diff --git a/setup.py b/setup.py index 55f8941f..241f68ee 100644 --- a/setup.py +++ b/setup.py @@ -1335,7 +1335,11 @@ def cpu_info(): else: try: txt = open("/proc/cpuinfo").read() - flags = set(re.search(r"^flags\s*:\s*(.*)$", txt, re.M).group(1).split()) + mflags = re.search(r"^flags\s*:\s*(.*)$", txt, re.M) + if mflags is None: + # aarch64 lists Features, not the x86 flags line. The normal engine is the native one. + return name, False, False + flags = set(mflags.group(1).split()) avx2 = "avx2" in flags avx512 = {"avx512f", "avx512bw", "avx512vl", "avx512_vnni", "avx512vbmi"} <= flags m = re.search(r"^model name\s*:\s*(.*)$", txt, re.M) @@ -1366,6 +1370,11 @@ def _cpuid_floor() -> str: return "" +def is_arm64() -> bool: + """An ARM64 CPU (DGX Spark / GB10, Grace): there is no AVX here, and no ready-made engine for it.""" + return platform.machine().lower() in ("aarch64", "arm64") + + def cpu_floor(avx2: bool) -> str: """The experimental older-CPU build this PC needs (#394 #595 #623): "" with AVX2 (the normal engine), "avx" (Sandy / Ivy Bridge, AMD Bulldozer), "none" (SSE4.2 + POPCNT: Nehalem, Westmere), or "unsupported". STRATA_ISA_FLOOR=avx| @@ -1373,6 +1382,8 @@ def cpu_floor(avx2: bool) -> str: forced = os.environ.get("STRATA_ISA_FLOOR", "").strip().lower() if forced in ("avx", "none"): return forced + if is_arm64(): + return "" if avx2: return "" if WIN: @@ -1457,7 +1468,21 @@ def gpus(): for line in s.strip().splitlines(): try: idx, name, mem, cc, drv = [x.strip() for x in line.split(",")] - found.append({"index": int(idx), "name": name, "vram_gb": float(mem) / 1024.0, "arch": cc.replace(".", ""), + try: + mem_gb = float(mem) / 1024.0 # nvidia-smi nounits: MiB + except ValueError: + if not is_arm64(): + raise # as before on x86: a card whose memory nvidia-smi does not report is not listed + # GB10 unified memory: memory.total is [N/A]. The pool is the machine RAM. + mem_gb = 0.0 + try: + for info in open("/proc/meminfo"): + if info.startswith("MemTotal:"): + mem_gb = int(info.split()[1]) / 1024.0 / 1024.0 + break + except OSError: + pass + found.append({"index": int(idx), "name": name, "vram_gb": mem_gb, "arch": cc.replace(".", ""), "driver": drv}) except ValueError: continue @@ -9014,7 +9039,7 @@ def main() -> int: pf = page_file_gb() if pf is not None and pf < 4: warn(page_file_advice(pf, page_file_setting())) - ok(f"CPU: {cpu} ({'AVX-512' if avx512 else 'AVX2' if avx2 else 'no AVX2'})") + ok(f"CPU: {cpu} ({'ARM64' if is_arm64() else 'AVX-512' if avx512 else 'AVX2' if avx2 else 'no AVX2'})") link = None if hip else pcie_link(int(gpu.get("index", 0))) if link is not None: line, problem = pcie_lines(link) @@ -9024,7 +9049,13 @@ def main() -> int: floor = cpu_floor(avx2) if floor == "unsupported": fail("this CPU has neither AVX2 nor SSE4.2; Strata needs at least SSE4.2 (Intel Nehalem, 2008, or newer)") - if not avx2: + if is_arm64(): + # DGX Spark / GB10 (#3622): no ready-made ARM64 engine; setup compiles it here. A warning, not a stop. + warn("ARM64 CPU: setup compiles the engine on this PC (once). The CPU's share of the experts runs on ggml-cpu's " + "NEON kernels. Tested: GGUF packs with i-quant experts (IQ3_S). The Q2_0 and Gyro packs need the x86 " + "kernels and do not run on this CPU. This port is new and has not been measured against an x86 PC") + a.build = True + elif not avx2: # #394 #595 #623: the ready-made engine is AVX2; an older CPU gets one compiled here, whose CPU experts run on # ggml-cpu's kernels for this CPU. Experimental: measured only on newer CPUs with the older path forced, and by # users on a few Xeons. A warning, not a stop. diff --git a/src/core/expert_source.cpp b/src/core/expert_source.cpp index b8ad8fb4..43ec2615 100644 --- a/src/core/expert_source.cpp +++ b/src/core/expert_source.cpp @@ -3184,6 +3184,10 @@ void RouterLookahead::run() { // one (kq_avx1.cpp) or, without AVX, the same sums in plain C++. From the Strata_Dirigo fork (rwkeyes): an // AVX-only Xeon E5-2687W died here (vpmovzxwd) on its first request. An estimate only (which experts to // prefetch); the order of the additions differs between the three, the output does not depend on it. +#if defined(__aarch64__) && !defined(STRATA_GPU_METAL) // Metal (Apple silicon) keeps the plain loop below: its bf16_rows_dot_multi is a stub + strata::kernels::cpu::bf16_rows_dot_multi(routers_[(size_t) layer].data(), (int) n_expert_, (int) n_embd_, + x_.data(), (int) nt, logits.data()); +#else if (strata::kernels::cpu::cpu_avx2_ok()) { strata::kernels::cpu::bf16_rows_dot_multi(routers_[(size_t) layer].data(), (int) n_expert_, (int) n_embd_, x_.data(), (int) nt, logits.data()); @@ -3209,6 +3213,7 @@ void RouterLookahead::run() { } } } +#endif const int kmax = stats_on_ ? (int) std::min(k_ + 10, n_expert_) : k_; const int ks[3] = {k_, (int) std::min(k_ + 4, n_expert_), kmax}; std::vector sm; diff --git a/src/core/graph.cpp b/src/core/graph.cpp index 2e9c7e70..da2c5032 100644 --- a/src/core/graph.cpp +++ b/src/core/graph.cpp @@ -1,7 +1,7 @@ // src/core/graph.cpp - P2.S5: the GraphRegistry implementation. #include "strata/core/graph.hpp" -#include +#include "strata/platform/cpu_relax.hpp" #include #include @@ -97,7 +97,7 @@ bool CapturedGraph::wait_ms(int timeout_ms) const { if (q == cudaSuccess) return true; if (q != cudaErrorNotReady) return false; // a real error, not "not finished" if (std::chrono::steady_clock::now() >= deadline) return false; - for (int i = 0; i < 64; ++i) _mm_pause(); + for (int i = 0; i < 64; ++i) strata_cpu_pause(); } } diff --git a/src/core/mtp.cpp b/src/core/mtp.cpp index 292a5d2b..af59e02f 100644 --- a/src/core/mtp.cpp +++ b/src/core/mtp.cpp @@ -44,11 +44,7 @@ #include #include #include -#if defined(_WIN32) -#include -#elif defined(__x86_64__) -#include -#endif +#include "strata/platform/cpu_relax.hpp" #ifdef STRATA_NATIVE_EXPERTS #include "ggml.h" // --mtp-q4: the head rows' and projections' formats, dequantized and requantized at load @@ -1641,10 +1637,11 @@ bool MtpDrafter::draft(int T, const int32_t* tokens, int64_t p, int a, int32_t* uint32_t spins = 0; while (((volatile int32_t*) h_out_)[j] < 0) { #if defined(_WIN32) || defined(__x86_64__) - _mm_pause(); + strata_cpu_pause(); #endif if ((++spins & 1023u) == 0 && cudaStreamQuery(cs_) != cudaErrorNotReady) break; } + strata_load_fence(); // the token was seen: its probability is read after it (a no-op on x86) drafts[j] = ((volatile int32_t*) h_out_)[j]; if (probs) probs[j] = ((volatile float*) h_prob_)[j]; prefetch_ple(drafts[j]); @@ -1677,10 +1674,11 @@ bool MtpDrafter::draft(int T, const int32_t* tokens, int64_t p, int a, int32_t* uint32_t spins = 0; while (((volatile int32_t*) h_out_)[j] < 0) { #if defined(_WIN32) || defined(__x86_64__) - _mm_pause(); + strata_cpu_pause(); #endif if ((++spins & 1023u) == 0 && cudaStreamQuery(cs_) != cudaErrorNotReady) break; } + strata_load_fence(); }; tl = Clock::now(); wait_step(0); diff --git a/src/core/session.cpp b/src/core/session.cpp index 0a25df79..42a7ff77 100644 --- a/src/core/session.cpp +++ b/src/core/session.cpp @@ -21,14 +21,9 @@ #include #include -// `_mm_pause` for the doorbell spin. Guarded because it is x86-only; a target without it still builds, the -// spin is just less polite to the pipeline. -#if defined(_MSC_VER) || defined(__x86_64__) || defined(__i386__) -#include -#define STRATA_SPIN_PAUSE() _mm_pause() -#else -#define STRATA_SPIN_PAUSE() ((void) 0) -#endif +// The doorbell spin. x86 pauses, aarch64 yields (include/strata/platform/cpu_relax.hpp). +#include "strata/platform/cpu_relax.hpp" +#define STRATA_SPIN_PAUSE() strata_cpu_pause() namespace strata::core { namespace { @@ -716,6 +711,7 @@ bool session_loop(const ModelGeometry& g, int64_t pos, int32_t pos_base, Session if (rang) break; // rung, and the graph is still running: the overlap STRATA_SPIN_PAUSE(); } + strata_load_fence(); // the ring was seen: the GPU's x_f / ids / weights are read after it (a no-op on x86) if (mid_graph) ++gr.rings_mid_graph; // the window, from just before the launch to the instant the ring was seen const auto t_ring = std::chrono::steady_clock::now(); @@ -968,11 +964,12 @@ bool session_run_token(const ModelGeometry& g, int64_t pos, int32_t pos_base, Se return false; } } + strata_load_fence(); // as above, for the token loop's wait const auto t1 = Clock::now(); progress_at("token: the CPU experts of layer", l); if (pool != nullptr) pool(user, s.db->h_x_f, s.db->h_ids, s.db->h_weights, g.n_embd, s.k, y_miss_host); std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); *flag = want; const auto t2 = Clock::now(); tg.ms_wait += std::chrono::duration(t1 - t0).count(); diff --git a/src/core/verify.cpp b/src/core/verify.cpp index 5e97b8f5..d5e41255 100644 --- a/src/core/verify.cpp +++ b/src/core/verify.cpp @@ -75,7 +75,7 @@ #include #include #include -#include +#include "strata/platform/cpu_relax.hpp" #if defined(STRATA_GPU_METAL) // The Metal backend's decode window (metal/src/engine/window.cpp): this file's init/run/commit hand over to it. @@ -517,7 +517,7 @@ bool Verifier::release_gpu_waits(int timeout_ms) { for (uint32_t* p : {h_flag_, h_flagA_, h_flagB_, hflag_in_h_}) // + an overlapped split's hand-off wait if (p != nullptr) *(volatile uint32_t*) p = UINT32_MAX; std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); const OnDevice on_device(device_); const Clock::time_point t0 = Clock::now(); for (cudaStream_t s : {cs_, copy_}) { @@ -2738,7 +2738,7 @@ void Verifier::stage_inputs(int T, const int32_t* tokens, int64_t pos0) { if (herr_h_ != nullptr) *(volatile uint32_t*) herr_h_ = 0; if (hdrop_h_ != nullptr) *(volatile uint32_t*) hdrop_h_ = (g_test_drop > 0 && windows + 1 == g_test_drop) ? 1u : 0u; std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); last_t_ = T; if (rs_on_ && T >= 1 && T <= 8) ++rs_win_[T]; last_pos0_ = pos0; @@ -2811,7 +2811,7 @@ bool Verifier::adapt_gate_wait(int32_t layer, std::string& err) { const Clock::time_point t0 = Clock::now(); uint32_t spins = 0; while (adapt_gate_->reached() < layer) { - _mm_pause(); + strata_cpu_pause(); if ((++spins & 1023u) != 0) continue; if (Clock::now() - t0 > std::chrono::seconds(20)) { err = "verify: the adaptive tier's copies for layer " + std::to_string(layer) + " never landed" + @@ -2917,7 +2917,7 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, const Clock::time_point tp = Clock::now(); if (!ss.ple.table->gather_batch(ple_rows, (size_t) T, h_ple_, err)) return false; std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); ms_host += ms_since(tp); } *(volatile uint32_t*) h_flag_ = 1; @@ -2992,7 +2992,7 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, const Clock::time_point tp = Clock::now(); if (!ss.ple.table->gather_batch(ple_rows, (size_t) T, h_ple_, err)) return false; std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); *flag = 1; ms_host += ms_since(tp); } else { @@ -3008,7 +3008,7 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, uint32_t spins = 0; progress_at("verify window: waiting for the GPU to reach layer", l); while (*seq < want) { - _mm_pause(); + strata_cpu_pause(); if ((++spins & 1023u) != 0) continue; const auto now = Clock::now(); if (now - last_flush > std::chrono::microseconds(2000)) { @@ -3044,6 +3044,7 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, return false; } } + strata_load_fence(); // the layer rang: the GPU's activations are read after it (a no-op on x86) const Clock::time_point b = Clock::now(); if (g_trace) trace_ev("RANG", k, l, (int64_t) std::chrono::duration_cast(b - a).count()); VDBG("layer %lld rang\n", (long long) l); @@ -3064,7 +3065,7 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, (int64_t) ms_since(b)); // aux: ms the CPU experts took progress_tick(); std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); if (h_bal_ != nullptr) { // set_pcie_balance: the step's plan, its CPU and its PCIe experts const size_t bi = (size_t) ((k / G) * 2 + grp); bal_ok_[bi] = *(volatile uint32_t*) h_flagA_ == want; @@ -3090,7 +3091,7 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, const Clock::time_point tp = Clock::now(); if (!ss.ple.table->gather_batch(ple_rows, (size_t) T, h_ple_, err)) return false; std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); ms_host += ms_since(tp); } if (!(test_stall && k + 1 == steps)) *flag = want; @@ -3331,7 +3332,7 @@ void Verifier::fetch_dma(void* ctx, const uint8_t* const* src, int n, size_t byt void Verifier::publish_plan(void* ctx) { Verifier* v = (Verifier*) ctx; - _mm_sfence(); + strata_store_fence(); *(volatile uint32_t*) v->h_flagA_ = v->cur_layer_ + 1; } @@ -3862,7 +3863,7 @@ bool Verifier::run_slot_rows(const int* rows, int S, const int32_t* tokens, cons if (ar_on()) { if (ss_->ple.ready() && ple_stage()) { std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); *flag = 1; } const cudaError_t se = cudaStreamSynchronize(cs_); @@ -3886,7 +3887,7 @@ bool Verifier::run_slot_rows(const int* rows, int S, const int32_t* tokens, cons uint32_t spins = 0; progress_at("verify batch: waiting for the GPU to reach layer", l); while (*seq < want) { - _mm_pause(); + strata_cpu_pause(); if ((++spins & 1023u) != 0) continue; const auto now = Clock::now(); if (now - last_flush > std::chrono::microseconds(2000)) { @@ -3903,6 +3904,7 @@ bool Verifier::run_slot_rows(const int* rows, int S, const int32_t* tokens, cons return false; } } + strata_load_fence(); // as above, for the batch loop's wait const Clock::time_point b = Clock::now(); cur_layer_ = want - 1; set_plan_slot(0); @@ -3914,7 +3916,7 @@ bool Verifier::run_slot_rows(const int* rows, int S, const int32_t* tokens, cons if (remote_opt_) remote_opt_->end(); progress_tick(); std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); if (*(volatile uint32_t*) h_flagA_ != want) { // the pool did not publish a plan: an empty one sink_.counts[0] = 0; sink_.counts[1] = 0; @@ -4039,7 +4041,7 @@ int Verifier::batch_serve(PoolMultiFn pool, void* user) { // run_slot_rows makes). Per-window (ar_on): with a loan in flight the doorbell graph needs the per-layer service. if (ar_on() && ss_->ple.ready() && ple_stage()) { std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); *h_flag_ = 1; } for (;;) { @@ -4059,7 +4061,7 @@ int Verifier::batch_serve(PoolMultiFn pool, void* user) { if (remote_opt_) remote_opt_->end(); progress_tick(); std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); if (*(volatile uint32_t*) h_flagA_ != want) { // the pool did not publish a plan: an empty one sink_.counts[0] = 0; sink_.counts[1] = 0; @@ -4354,7 +4356,7 @@ int Verifier::service(PoolMultiFn pool, void* user, std::string& err) { const Clock::time_point tp = Clock::now(); if (!ss.ple.table->gather_batch(pl_ple_rows_.data(), (size_t) T, h_ple_, err)) return false; std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); ms_host += ms_since(tp); return true; }; @@ -4405,7 +4407,7 @@ int Verifier::service(PoolMultiFn pool, void* user, std::string& err) { (int64_t) ms_since(b)); progress_tick(); std::atomic_thread_fence(std::memory_order_seq_cst); - _mm_sfence(); + strata_store_fence(); if (*(volatile uint32_t*) h_flagA_ != want) { // the pool did not publish a plan: an empty one sink_.counts[0] = 0; sink_.counts[1] = 0; @@ -4479,7 +4481,7 @@ bool Verifier::pl_commit_async(int n_keep, std::string& err) { const Clock::time_point t0 = Clock::now(); if (commit_live_) { // its graph reads the words below when it starts: the previous one has (a window ago) cudaError_t q; - while ((q = cudaEventQuery(ev_commit_)) == cudaErrorNotReady) _mm_pause(); + while ((q = cudaEventQuery(ev_commit_)) == cudaErrorNotReady) strata_cpu_pause(); if (q != cudaSuccess) { err = std::string("verify: commit: ") + cudaGetErrorString(q); return false; } } h_commit_[0] = n_keep; diff --git a/src/kernels/cpu/expert_layout.cpp b/src/kernels/cpu/expert_layout.cpp index 9d296c97..bcbedaa2 100644 --- a/src/kernels/cpu/expert_layout.cpp +++ b/src/kernels/cpu/expert_layout.cpp @@ -10,7 +10,7 @@ #if defined(_MSC_VER) #include #include -#else +#elif defined(__x86_64__) || defined(__i386__) #include #endif #include @@ -39,6 +39,38 @@ int cpu_isa_cap() { return cap; } +#if !defined(__x86_64__) && !defined(__i386__) && !defined(_M_X64) && !defined(_M_IX86) +bool cpu_avx512_ok() { return false; } +bool cpu_avx512bw_ok() { return false; } +bool cpu_avx2_ok() { return false; } +bool cpu_avx1_ok() { return false; } +bool cpu_sse42_ok() { return false; } +bool cpu_gather_fast() { return false; } +bool cpu_gather_fast_here() { return false; } +bool cpu_avxvnni_ok() { return false; } +std::string cpu_name() { +#if defined(__APPLE__) + return "unknown"; // as before: the Metal build's cpuid shim read all zeros +#else + std::ifstream f("/proc/cpuinfo"); + std::string line; + while (std::getline(f, line)) { + if (line.rfind("model name", 0) == 0 || line.rfind("Model", 0) == 0) { + const size_t c = line.find(':'); + if (c != std::string::npos) { + const size_t b = line.find_first_not_of(" \t", c + 1); + if (b != std::string::npos) return line.substr(b); + } + } + } +#if defined(__aarch64__) + return "aarch64"; +#else + return "unknown"; +#endif +#endif // __APPLE__ +} +#else bool cpu_avx512_ok() { static const bool ok = [] { if (const char* f = std::getenv("STRATA_FORCE_AVX2"); f != nullptr && f[0] == '1') return false; @@ -186,6 +218,8 @@ bool cpu_sse42_ok() { return ok; } +#endif // x86 cpu probes + const char* isa_floor_build() { #if defined(STRATA_ISA_FLOOR_AVX) return "avx"; @@ -196,6 +230,7 @@ const char* isa_floor_build() { #endif } +#if defined(__x86_64__) || defined(__i386__) || defined(_M_X64) || defined(_M_IX86) namespace { void cpuid_regs(unsigned leaf, unsigned sub, unsigned r[4]) { #if defined(_MSC_VER) @@ -207,6 +242,7 @@ void cpuid_regs(unsigned leaf, unsigned sub, unsigned r[4]) { #endif } } // namespace +#endif // x86 cpuid helper int iq256_gather_setting() { static const int s = [] { @@ -217,6 +253,7 @@ int iq256_gather_setting() { return s; } +#if defined(__x86_64__) || defined(__i386__) || defined(_M_X64) || defined(_M_IX86) bool cpu_gather_fast() { static const bool ok = [] { if (cpu_isa_cap() < 3) return false; // STRATA_FORCE_ISA: as on a CPU that stops at AVX2 @@ -295,6 +332,7 @@ std::string cpu_name() { const size_t b0 = name.find_first_not_of(' '), b1 = name.find_last_not_of(' '); return b0 == std::string::npos ? std::string("unknown") : name.substr(b0, b1 - b0 + 1); } +#endif // x86 cpuid probes // ---- the AVX-512 probe and the oracle flag (#795) ---- // @@ -322,6 +360,9 @@ const char* CpuFeatures::reason() const { CpuFeatures cpu_features() { CpuFeatures f; +#if !defined(__x86_64__) && !defined(__i386__) && !defined(_M_X64) && !defined(_M_IX86) + return f; // the AVX-512 kernels are x86-only: no AVX-512 on this architecture +#else int reg[4] = {0, 0, 0, 0}; #if defined(_MSC_VER) __cpuid(reg, 0); @@ -341,6 +382,7 @@ CpuFeatures cpu_features() { f.avx512_vnni = (ecx >> 11) & 1u; f.avx512_vbmi = (ecx >> 1) & 1u; return f; +#endif } void cpu_require_expert_support() { diff --git a/src/kernels/cpu/pool.cpp b/src/kernels/cpu/pool.cpp index cdcc8d79..2d362400 100644 --- a/src/kernels/cpu/pool.cpp +++ b/src/kernels/cpu/pool.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include "strata/platform/cpu_relax.hpp" #if defined(_MSC_VER) #include // __cpuidex, __rdtsc (STRATA_POOL_UMWAIT) #endif @@ -1010,7 +1010,7 @@ void ExpertPool::worker(int id) { } else #endif { - _mm_pause(); + strata_cpu_pause(); if ((++spins & 1023u) != 0) continue; } if (std::chrono::steady_clock::now() - parked_at < spin_before_sleep_) continue; @@ -1131,7 +1131,7 @@ void ExpertPool::wait_parked(const char* what) { uint32_t spins = 0; std::chrono::steady_clock::time_point t0{}; while (parked_count() != (uint32_t) n_) { - _mm_pause(); + strata_cpu_pause(); if ((++spins & 1023u) != 0) continue; const auto now = std::chrono::steady_clock::now(); if (spins == 1024u) t0 = now; @@ -1158,7 +1158,7 @@ void ExpertPool::wait_done(int n) { else d = done_.load(std::memory_order_acquire); if (d >= (uint32_t) n) return; // `>=`: never a wait that an overshoot outlives - _mm_pause(); + strata_cpu_pause(); if ((++spins & 1023u) != 0) continue; const auto now = std::chrono::steady_clock::now(); if (spins == 1024u || d != seen) { t0 = now; seen = d; } // progress restarts the clock @@ -1408,7 +1408,7 @@ void ExpertPool::run(ExpertJob* jobs, int n) { // ---- **THE HOST DRAINS TOO (R2.2), INSTEAD OF SPINNING ON `done_`.** // - // The loop below used to be `while (done_ != n) _mm_pause();`. The host is pinned to core 0 - the core + // The loop below used to be `while (done_ != n) strata_cpu_pause();`. The host is pinned to core 0 - the core // `physical_cores(true)` deliberately keeps the five workers off - so for the whole drain that core was // idle while five cores did six cores' worth of work. Measured before the change: 33.7 GB/s against // 5/6 x 44.14 = 36.8 for five workers and 44.14 for six. diff --git a/src/kernels/cpu/portable.cpp b/src/kernels/cpu/portable.cpp new file mode 100644 index 00000000..a8ff501c --- /dev/null +++ b/src/kernels/cpu/portable.cpp @@ -0,0 +1,214 @@ +// src/kernels/cpu/portable.cpp - the CPU kernel symbols on a non-x86 build (aarch64: DGX Spark / GB10, Grace). +// +// Replaces expert.cpp, q2_avx2.cpp, iq_avx2.cpp, iq_avx512.cpp, kq_avx2.cpp and kq_avx1.cpp, which are AVX-only. +// What a native (GGUF) pack needs from this file is the activation quantizer (expert.cpp's scalar loop, bitwise +// the same as its AVX-512 version) and the BF16 router dot (NEON). The i-quant / k-quant multi-token kernels +// report themselves unsupported, so native_expert.cpp takes ggml-cpu's vec_dot, which has NEON paths. The +// canonical Q2_0 pack (its own blob layout) is refused: those kernels are AVX-512 only. A native pack's Q2_0 +// rows are not routed here: q2_native_kernels() is false when cpu_avx2_ok() is, so they use ggml-cpu too. +// Local port of the approach in https://github.com/Niko1221/Strata/pull/1077 onto main e8ca9af. +#include "strata/kernels/cpu/expert.hpp" +#include "strata/kernels/cpu/iq_avx2.hpp" +#include "strata/kernels/cpu/iq_avx512.hpp" +#include "strata/kernels/cpu/kq_avx1.hpp" +#include "strata/kernels/cpu/kq_avx2.hpp" +#include "strata/kernels/cpu/iq_avx1.hpp" +#include "strata/kernels/cpu/nvfp4_avx2.hpp" +#include "strata/kernels/cpu/nvfp4_avx512.hpp" +#include "strata/kernels/cpu/tq_cpu.hpp" + +#include +#include +#include +#include +#include + +#if defined(__ARM_NEON) +#include +#endif + +namespace strata::kernels::cpu { +namespace { + +std::atomic oracle_q8_0{false}; + +[[noreturn]] void unsupported(const char* what) { + std::fprintf(stderr, + "strata: %s is x86-only (AVX2/AVX-512) and this build has no portable version.\n" + " On this CPU use a native GGUF pack (IQ2_XS, IQ3_S, IQ4_NL, Q4_K, ...).\n", + what); + std::abort(); +} + +inline float bf16f(uint16_t h) { + const uint32_t b = (uint32_t) h << 16; + float f; + std::memcpy(&f, &b, 4); + return f; +} + +} // namespace + + + + + +void act_quant_q8_1(const float* x, int n, ActQ& a) { + // The Q8_0 oracle contract (expert.cpp's quantize_oracle_q8_0) is an AVX rounding. The legacy scalar loop + // below is what the engine uses; the oracle flag would silently quantize differently here. + if (oracle_q8_0.load(std::memory_order_relaxed)) unsupported("the Q8_0 oracle activation quantizer"); + // expert.cpp's scalar loop: round half away from zero, clamp to +-127. + a.nchunks = n / QKA; + for (int k = 0; k < a.nchunks; ++k) { + const float* xb = x + k * QKA; + float amax = 0.f; + for (int j = 0; j < QKA; ++j) amax = std::fmax(amax, std::fabs(xb[j])); + const float s = amax > 0.f ? amax / 127.f : 0.f; + const float inv = s > 0.f ? 1.f / s : 0.f; + int32_t sum = 0; + int8_t* q = a.q + k * QKA; + for (int j = 0; j < QKA; ++j) { + const float t = xb[j] * inv; + const float r = t + (t >= 0.f ? 0.5f : -0.5f); + int v = (int) r; + v = v < -127 ? -127 : (v > 127 ? 127 : v); + q[j] = (int8_t) v; + sum += v; + } + a.scale[k] = s; + a.sum[k] = sum; + a.hx[k] = s * (float) sum; + } +} + +void act_quant_q8_1_avx2(const float* x, int n, ActQ& a) { act_quant_q8_1(x, n, a); } + +void s2_expert_vnni(const uint8_t*, const float*, float*, ExpertScratch&) { unsupported("the Q2_0 expert kernel"); } +void s2_expert_vnni_q(const uint8_t*, const ActQ&, float*, ExpertScratch&) { unsupported("the Q2_0 expert kernel"); } +void s2_expert_gu_rows(const uint8_t*, const ActQ&, float*, int, int) { unsupported("the Q2_0 expert kernel"); } +void s2_expert_down_rows(const uint8_t*, const ActQ&, float*, int, int) { unsupported("the Q2_0 expert kernel"); } +void s2_expert_gu_rows_multi(const uint8_t*, const ActQ* const*, int, float* const*, int, int) { + unsupported("the Q2_0 expert kernel"); +} +void s2_expert_down_rows_multi(const uint8_t*, const ActQ* const*, int, float* const*, int, int) { + unsupported("the Q2_0 expert kernel"); +} +void s2_expert_vnni_multi(const uint8_t*, const ActQ* const*, int, float* const*, ExpertScratchMulti&) { + unsupported("the Q2_0 expert kernel"); +} +void q2_0_gguf_rows_multi(const uint8_t*, size_t, int, const ActQ* const*, int, float* const*, int, int) { + unsupported("the Q2_0 GGUF row kernel"); +} +void q2_0_gguf_rows_multi_avx2(const uint8_t*, size_t, int, const ActQ* const*, int, float* const*, int, int) { + unsupported("the Q2_0 GGUF row kernel"); +} +bool q2_bitplane_enabled() { return false; } +void q2_0_gguf_rows_multi_avx2_legacy(const uint8_t*, size_t, int, const ActQ* const*, int, float* const*, int, int) { + unsupported("the Q2_0 GGUF row kernel"); +} +void q2_0_gguf_rows_multi_avx2_v(bool, const uint8_t*, size_t, int, const ActQ* const*, int, float* const*, int, int) { + unsupported("the Q2_0 GGUF row kernel"); +} +void s2_expert_scalar(const uint8_t*, const float*, float*, bool) { unsupported("the Q2_0 scalar oracle"); } + +void iq512_gu_rows(int, const uint8_t*, size_t, size_t, int, const void* const*, int, float* const*, int, int) { + unsupported("the AVX-512 i-quant kernel"); +} +void iq512_rows(int, const uint8_t*, size_t, int, const void* const*, int, float* const*, int, int) { + unsupported("the AVX-512 i-quant kernel"); +} +bool iq512_gather_on() noexcept { return false; } +bool iq512x_supported(int) noexcept { return false; } +void iq512x_rows(int, const uint8_t*, size_t, int, const void* const*, int, float* const*, int, int) { + unsupported("the AVX-512 exact i-quant kernel"); +} +int iq256_variant() noexcept { return 0; } +int iq256_variants() noexcept { return 0; } +void iq256_gu_rows_v(int, int, const uint8_t*, size_t, size_t, int, const void* const*, int, float* const*, int, int) { + unsupported("the AVX-2 i-quant kernel"); +} +void iq256_rows_v(int, int, const uint8_t*, size_t, int, const void* const*, int, float* const*, int, int) { + unsupported("the AVX-2 i-quant kernel"); +} +void iq256_gu_rows(int, const uint8_t*, size_t, size_t, int, const void* const*, int, float* const*, int, int) { + unsupported("the AVX-2 i-quant kernel"); +} +void iq256_rows(int, const uint8_t*, size_t, int, const void* const*, int, float* const*, int, int) { + unsupported("the AVX-2 i-quant kernel"); +} +void iq4nl256_down_rows_v(int, const uint8_t*, size_t, int, const void* const*, int, float* const*, int, int) { + unsupported("the AVX-2 IQ4_NL kernel"); +} +void iq4nl256_down_rows(const uint8_t*, size_t, int, const void* const*, int, float* const*, int, int) { + unsupported("the AVX-2 IQ4_NL kernel"); +} +void q8k_quant_avx2(const float*, void*, int64_t) { unsupported("the AVX-2 Q8_K quantizer"); } +void kq256_gu_rows(int, const uint8_t*, size_t, size_t, int, const void* const*, int, float* const*, int, int) { + unsupported("the AVX-2 k-quant kernel"); +} +void kq256_rows(int, const uint8_t*, size_t, int, const void* const*, int, float* const*, int, int) { + unsupported("the AVX-2 k-quant kernel"); +} +void bf16_rows_dot_multi_avx1(const uint16_t*, int, int, const float*, int, float*) { + unsupported("the AVX1 router dot"); +} + +void bf16_rows_dot(const uint16_t* w, int rows, int cols, const float* x, float* out) { + bf16_rows_dot_multi(w, rows, cols, x, 1, out); +} + +void bf16_rows_dot_multi(const uint16_t* w, int rows, int cols, const float* x, int nt, float* out) { + for (int r = 0; r < rows; ++r) { + const uint16_t* wr = w + (size_t) r * (size_t) cols; + for (int t = 0; t < nt; ++t) { + const float* xt = x + (size_t) t * cols; + int c = 0; + float s = 0.f; +#if defined(__ARM_NEON) + float32x4_t acc0 = vdupq_n_f32(0.f), acc1 = vdupq_n_f32(0.f); + for (; c + 8 <= cols; c += 8) { + const uint16x8_t h = vld1q_u16(wr + c); + const float32x4_t w0 = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(h), 16)); + const float32x4_t w1 = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(h), 16)); + acc0 = vfmaq_f32(acc0, w0, vld1q_f32(xt + c)); + acc1 = vfmaq_f32(acc1, w1, vld1q_f32(xt + c + 4)); + } + s = vaddvq_f32(vaddq_f32(acc0, acc1)); +#endif + for (; c < cols; ++c) s += bf16f(wr[c]) * xt[c]; + out[(size_t) t * rows + r] = s; + } + } +} + + +// aarch64 has none of the x86 kernels: these names are referenced by shared code but reached only behind +// the x86 feature checks (all of them report unsupported here). +[[noreturn]] static void cpu_kernel_unreachable(const char* name) { + std::fprintf(stderr, "strata: %s is an x86 kernel and cannot run on this CPU\n", name); + std::abort(); +} +void act_quant_q8_1_avx1(const float* x, int n, ActQ& a) { cpu_kernel_unreachable("act_quant_q8_1_avx1"); } +void iq128_down_rows(int ggml_type, const uint8_t* w, size_t row_bytes, int n, const void* const* hq, int nt, float* const* out, int r0, int r1) { cpu_kernel_unreachable("iq128_down_rows"); } +bool iq128_down_supported(int ggml_type) noexcept { return false; } +void iq128_gu_rows(int ggml_type, const uint8_t* blob, size_t gu_row, size_t up_off, int n, const void* const* act, int nt, float* const* ff, int r0, int r1) { cpu_kernel_unreachable("iq128_gu_rows"); } +bool iq128_supported(int ggml_type) noexcept { return false; } +void iq256_gu_rows_exact_one(int ggml_type, const uint8_t* blob, size_t gu_row, size_t up_off, int n, const void* act, float* ff, int r0, int r1) { cpu_kernel_unreachable("iq256_gu_rows_exact_one"); } +bool nvfp4_256_fits(int n) { return false; } +void nvfp4_256_gu_rows(const uint8_t* blob, size_t gu_row, size_t up_off, int n, const void* const* act, int nt, float* const* ff, int r0, int r1, float s_gate, float s_up) { cpu_kernel_unreachable("nvfp4_256_gu_rows"); } +void nvfp4_256_rows(const uint8_t* w, size_t row_bytes, int n, const void* const* act, int nt, float* const* out, int r0, int r1, float scale) { cpu_kernel_unreachable("nvfp4_256_rows"); } +bool nvfp4_512_fits(int n) { return false; } +void nvfp4_512_gu_rows(const uint8_t* blob, size_t gu_row, size_t up_off, int n, const void* const* act, int nt, float* const* ff, int r0, int r1, float s_gate, float s_up) { cpu_kernel_unreachable("nvfp4_512_gu_rows"); } +void nvfp4_512_rows(const uint8_t* w, size_t row_bytes, int n, const void* const* act, int nt, float* const* out, int r0, int r1, float scale) { cpu_kernel_unreachable("nvfp4_512_rows"); } +void q2_0_gguf_rows_multi_avx1_legacy(const uint8_t* w, size_t row_bytes, int nblocks, const ActQ* const* a, int nt, float* const* out, int r0, int r1) { cpu_kernel_unreachable("q2_0_gguf_rows_multi_avx1_legacy"); } +void q2_0_gguf_rows_multi_avx1(const uint8_t* w, size_t row_bytes, int nblocks, const ActQ* const* a, int nt, float* const* out, int r0, int r1) { cpu_kernel_unreachable("q2_0_gguf_rows_multi_avx1"); } +void q8k_quant_avx1(const float* x, void* y, int64_t n) { cpu_kernel_unreachable("q8k_quant_avx1"); } +void tq_gu_rows(TqIsa isa, int ggml_type, const uint8_t* blob, size_t gu_row, size_t up_off, int n, const void* const* act, int nt, float* const* ff, int r0, int r1) { cpu_kernel_unreachable("tq_gu_rows"); } +TqIsa tq_isa() { return TqIsa::scalar; } +int tq_kbits(int ggml_type) noexcept { return 0; } +bool tq_kernels_enabled() { return false; } +void tq_rows(TqIsa isa, int ggml_type, const uint8_t* w, size_t row_bytes, int n, const void* const* act, int nt, float* const* out, int r0, int r1) { cpu_kernel_unreachable("tq_rows"); } + +bool tq_supported(int) noexcept { return false; } + +} // namespace strata::kernels::cpu diff --git a/src/program/generate.cpp b/src/program/generate.cpp index 75d404ac..63ae3e34 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -3447,6 +3447,7 @@ int main(int argc, char** argv) { // compiled `/arch:AVX512`, so on a CPU without those features it does not fail - it executes an illegal // instruction at some unpredictable token. Refusing at second zero is the whole point of P2.S3's check. strata::kernels::cpu::expert_set_oracle_q8_0(o.cpu_oracle_q8_0); +#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86) || defined(STRATA_GPU_METAL) // ... and nothing runs on a CPU without AVX2: every CPU expert kernel is AVX2 at least (the AVX-512 ones are // chosen above it), and so is ggml-cpu in the release build, which the native pack's layout load initializes // next. Refused here, by name, rather than an illegal instruction in the first expert. @@ -3484,6 +3485,14 @@ int main(int argc, char** argv) { "and the normal build is faster on it\n", isa_floor); } +#else + if (o.cpu_oracle_q8_0) { + std::fprintf(stderr, "strata generate: --cpu-oracle-q8-0 is the x86 Q8_0 activation contract; this CPU uses the " + "scalar quantizer\n"); + return 2; + } +#endif + std::string err; if (!o.native_head_gguf.empty() && !o.stream_token) { std::fprintf(stderr, "--native-head-gguf requires --stream-token\n"); diff --git a/tests/terminal/test_setup_arm64.py b/tests/terminal/test_setup_arm64.py new file mode 100644 index 00000000..7f51056a --- /dev/null +++ b/tests/terminal/test_setup_arm64.py @@ -0,0 +1,68 @@ +"""setup.py on an ARM64 machine (DGX Spark / GB10, #3622): the GPU line, the CPU probe and the older-CPU floor. +x86 machines must keep their old answers: a card whose memory nvidia-smi does not report stays unlisted there.""" + +from __future__ import annotations + +import os +import sys +import unittest +from pathlib import Path +from unittest import mock + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT)) + +import setup # noqa: E402 + +GB10 = "0, NVIDIA GB10, [N/A], 12.1, 580.95.05\n" +MEMINFO = "MemTotal: 130000000 kB\nMemFree: 1000000 kB\n" +ARM_CPUINFO = ("processor\t: 0\nBogoMIPS\t: 2000.00\nFeatures\t: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics\n" + "CPU implementer\t: 0x41\n") + + +class UnifiedMemoryCard(unittest.TestCase): + def test_a_card_without_a_memory_figure_is_unlisted_on_x86_as_before(self): + with mock.patch.object(setup, "out", return_value=GB10), mock.patch.object(setup, "is_arm64", return_value=False): + self.assertEqual(setup.gpus(), []) + + def test_gb10_takes_the_pool_size_from_the_machine_ram(self): + with mock.patch.object(setup, "out", return_value=GB10), mock.patch.object(setup, "is_arm64", return_value=True), \ + mock.patch("builtins.open", mock.mock_open(read_data=MEMINFO)): + found = setup.gpus() + self.assertEqual(len(found), 1) + self.assertEqual(found[0]["name"], "NVIDIA GB10") + self.assertEqual(found[0]["arch"], "121") + self.assertAlmostEqual(found[0]["vram_gb"], 130000000 / 1024.0 / 1024.0, places=3) + + def test_a_normal_card_is_unchanged(self): + with mock.patch.object(setup, "out", return_value="0, NVIDIA GeForce RTX 5070, 12227, 12.0, 580.1\n"): + found = setup.gpus() + self.assertAlmostEqual(found[0]["vram_gb"], 12227 / 1024.0, places=6) + + +class ArmCpu(unittest.TestCase): + def test_an_arm_cpuinfo_without_the_x86_flags_line_does_not_crash(self): + with mock.patch.object(setup, "WIN", False), mock.patch("builtins.open", mock.mock_open(read_data=ARM_CPUINFO)): + name, avx2, avx512 = setup.cpu_info() + self.assertFalse(avx2) + self.assertFalse(avx512) + + def test_an_x86_cpuinfo_is_read_as_before(self): + text = "flags\t\t: fpu avx avx2 avx512f avx512bw avx512vl avx512_vnni avx512vbmi\nmodel name\t: Test CPU\n" + with mock.patch.object(setup, "WIN", False), mock.patch("builtins.open", mock.mock_open(read_data=text)): + name, avx2, avx512 = setup.cpu_info() + self.assertEqual((name, avx2, avx512), ("Test CPU", True, True)) + + def test_no_older_cpu_floor_on_arm(self): + with mock.patch.dict(os.environ, {"STRATA_ISA_FLOOR": ""}), mock.patch.object(setup, "is_arm64", return_value=True): + self.assertEqual(setup.cpu_floor(False), "") + + def test_the_floor_on_x86_without_avx2_is_still_asked_for(self): + with mock.patch.dict(os.environ, {"STRATA_ISA_FLOOR": ""}), mock.patch.object(setup, "is_arm64", return_value=False), \ + mock.patch.object(setup, "WIN", False), \ + mock.patch("builtins.open", mock.mock_open(read_data="flags\t: fpu sse4_2 popcnt avx\n")): + self.assertEqual(setup.cpu_floor(False), "avx") + + +if __name__ == "__main__": + unittest.main() From 38b06316efff96ea3b6c1d53a1cd23175148f5cb Mon Sep 17 00:00:00 2001 From: Kevin Read <1876063+unverbraucht@users.noreply.github.com> Date: Tue, 27 Oct 2026 01:57:57 +0100 Subject: [PATCH 2/2] Gyro-L: read the pair-trellis types TQP5/TQP6/TQP7 (ggml 147-149) on the CPU, CUDA and the MMQ prompt path (#3806) Gyro-L (agentionai's 50 GiB file) stores its experts as a pair trellis: 128 values per block, a fp16 scale and 64 steps of K bits (K = 5, 6, 7), two weights per step taken from the first pair of the TQ2_T codebook entry the step's state selects. The engine refused the file because it did not know the three type ids. This adds them next to the existing trellis types 144-146 without changing those. What changed: - gguf_reader.hpp / tools/gguf_reader.py: names and block sizes (42 / 50 / 58 bytes). - tq_cpu.*, tq_states.hpp, tq_avx2.cpp: tq_supported / tq_kbits / tq_block_bytes for the new ids (the bytes per block are 2 + 8K for the pair layout, 2 + 4K for the others), a scalar pair kernel (the reference) and an AVX2 pair kernel that is bit for bit equal to it; AVX-512 uses the AVX2 pair kernel. - iq_kernels.cu: TqInfo, the lane decode (tqp_states8, a byte window), the dot product, the decode-once Split, the fp16 dequantizer and the format lists for 147-149. They are reached only by a pack that holds these types. - moe_mmq.cu, strata_ggml_tq_overlay.cmake, third_party/ggml-tq/tq.cuh, ggml-common-tq.h: the MMQ prompt path for the three types (overlay enumerators, traits, tile loader branch, instances). - native_expert_parity: names tqp5-7 and skips its synthetic check, with a message, on a ggml that has no trellis traits. - docs/GYRO.md: the "not supported" sentence now says how Gyro-L is read. Changes to the contributor's patch: - tq_cpu.cpp called the AVX2 pair kernel without the x86 guard the rest of the file has, which breaks the Metal build (it compiles tq_cpu.cpp without tq_avx2.cpp): the pair rows go through one helper that falls back to the scalar kernel. - moe_mmq.cu named GGML_TYPE_TQP5-7 under STRATA_MMQ_TQ, but the public agentionai llama.cpp has the older trellis types and not these: a new STRATA_MMQ_TQP (set for the overlay and for a tree that has the tqp instances) guards them, so that tree still builds. The HIP overlay lists get the three instances too, and ggml_tq_overlay_test checks them (and sets the CMake policy its IN_LIST needs). - native_expert_parity: the skip test compared a type-name pointer to null, but ggml_type_name of an id past the table reads past it (the test crashed on an upstream ggml); it compares against GGML_TYPE_COUNT. - docs/GYRO.md keeps the plain description and leaves out the contributor's measured speeds (we have not repeated them). Tests added: src/kernels/cpu/tqp_cpu_test.cpp (no ggml, no model: the scalar kernel against a bit-by-bit reference from the format definition, every ISA equal to scalar for 1-8 tokens and for gate/up rows, the type predicates for ids 0-299 against the old formulas), the pair types in tq_pascal_bench (fp16 dequantizer bit for bit against the definition, decode kernels within 3e-2, decode-once == per-column), tools/test_gguf_tqp.py, and native_expert_parity tqp5/6/7 registered when the ggml tree has the tqp instances. Checked (no model): tqp_cpu_test passes on x86 (AVX2 on z440; AVX2 and AVX-512 on the owner PC with MSVC); tq_pascal_bench on an RTX 5070 (sm_120, MSVC build): 0 failures, the TQP dequantizer matches the definition in all bits, and the checksum lines of TQ2_T/TQK6/TQK7 are byte-identical to the build without this change. Builds: CUDA sm_86 (gcc), sm_60 (CUDA 12.9), sm_120 (MSVC), SYCL, and the CPU files for aarch64. Not covered: the contributor's own parity runs need a ggml tree that carries the types (not public here), and the 50 GiB model run (load check, coherent answer, 32K prompt, MTP, Gyro-S/M identity) is still to be done. Metal, SYCL and HIP have no kernels for the types; HIP was read, not built. --- CMakeLists.txt | 33 +++- cmake/ggml_tq_overlay_test.cmake | 13 +- cmake/strata_ggml_tq_overlay.cmake | 17 +- docs/GYRO.md | 6 +- include/strata/artifact/gguf_reader.hpp | 19 ++ include/strata/kernels/cpu/tq_cpu.hpp | 14 +- src/kernels/cpu/tq_avx2.cpp | 71 ++++++++ src/kernels/cpu/tq_cpu.cpp | 82 ++++++++- src/kernels/cpu/tq_states.hpp | 20 +++ src/kernels/cpu/tqp_cpu_test.cpp | 229 ++++++++++++++++++++++++ src/kernels/cuda/iq_kernels.cu | 138 +++++++++++++- src/kernels/native_expert_parity.cpp | 27 +++ src/kernels/tq_pascal_bench.cpp | 141 ++++++++++++++- src/prefill/moe_mmq.cu | 18 +- third_party/ggml-tq/tq.cuh | 74 +++++++- third_party/ggml/ggml-common-tq.h | 20 +++ tools/gguf_reader.py | 4 + tools/test_gguf_tqp.py | 34 ++++ 18 files changed, 921 insertions(+), 39 deletions(-) create mode 100644 src/kernels/cpu/tqp_cpu_test.cpp create mode 100644 tools/test_gguf_tqp.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 57c6b411..6fdd7d94 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1149,7 +1149,7 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) # the Gyro trellis decode kernels (TQ2_T / TQK6 / TQK7) on synthetic blocks: decode-once vs per-column bitwise, output # checksums to compare two builds byte for byte, --bench for timings (GPU, no model) add_executable(tq_pascal_bench src/kernels/tq_pascal_bench.cpp) - target_include_directories(tq_pascal_bench PRIVATE ${_strata_gpu_include_directories}) + target_include_directories(tq_pascal_bench PRIVATE ${_strata_gpu_include_directories} ${CMAKE_CURRENT_SOURCE_DIR}/third_party/ggml) target_link_libraries(tq_pascal_bench PRIVATE strata_kernels ${_strata_gpu_runtime_target}) add_test(NAME tq_pascal_bench COMMAND tq_pascal_bench) set_tests_properties(iq_multi_parity PROPERTIES ENVIRONMENT "STRATA_KQ_SPLIT=1") # the Q4_K / Q5_K / Q8_0 Split path against the old kernels @@ -1881,6 +1881,11 @@ if(STRATA_NATIVE_EXPERTS) # tree leaves them to the FP16 prompt path if(STRATA_GGML_DIR AND EXISTS "${STRATA_GGML_DIR}/ggml/src/ggml-cuda/template-instances/mmq-instance-tqk6.cu") list(APPEND _strata_mmq_types tq2_t tqk6 tqk7) + # Gyro-L's pair trellis (types 147-149): the same fork's ggml has their MMQ instances + if(EXISTS "${STRATA_GGML_DIR}/ggml/src/ggml-cuda/template-instances/mmq-instance-tqp5.cu") + list(APPEND _strata_mmq_types tqp5 tqp6 tqp7) + set(_strata_mmq_tqp ON) # moe_mmq.cu names GGML_TYPE_TQP5-7 only where the ggml has them (STRATA_MMQ_TQP) + endif() set(_strata_mmq_tq ON) elseif(STRATA_MMQ_TQ_OVERLAY AND (STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)) # an upstream llama.cpp: the trellis MMQ instances come from an overlay of the pinned ggml-cuda headers written at @@ -1897,6 +1902,7 @@ if(STRATA_NATIVE_EXPERTS) configure_file(third_party/ggml/ggml-common-tq.h "${_strata_tq_overlay_dir}/ggml-common-tq.h" COPYONLY) configure_file(src/kernels/cuda/tq_lut_i8.cuh "${_strata_tq_overlay_dir}/tq_lut_i8.cuh" COPYONLY) set(_strata_mmq_tq ON) + set(_strata_mmq_tqp ON) # the overlay adds the pair-trellis enumerators and instances set(_strata_mmq_tq_overlay ON) message(STATUS "trellis MMQ (Gyro prompts): overlay on the pinned ggml-cuda in ${_strata_tq_overlay_dir}") if(STRATA_BUILD_TESTS) @@ -1946,7 +1952,7 @@ if(STRATA_NATIVE_EXPERTS) target_sources(strata_mmq PRIVATE ${_strata_mmq}/template-instances/mmq-instance-${_t}.cu) endforeach() if(_strata_mmq_tq_overlay) - foreach(_t tq2_t tqk6 tqk7) + foreach(_t tq2_t tqk6 tqk7 tqp5 tqp6 tqp7) target_sources(strata_mmq PRIVATE ${_strata_tq_overlay_dir}/template-instances/mmq-instance-${_t}.cu) endforeach() # first on the include path: the patched mmq.cuh / ggml.h shadow the pinned ones for this target only @@ -2064,6 +2070,9 @@ if(STRATA_NATIVE_EXPERTS) if(_strata_mmq_tq) target_compile_definitions(strata_mmq PRIVATE STRATA_MMQ_TQ=1) endif() + if(_strata_mmq_tqp) + target_compile_definitions(strata_mmq PRIVATE STRATA_MMQ_TQP=1) + endif() target_link_libraries(strata_prefill PUBLIC strata_mmq) target_compile_definitions(strata_prefill PUBLIC STRATA_PREFILL_MMQ=1 STRATA_PREFILL_FUSED=1) if(STRATA_BUILD_TESTS) @@ -2141,7 +2150,7 @@ if(STRATA_NATIVE_EXPERTS) list(APPEND _strata_mmq_sources "${_strata_mmq}/template-instances/mmq-instance-${_t}.cu") endforeach() if(_strata_mmq_tq_overlay) - foreach(_t tq2_t tqk6 tqk7) + foreach(_t tq2_t tqk6 tqk7 tqp5 tqp6 tqp7) list(APPEND _strata_mmq_sources "${_strata_tq_overlay_dir}/template-instances/mmq-instance-${_t}.cu") endforeach() endif() @@ -2166,6 +2175,9 @@ if(STRATA_NATIVE_EXPERTS) if(_strata_mmq_tq) target_compile_definitions(strata_mmq PRIVATE STRATA_MMQ_TQ=1) endif() + if(_strata_mmq_tqp) + target_compile_definitions(strata_mmq PRIVATE STRATA_MMQ_TQP=1) + endif() target_link_libraries(strata_mmq PUBLIC ggml-base hip::host roc::hipblas) target_link_libraries(strata_prefill PUBLIC strata_mmq) target_compile_definitions(strata_prefill PUBLIC STRATA_PREFILL_MMQ=1 STRATA_PREFILL_FUSED=1) @@ -2257,7 +2269,11 @@ if(STRATA_NATIVE_EXPERTS) # APR: the trellis expert types and the Hadamard activation rotation, when ggml comes from the agention fork # (the only tree with these types; marked by its ggml-cuda/tq.cuh) if(STRATA_GGML_DIR AND EXISTS "${STRATA_GGML_DIR}/ggml/src/ggml-cuda/tq.cuh") - foreach(_pair tqk6/tqk7 tq2_t/tq2_t) + set(_strata_tq_pairs tqk6/tqk7 tq2_t/tq2_t) + if(EXISTS "${STRATA_GGML_DIR}/ggml/src/ggml-cuda/template-instances/mmq-instance-tqp5.cu") # Gyro-L's pair trellis + list(APPEND _strata_tq_pairs tqp5/tqp5 tqp6/tqp6 tqp7/tqp7) + endif() + foreach(_pair ${_strata_tq_pairs}) string(REPLACE "/" "_" _name "${_pair}") add_test(NAME native_expert_parity_${_name} COMMAND native_expert_parity --synthetic ${_pair}) add_test(NAME native_expert_parity_hadamard_${_name} COMMAND native_expert_parity --synthetic-hadamard ${_pair}) @@ -2331,6 +2347,13 @@ if(STRATA_BUILD_TESTS) add_test(NAME tq_cpu_parity COMMAND tq_cpu_parity) add_test(NAME tq_cpu_parity_avx2 COMMAND ${CMAKE_COMMAND} -E env STRATA_FORCE_AVX2=1 $) set_tests_properties(tq_cpu_parity tq_cpu_parity_avx2 PROPERTIES TIMEOUT 300 SKIP_RETURN_CODE 77) + # Gyro-L's pair trellis (types 147-149) and the older trellis rows on synthetic blocks, no ggml and no model: the scalar + # kernel against a bit-by-bit reference, every ISA bit for bit equal to scalar, the type predicates unchanged + add_executable(tqp_cpu_test src/kernels/cpu/tqp_cpu_test.cpp) + target_link_libraries(tqp_cpu_test PRIVATE strata_kernels_cpu) + add_test(NAME tqp_cpu_test COMMAND tqp_cpu_test) + add_test(NAME tqp_cpu_test_avx2 COMMAND ${CMAKE_COMMAND} -E env STRATA_FORCE_AVX2=1 $) + set_tests_properties(tqp_cpu_test tqp_cpu_test_avx2 PROPERTIES TIMEOUT 120) endif() add_executable(expert_parity src/kernels/cpu/expert_parity.cpp) @@ -2466,7 +2489,7 @@ if(STRATA_HIP_GFX906) # f16kq::covers() answers false and the int8 kernels run, as on every other card) if(_strata_mmq_tq_overlay) set(_strata_tq_hip ${_strata_tq_overlay_dir}/quantize.cu) - foreach(_t tq2_t tqk6 tqk7) + foreach(_t tq2_t tqk6 tqk7 tqp5 tqp6 tqp7) list(APPEND _strata_tq_hip ${_strata_tq_overlay_dir}/template-instances/mmq-instance-${_t}.cu) endforeach() list(APPEND _strata_cu ${_strata_tq_hip}) diff --git a/cmake/ggml_tq_overlay_test.cmake b/cmake/ggml_tq_overlay_test.cmake index e56604c8..02f8eb16 100644 --- a/cmake/ggml_tq_overlay_test.cmake +++ b/cmake/ggml_tq_overlay_test.cmake @@ -1,6 +1,7 @@ # cmake -DLLAMA= -DOUT= -P cmake/ggml_tq_overlay_test.cmake # The trellis MMQ overlay on a pinned llama.cpp: it generates, every copied header is the pinned one except the patched # ones, and the patched ones only gain lines (TQ routing, three externs, config copies of IQ2_XS's). +cmake_policy(SET CMP0057 NEW) # script mode starts with the old policies: IN_LIST below needs this one get_filename_component(CMAKE_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/.." ABSOLUTE) include(${CMAKE_CURRENT_LIST_DIR}/strata_ggml_tq_overlay.cmake) file(REMOVE_RECURSE "${OUT}") @@ -24,7 +25,8 @@ foreach(_f IN LISTS _cuh) endforeach() file(READ "${OUT}/mmq.cuh" _m) foreach(_pat "case GGML_TYPE_TQ2_T: return MMQ_DP4A_TXS_Q8_0_16;" "ggml_cuda_mmq_load_tiles_tq,\n ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a" - "extern DECL_MMQ_CASE(GGML_TYPE_TQK7);") + "extern DECL_MMQ_CASE(GGML_TYPE_TQK7);" "extern DECL_MMQ_CASE(GGML_TYPE_TQP5);" "extern DECL_MMQ_CASE(GGML_TYPE_TQP7);" + "case GGML_TYPE_TQP6: return MMQ_DP4A_TXS_Q8_0_16;") string(FIND "${_m}" "${_pat}" _p) if(_p LESS 0) message(FATAL_ERROR "mmq.cuh lacks: ${_pat}") @@ -35,6 +37,15 @@ string(FIND "${_g}" "GGML_TYPE_TQK7 = 146," _p) if(_p LESS 0) message(FATAL_ERROR "ggml.h lacks the trellis enumerators") endif() +string(FIND "${_g}" "GGML_TYPE_TQP7 = 149," _p) # Gyro-L's pair trellis +if(_p LESS 0) + message(FATAL_ERROR "ggml.h lacks the pair-trellis enumerators") +endif() +foreach(_t tqp5 tqp6 tqp7) + if(NOT EXISTS "${OUT}/template-instances/mmq-instance-${_t}.cu") + message(FATAL_ERROR "the overlay has no MMQ instance for ${_t}") + endif() +endforeach() foreach(_c IN ITEMS ampere rdna4) file(READ "${CU}/mmq-config-${_c}.cuh" _o) file(READ "${OUT}/mmq-config-${_c}.cuh" _n2) diff --git a/cmake/strata_ggml_tq_overlay.cmake b/cmake/strata_ggml_tq_overlay.cmake index df24af25..a29ae4be 100644 --- a/cmake/strata_ggml_tq_overlay.cmake +++ b/cmake/strata_ggml_tq_overlay.cmake @@ -79,6 +79,9 @@ function(strata_ggml_tq_overlay LLAMA OUT) " GGML_TYPE_TQ2_T = 144, // Agention trellis types (Strata overlay; the fork's ids) GGML_TYPE_TQK6 = 145, GGML_TYPE_TQK7 = 146, + GGML_TYPE_TQP5 = 147, // Gyro-L pair trellis (the same fork's 147-149) + GGML_TYPE_TQP6 = 148, + GGML_TYPE_TQP7 = 149, ") file(WRITE "${OUT}/include/ggml.h" "${_h}") @@ -99,15 +102,15 @@ function(strata_ggml_tq_overlay LLAMA OUT) set(_tq_fail "the MMQ helper signatures of this ggml-cuda are not the ones the overlay knows (ggml_cuda_mmq_get_nthreads)") endif() _tq_edit(_m AFTER " case GGML_TYPE_IQ2_XXS:\n case GGML_TYPE_IQ2_XS:\n" - " case GGML_TYPE_TQ2_T:\n case GGML_TYPE_TQK6:\n case GGML_TYPE_TQK7:\n") + " case GGML_TYPE_TQ2_T:\n case GGML_TYPE_TQK6:\n case GGML_TYPE_TQK7:\n case GGML_TYPE_TQP5:\n case GGML_TYPE_TQP6:\n case GGML_TYPE_TQP7:\n") _tq_edit(_m AFTER " case GGML_TYPE_IQ2_XS: return MMQ_DP4A_TXS_Q8_0_16;\n" - " case GGML_TYPE_TQ2_T: return MMQ_DP4A_TXS_Q8_0_16;\n case GGML_TYPE_TQK6: return MMQ_DP4A_TXS_Q8_0_16;\n case GGML_TYPE_TQK7: return MMQ_DP4A_TXS_Q8_0_16;\n") + " case GGML_TYPE_TQ2_T: return MMQ_DP4A_TXS_Q8_0_16;\n case GGML_TYPE_TQK6: return MMQ_DP4A_TXS_Q8_0_16;\n case GGML_TYPE_TQK7: return MMQ_DP4A_TXS_Q8_0_16;\n case GGML_TYPE_TQP5: return MMQ_DP4A_TXS_Q8_0_16;\n case GGML_TYPE_TQP6: return MMQ_DP4A_TXS_Q8_0_16;\n case GGML_TYPE_TQP7: return MMQ_DP4A_TXS_Q8_0_16;\n") _tq_edit(_m BEFORE " case GGML_TYPE_IQ3_XXS:\n return ggml_cuda_mmq_util_funcs(\n VDR_IQ3_XXS_Q8_1_MMQ," - " case GGML_TYPE_TQ2_T:\n case GGML_TYPE_TQK6:\n case GGML_TYPE_TQK7:\n return ggml_cuda_mmq_util_funcs(\n VDR_IQ2_XS_Q8_1_MMQ,\n ggml_cuda_mmq_load_tiles_tq,\n ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a,\n ggml_cuda_mmq_write_back_dp4a);\n") + " case GGML_TYPE_TQ2_T:\n case GGML_TYPE_TQK6:\n case GGML_TYPE_TQK7:\n case GGML_TYPE_TQP5:\n case GGML_TYPE_TQP6:\n case GGML_TYPE_TQP7:\n return ggml_cuda_mmq_util_funcs(\n VDR_IQ2_XS_Q8_1_MMQ,\n ggml_cuda_mmq_load_tiles_tq,\n ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a,\n ggml_cuda_mmq_write_back_dp4a);\n") _tq_edit(_m BEFORE " case GGML_TYPE_IQ3_XXS:\n return ggml_cuda_mmq_util_funcs(\n -1,\n ggml_cuda_mmq_load_tiles_iq3_xxs," - " case GGML_TYPE_TQ2_T:\n case GGML_TYPE_TQK6:\n case GGML_TYPE_TQK7:\n return ggml_cuda_mmq_util_funcs(\n -1,\n ggml_cuda_mmq_load_tiles_tq,\n ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma,\n ggml_cuda_mmq_write_back_mma);\n") + " case GGML_TYPE_TQ2_T:\n case GGML_TYPE_TQK6:\n case GGML_TYPE_TQK7:\n case GGML_TYPE_TQP5:\n case GGML_TYPE_TQP6:\n case GGML_TYPE_TQP7:\n return ggml_cuda_mmq_util_funcs(\n -1,\n ggml_cuda_mmq_load_tiles_tq,\n ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma,\n ggml_cuda_mmq_write_back_mma);\n") _tq_edit(_m AFTER "extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS);\n" - "extern DECL_MMQ_CASE(GGML_TYPE_TQ2_T);\nextern DECL_MMQ_CASE(GGML_TYPE_TQK6);\nextern DECL_MMQ_CASE(GGML_TYPE_TQK7);\n") + "extern DECL_MMQ_CASE(GGML_TYPE_TQ2_T);\nextern DECL_MMQ_CASE(GGML_TYPE_TQK6);\nextern DECL_MMQ_CASE(GGML_TYPE_TQK7);\nextern DECL_MMQ_CASE(GGML_TYPE_TQP5);\nextern DECL_MMQ_CASE(GGML_TYPE_TQP6);\nextern DECL_MMQ_CASE(GGML_TYPE_TQP7);\n") file(WRITE "${OUT}/mmq.cuh" "${_m}") # ---- mmq-load-tiles.cuh: the loader, at the end ------------------------------------------------------------------- @@ -131,7 +134,7 @@ function(strata_ggml_tq_overlay LLAMA OUT) list(REMOVE_DUPLICATES _lines) foreach(_ln IN LISTS _lines) set(_add "") - foreach(_ty TQ2_T TQK6 TQK7) + foreach(_ty TQ2_T TQK6 TQK7 TQP5 TQP6 TQP7) string(REPLACE "GGML_TYPE_IQ2_XS," "GGML_TYPE_${_ty}," _x "${_ln}") string(APPEND _add "${_x}") endforeach() @@ -142,7 +145,7 @@ function(strata_ggml_tq_overlay LLAMA OUT) endforeach() # ---- template instances ------------------------------------------------------------------------------------------- - foreach(_ty tq2_t tqk6 tqk7) + foreach(_ty tq2_t tqk6 tqk7 tqp5 tqp6 tqp7) string(TOUPPER "${_ty}" _TY) file(WRITE "${OUT}/template-instances/mmq-instance-${_ty}.cu" "// Strata trellis overlay (cmake/strata_ggml_tq_overlay.cmake)\n\n#include \"../mmq.cuh\"\n\nDECL_MMQ_CASE(GGML_TYPE_${_TY});\n") diff --git a/docs/GYRO.md b/docs/GYRO.md index 6c390e1b..6ab2bd22 100644 --- a/docs/GYRO.md +++ b/docs/GYRO.md @@ -39,8 +39,10 @@ ready-made engine is too old for the pack, setup replaces it or the app offers U About half of each file is the n-gram (PLE) table. It stays on the SSD and is read row by row for every token, so keep the file on a fast local drive. The memory figure is what the weights take once loaded. -Gyro-L (agentionai's 50 GiB file) is not supported. It uses another trellis family (types 147 to 149) that the engine -does not have. +Gyro-L (agentionai's 50 GiB file) uses another trellis family, the pair trellis (types 147 to 149, TQP5 to TQP7). The +engine reads it with the same pack step and expert kernels as the other Gyro files, on the CPU and on CUDA cards. The +prompt path uses the trellis MMQ build where the build has it; Metal and SYCL have no kernels for it, so they use the CPU +rows and the FP16 prompt path. We have not run the 50 GiB file ourselves yet. **Pictures.** Gyro reads pictures with the original model's picture reader file, the same one every Qwen3.8-Flash-Next model uses. Setup downloads it with the model. The reader opens only the vocabulary and the header of diff --git a/include/strata/artifact/gguf_reader.hpp b/include/strata/artifact/gguf_reader.hpp index 5b67f54b..5800fe94 100644 --- a/include/strata/artifact/gguf_reader.hpp +++ b/include/strata/artifact/gguf_reader.hpp @@ -124,6 +124,12 @@ inline const char* ggml_type_name(uint32_t t) { return "TQK6"; case 146: return "TQK7"; + case 147: + return "TQP5"; + case 148: + return "TQP6"; + case 149: + return "TQP7"; default: return "?"; } @@ -252,6 +258,19 @@ inline bool block_geometry(uint32_t t, int& elems, int& bytes) { elems = 128; bytes = 30; return true; + // Gyro-L pair trellis: 128 values = fp16 scale + 64 steps of K bits, 2 weights per step (8*K bytes) + case 147: + elems = 128; + bytes = 42; + return true; + case 148: + elems = 128; + bytes = 50; + return true; + case 149: + elems = 128; + bytes = 58; + return true; default: return false; } diff --git a/include/strata/kernels/cpu/tq_cpu.hpp b/include/strata/kernels/cpu/tq_cpu.hpp index 7555fccf..3eaaf533 100644 --- a/include/strata/kernels/cpu/tq_cpu.hpp +++ b/include/strata/kernels/cpu/tq_cpu.hpp @@ -44,7 +44,12 @@ bool tq_supported(int ggml_type) noexcept; int tq_kbits(int ggml_type) noexcept; /// Values per block of the trellis types, and bytes per block (2 + 4 * bits: 34, 26, 30, 22, 18). constexpr int kTqBlock = 128; -inline size_t tq_block_bytes(int ggml_type) noexcept { return (size_t) (2 + 4 * tq_kbits(ggml_type)); } +// TQ2_T / TQK6 / TQK7: 4 weights per step (2 + 4K bytes per 128-value block). TQP5/6/7 (147-149, Gyro-L) are the +// pair layout: 2 weights per step, 8K bytes per block (2 + 8K). +inline bool tq_pair_layout(int ggml_type) noexcept { return ggml_type >= 147 && ggml_type <= 149; } +inline size_t tq_block_bytes(int ggml_type) noexcept { + return (size_t) (2 + (tq_pair_layout(ggml_type) ? 8 : 4) * tq_kbits(ggml_type)); +} /// The best ISA this CPU runs (capped by STRATA_TQ_ISA). TqIsa tq_isa(); /// Whether this CPU can run `isa` at all (scalar: always). @@ -88,6 +93,13 @@ void tq_rows_scalar(int kb, const int16_t* tab, const uint8_t* w, size_t row_byt float* const* out, int r0, int r1); void tq_rows_avx2(int kb, const int16_t* tab, const uint8_t* w, size_t row_bytes, int n, const TqAct* const* a, int nt, float* const* out, int r0, int r1); +/// The TQP pair layout (types 147-149) on AVX2. AVX-512 also takes this kernel: the pair layout has no +/// AVX-512 form yet (no AVX-512 machine has verified one), and the AVX2 path is exact. +void tq_rows_avx2_tqp(int kb, const uint8_t* w, size_t row_bytes, int n, const TqAct* const* a, int nt, + float* const* out, int r0, int r1); +/// The scalar reference of the pair layout (any CPU). +void tq_rows_scalar_tqp(int kb, const uint8_t* w, size_t row_bytes, int n, const TqAct* const* a, int nt, + float* const* out, int r0, int r1); void tq_rows_avx512(int kb, const int16_t* tab, const uint8_t* w, size_t row_bytes, int n, const TqAct* const* a, int nt, float* const* out, int r0, int r1); } // namespace tq_detail diff --git a/src/kernels/cpu/tq_avx2.cpp b/src/kernels/cpu/tq_avx2.cpp index 9a43393f..dbc6b5aa 100644 --- a/src/kernels/cpu/tq_avx2.cpp +++ b/src/kernels/cpu/tq_avx2.cpp @@ -75,8 +75,79 @@ void rows_t(const int16_t* table, const uint8_t* w, size_t row_bytes, int n, con } } +// TQP (147-149): 64 pair-steps per 128-value block. A step's two weights are the FIRST pair of its state +// table row, so 8 states make one YMM of 16 int16 weights, and a 32-value group (16 steps) is two YMM of +// weights against two YMM of activations: VPMADDWD, one exact int32 sum per group, then the four group sums +// go into the same per-token __m128 the other kernels use, so the float part is the scalar kernel's. +inline __m256i pairs8(const int16_t* tab, uint32_t s0, uint32_t s1, uint32_t s2, uint32_t s3, uint32_t s4, + uint32_t s5, uint32_t s6, uint32_t s7) { + const __m128i a0 = _mm_unpacklo_epi32(_mm_loadl_epi64((const __m128i*) (tab + 4 * s0)), + _mm_loadl_epi64((const __m128i*) (tab + 4 * s1))); + const __m128i a1 = _mm_unpacklo_epi32(_mm_loadl_epi64((const __m128i*) (tab + 4 * s2)), + _mm_loadl_epi64((const __m128i*) (tab + 4 * s3))); + const __m128i a2 = _mm_unpacklo_epi32(_mm_loadl_epi64((const __m128i*) (tab + 4 * s4)), + _mm_loadl_epi64((const __m128i*) (tab + 4 * s5))); + const __m128i a3 = _mm_unpacklo_epi32(_mm_loadl_epi64((const __m128i*) (tab + 4 * s6)), + _mm_loadl_epi64((const __m128i*) (tab + 4 * s7))); + return _mm256_set_m128i(_mm_unpacklo_epi64(a2, a3), _mm_unpacklo_epi64(a0, a1)); +} + +template +void rows_tqp_t(const uint8_t* w, size_t row_bytes, int n, const TqAct* const* a, int nt, float* const* out, int r0, + int r1) { + const int nb = n / 128; + constexpr size_t bsz = 2 + 8 * KB; + const int16_t* tab = tq_state_table(); + for (int r = r0; r < r1; ++r) { + const uint8_t* x = w + (size_t) r * row_bytes; + __m128 acc[kTqMaxTokens]; + for (int t = 0; t < nt; ++t) acc[t] = _mm_setzero_ps(); + for (int i = 0; i < nb; ++i) { + const uint8_t* blk = x + (size_t) i * bsz; + uint32_t st[64]; + tq_states_tqp(blk + 2, st); + __m256i wv[8]; // wv[2g], wv[2g+1]: the weights of the 32-value group g +#pragma GCC unroll 4 + for (int g = 0; g < 4; ++g) { + const uint32_t* s = st + 16 * g; + wv[2 * g] = pairs8(tab, s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7]); + wv[2 * g + 1] = pairs8(tab, s[8], s[9], s[10], s[11], s[12], s[13], s[14], s[15]); + } + const __m128 dx = _mm_set1_ps(tq_half_to_float(tq_ld16(blk))); + for (int t = 0; t < nt; ++t) { + const __m256i* q = (const __m256i*) (a[t]->q + 128 * i); + int32_t sum[4]; +#pragma GCC unroll 4 + for (int g = 0; g < 4; ++g) { + const __m256i p = _mm256_add_epi32(_mm256_madd_epi16(wv[2 * g], _mm256_load_si256(q + 2 * g)), + _mm256_madd_epi16(wv[2 * g + 1], _mm256_load_si256(q + 2 * g + 1))); + __m128i s = _mm_add_epi32(_mm256_castsi256_si128(p), _mm256_extracti128_si256(p, 1)); + s = _mm_hadd_epi32(s, s); + s = _mm_hadd_epi32(s, s); + sum[g] = _mm_cvtsi128_si32(s); + } + const __m128 sc = _mm_mul_ps(dx, _mm_loadu_ps(a[t]->d + 4 * i)); + const __m128i s4 = _mm_set_epi32(sum[3], sum[2], sum[1], sum[0]); + acc[t] = _mm_fmadd_ps(_mm_cvtepi32_ps(s4), sc, acc[t]); + } + } + for (int t = 0; t < nt; ++t) { + alignas(16) float v[4]; + _mm_store_ps(v, acc[t]); + out[t][r] = ((v[0] + v[1]) + (v[2] + v[3])) * (1.0f / 8192.0f); + } + } +} + } // namespace +void tq_rows_avx2_tqp(int kb, const uint8_t* w, size_t row_bytes, int n, const TqAct* const* a, int nt, + float* const* out, int r0, int r1) { + if (kb == 5) rows_tqp_t<5>(w, row_bytes, n, a, nt, out, r0, r1); + else if (kb == 6) rows_tqp_t<6>(w, row_bytes, n, a, nt, out, r0, r1); + else rows_tqp_t<7>(w, row_bytes, n, a, nt, out, r0, r1); +} + void tq_rows_avx2(int kb, const int16_t* table, const uint8_t* w, size_t row_bytes, int n, const TqAct* const* a, int nt, float* const* out, int r0, int r1) { if (kb == 8) rows_t<8>(table, w, row_bytes, n, a, nt, out, r0, r1); diff --git a/src/kernels/cpu/tq_cpu.cpp b/src/kernels/cpu/tq_cpu.cpp index 22a874bc..9c898cdf 100644 --- a/src/kernels/cpu/tq_cpu.cpp +++ b/src/kernels/cpu/tq_cpu.cpp @@ -86,6 +86,41 @@ void rows_scalar_t(const int16_t* tab, const uint8_t* w, size_t row_bytes, int n } } +// TQP (Gyro-L): 64 pair-trellis steps of 2 weights; a 32-value group is 16 steps. The table row holds the +// TQ2_T state's two pairs; TQP uses the first pair only (weight 2t+j = lut[2*hyb(2*s)+j]). +template +void rows_scalar_tqp_t(const uint8_t* w, size_t row_bytes, int n, const tq_detail::TqAct* const* a, int nt, + float* const* out, int r0, int r1) { + const int nb = n / 128; + constexpr size_t bsz = 2 + 8 * KB; + const int16_t* tab = tq_state_table(); + for (int r = r0; r < r1; ++r) { + const uint8_t* x = w + (size_t) r * row_bytes; + float acc[kTqMaxTokens][4] = {}; + for (int i = 0; i < nb; ++i) { + const uint8_t* blk = x + (size_t) i * bsz; + const float dx = tq_detail::tq_half_to_float(tq_detail::tq_ld16(blk)); + uint32_t st[64]; + tq_detail::tq_states_tqp(blk + 2, st); + for (int t = 0; t < nt; ++t) { + const int16_t* q = a[t]->q + 128 * i; + for (int k = 0; k < 4; ++k) { + int32_t s = 0; + for (int j = 0; j < 16; ++j) { // step 16k+j: two weights + const int16_t* wv = tab + 4 * st[16 * k + j]; + const int16_t* qv = q + 32 * k + 2 * j; + s += wv[0] * qv[0] + wv[1] * qv[1]; + } + const float sc = dx * a[t]->d[4 * i + k]; + acc[t][k] = std::fma((float) s, sc, acc[t][k]); + } + } + } + for (int t = 0; t < nt; ++t) + out[t][r] = ((acc[t][0] + acc[t][1]) + (acc[t][2] + acc[t][3])) * (1.0f / 8192.0f); + } +} + } // namespace const int16_t* tq_state_table() { @@ -103,11 +138,18 @@ const int16_t* tq_state_table_computed() { const int16_t* tq_state_table_for(int t) { return (t == kTqTypeK5X || t == kTqTypeK4X) ? tq_state_table_computed() : tq_state_table(); } bool tq_supported(int t) noexcept { - return t == 144 || t == 145 || t == 146 || t == kTqTypeK5X || t == kTqTypeK4X || t == kTqTypeK5T || t == kTqTypeK4T; + return t == 144 || t == 145 || t == 146 || t == 147 || t == 148 || t == 149 || + t == kTqTypeK5X || t == kTqTypeK4X || t == kTqTypeK5T || t == kTqTypeK4T; } int tq_kbits(int t) noexcept { - return t == 144 ? 8 : t == 145 ? 6 : t == 146 ? 7 : (t == kTqTypeK5X || t == kTqTypeK5T) ? 5 : (t == kTqTypeK4X || t == kTqTypeK4T) ? 4 : 0; + if (t == 148) return 6; // TQP6 + if (t == 144) return 8; + if (t == 146 || t == 149) return 7; + if (t == 145) return 6; + if (t == 147 || t == kTqTypeK5X || t == kTqTypeK5T) return 5; + if (t == kTqTypeK4X || t == kTqTypeK4T) return 4; + return 0; } bool tq_isa_ok(TqIsa isa) { @@ -166,6 +208,13 @@ void tq_rows_scalar(int kb, const int16_t* tab, const uint8_t* w, size_t row_byt else rows_scalar_t<7>(tab, w, row_bytes, n, a, nt, out, r0, r1); } +void tq_rows_scalar_tqp(int kb, const uint8_t* w, size_t row_bytes, int n, const TqAct* const* a, int nt, + float* const* out, int r0, int r1) { + if (kb == 5) rows_scalar_tqp_t<5>(w, row_bytes, n, a, nt, out, r0, r1); + else if (kb == 6) rows_scalar_tqp_t<6>(w, row_bytes, n, a, nt, out, r0, r1); + else rows_scalar_tqp_t<7>(w, row_bytes, n, a, nt, out, r0, r1); +} + } // namespace tq_detail namespace { @@ -197,6 +246,21 @@ void rows_on(TqIsa isa, int kb, const int16_t* tab, const uint8_t* w, size_t row } } +// The pair layout (types 147-149): AVX2 and AVX-512 both take the AVX2 pair kernel; scalar, and any build without the +// x86 units (arm64: the Metal build), the scalar one. +void rows_pair_on(TqIsa isa, int kb, const uint8_t* w, size_t row_bytes, int n, const tq_detail::TqAct* const* a, int nt, + float* const* out, int r0, int r1) { +#if defined(STRATA_TQ_X86) + if (isa != TqIsa::scalar) { + tq_detail::tq_rows_avx2_tqp(kb, w, row_bytes, n, a, nt, out, r0, r1); + return; + } +#else + (void) isa; +#endif + tq_detail::tq_rows_scalar_tqp(kb, w, row_bytes, n, a, nt, out, r0, r1); +} + } // namespace void tq_rows(TqIsa isa, int ggml_type, const uint8_t* w, size_t row_bytes, int n, const void* const* act, int nt, @@ -204,6 +268,11 @@ void tq_rows(TqIsa isa, int ggml_type, const uint8_t* w, size_t row_bytes, int n if (!tq_isa_ok(isa)) std::abort(); // a caller bug: asked for an ISA this CPU does not have const int16_t* tab = tq_state_table_for(ggml_type); const ActSet& s = unpack_all(act, n, nt); + if (ggml_type >= 147 && ggml_type <= 149) { // Gyro-L pair trellis: the AVX2 pair kernel (scalar is the reference) + const int kb = tq_kbits(ggml_type); + rows_pair_on(isa, kb, w, row_bytes, n, s.ptr, nt, out, r0, r1); + return; + } rows_on(isa, tq_kbits(ggml_type), tab, w, row_bytes, n, s.ptr, nt, out, r0, r1); } @@ -222,8 +291,13 @@ void tq_gu_rows(TqIsa isa, int ggml_type, const uint8_t* blob, size_t gu_row, si up[t] = ubuf.data() + (size_t) t * nr; } const int kb = tq_kbits(ggml_type); - rows_on(isa, kb, tab, blob, gu_row, n, s.ptr, nt, gp, r0, r1); - rows_on(isa, kb, tab, blob + up_off, gu_row, n, s.ptr, nt, up, r0, r1); + const bool pair = ggml_type >= 147 && ggml_type <= 149; // Gyro-L pair trellis: its own kernels + auto rows = [&](const uint8_t* w, float* const* o) { + if (pair) rows_pair_on(isa, kb, w, gu_row, n, s.ptr, nt, o, r0, r1); + else rows_on(isa, kb, tab, w, gu_row, n, s.ptr, nt, o, r0, r1); + }; + rows(blob, gp); + rows(blob + up_off, up); for (int t = 0; t < nt; ++t) for (int r = r0; r < r1; ++r) { const float g = gp[t][r], u = up[t][r]; diff --git a/src/kernels/cpu/tq_states.hpp b/src/kernels/cpu/tq_states.hpp index bb4eaa7a..8bdbfb66 100644 --- a/src/kernels/cpu/tq_states.hpp +++ b/src/kernels/cpu/tq_states.hpp @@ -60,6 +60,26 @@ inline void tq_states(const uint8_t* qs, uint32_t* st) { } } +/// TQP (Gyro-L): 64 pair-trellis steps; state t = the 16 stream bits from bit (63 - t)*KB, LSB first, +/// circular over 8*KB bytes (tqp_state in the fork's ggml-common.h). +template +inline uint32_t tq_state_tqp(const uint8_t* qs, int t) { + constexpr int NB = 8 * KB; + const uint32_t off = (uint32_t) (63 - t) * (uint32_t) KB; + const uint32_t i0 = off >> 3; + const uint32_t i1 = (i0 + 1) % (uint32_t) NB; + const uint32_t i2 = (i0 + 2) % (uint32_t) NB; + const uint32_t w = (uint32_t) qs[i0] | ((uint32_t) qs[i1] << 8) | ((uint32_t) qs[i2] << 16); + return (w >> (off & 7)) & kTqStateMask; +} +template +inline void tq_states_tqp(const uint8_t* qs, uint32_t* st) { +#if defined(__GNUC__) +#pragma GCC unroll 64 +#endif + for (int t = 0; t < 64; ++t) st[t] = tq_state_tqp(qs, t); +} + /// fp16 bits -> float, exact (the block scales and the Q8_0 activation scales). inline float tq_half_to_float(uint16_t h) { const uint32_t sign = (uint32_t) (h & 0x8000u) << 16; diff --git a/src/kernels/cpu/tqp_cpu_test.cpp b/src/kernels/cpu/tqp_cpu_test.cpp new file mode 100644 index 00000000..a8da31b2 --- /dev/null +++ b/src/kernels/cpu/tqp_cpu_test.cpp @@ -0,0 +1,229 @@ +// src/kernels/cpu/tqp_cpu_test.cpp - the trellis CPU rows (tq_cpu.hpp) on synthetic blocks, with no ggml and no model: +// +// * the pair-trellis types TQP5 / TQP6 / TQP7 (ggml 147-149, Gyro-L): the scalar kernel against a reference written +// here from the format definition (state t = the 16 stream bits that start at bit (63 - t) * K, LSB first, +// circular over 64 * K bits; the step's two weights are the first pair of the state's codebook entry), summed in +// double; then every ISA this CPU has (AVX2, AVX-512 takes the AVX2 pair kernel) bit for bit equal to the scalar +// kernel, for groups of 1..8 tokens, one row range or two, and for the fused gate/up rows; +// * the older trellis types (144-146): AVX2 / AVX-512 bit for bit equal to the scalar kernel (the pair types were +// added next to them: this shows they still agree), and their block sizes and K unchanged; +// * the type predicates (tq_supported, tq_kbits, tq_block_bytes) for every id 0..299: the answers the code gave +// before the pair types, plus 147-149. +// +// Random stream bytes are valid codes (any 15-bit state indexes the codebook). CPU only. Exit 0 = pass. +#include "strata/kernels/cpu/tq_cpu.hpp" + +#include "tq_states.hpp" + +#include +#include +#include +#include +#include +#include + +using namespace strata::kernels::cpu; + +namespace { + +int g_fail = 0; + +void fail(const char* what, int a = 0, int b = 0) { + std::printf(" FAIL: %s %d %d\n", what, a, b); + ++g_fail; +} + +// ---- the type predicates as they were before the pair types (the old formulas, copied) +bool old_supported(int t) { + return t == 144 || t == 145 || t == 146 || t == kTqTypeK5X || t == kTqTypeK4X || t == kTqTypeK5T || t == kTqTypeK4T; +} +int old_kbits(int t) { + return t == 144 ? 8 : t == 145 ? 6 : t == 146 ? 7 : (t == kTqTypeK5X || t == kTqTypeK5T) ? 5 : (t == kTqTypeK4X || t == kTqTypeK4T) ? 4 : 0; +} + +void check_predicates() { + for (int t = 0; t < 300; ++t) { + const bool pair = t >= 147 && t <= 149; + if (tq_supported(t) != (old_supported(t) || pair)) fail("tq_supported", t); + const int want_k = pair ? (t == 147 ? 5 : t == 148 ? 6 : 7) : old_kbits(t); + if (tq_kbits(t) != want_k) fail("tq_kbits", t, tq_kbits(t)); + const size_t want_b = pair ? (size_t) (2 + 8 * want_k) : (size_t) (2 + 4 * want_k); + if (tq_block_bytes(t) != want_b) fail("tq_block_bytes", t, (int) tq_block_bytes(t)); + } + if (tq_block_bytes(144) != 34 || tq_block_bytes(145) != 26 || tq_block_bytes(146) != 30) fail("sizes of 144-146"); + if (tq_block_bytes(147) != 42 || tq_block_bytes(148) != 50 || tq_block_bytes(149) != 58) fail("sizes of 147-149"); + std::printf("predicates (ids 0..299): %s\n", g_fail ? "FAIL" : "ok"); +} + +// ---- synthetic data +uint16_t random_half(std::mt19937& rng) { + std::uniform_int_distribution ex(2, 8), man(0, 1023), sgn(0, 3); + return (uint16_t) ((sgn(rng) == 0 ? 0x8000 : 0) | (ex(rng) << 10) | man(rng)); // 2^-13 .. 2^-6 +} + +std::vector random_rows(int type, int rows, int n, std::mt19937& rng, size_t& row_bytes) { + const size_t bs = tq_block_bytes(type); + row_bytes = bs * (size_t) (n / 128); + std::vector w(row_bytes * (size_t) rows); + std::uniform_int_distribution byte(0, 255); + for (auto& b : w) b = (uint8_t) byte(rng); + for (size_t o = 0; o < w.size(); o += bs) { + const uint16_t h = random_half(rng); + std::memcpy(&w[o], &h, 2); + } + return w; +} + +// n values as Q8_0 blocks (fp16 d, 32 x int8): 34 bytes per 32 values +std::vector random_q8_0(int n, std::mt19937& rng, std::vector& d, std::vector& q) { + std::vector a((size_t) (n / 32) * 34); + std::uniform_int_distribution val(-127, 127), ex(8, 13), man(0, 1023); + d.assign((size_t) n / 32, 0.f); + q.assign((size_t) n, 0); + for (int b = 0; b < n / 32; ++b) { + const uint16_t h = (uint16_t) ((ex(rng) << 10) | man(rng)); + std::memcpy(&a[(size_t) b * 34], &h, 2); + d[(size_t) b] = strata::kernels::cpu::tq_detail::tq_half_to_float(h); + for (int j = 0; j < 32; ++j) { + const int v = val(rng); + q[(size_t) b * 32 + j] = v; + a[(size_t) b * 34 + 2 + j] = (uint8_t) (int8_t) v; + } + } + return a; +} + +// ---- the reference for the pair types, from the format definition (bit by bit, not the window code of the kernels) +uint32_t ref_state(const uint8_t* qs, int K, int t) { + const int nbits = 64 * K; + uint32_t s = 0; + for (int b = 0; b < 16; ++b) { + const int pos = ((63 - t) * K + b) % nbits; + s |= (uint32_t) ((qs[pos >> 3] >> (pos & 7)) & 1) << b; + } + return s & 0x7FFFu; +} + +// w . a in double; `bound` = the sum of the absolute terms (the scale of the float rounding allowed) +double ref_dot(int type, const uint8_t* row, int n, const std::vector& d, const std::vector& q, double& bound) { + const int K = tq_kbits(type); + const int16_t* tab = tq_state_table(); + const size_t bs = tq_block_bytes(type); + double sum = 0, abs_sum = 0; + for (int i = 0; i < n / 128; ++i) { + const uint8_t* blk = row + (size_t) i * bs; + uint16_t hd; + std::memcpy(&hd, blk, 2); + const double dw = strata::kernels::cpu::tq_detail::tq_half_to_float(hd); + uint32_t st[64]; + for (int t = 0; t < 64; ++t) st[t] = ref_state(blk + 2, K, t); + for (int k = 0; k < 4; ++k) { + long long s = 0, sa = 0; + for (int j = 0; j < 32; ++j) { + const int v = 32 * k + j; // weight v of the block: step v / 2, element v & 1 of the first pair + const int w16 = tab[4 * st[v >> 1] + (v & 1)]; + s += (long long) w16 * q[(size_t) 128 * i + v]; + sa += (long long) std::abs(w16) * std::abs(q[(size_t) 128 * i + v]); + } + const double sc = dw * d[(size_t) 4 * i + k]; + sum += sc * (double) s; + abs_sum += std::fabs(sc) * (double) sa; + } + } + bound = abs_sum / 8192.0; + return sum / 8192.0; +} + +bool same_bits(const float* a, const float* b, size_t n) { return std::memcmp(a, b, n * sizeof(float)) == 0; } + +constexpr int NT = kTqMaxTokens; + +struct Case { + int type, rows, n; + std::vector w; + size_t rb = 0; + std::vector> act; + std::vector> d; + std::vector> q; + const void* ap[NT]; + Case(int type_, int rows_, int n_, uint32_t seed) : type(type_), rows(rows_), n(n_) { + std::mt19937 rng(seed); + w = random_rows(type, rows, n, rng, rb); + act.resize(NT); d.resize(NT); q.resize(NT); + for (int t = 0; t < NT; ++t) { + act[t] = random_q8_0(n, rng, d[t], q[t]); + ap[t] = act[t].data(); + } + } + void run(TqIsa isa, int nt, const void* const* a, std::vector& o) const { + o.assign((size_t) nt * rows, 0.f); + float* op[NT]; + for (int t = 0; t < nt; ++t) op[t] = o.data() + (size_t) t * rows; + tq_rows(isa, type, w.data(), rb, n, a, nt, op, 0, rows / 3); // two row ranges, as the pool splits them + tq_rows(isa, type, w.data(), rb, n, a, nt, op, rows / 3, rows); + } +}; + +void check_type(int type, bool reference) { + Case c(type, 37, 1280, 1000u + (uint32_t) type); + std::vector ref, got; + c.run(TqIsa::scalar, NT, c.ap, ref); + int bad = 0; + if (reference) { // the pair types: the scalar kernel against the definition + double worst = 0; + for (int t = 0; t < NT; ++t) + for (int r = 0; r < c.rows; ++r) { + double bound = 0; + const double want = ref_dot(type, c.w.data() + (size_t) r * c.rb, c.n, c.d[t], c.q[t], bound); + const double err = std::fabs((double) ref[(size_t) t * c.rows + r] - want); + worst = std::fmax(worst, bound > 0 ? err / bound : err); + if (err > 2e-5 * bound + 1e-30) ++bad; + } + std::printf("type %d scalar vs the definition (double): %d of %d outside 2e-5 of the term sum, worst %.2e\n", type, bad, + NT * c.rows, worst); + if (bad) fail("scalar kernel differs from the definition", type); + } + for (TqIsa isa : {TqIsa::avx2, TqIsa::avx512}) { + if (!tq_isa_ok(isa)) { std::printf("type %d %-6s not on this CPU: skipped\n", type, tq_isa_name(isa)); continue; } + bool ok = true; + for (int nt = 1; nt <= NT; ++nt) + for (int t0 = 0; t0 + nt <= NT; t0 += nt) { + c.run(isa, nt, c.ap + t0, got); + ok = ok && same_bits(got.data(), ref.data() + (size_t) t0 * c.rows, (size_t) nt * c.rows); + } + std::printf("type %d %-6s groups of 1..%d tokens vs the scalar kernel: %s\n", type, tq_isa_name(isa), NT, + ok ? "bitwise equal" : "DIFFERENT"); + if (!ok) fail("ISA differs from scalar", type); + } + // the fused gate/up rows: gate rows at the blob, up rows after them + { + std::vector blob = c.w; + std::mt19937 rng(77u + (uint32_t) type); + size_t rb2 = 0; + const auto up = random_rows(type, c.rows, c.n, rng, rb2); + blob.insert(blob.end(), up.begin(), up.end()); + const size_t up_off = c.w.size(); + std::vector ff0((size_t) NT * c.rows), ff1((size_t) NT * c.rows); + float* p0[NT]; float* p1[NT]; + for (int t = 0; t < NT; ++t) { p0[t] = ff0.data() + (size_t) t * c.rows; p1[t] = ff1.data() + (size_t) t * c.rows; } + tq_gu_rows(TqIsa::scalar, type, blob.data(), c.rb, up_off, c.n, c.ap, NT, p0, 0, c.rows); + bool ok = true; + for (TqIsa isa : {TqIsa::avx2, TqIsa::avx512}) { + if (!tq_isa_ok(isa)) continue; + tq_gu_rows(isa, type, blob.data(), c.rb, up_off, c.n, c.ap, NT, p1, 0, c.rows); + ok = ok && same_bits(ff0.data(), ff1.data(), ff0.size()); + } + std::printf("type %d gate/up rows, every ISA vs scalar: %s\n", type, ok ? "bitwise equal" : "DIFFERENT"); + if (!ok) fail("gate/up rows differ", type); + } +} + +} // namespace + +int main() { + check_predicates(); + for (int type : {147, 148, 149}) check_type(type, true); + for (int type : {144, 145, 146}) check_type(type, false); + std::printf("tqp_cpu_test: %d failures\n", g_fail); + return g_fail ? 1 : 0; +} diff --git a/src/kernels/cuda/iq_kernels.cu b/src/kernels/cuda/iq_kernels.cu index d2f2e34e..f81efd69 100644 --- a/src/kernels/cuda/iq_kernels.cu +++ b/src/kernels/cuda/iq_kernels.cu @@ -616,6 +616,10 @@ template struct TqInfo; template<> struct TqInfo<144> { using block_t = block_tq2_t; }; // byte-aligned, no K template<> struct TqInfo<145> { using block_t = block_tqk6; static constexpr int K = 6; }; template<> struct TqInfo<146> { using block_t = block_tqk7; static constexpr int K = 7; }; +// Gyro-L: pair trellis (64 steps x K bits, 2 weights per step), the TQ2_T codebook's first pair +template<> struct TqInfo<147> { using block_t = block_tqp5; static constexpr int K = 5; }; +template<> struct TqInfo<148> { using block_t = block_tqp6; static constexpr int K = 6; }; +template<> struct TqInfo<149> { using block_t = block_tqp7; static constexpr int K = 7; }; __device__ __forceinline__ float tq_h2f(uint32_t bits) { return __half2float(__ushort_as_half((unsigned short) bits)); } // the four unscaled weights of the step with state s: (a.x, a.y, b.x, b.y) __device__ __forceinline__ void tq_step(uint32_t s, float2& a, float2& b) { @@ -716,6 +720,71 @@ __device__ __forceinline__ float vec_dot_tq_q8_1(const void* __restrict__ vbq, c return tq_apply(tq_load(vbq, kbx, iqs), bq8_1, iqs); } +// ---- TQP (Gyro-L): 64 trellis steps of 2 weights; a lane is 8 steps (16 values). +// State of step t: the 16 stream bits from bit (63 - t)*K, LSB-first, circular over 64*K bits +// (tqp_state in the fork's ggml-common.h). The step's 2 weights are tq2t_lut_f16[2*hyb(2*s)+j], +// i.e. the FIRST pair of TQ2_T state s = tq_lut_i8[x >> 21] (tq_lut_i8 holds the pair as two int8). +// A lane's 8 states are the 16-bit fields at bit i*K of the bit stream starting at its 8th-to-last +// step, so the whole lane needs K + 2 bytes (7, 8 or 9 for K = 5, 6, 7): one byte-built 64-bit +// window plus one byte for K = 7's first state (which ends at bit 65), then eight shifts. That is +// 7-9 loads per lane instead of 24. The bytes are loaded (not one u64) because a block's qs array +// is only 2-byte aligned (block strides are 42/50/58 bytes). +template +__device__ __forceinline__ void tqp_states8(const uint8_t* __restrict__ qs, int lane, uint32_t s[8]) { + constexpr int K = TqInfo::K, NB = 8 * K, NW = K + 2 < 8 ? K + 2 : 8; + const int pidx = (7 - lane) * K; // the lane's first byte, always < NB + uint64_t w = 0; // bits pidx*8 .. pidx*8+63 +#pragma unroll + for (int j = 0; j < NW; ++j) { + const int idx0 = pidx + j; + w |= (uint64_t) qs[idx0 < NB ? idx0 : idx0 - NB] << (8 * j); + } + const int i8 = pidx + 8; + const uint32_t hi = K == 7 ? (uint32_t) qs[i8 < NB ? i8 : i8 - NB] : 0u; // the 65th bit +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int sh = (7 - i) * K; // 0 .. 7K (<= 49) + const uint32_t lo = (uint32_t) (w >> sh); + const uint32_t top = sh > 48 ? (uint32_t) (hi << (64 - sh)) : 0u; + s[i] = (lo | top) & 0xFFFFu; + } +} +// The per-state form: 3 byte loads and a shift, no window. Used by dq_tqp (the f16 expert uploads), which +// benches faster with it than with the window (the window's register pressure costs the copy kernels). +template +__device__ __forceinline__ void tqp_states8_cheap(const uint8_t* __restrict__ qs, int lane, uint32_t s[8]) { + constexpr int K = TqInfo::K, NB = 8 * K; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t off = (uint32_t) (63 - (8 * lane + i)) * (uint32_t) K; + const uint32_t i0 = off >> 3; + const uint32_t i1 = i0 + 1 < (uint32_t) NB ? i0 + 1 : 0; + const uint32_t i2 = i1 + 1 < (uint32_t) NB ? i1 + 1 : 0; + const uint32_t w = (uint32_t) qs[i0] | ((uint32_t) qs[i1] << 8) | ((uint32_t) qs[i2] << 16); + s[i] = (w >> (off & 7)) & 0xFFFFu; + } +} +template +__device__ __forceinline__ TqLane tq_load_tqp(const void* __restrict__ vbq, int kbx, int iqs) { + const auto* b = (const typename TqInfo::block_t*) vbq + kbx; + uint32_t s[8]; + tqp_states8(b->qs, iqs / 4, s); + TqLane r; +#pragma unroll + for (int i = 0; i < 4; ++i) { // two steps (2+2 weights) per int + const uint32_t x0 = (s[2 * i] & TQ_STATE_MASK) * 0x9e3779b1u; + const uint32_t x1 = (s[2 * i + 1] & TQ_STATE_MASK) * 0x9e3779b1u; + r.q[i] = (int) ((uint32_t) tq_lut_i8[x0 >> 21] | ((uint32_t) tq_lut_i8[x1 >> 21] << 16)); + } + r.d = tq_h2f(b->d); + return r; +} +template +__device__ __forceinline__ float vec_dot_tqp_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + return tq_apply(tq_load_tqp(vbq, kbx, iqs), bq8_1, iqs); +} + // ---------------------------------------------------------------- the formats // qk = values per block, ipb = dot calls per block (qi / vdr), step = the iqs stride between calls. template struct Fmt; @@ -760,21 +829,28 @@ template<> struct Fmt<145> { static constexpr int qk = 128, ipb = 8, step = 4; __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_tq_q8_1<145>(v, y, kbx, iqs); } }; template<> struct Fmt<146> { static constexpr int qk = 128, ipb = 8, step = 4; __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_tq_q8_1<146>(v, y, kbx, iqs); } }; +// Gyro-L pair trellis: 8 lanes of 8 steps (16 values) per 128-value block +#define STRATA_TQP_FMT(T) template<> struct Fmt { static constexpr int qk = 128, ipb = 8, step = 4; \ + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_tqp_q8_1(v, y, kbx, iqs); } }; +STRATA_TQP_FMT(147) +STRATA_TQP_FMT(148) +STRATA_TQP_FMT(149) +#undef STRATA_TQP_FMT // The formats of each role, one list each so a type cannot be in one switch and missing from another. Every // entry is a kernel template for each CUDA architecture of the build, hence two lists rather than one. #ifdef STRATA_Q6K_EXPERTS // opt-in build (-DSTRATA_Q6K_EXPERTS=ON): one more instance per kernel, loaded at start -#define STRATA_GU_FMTS(X) X(16) X(17) X(18) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(14) X(6) X(2) X(3) X(8) X(40) X(144) X(145) +#define STRATA_GU_FMTS(X) X(16) X(17) X(18) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(14) X(6) X(2) X(3) X(8) X(40) X(144) X(145) X(147) X(148) X(149) #else -#define STRATA_GU_FMTS(X) X(16) X(17) X(18) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(6) X(2) X(3) X(8) X(40) X(144) X(145) +#define STRATA_GU_FMTS(X) X(16) X(17) X(18) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(6) X(2) X(3) X(8) X(40) X(144) X(145) X(147) X(148) X(149) #endif // Q5_K / Q6_K downs: Qwen3.6-35B-A3B (FF 512, whole 256-blocks per row); Flash-Next's FF 640 never selects them #ifdef STRATA_Q6K_EXPERTS -#define STRATA_D_FMTS(X) X(20) X(23) X(42) X(7) X(6) X(2) X(3) X(8) X(13) X(14) X(40) X(144) X(146) +#define STRATA_D_FMTS(X) X(20) X(23) X(42) X(7) X(6) X(2) X(3) X(8) X(13) X(14) X(40) X(144) X(146) X(147) X(148) X(149) #else -#define STRATA_D_FMTS(X) X(20) X(23) X(42) X(7) X(6) X(2) X(3) X(8) X(13) X(40) X(144) X(146) +#define STRATA_D_FMTS(X) X(20) X(23) X(42) X(7) X(6) X(2) X(3) X(8) X(13) X(40) X(144) X(146) X(147) X(148) X(149) #endif -#define STRATA_MMVQ_FMTS(X) X(16) X(17) X(18) X(20) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(7) X(6) X(2) X(3) X(8) X(40) X(144) X(145) X(146) +#define STRATA_MMVQ_FMTS(X) X(16) X(17) X(18) X(20) X(21) X(22) X(23) X(29) X(42) X(12) X(13) X(7) X(6) X(2) X(3) X(8) X(40) X(144) X(145) X(146) X(147) X(148) X(149) template<> struct Fmt<40> { static constexpr int qk = 64, ipb = 2, step = 4; // NVFP4: 2 calls of 32 values __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_nvfp4_q8_1(v, y, kbx, iqs); } }; @@ -1390,6 +1466,23 @@ template struct SplitTq { template<> struct Split<144> : SplitTq<144> {}; template<> struct Split<145> : SplitTq<145> {}; template<> struct Split<146> : SplitTq<146> {}; +// Gyro-L TQP (147-149): the pair layout needs tq_load_tqp, not tq_load, so its split is its own trait. +// Until this existed the launchers took the per-entry mmvq/native kernels for TQP, and row_dot_multi's +// callers re-decoded the lane once per activation. +template struct SplitTqp { + using W = TqLane; + template // (no grid table to stage: the codebook is read where it is) + __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ = nullptr) { + return tq_load_tqp(vbq, kbx, iqs); + } + __device__ static float apply(const W& r, const block_q8_1* __restrict__ bq8_1, int iqs) { return tq_apply(r, bq8_1, iqs); } +}; +template<> inline constexpr bool kSplit<147> = true; +template<> inline constexpr bool kSplit<148> = true; +template<> inline constexpr bool kSplit<149> = true; +template<> struct Split<147> : SplitTqp<147> {}; +template<> struct Split<148> : SplitTqp<148> {}; +template<> struct Split<149> : SplitTqp<149> {}; // One row against the n <= NC activations x + off[0..n) (n >= 1, warp-uniform; offsets in q8_1 blocks, 32-bit to // spare registers), the whole warp. Per activation this is row_dot: the same calls k, lane-strided the same way, @@ -3170,6 +3263,31 @@ __device__ void dq_tq(const void* vx, int64_t ibs, dst_t* yy, int tid) { } } +// TQP (Gyro-L): pair trellis; the step's 2 weights are the first pair of TQ2_T state s. +__device__ __forceinline__ float2 tqp_pair(uint32_t s) { + const uint32_t* lut = (const uint32_t*) tq_lut_f16; + const uint32_t x = (s & TQ_STATE_MASK) * 0x9e3779b1u; + const uint32_t pa = lut[x >> 21]; + return make_float2(tq_h2f(pa & 0xFFFFu), tq_h2f(pa >> 16)); +} +// one 256-value call = two 128-blocks; thread tid decodes block tid/16, lane (tid>>1)&7, steps +// 8*l + 4*half .. +3 (8 values, 2 per step) +template +__device__ void dq_tqp(const void* vx, int64_t ibs, dst_t* yy, int tid) { + const int blk = tid >> 4, l = (tid >> 1) & 7, half = tid & 1; + const auto* b = (const typename TqInfo::block_t*) vx + 2 * ibs + blk; + const float d = tq_h2f(b->d); + uint32_t s[8]; + tqp_states8_cheap(b->qs, l, s); + dst_t* y = yy + blk * 128 + 16 * l + 8 * half; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 a = tqp_pair(s[4 * half + i]); + y[2 * i + 0] = cvt(d * a.x); + y[2 * i + 1] = cvt(d * a.y); + } +} + // Every type below must also be in is_iq() (BF16: embed_type_supported): the host entry points refuse the others, // so the default is unreachable. template @@ -3198,6 +3316,9 @@ __device__ __forceinline__ void dq_dispatch(int ty, const void* vx, int64_t ibs, case 144: dq_tq<144>(vx, ibs, y, tid); break; case 145: dq_tq<145>(vx, ibs, y, tid); break; case 146: dq_tq<146>(vx, ibs, y, tid); break; + case 147: dq_tqp<147>(vx, ibs, y, tid); break; + case 148: dq_tqp<148>(vx, ibs, y, tid); break; + case 149: dq_tqp<149>(vx, ibs, y, tid); break; default: break; } } @@ -3239,9 +3360,9 @@ __global__ void dequant_gu_kernel(int ty, const void* __restrict__ gate, const v bool is_iq(int t) { return t == 16 || t == 17 || t == 18 || t == 20 || t == 21 || t == 22 || t == 23 || t == 29 || t == 42 || t == 11 || #ifdef STRATA_Q6K_EXPERTS - t == 12 || t == 13 || t == 14 || t == 7 || t == 6 || t == 2 || t == 3 || t == 8 || t == 40 || t == 144 || t == 145 || t == 146; + t == 12 || t == 13 || t == 14 || t == 7 || t == 6 || t == 2 || t == 3 || t == 8 || t == 40 || t == 144 || t == 145 || t == 146 || t == 147 || t == 148 || t == 149; #else - t == 12 || t == 13 || t == 7 || t == 6 || t == 2 || t == 3 || t == 8 || t == 40 || t == 144 || t == 145 || t == 146; + t == 12 || t == 13 || t == 7 || t == 6 || t == 2 || t == 3 || t == 8 || t == 40 || t == 144 || t == 145 || t == 146 || t == 147 || t == 148 || t == 149; #endif } // values per block of the types the grouped expert kernels take (0 = none) @@ -4202,6 +4323,9 @@ size_t iq_row_bytes(int t, int64_t n) noexcept { case 144: return (size_t) (n / 128) * sizeof(block_tq2_t); // APR trellis types case 145: return (size_t) (n / 128) * sizeof(block_tqk6); case 146: return (size_t) (n / 128) * sizeof(block_tqk7); + case 147: return (size_t) (n / 128) * sizeof(block_tqp5); + case 148: return (size_t) (n / 128) * sizeof(block_tqp6); + case 149: return (size_t) (n / 128) * sizeof(block_tqp7); default: return 0; } } diff --git a/src/kernels/native_expert_parity.cpp b/src/kernels/native_expert_parity.cpp index ffe8a6ca..a49a8515 100644 --- a/src/kernels/native_expert_parity.cpp +++ b/src/kernels/native_expert_parity.cpp @@ -253,6 +253,19 @@ int check_blob(const cpu::NativeFmt& f, const std::vector& blob, int se std::printf(" | %s %s", cpu::tq_isa_name(isa), same ? "bitwise equal" : "DIFFERENT"); if (!same) ++failures; } + { // time: the scalar kernel against the best SIMD one this CPU has, same rows and tokens + const int it = 30; + const auto t0 = std::chrono::steady_clock::now(); + for (int k = 0; k < it; ++k) cpu::tq_rows(cpu::TqIsa::scalar, type, w, rb, n, acts, NT, rp, 0, rows); + const auto t1 = std::chrono::steady_clock::now(); + const cpu::TqIsa best = cpu::tq_isa(); + for (int k = 0; k < it; ++k) cpu::tq_rows(best, type, w, rb, n, acts, NT, gp, 0, rows); + const auto t2 = std::chrono::steady_clock::now(); + const double us0 = std::chrono::duration(t1 - t0).count() / it; + const double us1 = std::chrono::duration(t2 - t1).count() / it; + std::printf(" | %s %.0f us, scalar %.0f us (%.2fx)", cpu::tq_isa_name(best), us1, us0, + us1 > 0 ? us0 / us1 : 0.0); + } std::printf("\n"); } // UD-Q4_K_XL's formats: the multi-token AVX2 kernels (kq_avx2.cpp) against ggml-cpu's vec_dot per token, @@ -482,6 +495,11 @@ int check_blob(const cpu::NativeFmt& f, const std::vector& blob, int se int type_of(std::string name) { for (auto& c : name) c = (char) std::tolower((unsigned char) c); + // The trellis types sit above GGML_TYPE_COUNT on an upstream ggml (the overlay adds them to ggml-cuda, not to the + // type table), so ggml_type_name cannot resolve them: name them here, as the engine's own type ids do. + static const struct { const char* n; int t; } trellis[] = {{"tq2_t", 144}, {"tqk6", 145}, {"tqk7", 146}, + {"tqp5", 147}, {"tqp6", 148}, {"tqp7", 149}}; + for (const auto& e : trellis) if (name == e.n) return e.t; for (int t = 0; t < GGML_TYPE_COUNT; ++t) { const char* n = ggml_type_name((ggml_type) t); if (n == nullptr) continue; @@ -837,6 +855,15 @@ int main(int argc, char** argv) { const int dn = slash == std::string::npos ? -1 : type_of(arg.substr(slash + 1)); cpu::NativeFmt f; std::string err; + // The synthetic check quantizes the blob with ggml and compares against ggml's to_float. An upstream + // ggml (the MMQ overlay) has no traits for the trellis types (they sit above GGML_TYPE_COUNT), so the + // check needs a ggml that carries them (a build with -DSTRATA_GGML_DIR=): skip it + // here with a message instead of reading past the type table. + if (gu >= (int) GGML_TYPE_COUNT || dn >= (int) GGML_TYPE_COUNT) { // (ggml_type_name of an id past the table is not a null pointer: it reads past it) + std::printf("%s: skipped (this ggml has no trellis type traits: the check needs STRATA_GGML_DIR=)\n", arg.c_str()); + continue; + } if (gu < 0 || dn < 0 || !cpu::native_fmt(gu, dn, H, FF, f, err) || !strata::kernels::native_expert_supported(gu, dn, H, FF)) { std::printf("%s: %s\n", arg.c_str(), err.empty() ? "not a pair the GPU expert kernels take" : err.c_str()); diff --git a/src/kernels/tq_pascal_bench.cpp b/src/kernels/tq_pascal_bench.cpp index bbed71ac..1c39d827 100644 --- a/src/kernels/tq_pascal_bench.cpp +++ b/src/kernels/tq_pascal_bench.cpp @@ -1,4 +1,4 @@ -// src/kernels/tq_pascal_bench.cpp - the Gyro trellis decode kernels (TQ2_T 144, TQK6 145, TQK7 146) on synthetic blocks. +// src/kernels/tq_pascal_bench.cpp - the Gyro trellis decode kernels (TQ2_T 144, TQK6 145, TQK7 146 and Gyro-L's TQP5-7 147-149) on synthetic blocks. // // build/tq_pascal_bench checks: the decode-once kernels against the per-column ones (bitwise), and a // checksum of every output, so two builds can be compared byte for byte @@ -11,6 +11,8 @@ #include "strata/kernels/iq_kernels.hpp" #include "strata/kernels/native_mmvq.hpp" +#include "ggml-common-tq.h" // tq_lut_f16: the codebook, as host data + #include #include @@ -45,7 +47,10 @@ T* dalloc(size_t n) { return p; } -const char* name_of(int t) { return t == 144 ? "TQ2_T" : t == 145 ? "TQK6" : t == 146 ? "TQK7" : t == 42 ? "Q2_0" : "?"; } +const char* name_of(int t) { + return t == 144 ? "TQ2_T" : t == 145 ? "TQK6" : t == 146 ? "TQK7" : t == 147 ? "TQP5" : t == 148 ? "TQP6" : t == 149 ? "TQP7" + : t == 42 ? "Q2_0" : "?"; +} uint64_t fnv(const void* p, size_t n, uint64_t h = 1469598103934665603ull) { const uint8_t* b = (const uint8_t*) p; @@ -206,6 +211,124 @@ void check_grouped(int gu, int dt, int64_t H, int64_t FF, cudaStream_t s, std::m if (!ok) { std::printf(" %zu of %zu floats differ%s\n", diff, a.size(), finite ? "" : ", non-finite outputs"); ++g_fail; } } +// ---- Gyro-L pair trellis (TQP5/6/7, ggml 147-149) against a reference written from the format definition (no ggml): +// state t of a block = the 16 stream bits that start at bit (63 - t) * K, LSB first, circular over 64 * K bits (15 bits +// are used); the step's two weights are the first pair of the state's codebook entry, times the block's fp16 scale. +uint32_t tqp_ref_state(const uint8_t* qs, int K, int t) { + const int nbits = 64 * K; + uint32_t s = 0; + for (int b = 0; b < 16; ++b) { + const int pos = ((63 - t) * K + b) % nbits; + s |= (uint32_t) ((qs[pos >> 3] >> (pos & 7)) & 1) << b; + } + return s & 0x7FFFu; +} + +float half_to_f(uint16_t h) { // normal and subnormal fp16 (no inf / nan in these tests) + const int e = (h >> 10) & 31, m = h & 1023; + const float v = e == 0 ? std::ldexp((float) m, -24) : std::ldexp((float) (1024 + m), e - 25); + return (h & 0x8000) ? -v : v; +} + +uint16_t f_to_half_rn(float f) { // round to nearest even, the GPU's __float2half_rn + uint32_t x; + std::memcpy(&x, &f, 4); + const uint32_t sign = (x >> 16) & 0x8000u; + const int e = (int) ((x >> 23) & 0xFF) - 127 + 15; + uint32_t m = x & 0x7FFFFFu; + if (e >= 31) return (uint16_t) (sign | 0x7C00u); + if (e <= 0) { + if (e < -10) return (uint16_t) sign; + m |= 0x800000u; + const int shift = 14 - e; + uint32_t half = m >> shift; + const uint32_t rem = m & ((1u << shift) - 1), mid = 1u << (shift - 1); + if (rem > mid || (rem == mid && (half & 1u))) ++half; + return (uint16_t) (sign | half); + } + uint32_t half = ((uint32_t) e << 10) | (m >> 13); + const uint32_t rem = m & 0x1FFFu; + if (rem > 0x1000u || (rem == 0x1000u && (half & 1u))) ++half; + return (uint16_t) (sign | half); +} + +// the weights of `rows` rows of `n` values, as floats (d * codebook, unrounded) +std::vector tqp_ref_weights(int t, const std::vector& w, int64_t rows, int64_t n) { + const int K = t - 142; // 147 -> 5, 148 -> 6, 149 -> 7 + const size_t bs = k::iq_row_bytes(t, 128); + std::vector out((size_t) (rows * n)); + for (size_t b = 0; b < (size_t) (rows * n / 128); ++b) { + const uint8_t* blk = &w[b * bs]; + uint16_t hd; + std::memcpy(&hd, blk, 2); + const float d = half_to_f(hd); + for (int step = 0; step < 64; ++step) { + const uint32_t s = tqp_ref_state(blk + 2, K, step); + const uint32_t x = (s & 0x7FFFu) * 0x9e3779b1u; + const uint32_t pt = x >> 21; + out[b * 128 + 2 * step + 0] = d * half_to_f(tq_lut_f16[2 * pt + 0]); + out[b * 128 + 2 * step + 1] = d * half_to_f(tq_lut_f16[2 * pt + 1]); + } + } + return out; +} + +void check_tqp_reference(int t, int n_in, int n_out, cudaStream_t s, std::mt19937& rng) { + const auto w = random_rows(t, n_out, n_in, rng); + const auto ref = tqp_ref_weights(t, w, n_out, n_in); + uint8_t* dw = dalloc(w.size()); + ck(cudaMemcpy(dw, w.data(), w.size(), cudaMemcpyHostToDevice), "w"); + // (1) the prompt path's fp16 dequantizer: bit for bit the fp16 of the reference weights + const int64_t nv = (int64_t) n_out * n_in; + uint16_t* dh = dalloc((size_t) nv); + k::iq_dequant_f16(t, dw, nv, dh, s); + ck(cudaStreamSynchronize(s), "sync"); + std::vector h((size_t) nv); + ck(cudaMemcpy(h.data(), dh, h.size() * 2, cudaMemcpyDeviceToHost), "h"); + size_t bad = 0; + for (size_t i = 0; i < h.size(); ++i) bad += h[i] != f_to_half_rn(ref[i]); + std::printf("tqp ref %-6s dequant_f16 %5d x %5d: %zu of %zu fp16 values differ from the definition %s\n", name_of(t), n_out, n_in, + bad, h.size(), bad ? "FAIL" : "bitwise equal"); + g_fail += bad != 0; + // (2) the decode kernels (int8 codebook, q8_1 activations): relative L1 error against the float weights . the + // dequantized q8_1 activations, in double, below the 3e-2 every expert type gets in native_expert_parity + const int nc = 3; + const auto x = random_x((size_t) nc * n_in, rng); + float* dx = dalloc(x.size()); + ck(cudaMemcpy(dx, x.data(), x.size() * 4, cudaMemcpyHostToDevice), "x"); + uint8_t* dxq = dalloc((size_t) nc * (n_in / 32) * 36); + k::quantize_q8_1_rows(dx, nc, n_in, dxq, s); + float* dy = dalloc((size_t) nc * n_out); + k::iq_mmvq(t, dw, dxq, dy, n_in, n_out, nc, s); + ck(cudaStreamSynchronize(s), "sync"); + std::vector y((size_t) nc * n_out); + ck(cudaMemcpy(y.data(), dy, y.size() * 4, cudaMemcpyDeviceToHost), "y"); + std::vector xq((size_t) nc * (n_in / 32) * 36); + ck(cudaMemcpy(xq.data(), dxq, xq.size(), cudaMemcpyDeviceToHost), "xq"); + double l1 = 0, l1ref = 0; + for (int c = 0; c < nc; ++c) { + std::vector xd((size_t) n_in); + for (int b = 0; b < n_in / 32; ++b) { + const uint8_t* blk = &xq[((size_t) c * (n_in / 32) + b) * 36]; + uint16_t hd; + std::memcpy(&hd, blk, 2); + const double d = half_to_f(hd); + for (int j = 0; j < 32; ++j) xd[(size_t) b * 32 + j] = d * (double) (int8_t) blk[4 + j]; + } + for (int r = 0; r < n_out; ++r) { + double acc = 0; + for (int i = 0; i < n_in; ++i) acc += (double) ref[(size_t) r * n_in + i] * xd[(size_t) i]; + l1 += std::fabs((double) y[(size_t) c * n_out + r] - acc); + l1ref += std::fabs(acc); + } + } + const double rel = l1 / l1ref; + std::printf("tqp ref %-6s mmvq %5d x %5d x %d cols: relative L1 error %.2e (bound 3e-2) %s\n", name_of(t), n_out, n_in, nc, rel, + rel < 3e-2 ? "ok" : "FAIL"); + g_fail += !(rel < 3e-2); + cudaFree(dw); cudaFree(dh); cudaFree(dx); cudaFree(dxq); cudaFree(dy); +} + float time_ms(cudaStream_t s, int it, const auto& fn) { cudaEvent_t e0, e1; cudaEventCreate(&e0); @@ -286,6 +409,7 @@ int main(int argc, char** argv) { cudaGetDeviceProperties(&prop, 0); std::printf("device %s cc %d.%d\n", prop.name, prop.major, prop.minor); std::mt19937 rng(18); + // Gyro-S / Gyro-M types first, on the generator they always had: their "SUM" lines are the same bits as before the pair types for (int t : {144, 145, 146}) { check_mmvq(t, 2560, 67, s, rng); check_mmvq(t, 1024, 5, s, rng); @@ -299,6 +423,19 @@ int main(int argc, char** argv) { check_grouped(gu, 144, 2560, 512, s, rng); check_grouped(gu, 146, 2048, 256, s, rng); } + // Gyro-L's pair trellis (own generator): decode-once == per-column, then a reference from the format definition + // (fp16 dequantizer bit for bit, decode kernels to 3e-2) + std::mt19937 rng_p(19); + for (int t : {147, 148, 149}) { + check_mmvq(t, 2560, 67, s, rng_p); + check_mmvq(t, 1024, 5, s, rng_p); + check_tqp_reference(t, 2560, 67, s, rng_p); + check_tqp_reference(t, 1024, 5, s, rng_p); + check_grouped(t, t, 2560, 640, s, rng_p); // the pair trellis in both roles, as its file has it (tqp6 / tqp6) + check_grouped(t, t, 5120, 640, s, rng_p); + } + check_grouped(148, 146, 2560, 640, s, rng_p); + check_grouped(145, 148, 2560, 640, s, rng_p); if (do_bench) bench(s, rng); std::printf("tq_pascal_bench: %d failures\n", g_fail); cudaStreamDestroy(s); diff --git a/src/prefill/moe_mmq.cu b/src/prefill/moe_mmq.cu index e7ff5d45..96098dfa 100644 --- a/src/prefill/moe_mmq.cu +++ b/src/prefill/moe_mmq.cu @@ -21,11 +21,15 @@ namespace strata::prefill::mmq { namespace { // The Gyro trellis types (ids 144-146) are outside the type table of an upstream ggml (indexed by the id): their block // geometry is known here (128 weights; 34 / 26 / 30 bytes), everything else asks ggml. -inline bool is_tq(int t) { return t >= 144 && t <= 146; } +inline bool is_tq(int t) { return t >= 144 && t <= 149; } // 144-146 TQ/TQK, 147-149 the pair trellis (Gyro-L) inline int64_t type_blck(int t) { return is_tq(t) ? 128 : (int64_t) ggml_blck_size((ggml_type) t); } -inline size_t type_bytes(int t) { return t == 144 ? 34 : t == 145 ? 26 : t == 146 ? 30 : ggml_type_size((ggml_type) t); } +inline size_t type_bytes(int t) { + return t == 144 ? 34 : t == 145 ? 26 : t == 146 ? 30 : t == 147 ? 42 : t == 148 ? 50 : t == 149 ? 58 + : ggml_type_size((ggml_type) t); +} inline const char* type_label(int t) { - return t == 144 ? "tq2_t" : t == 145 ? "tqk6" : t == 146 ? "tqk7" : ggml_type_name((ggml_type) t); + return t == 144 ? "tq2_t" : t == 145 ? "tqk6" : t == 146 ? "tqk7" + : t == 147 ? "tqp5" : t == 148 ? "tqp6" : t == 149 ? "tqp7" : ggml_type_name((ggml_type) t); } } // namespace #if defined(STRATA_KQ_DSO) @@ -275,6 +279,9 @@ bool supported(int t) { case GGML_TYPE_Q8_0: // the draft layer's dense matrices (E-9) #ifdef STRATA_MMQ_TQ case GGML_TYPE_TQ2_T: case GGML_TYPE_TQK6: case GGML_TYPE_TQK7: // APR experts (agention fork's MMQ) +#endif +#ifdef STRATA_MMQ_TQP // a ggml with the pair-trellis types (the overlay, or an agention tree that has their instances) + case GGML_TYPE_TQP5: case GGML_TYPE_TQP6: case GGML_TYPE_TQP7: // Gyro-L pair trellis (types 147-149) #endif return true; #ifdef STRATA_MMQ_KQUANTS @@ -462,6 +469,11 @@ void Context::run(const Product& p, void* stream) { case GGML_TYPE_TQK6: mul_mat_q_case(ctx, a, s); break; case GGML_TYPE_TQK7: mul_mat_q_case(ctx, a, s); break; #endif +#ifdef STRATA_MMQ_TQP + case GGML_TYPE_TQP5: mul_mat_q_case(ctx, a, s); break; + case GGML_TYPE_TQP6: mul_mat_q_case(ctx, a, s); break; + case GGML_TYPE_TQP7: mul_mat_q_case(ctx, a, s); break; +#endif #if defined(STRATA_MMQ_KQUANTS) && defined(STRATA_KQ_DSO) case GGML_TYPE_Q4_K: case GGML_TYPE_Q5_K: case GGML_TYPE_Q6_K: case GGML_TYPE_Q5_1: { const int e = kq::run((int) t, &a, s); diff --git a/third_party/ggml-tq/tq.cuh b/third_party/ggml-tq/tq.cuh index 9e9e8a26..a20fe3b2 100644 --- a/third_party/ggml-tq/tq.cuh +++ b/third_party/ggml-tq/tq.cuh @@ -15,6 +15,7 @@ #define TQ_LUT_POINTS 2048 #define TQ_LANES 8 // lanes per 128-weight block #define TQ_STEPS_LANE 4 // trellis steps per lane +#define TQ_PAIRS_LANE 8 // codebook pairs per lane (the TQP layout: 2 weights per step) #define QR_TQ 4 #define QI_TQ (QK_TQK/QR_TQ) @@ -44,14 +45,34 @@ struct ggml_cuda_type_traits { static constexpr int bs = sizeof(block_tqk7); }; +// Gyro-L pair trellis (types 147-149): the same lane packing (16 values = 4 ints of 4 int8), 64 steps of +// 2 weights per 128-value block. +template<> struct ggml_cuda_type_traits { + static constexpr int qk = QK_TQP; static constexpr int qr = QR_TQ; static constexpr int qi = QI_TQ; + static constexpr int bs = sizeof(block_tqp5); }; +template<> struct ggml_cuda_type_traits { + static constexpr int qk = QK_TQP; static constexpr int qr = QR_TQ; static constexpr int qi = QI_TQ; + static constexpr int bs = sizeof(block_tqp6); }; +template<> struct ggml_cuda_type_traits { + static constexpr int qk = QK_TQP; static constexpr int qr = QR_TQ; static constexpr int qi = QI_TQ; + static constexpr int bs = sizeof(block_tqp7); }; + +static constexpr __host__ __device__ bool ggml_cuda_type_is_tqp(ggml_type type) { + return type == GGML_TYPE_TQP5 || type == GGML_TYPE_TQP6 || type == GGML_TYPE_TQP7; +} + static constexpr __host__ __device__ bool ggml_cuda_type_is_tq(ggml_type type) { - return type == GGML_TYPE_TQ2_T || type == GGML_TYPE_TQK6 || type == GGML_TYPE_TQK7; + return type == GGML_TYPE_TQ2_T || type == GGML_TYPE_TQK6 || type == GGML_TYPE_TQK7 || ggml_cuda_type_is_tqp(type); } template struct tq_block_info; template <> struct tq_block_info { typedef block_tq2_t block_t; static constexpr int K = 8; }; template <> struct tq_block_info { typedef block_tqk6 block_t; static constexpr int K = 6; }; template <> struct tq_block_info { typedef block_tqk7 block_t; static constexpr int K = 7; }; +template <> struct tq_block_info { typedef block_tqp5 block_t; static constexpr int K = 5; }; +template <> struct tq_block_info { typedef block_tqp6 block_t; static constexpr int K = 6; }; +template <> struct tq_block_info { typedef block_tqp7 block_t; static constexpr int K = 7; }; +static_assert(QK_TQP == QK_TQK, "trellis block sizes differ"); // States of steps 4l..4l+3 of one block. qs is the block's qs array, 2-byte aligned // (every trellis block has an even size and qs sits right after the fp16 scale). @@ -88,6 +109,34 @@ static __device__ __forceinline__ void tq_states4(const uint8_t * qs, const int } } +// TQP: the 15-bit states of steps 8l..8l+7 (s[j] = step 8l + j). Step t's window starts at stream bit +// (63 - t)*K, so the lane's steps lie in the bits from S = (56 - 8l)*K: step 8l + j at S + (7 - j)*K. +// S = 8*(7 - l)*K, so r = S & 15 is 0 or 8 and r + 7*K + 15 <= 72: five u16 words (wrapped modulo the +// 4*K words of the stream) cover it, the fifth only when r + 7*K + 15 > 64 (never for K = 5). +template +static __device__ __forceinline__ void tqp_states8(const uint8_t * qs, const int l, uint32_t s[TQ_PAIRS_LANE]) { + constexpr int K = tq_block_info::K; + constexpr int nwords = 4*K; + const uint16_t * qs16 = (const uint16_t *) qs; + const int S = (56 - 8*l)*K; + const int wb = S >> 4; + const int r = S & 15; + const int wi1 = wb + 1 >= nwords ? wb + 1 - nwords : wb + 1; + const int wi2 = wb + 2 >= nwords ? wb + 2 - nwords : wb + 2; + const int wi3 = wb + 3 >= nwords ? wb + 3 - nwords : wb + 3; + uint64_t v = (uint64_t) qs16[wb] | ((uint64_t) qs16[wi1] << 16) | + ((uint64_t) qs16[wi2] << 32) | ((uint64_t) qs16[wi3] << 48); + v >>= r; + if constexpr (8 + 7*K + TQ_STATE_BITS > 64) { + const int wi4 = wb + 4 >= nwords ? wb + 4 - nwords : wb + 4; + v |= r ? (uint64_t) qs16[wi4] << (64 - r) : 0; + } +#pragma unroll + for (int j = 0; j < TQ_PAIRS_LANE; ++j) { + s[j] = (uint32_t) (v >> ((7 - j)*K)) & TQ_STATE_MASK; + } +} + // MMQ (batched) path: lane l of block x as 16 int8 values (4 packed ints, weight order) with one float scale per // lane, i.e. weights ~= scale * q. The int8 codebook (one global scale) is already int8: no decode to float. template @@ -95,13 +144,24 @@ static __device__ __forceinline__ float tq_lane_to_q8(const void * x, const int typedef typename tq_block_info::block_t block_t; const block_t * b = (const block_t *) x; - uint32_t s[TQ_STEPS_LANE]; - tq_states4(b->qs, l, s); - + if constexpr (ggml_cuda_type_is_tqp(type)) { + // 8 pairs of 2 weights, in weight order: two pairs per int, as the TQ path packs two steps + uint32_t s[TQ_PAIRS_LANE]; + tqp_states8(b->qs, l, s); #pragma unroll - for (int i = 0; i < TQ_STEPS_LANE; ++i) { - const uint32_t xh = (s[i] & TQ_STATE_MASK) * 0x9e3779b1u; - q[i] = (int) ((uint32_t) tq_lut_i8[xh >> 21] | ((uint32_t) tq_lut_i8[(xh >> 10) & 2047u] << 16)); + for (int i = 0; i < TQ_STEPS_LANE; ++i) { + const uint32_t p0 = tq_lut_i8[((s[2 * i] & TQ_STATE_MASK) * 0x9e3779b1u) >> 21]; + const uint32_t p1 = tq_lut_i8[((s[2 * i + 1] & TQ_STATE_MASK) * 0x9e3779b1u) >> 21]; + q[i] = (int) (p0 | (p1 << 16)); + } + } else { + uint32_t s[TQ_STEPS_LANE]; + tq_states4(b->qs, l, s); +#pragma unroll + for (int i = 0; i < TQ_STEPS_LANE; ++i) { + const uint32_t xh = (s[i] & TQ_STATE_MASK) * 0x9e3779b1u; + q[i] = (int) ((uint32_t) tq_lut_i8[xh >> 21] | ((uint32_t) tq_lut_i8[(xh >> 10) & 2047u] << 16)); + } } // Strata's block structs hold the fp16 scale as uint16_t (the fork's as ggml_half): reinterpret, do not convert return __half2float(__ushort_as_half(b->d)) * TQ_LUT_I8_SCALE; diff --git a/third_party/ggml/ggml-common-tq.h b/third_party/ggml/ggml-common-tq.h index 1ce3773b..6e00a97d 100644 --- a/third_party/ggml/ggml-common-tq.h +++ b/third_party/ggml/ggml-common-tq.h @@ -32,6 +32,26 @@ typedef struct { } block_tqk7; static_assert(sizeof(block_tqk7) == 30, "wrong tqk7 block size/padding"); +// Gyro-L pair trellis (agentionai llama.cpp release/gyro-l): 64 steps of K bits, 2 weights per step +// (the TQ2_T codebook's first pair). Defined in the fork's ggml-common.h; repeated here so Strata's +// device and CPU kernels compile without GGML_COMMON_IMPL. +#define QK_TQP 128 +typedef struct { + uint16_t d; + uint8_t qs[8*5]; // 64 steps x 5 bits +} block_tqp5; +static_assert(sizeof(block_tqp5) == 42, "wrong tqp5 block size/padding"); +typedef struct { + uint16_t d; + uint8_t qs[8*6]; // 64 steps x 6 bits +} block_tqp6; +static_assert(sizeof(block_tqp6) == 50, "wrong tqp6 block size/padding"); +typedef struct { + uint16_t d; + uint8_t qs[8*7]; // 64 steps x 7 bits +} block_tqp7; +static_assert(sizeof(block_tqp7) == 58, "wrong tqp7 block size/padding"); + // device memory in a CUDA / HIP translation unit, ordinary host data everywhere else #if defined(__CUDACC__) || defined(__HIPCC__) #define STRATA_TQ_LUT_SPACE __device__ diff --git a/tools/gguf_reader.py b/tools/gguf_reader.py index 4da9d2d2..1c65a288 100644 --- a/tools/gguf_reader.py +++ b/tools/gguf_reader.py @@ -27,6 +27,8 @@ GGML_TYPES: dict[int, str] = { 39: "MXFP4", 40: "NVFP4", 41: "Q4_0_8_8", 42: "Q2_0", # Agention's llama.cpp fork (agentionai/llama.cpp): the APR trellis expert types (128-value blocks) 143: "PTQ1_0", 144: "TQ2_T", 145: "TQK6", 146: "TQK7", + # pair-trellis family added for Gyro-L (agentionai llama.cpp release/gyro-l) + 147: "TQP5", 148: "TQP6", 149: "TQP7", } # GGUF metadata value type ids @@ -49,6 +51,8 @@ BLOCK_GEOMETRY: dict[str, tuple[int, int]] = { "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), + # TQP5/6/7: fp16 scale + 8*K bytes of a 64-step x K-bit circular stream (K = 5/6/7) + "TQP5": (128, 42), "TQP6": (128, 50), "TQP7": (128, 58), } diff --git a/tools/test_gguf_tqp.py b/tools/test_gguf_tqp.py new file mode 100644 index 00000000..338eda96 --- /dev/null +++ b/tools/test_gguf_tqp.py @@ -0,0 +1,34 @@ +"""The pair-trellis ggml types (147-149: TQP5 / TQP6 / TQP7, Gyro-L): the Python reader and the C++ reader agree on +their names and block sizes, and the sizes follow the layout (fp16 scale + 8 * K bytes per 128 values).""" + +import re +import sys +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "tools")) + +import gguf_reader # noqa: E402 + + +class PairTrellisTypes(unittest.TestCase): + def test_python_reader_names_and_sizes(self): + for t, name, k in ((147, "TQP5", 5), (148, "TQP6", 6), (149, "TQP7", 7)): + self.assertEqual(gguf_reader.GGML_TYPES[t], name) + self.assertEqual(gguf_reader.BLOCK_GEOMETRY[name], (128, 2 + 8 * k)) + + def test_older_trellis_types_unchanged(self): + self.assertEqual(gguf_reader.BLOCK_GEOMETRY["TQ2_T"], (128, 34)) + self.assertEqual(gguf_reader.BLOCK_GEOMETRY["TQK6"], (128, 26)) + self.assertEqual(gguf_reader.BLOCK_GEOMETRY["TQK7"], (128, 30)) + + def test_cpp_reader_has_the_same_numbers(self): + text = (ROOT / "include" / "strata" / "artifact" / "gguf_reader.hpp").read_text(encoding="utf-8") + for t, name, k in ((147, "TQP5", 5), (148, "TQP6", 6), (149, "TQP7", 7)): + self.assertRegex(text, r"case %d:\s*return \"%s\";" % (t, name)) + self.assertRegex(text, r"case %d:\s*elems = 128;\s*bytes = %d;" % (t, 2 + 8 * k)) + + +if __name__ == "__main__": + unittest.main()