diff --git a/CMakeLists.txt b/CMakeLists.txt index f5253583..acc55093 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) @@ -1192,7 +1203,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 @@ -1732,8 +1743,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). @@ -1845,11 +1861,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 @@ -1966,6 +1984,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 @@ -1982,6 +2005,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) @@ -2037,7 +2061,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 @@ -2155,6 +2179,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_EXPERIMENTAL_SM60 AND NOT STRATA_HIP_GFX906) @@ -2264,7 +2291,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() @@ -2289,6 +2316,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) @@ -2394,7 +2424,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}) @@ -2468,6 +2502,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) @@ -2607,7 +2648,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 ca637e5f..26447c83 100644 --- a/cmake/ggml_tq_overlay_test.cmake +++ b/cmake/ggml_tq_overlay_test.cmake @@ -26,7 +26,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}") @@ -37,6 +38,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/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 286c0b2c..d3cda7ba 100644 --- a/setup.py +++ b/setup.py @@ -1336,7 +1336,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) @@ -1367,6 +1371,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| @@ -1374,6 +1383,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: @@ -1458,7 +1469,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 @@ -9082,7 +9107,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) @@ -9092,7 +9117,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 b1e559b7..1b5b1155 100644 --- a/src/core/expert_source.cpp +++ b/src/core/expert_source.cpp @@ -3185,6 +3185,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()); @@ -3210,6 +3214,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 b9a5f55c..0aa4c505 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 @@ -1689,10 +1685,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]); @@ -1725,10 +1722,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 42bdc889..920b55c1 100644 --- a/src/core/verify.cpp +++ b/src/core/verify.cpp @@ -79,7 +79,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. @@ -609,7 +609,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_}) { @@ -2899,7 +2899,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; @@ -2977,7 +2977,7 @@ bool Verifier::adapt_gate_wait(int32_t layer, std::string& err) { progress_at("verify window: waiting for the adaptive tier's copies of layer", layer); announced = true; } - _mm_pause(); + strata_cpu_pause(); if ((++spins & 1023u) != 0) continue; if (Clock::now() - t0 > std::chrono::seconds(20)) { // the release drains the window's own streams, never the tier's copy stream: say so, and hand out the @@ -3115,7 +3115,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; @@ -3190,7 +3190,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 { @@ -3206,7 +3206,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)) { @@ -3242,6 +3242,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); @@ -3262,7 +3263,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; @@ -3288,7 +3289,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; @@ -3544,7 +3545,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; } @@ -4172,7 +4173,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_); @@ -4196,7 +4197,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)) { @@ -4213,6 +4214,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); @@ -4224,7 +4226,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; @@ -4349,7 +4351,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 (;;) { @@ -4369,7 +4371,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; @@ -4671,7 +4673,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; }; @@ -4722,7 +4724,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; @@ -4796,7 +4798,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 ee1730b6..bca01c0b 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 d7355721..a43b5c6e 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 @@ -1415,7 +1415,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/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 50f0a839..8262f566 100644 --- a/src/kernels/cuda/iq_kernels.cu +++ b/src/kernels/cuda/iq_kernels.cu @@ -617,6 +617,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) { @@ -717,6 +721,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; @@ -761,21 +830,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); } }; @@ -1391,6 +1467,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, @@ -3171,6 +3264,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 @@ -3199,6 +3317,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; } } @@ -3240,9 +3361,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) @@ -5209,6 +5330,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 e832cffe..8a53f122 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) @@ -323,6 +327,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 @@ -510,6 +517,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/src/program/generate.cpp b/src/program/generate.cpp index 186b23f3..b6b85275 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -3908,6 +3908,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. @@ -3945,6 +3946,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() 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 de30d3b9..5a5a05eb 100644 --- a/tools/gguf_reader.py +++ b/tools/gguf_reader.py @@ -28,6 +28,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 @@ -50,6 +52,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()