diff --git a/.gitignore b/.gitignore index 16f92c6b471..eca1ac7dff4 100644 --- a/.gitignore +++ b/.gitignore @@ -34,6 +34,8 @@ td.log /run-*.sh /strata-*.json /strata-*.log +# the web app's settings shared with other apps, written next to the config the server runs (examples/) +*.shared-settings.json /build-strata.bat # image encoder (setup.py --vision): build output and the generated build script diff --git a/CMakeLists.txt b/CMakeLists.txt index 74ae6916a0c..75364617d9a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -4,18 +4,12 @@ # `strata-pack` are added when their sources exist - a target with no sources would make the build graph # claim something the tree does not have, which is worse than an absent target. # -# CUDA targets arrive in Phase 2. The kernels need sm_80 or newer (tf32 mma, bf16 math) and the guard -# below refuses anything older: silently building for a pre-Ampere arch would produce a binary that runs -# and is wrong about its own performance, which is the worst of both outcomes. +# CUDA targets arrive in Phase 2 and this build REJECTS any architecture other than 120: the engine is +# written against sm_120 (RTX 5000 / Blackwell) and silently building for an older arch would produce a +# binary that runs and is wrong about its own performance, which is the worst of both outcomes. cmake_minimum_required(VERSION 3.24) project(strata VERSION 0.1.27 LANGUAGES CXX) add_compile_definitions(STRATA_VERSION="${PROJECT_VERSION}") -# issue #31: release builds keep a PDB beside the exe (not shipped in it), so a minidump from a stalled engine can -# be read back to the exact source line -if(MSVC AND CMAKE_BUILD_TYPE STREQUAL "Release") - add_compile_options($<$:/Zi>) - add_link_options(/DEBUG /OPT:REF /OPT:ICF) -endif() # The prebuilt release (setup.py downloads it, so users need no compiler): ggml's CPU backend for any AVX2 CPU # instead of this machine's (Strata's own AVX-512 kernels are per-file and chosen at run time either way), and @@ -33,17 +27,6 @@ if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) endif() option(STRATA_ENABLE_CUDA "Build the CUDA targets (Phase 2+)" OFF) -option(STRATA_ENABLE_HIP "Build the HIP targets for gfx1100 wave32 hardware" OFF) -option(STRATA_PREFILL_MMQ "Build the optional GGML MMQ prefill path for HIP" OFF) -if(STRATA_PREFILL_MMQ AND NOT STRATA_ENABLE_HIP) - message(FATAL_ERROR "STRATA_PREFILL_MMQ requires STRATA_ENABLE_HIP=ON") -endif() -if(STRATA_ENABLE_CUDA AND STRATA_ENABLE_HIP) - message(FATAL_ERROR "STRATA_ENABLE_CUDA and STRATA_ENABLE_HIP are separate backends; enable only one") -endif() -if(STRATA_ENABLE_HIP) - include(cmake/hip_backend.cmake) -endif() # The test suite and the micro benchmarks live in tests/ and bench/micro/, which the published source leaves out. set(_strata_tests_default OFF) if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/tests/CMakeLists.txt AND EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/bench/micro) @@ -62,26 +45,18 @@ if(STRATA_ENABLE_CUDA) endif() foreach(_a IN LISTS CMAKE_CUDA_ARCHITECTURES) string(REGEX REPLACE "-.*$" "" _base "${_a}") # 120-real / 120-virtual are still 120 - # Turing port: developed and measured on sm_120 (RTX 50); the tf32 mma in the QSA scorer has a - # portable fp32-FMA fallback, so RTX 20 (75), 30 (86), RTX 40 (89) and RTX 50 (120) build too. - # Pre-sm_75 is refused. + # Plan v0.3 P6: developed and measured on sm_120 (RTX 50); the kernels need sm_80 or newer (tf32 mma in + # the attention scorer, bf16 math), so RTX 30 (86) and RTX 40 (89) build too. Older cards are refused. if(_base STREQUAL "native" OR _base STREQUAL "all" OR _base STREQUAL "all-major") continue() endif() - if(_base LESS 75) + if(_base LESS 80) message(FATAL_ERROR - "Strata needs an NVIDIA GPU of compute capability 7.5 or newer (RTX 20 / 30 / 40 / 50 series), but " + "Strata needs an NVIDIA GPU of compute capability 8.0 or newer (RTX 30 / 40 / 50 series), but " "CMAKE_CUDA_ARCHITECTURES is '${CMAKE_CUDA_ARCHITECTURES}'.") endif() endforeach() message(STATUS "Strata: CUDA enabled, arch ${CMAKE_CUDA_ARCHITECTURES}") - set(_strata_gpu_runtime_target CUDA::cudart) - set(_strata_gpu_blas_target CUDA::cublas) - set(_strata_gpu_include_directories ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) -elseif(STRATA_ENABLE_HIP) - set(_strata_gpu_runtime_target strata_hip_runtime) - set(_strata_gpu_blas_target roc::hipblas) - set(_strata_gpu_include_directories ${STRATA_HIP_COMPAT_INCLUDE_DIR}) endif() # Warnings live on an INTERFACE target the project's own code links, NOT in add_compile_options: applied @@ -98,13 +73,14 @@ target_compile_options(strata_warnings INTERFACE $<$,$>>:-Wall> $<$,$>>:-Wextra> $<$,$>:-Xcompiler=/W4> - $<$:-Wall> - $<$:-Wextra>) + # C4211 comes from the host stubs nvcc generates for kernels with internal linkage (anonymous namespaces). + $<$,$>:-Xcompiler=/wd4211>) +# fopen, sscanf and getenv are standard C; MSVC deprecates them in favour of its own _s variants. +target_compile_definitions(strata_warnings INTERFACE $<$:_CRT_SECURE_NO_WARNINGS>) if(STRATA_WERROR) target_compile_options(strata_warnings INTERFACE $<$,$>:/WX> - $<$,$>>:-Werror> - $<$:-Werror>) + $<$,$>>:-Werror>) endif() # ---- strata_artifact ------------------------------------------------------- @@ -122,12 +98,6 @@ add_executable(strata-dequant src/artifact/dequant.cpp) target_link_libraries(strata-dequant PRIVATE strata_artifact) # this file is now nothing but the CLI entry point, and its inherited guard spells that as a "selftest" target_compile_definitions(strata-dequant PRIVATE STRATA_DEQUANT_SELFTEST) -if(STRATA_BUILD_TESTS) - # the reader refuses a duplicate tensor name at open, naming it (a synthetic GGUF: no model, no GPU) - add_executable(gguf_reader_test src/artifact/gguf_reader_test.cpp) - target_link_libraries(gguf_reader_test PRIVATE strata_artifact) - add_test(NAME gguf_reader_test COMMAND gguf_reader_test) -endif() # ---- strata_plan ---------------------------------------------------------- # Also header-only. A "planner" that hides state would be one whose decision is not inspectable, and the @@ -139,16 +109,17 @@ target_link_libraries(strata_plan INTERFACE strata_warnings) add_executable(strata-plan src/plan/plan_main.cpp) target_link_libraries(strata-plan PRIVATE strata_plan) -# ---- strata_core (GPU runtime) --------------------------------------------- -# Only built when a GPU backend is on, so the Phase 1 targets keep working without a toolkit. -if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) +# ---- strata_core (CUDA, Phase 2) ------------------------------------------ +# Only built when CUDA is on, so the Phase 1 targets keep working on a machine with no toolkit. The arch +# guard above has already refused anything but sm_120 by the time this runs. +if(STRATA_ENABLE_CUDA) add_library(strata_core STATIC src/core/device.cu src/core/pinned.cu src/platform/memory.cpp src/core/graph.cpp src/core/weights.cpp src/core/layout.cpp) target_include_directories(strata_core PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include) - # The selected runtime is PUBLIC because graph.cpp is a HOST translation unit that calls cudaGraphLaunch and + # CUDA::cudart is PUBLIC because graph.cpp is a HOST translation unit that calls cudaGraphLaunch and # cudaEventQuery: the .cu sources get cudart through the CUDA language rules, a .cpp does not, and the # linker error ("unresolved external symbol cudaGraphInstantiate") names the function rather than the cause. - target_link_libraries(strata_core PUBLIC strata_plan strata_warnings ${_strata_gpu_runtime_target}) + target_link_libraries(strata_core PUBLIC strata_plan strata_warnings CUDA::cudart) add_executable(strata-device src/core/device_main.cpp) target_link_libraries(strata-device PRIVATE strata_core) @@ -157,13 +128,13 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) target_link_libraries(strata-load PRIVATE strata_core strata_kernels) add_executable(strata-concurrent src/core/concurrent_main.cpp) - target_include_directories(strata-concurrent PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(strata-concurrent PRIVATE strata_core strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(strata-concurrent PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(strata-concurrent PRIVATE strata_core strata_kernels CUDA::cudart) add_executable(strata-overlap src/core/overlap_main.cpp) # a .cpp that calls cudaMemcpyAsync itself needs the toolkit headers and cudart explicitly - target_include_directories(strata-overlap PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(strata-overlap PRIVATE strata_core strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(strata-overlap PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(strata-overlap PRIVATE strata_core strata_kernels CUDA::cudart) # The load test needs the 34 GB artifact, so it is registered ONLY where that artifact exists. A test that # silently skips is a test that stops running; an absent registration is visible in the configure output. @@ -211,12 +182,8 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) src/kernels/cuda/native_qsa_score.cu src/kernels/cuda/ple.cu src/kernels/cuda/kv_q8.cu - src/kernels/cuda/kv_stream.cu - src/kernels/cuda/kv_q4.cu src/kernels/cuda/fused_gr.cu - src/kernels/cuda/cvec.cu src/kernels/cuda/qsa_decode_attn.cu - src/kernels/cuda/qsa_prompt_attn.cu src/kernels/cuda/fused_gdn.cu src/kernels/cuda/dequant_bf16.cu src/kernels/cuda/qsa_select.cu @@ -226,48 +193,33 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) src/ngram/ple_reader.cpp src/platform/direct_file.cpp) target_link_libraries(strata_kernels PUBLIC strata_core strata_artifact) - # perf-review D-1: the tensor-core prompt attention vs the FP32 kernel and FP64 (GPU, synthetic, no model); with - # the tests, or alone with -DSTRATA_PARITY_PROMPT_ATTN=ON - if(STRATA_BUILD_TESTS OR STRATA_PARITY_PROMPT_ATTN) - add_executable(qsa_prompt_attn_parity src/kernels/qsa_prompt_attn_parity.cpp) - target_include_directories(qsa_prompt_attn_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) - target_link_libraries(qsa_prompt_attn_parity PRIVATE strata_kernels CUDA::cudart) - endif() # plan v0.3 P6: llama.cpp's block layouts and codebook grids for the i-quant formats (MIT, third_party/ggml) target_include_directories(strata_kernels PRIVATE third_party/ggml) # Match the pinned llama.cpp CUDA arithmetic contract for this adapter only. - if(STRATA_ENABLE_CUDA) - set_source_files_properties(src/kernels/cuda/native_mmvq.cu src/kernels/cuda/native_bf16.cu - src/kernels/cuda/native_gr_norm.cu - src/kernels/cuda/native_moe.cu - src/kernels/cuda/native_gdn.cu - src/kernels/cuda/native_gr_postops.cu - src/kernels/cuda/native_router.cu - src/kernels/cuda/native_gdn_preprocess.cu - src/kernels/cuda/native_qsa.cu - src/kernels/cuda/native_rope.cu - src/kernels/cuda/native_ple_postops.cu - src/kernels/cuda/native_flash_attn.cu - src/kernels/cuda/native_qsa_indexer.cu - src/kernels/cuda/native_qsa_score.cu - PROPERTIES COMPILE_OPTIONS "--use_fast_math") - elseif(STRATA_ENABLE_HIP) - # Keep Q8_K's multiply and magic-bias rounding as separate FP32 operations on HIP; contraction to an - # FMA changes tie results. The embedding gather also relies on separately rounded operations. - set_source_files_properties(src/kernels/cuda/elementwise.cu src/kernels/cuda/quantize_act.cu - PROPERTIES COMPILE_OPTIONS "-ffp-contract=off") - endif() + set_source_files_properties(src/kernels/cuda/native_mmvq.cu src/kernels/cuda/native_bf16.cu + src/kernels/cuda/native_gr_norm.cu + src/kernels/cuda/native_moe.cu + src/kernels/cuda/native_gdn.cu + src/kernels/cuda/native_gr_postops.cu + src/kernels/cuda/native_router.cu + src/kernels/cuda/native_gdn_preprocess.cu + src/kernels/cuda/native_qsa.cu + src/kernels/cuda/native_rope.cu + src/kernels/cuda/native_ple_postops.cu + src/kernels/cuda/native_flash_attn.cu + src/kernels/cuda/native_qsa_indexer.cu + src/kernels/cuda/native_qsa_score.cu + PROPERTIES COMPILE_OPTIONS "--use_fast_math") set(STRATA_ORACLE_SOURCE_DIR "" CACHE PATH "Pinned llama.cpp source for independent CUDA parity tests") set(STRATA_ORACLE_BUILD_DIR "" CACHE PATH "Matching existing shared llama.cpp build") - # This oracle is a CUDA ggml build; do not pull CUDA libraries into a HIP build. - if(STRATA_ENABLE_CUDA AND STRATA_ORACLE_SOURCE_DIR AND STRATA_ORACLE_BUILD_DIR) + if(STRATA_ORACLE_SOURCE_DIR AND STRATA_ORACLE_BUILD_DIR) foreach(parity IN ITEMS native_mmvq native_bf16 native_shared native_gr_norm native_moe native_gdn native_gr_postops native_router native_gdn_preprocess native_gdn_pipeline native_qsa native_rope native_ple_postops native_flash_attn native_qsa_indexer native_qsa_score) add_executable(${parity}_parity bench/micro/${parity}_parity.cpp) target_include_directories(${parity}_parity PRIVATE - "${STRATA_ORACLE_SOURCE_DIR}/ggml/include" ${_strata_gpu_include_directories}) + "${STRATA_ORACLE_SOURCE_DIR}/ggml/include" ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) target_compile_definitions(${parity}_parity PRIVATE GGML_SHARED GGML_BACKEND_SHARED) - target_link_libraries(${parity}_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_link_libraries(${parity}_parity PRIVATE strata_kernels CUDA::cudart) foreach(component IN ITEMS ggml ggml-base ggml-cpu ggml-cuda) string(REPLACE "-" "_" variable "STRATA_ORACLE_${component}") find_library(${variable} NAMES ${component} PATHS @@ -288,166 +240,105 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) # already links `strata_core` - putting `gdn_layer` in core would make the two mutually dependent. A third # library is the honest shape of the dependency graph: core knows about devices and weights, kernels know # about arithmetic, and this one knows the ORDER the arithmetic happens in. - add_library(strata_engine STATIC src/core/layer.cpp src/core/session.cpp src/core/expert_source.cpp src/core/remote_experts.cpp - src/core/expert_cache.cpp src/core/native_head.cpp src/core/native_dense.cpp src/core/verify.cpp src/core/mtp.cpp) + add_library(strata_engine STATIC src/core/layer.cpp src/core/session.cpp src/core/expert_source.cpp + src/core/expert_cache.cpp src/core/native_head.cpp src/core/native_dense.cpp src/core/verify.cpp src/core/mtp.cpp + src/core/prompt_cache.cpp src/core/adaptive_tier.cpp src/core/second_gpu.cpp + src/core/split_head.cpp) target_include_directories(strata_engine PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include) - target_link_libraries(strata_engine PUBLIC strata_core strata_kernels strata_kernels_cpu ${_strata_gpu_runtime_target}) - if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/core/file_expert_source_test.cpp") - enable_testing() - add_executable(file_expert_source_test tests/core/file_expert_source_test.cpp) - target_link_libraries(file_expert_source_test PRIVATE strata_engine) - add_test(NAME file_expert_source_test COMMAND file_expert_source_test) - set_tests_properties(file_expert_source_test PROPERTIES TIMEOUT 60) - endif() + target_link_libraries(strata_engine PUBLIC strata_core strata_kernels strata_kernels_cpu CUDA::cudart) # `strata generate` (P2.S6). The FIRST user-facing program: everything under it is a component, and this is # what composes them into a token. No tokenizer - the prompt arrives as ids - because Gate C1 compares logits # against llama.cpp on the SAME ids and a tokenizer on one side of that comparison is a second variable. # plan v0.3 P5: batched prompt processing (tensor-core GEMMs through cuBLAS). - add_library(strata_prefill STATIC src/prefill/gemm.cu src/prefill/kernels.cu src/prefill/prefill.cpp) - target_include_directories(strata_prefill PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include ${_strata_gpu_include_directories}) - target_link_libraries(strata_prefill PUBLIC strata_engine ${_strata_gpu_blas_target} ${_strata_gpu_runtime_target}) - if(STRATA_ENABLE_HIP AND STRATA_HIPBLASLT_AVAILABLE) - target_compile_definitions(strata_prefill PRIVATE STRATA_HIPBLASLT_AVAILABLE=1) - target_link_libraries(strata_prefill PRIVATE roc::hipblaslt) - - add_executable(tune_hipblaslt tools/hip/tune_hipblaslt.cpp) - target_link_libraries(tune_hipblaslt PRIVATE hip::host roc::hipblas roc::hipblaslt) - endif() + add_library(strata_prefill STATIC src/prefill/gemm.cu src/prefill/kernels.cu src/prefill/experts.cpp + src/prefill/prefill.cpp) + target_include_directories(strata_prefill PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(strata_prefill PUBLIC strata_engine CUDA::cublas CUDA::cudart) add_executable(strata src/program/generate.cpp) - target_include_directories(strata PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(strata PRIVATE strata_engine strata_prefill strata_spec ${_strata_gpu_runtime_target}) - - if(STRATA_ENABLE_HIP) - enable_testing() - set(_strata_hip_tests) - add_executable(hip_expert_cache_staging tests/hip/expert_cache_staging.cpp) - target_link_libraries(hip_expert_cache_staging PRIVATE strata_engine) - add_test(NAME hip_expert_cache_staging COMMAND hip_expert_cache_staging) - list(APPEND _strata_hip_tests hip_expert_cache_staging) - if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/hip/intrinsics.cpp") - add_executable(hip_intrinsics tests/hip/intrinsics.cpp) - target_link_libraries(hip_intrinsics PRIVATE strata_hip_runtime) - add_test(NAME hip_intrinsics COMMAND hip_intrinsics) - list(APPEND _strata_hip_tests hip_intrinsics) - endif() - if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/hip/handoff.cpp") - add_executable(hip_handoff tests/hip/handoff.cpp) - target_link_libraries(hip_handoff PRIVATE strata_kernels strata_hip_runtime) - add_test(NAME hip_handoff COMMAND hip_handoff) - list(APPEND _strata_hip_tests hip_handoff) - endif() - if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/hip/native_qsa_score.cpp") - add_executable(hip_native_qsa_score tests/hip/native_qsa_score.cpp) - target_link_libraries(hip_native_qsa_score PRIVATE strata_kernels strata_hip_runtime) - add_test(NAME hip_native_qsa_score COMMAND hip_native_qsa_score) - list(APPEND _strata_hip_tests hip_native_qsa_score) - endif() - if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/hip/prefill_native_batch.cpp") - add_executable(hip_prefill_native_batch tests/hip/prefill_native_batch.cpp) - target_link_libraries(hip_prefill_native_batch PRIVATE strata_kernels strata_hip_runtime) - add_test(NAME hip_prefill_native_batch COMMAND hip_prefill_native_batch) - list(APPEND _strata_hip_tests hip_prefill_native_batch) - endif() - if(STRATA_HIPBLASLT_AVAILABLE AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/hip/prefill_hipblaslt_gemm.cpp") - add_executable(hip_prefill_hipblaslt_gemm tests/hip/prefill_hipblaslt_gemm.cpp) - target_include_directories(hip_prefill_hipblaslt_gemm PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src/prefill") - target_link_libraries(hip_prefill_hipblaslt_gemm PRIVATE strata_prefill strata_hip_runtime roc::hipblaslt roc::hipblas) - add_test(NAME hip_prefill_hipblaslt_gemm COMMAND hip_prefill_hipblaslt_gemm) - set_tests_properties(hip_prefill_hipblaslt_gemm PROPERTIES TIMEOUT 120 SKIP_RETURN_CODE 77) - list(APPEND _strata_hip_tests hip_prefill_hipblaslt_gemm) - endif() - if(_strata_hip_tests) - set_tests_properties(${_strata_hip_tests} PROPERTIES TIMEOUT 60) - endif() - if(TARGET hip_prefill_native_batch) - set_tests_properties(hip_prefill_native_batch PROPERTIES TIMEOUT 120) - endif() - if(TARGET hip_prefill_hipblaslt_gemm) - set_tests_properties(hip_prefill_hipblaslt_gemm PROPERTIES TIMEOUT 120 SKIP_RETURN_CODE 77) - endif() - endif() + target_include_directories(strata PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(strata PRIVATE strata_engine strata_prefill strata_spec CUDA::cudart) add_executable(dequant_s2_parity src/kernels/dequant_s2_parity.cpp) # The tests are HOST C++ files that call cudaMalloc/cudaMemcpy themselves, so they need the toolkit headers # and cudart explicitly: enabling the CUDA language gives those to .cu sources automatically, but a .cpp # target is just a C++ target and gets neither. - target_include_directories(dequant_s2_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(dequant_s2_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(dequant_s2_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(dequant_s2_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME dequant_s2_parity COMMAND dequant_s2_parity --selftest) add_executable(s2_gemv_parity src/kernels/s2_gemv_parity.cpp) - target_include_directories(s2_gemv_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(s2_gemv_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(s2_gemv_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(s2_gemv_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME s2_gemv_parity COMMAND s2_gemv_parity --selftest) add_executable(iq_parity src/kernels/iq_parity.cpp) - target_include_directories(iq_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(iq_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(iq_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(iq_parity PRIVATE strata_kernels CUDA::cudart) add_executable(shared_expert_parity src/kernels/shared_expert_parity.cpp) - target_include_directories(shared_expert_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(shared_expert_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(shared_expert_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(shared_expert_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME shared_expert_parity COMMAND shared_expert_parity --selftest) add_executable(gr_parity src/kernels/gr_parity.cpp) - target_include_directories(gr_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(gr_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(gr_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(gr_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME gr_parity COMMAND gr_parity --selftest) add_executable(gdn_parity src/kernels/gdn_parity.cpp) - target_include_directories(gdn_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(gdn_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(gdn_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(gdn_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME gdn_parity COMMAND gdn_parity --selftest) add_executable(s2_gemv_q8_parity src/kernels/s2_gemv_q8_parity.cpp) - target_include_directories(s2_gemv_q8_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(s2_gemv_q8_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(s2_gemv_q8_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(s2_gemv_q8_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME s2_gemv_q8_parity COMMAND s2_gemv_q8_parity --selftest) add_executable(sampler_parity src/kernels/sampler_parity.cpp) - target_include_directories(sampler_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(sampler_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(sampler_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(sampler_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME sampler_parity COMMAND sampler_parity --selftest) add_executable(rope_parity src/kernels/rope_parity.cpp) - target_include_directories(rope_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(rope_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(rope_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(rope_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME rope_parity COMMAND rope_parity --selftest) add_executable(quantize_act_parity src/kernels/quantize_act_parity.cpp) - target_include_directories(quantize_act_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(quantize_act_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(quantize_act_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(quantize_act_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME quantize_act_parity COMMAND quantize_act_parity --selftest) add_executable(router_top10_parity src/kernels/router_top10_parity.cpp) - target_include_directories(router_top10_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(router_top10_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(router_top10_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(router_top10_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME router_top10_parity COMMAND router_top10_parity --selftest) add_executable(s_gemv_parity src/kernels/s_gemv_parity.cpp) - target_include_directories(s_gemv_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(s_gemv_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(s_gemv_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(s_gemv_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME s_gemv_parity COMMAND s_gemv_parity --selftest) add_executable(elementwise_parity src/kernels/elementwise_parity.cpp) - target_include_directories(elementwise_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(elementwise_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(elementwise_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(elementwise_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME elementwise_parity COMMAND elementwise_parity --selftest) add_executable(bf16_gemv_parity src/kernels/bf16_gemv_parity.cpp) - target_include_directories(bf16_gemv_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(bf16_gemv_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(bf16_gemv_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(bf16_gemv_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME bf16_gemv_parity COMMAND bf16_gemv_parity --selftest) add_executable(s_gemv_q8k_parity src/kernels/s_gemv_q8k_parity.cpp) - target_include_directories(s_gemv_q8k_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(s_gemv_q8k_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(s_gemv_q8k_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(s_gemv_q8k_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME s_gemv_q8k_parity COMMAND s_gemv_q8k_parity --selftest) add_executable(qsa_parity src/kernels/qsa_parity.cpp) - target_include_directories(qsa_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(qsa_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(qsa_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(qsa_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME qsa_parity COMMAND qsa_parity --selftest) # ---- P2.S4: the PLE n-gram hash, the IQ4_NL table read, and the PLE block ---- @@ -459,8 +350,8 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) # INSIDE the test. A test that passes by finding nothing to compare is the failure mode this project keeps # re-learning. add_executable(ple_parity src/kernels/ple_parity.cpp) - target_include_directories(ple_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(ple_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(ple_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(ple_parity PRIVATE strata_kernels CUDA::cudart) set(STRATA_PLE_FIXTURE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/bench/micro" CACHE PATH "Directory containing required ple_in.bin and ple_out.bin parity fixtures") add_test(NAME ple_parity COMMAND ple_parity --selftest @@ -469,86 +360,72 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) # Plan v0.3 P7: INT8 KV storage vs a host reference and vs the FP16 path (GPU, synthetic, no model). add_executable(kv_q8_parity src/kernels/kv_q8_parity.cpp) - target_include_directories(kv_q8_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(kv_q8_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(kv_q8_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(kv_q8_parity PRIVATE strata_kernels CUDA::cudart) add_test(NAME kv_q8_parity COMMAND kv_q8_parity) - # KV streaming (docs/kv-streaming-design.md): streamed vs resident attention, bitwise (GPU, synthetic, no model). - add_executable(kv_stream_parity src/kernels/kv_stream_parity.cpp) - target_include_directories(kv_stream_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(kv_stream_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) - add_test(NAME kv_stream_parity COMMAND kv_stream_parity) - - # PR #120: the hybrid K8V4 KV (INT8 K + rotated Q4_0 V) vs host references (GPU, synthetic, no model). - add_executable(kv_hybrid_parity src/kernels/kv_hybrid_parity.cpp) - target_include_directories(kv_hybrid_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(kv_hybrid_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) - add_test(NAME kv_hybrid_parity COMMAND kv_hybrid_parity) - # PR #21: Q4_0 KV with Walsh-Hadamard rotation vs a host reference (GPU, synthetic, no model). - add_executable(kv_q4_parity src/kernels/kv_q4_parity.cpp) - target_include_directories(kv_q4_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(kv_q4_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) - add_test(NAME kv_q4_parity COMMAND kv_q4_parity) - add_executable(cvec_parity src/kernels/cvec_parity.cpp) - target_include_directories(cvec_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(cvec_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) - add_test(NAME cvec_parity COMMAND cvec_parity) + # The verify window's batched kernels vs the per-token kernel sequences they replace, bitwise (synthetic data). + add_executable(verify_parity src/kernels/verify_parity.cpp) + target_include_directories(verify_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(verify_parity PRIVATE strata_kernels strata_prefill strata_kernels_cpu CUDA::cudart) + add_test(NAME verify_parity COMMAND verify_parity --selftest) + + # The QSA block scores and top-k against the kernels they replaced, kept in the test: bitwise (synthetic data). + add_executable(qsa_select_parity src/kernels/qsa_select_parity.cu) + target_link_libraries(qsa_select_parity PRIVATE strata_kernels CUDA::cudart) + add_test(NAME qsa_select_parity COMMAND qsa_select_parity) + # The MMVQ of 2..8 tokens from interleaved activations against native_mmvq: bitwise (synthetic matrices). + add_executable(mmvq_il_parity src/kernels/mmvq_il_parity.cu) + target_link_libraries(mmvq_il_parity PRIVATE strata_kernels CUDA::cudart) + add_test(NAME mmvq_il_parity COMMAND mmvq_il_parity) # Plan v0.3 P3: native MMVQ for 1..8 tokens (bitwise selftest + --bench timing); synthetic data, no model. if(STRATA_BUILD_TESTS) add_executable(native_mmvq_multi bench/micro/native_mmvq_multi.cpp) - target_include_directories(native_mmvq_multi PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(native_mmvq_multi PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(native_mmvq_multi PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(native_mmvq_multi PRIVATE strata_kernels CUDA::cudart) add_test(NAME native_mmvq_multi COMMAND native_mmvq_multi --selftest) endif() # Plan v0.3 P0.1: working-set lock of a host region (CPU only). add_executable(platform_memory_test src/platform/memory_test.cpp) - target_link_libraries(platform_memory_test PRIVATE strata_core ${_strata_gpu_runtime_target}) + target_link_libraries(platform_memory_test PRIVATE strata_core CUDA::cudart) add_test(NAME platform_memory_test COMMAND platform_memory_test) # Plan v0.3 P2: the n-gram table read straight from the SSD. The selftest needs only the CPU and a temp file. add_executable(ple_reader_test src/ngram/ple_reader_test.cpp) - target_include_directories(ple_reader_test PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(ple_reader_test PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + target_include_directories(ple_reader_test PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(ple_reader_test PRIVATE strata_kernels CUDA::cudart) add_test(NAME ple_reader_selftest COMMAND ple_reader_test --selftest) - # The GPU test runs ON THE TARGET PC. `--selftest` allocates + # The GPU test is registered only when CUDA is on, and it runs ON THE TARGET PC. `--selftest` allocates # and writes real VRAM, so a machine without a device fails this test rather than skipping it silently - # which is correct: an engine that targets one GPU should not report a pass for having no GPU. - if(STRATA_ENABLE_CUDA) - add_test(NAME cuda_device_selftest COMMAND strata-device --selftest) - else() - add_test(NAME hip_device_selftest COMMAND strata-device --selftest) - endif() + add_test(NAME cuda_device_selftest COMMAND strata-device --selftest) # ---- P2.S5: the GraphRegistry -------------------------------------------------- # A .cu because the test needs real kernels to capture; it links `strata_core` and not `strata_kernels`, # because what it exercises is the capture machinery and not any kernel in it. - if(STRATA_BUILD_TESTS AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/bench/micro/graph_registry_test.cu") - add_executable(graph_registry_test bench/micro/graph_registry_test.cu) - target_include_directories(graph_registry_test PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(graph_registry_test PRIVATE strata_core ${_strata_gpu_runtime_target}) - add_test(NAME graph_registry_test COMMAND graph_registry_test) - endif() + if(STRATA_BUILD_TESTS) + add_executable(graph_registry_test bench/micro/graph_registry_test.cu) + target_include_directories(graph_registry_test PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(graph_registry_test PRIVATE strata_core CUDA::cudart) + add_test(NAME graph_registry_test COMMAND graph_registry_test) # ---- P2.S5's GATE: is a device<->pinned-memory handoff graph-capturable? -------------------- # The doorbell protocol puts the D2H of `x_f` and the miss list INSIDE G[l] - that is the whole overlap. The # phase doc says to verify this "before building on it", and this project has already lost two rounds to # building first (a memory-only spin that DEADLOCKS on Windows, and a `cudaEventRecord` inside a capture that # is SILENTLY DROPPED). It links only cudart: it is about the capture machinery, not about any kernel. - if(STRATA_ENABLE_CUDA AND STRATA_BUILD_TESTS AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/bench/micro/pinned_capture.cu") - add_executable(pinned_capture bench/micro/pinned_capture.cu) - target_include_directories(pinned_capture PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(pinned_capture PRIVATE ${_strata_gpu_runtime_target}) - add_test(NAME pinned_capture COMMAND pinned_capture) + add_executable(pinned_capture bench/micro/pinned_capture.cu) + target_include_directories(pinned_capture PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(pinned_capture PRIVATE CUDA::cudart) + add_test(NAME pinned_capture COMMAND pinned_capture) endif() # Tells tests/CMakeLists.txt to add the CUDA-dependent Catch2 target. It cannot be declared there # unconditionally: tests/ is added at the end of this file, but the flag has to be known before it runs. - if(STRATA_ENABLE_CUDA) - set(STRATA_CUDA_TESTS ON) - endif() + set(STRATA_CUDA_TESTS ON) endif() # ---- strata_kernels_cpu (P2.S3) ---- @@ -563,22 +440,15 @@ endif() # worse opt-out. `cpu_require_expert_support()` therefore refuses at startup on a CPU without the # features, and the fallback exists for tests, exactly as P2.S3 specifies. # Plan v0.3 P6: speculation components that need no GPU. -add_library(strata_spec STATIC src/spec/suffix_drafter.cpp src/spec/controller.cpp src/spec/draft_policy.cpp) +add_library(strata_spec STATIC src/spec/suffix_drafter.cpp src/spec/controller.cpp) target_include_directories(strata_spec PUBLIC include) +add_executable(suffix_drafter_test src/spec/suffix_drafter_test.cpp) +target_link_libraries(suffix_drafter_test PRIVATE strata_spec) +add_test(NAME suffix_drafter_test COMMAND suffix_drafter_test) if(STRATA_BUILD_TESTS) - add_executable(suffix_drafter_test src/spec/suffix_drafter_test.cpp) - target_link_libraries(suffix_drafter_test PRIVATE strata_spec) - add_test(NAME suffix_drafter_test COMMAND suffix_drafter_test) add_executable(controller_test src/spec/controller_test.cpp) target_link_libraries(controller_test PRIVATE strata_spec) add_test(NAME controller_test COMMAND controller_test) - add_executable(draft_policy_test src/spec/draft_policy_test.cpp) - target_link_libraries(draft_policy_test PRIVATE strata_spec) - add_test(NAME draft_policy_test COMMAND draft_policy_test) - # the serve loop's conversation-cache retention policy (header-only, no library to link) - add_executable(conv_cache_test src/program/conv_cache_test.cpp) - target_include_directories(conv_cache_test PRIVATE include) - add_test(NAME conv_cache_test COMMAND conv_cache_test) endif() add_library(strata_kernels_cpu STATIC src/kernels/cpu/expert.cpp src/kernels/cpu/pool.cpp @@ -597,6 +467,10 @@ else() endif() find_package(Threads REQUIRED) target_link_libraries(strata_kernels_cpu PUBLIC Threads::Threads) +# The expert pool against serial runs; `--stress` needs no pack (synthetic experts, idle gaps that put the workers +# to sleep). +add_executable(pool_test src/kernels/cpu/pool_test.cpp) +target_link_libraries(pool_test PRIVATE strata_kernels_cpu) # ---- plan v0.3 P6: the i-quant models' CPU experts (IQ2_XS / IQ3_XXS files) ---- # @@ -605,9 +479,6 @@ target_link_libraries(strata_kernels_cpu PUBLIC Threads::Threads) # with -DSTRATA_GGML_DIR=..., otherwise the pinned upstream commit (third_party/ggml/VERSION.txt). option(STRATA_NATIVE_EXPERTS "Native (IQ) expert support through ggml-cpu" ON) set(STRATA_GGML_DIR "" CACHE PATH "llama.cpp checkout to take ggml from (empty: fetch the pinned commit)") -if(STRATA_PREFILL_MMQ AND NOT STRATA_NATIVE_EXPERTS) - message(FATAL_ERROR "STRATA_PREFILL_MMQ requires STRATA_NATIVE_EXPERTS=ON for GGML sources") -endif() if(STRATA_NATIVE_EXPERTS) enable_language(C) set(BUILD_SHARED_LIBS OFF CACHE BOOL "" FORCE) @@ -655,9 +526,9 @@ if(STRATA_NATIVE_EXPERTS) target_link_libraries(strata_kernels_cpu PUBLIC ggml-cpu ggml-base) target_compile_definitions(strata_kernels_cpu PUBLIC STRATA_NATIVE_EXPERTS=1) if(STRATA_ENABLE_CUDA) - # prompt-speed plan step 2b: llama.cpp's MMQ kernels (ggml-cuda, MIT) for the prompt path's experts, compiled - # from the same checkout as ggml above - only mmq.cuh's instances for the packs' types and the quantizer, with - # the few host symbols of ggml-cuda.cu they use in src/prefill/ggml_cuda_host.cu + # llama.cpp's MMQ kernels (ggml-cuda, MIT) for the prompt path's experts, from the same checkout as ggml: mmq.cuh's + # instances for the expert types of ISTA-DASLab's and unsloth's files and the quantizer, with the few host + # symbols of ggml-cuda.cu they use in src/prefill/ggml_cuda_host.cu if(STRATA_GGML_DIR) set(_strata_gg ${STRATA_GGML_DIR}/ggml) else() @@ -665,7 +536,7 @@ if(STRATA_NATIVE_EXPERTS) endif() set(_strata_mmq ${_strata_gg}/src/ggml-cuda) add_library(strata_mmq STATIC src/prefill/moe_mmq.cu src/prefill/ggml_cuda_host.cu ${_strata_mmq}/quantize.cu) - foreach(_t q2_0 iq2_xxs iq2_xs iq2_s iq3_xxs iq3_s iq4_nl iq4_xs q8_0) # q8_0: the drafter's matrices (E-9) + foreach(_t q2_0 iq1_s iq2_xxs iq2_xs iq2_s iq3_xxs iq3_s iq4_nl iq4_xs q4_k q5_k q6_k q5_1 q8_0) target_sources(strata_mmq PRIVATE ${_strata_mmq}/template-instances/mmq-instance-${_t}.cu) endforeach() target_include_directories(strata_mmq PRIVATE ${_strata_gg}/include ${_strata_gg}/src ${_strata_mmq} @@ -674,71 +545,30 @@ if(STRATA_NATIVE_EXPERTS) if(MSVC) target_compile_options(strata_mmq PRIVATE $<$:-Xcompiler=/bigobj>) endif() - target_link_libraries(strata_mmq PUBLIC ggml-base ${_strata_gpu_runtime_target}) - target_link_libraries(strata_prefill PUBLIC strata_mmq) - target_compile_definitions(strata_prefill PUBLIC STRATA_PREFILL_MMQ=1) - elseif(STRATA_PREFILL_MMQ) - # Compile only the pinned GGML MMQ and activation-quantizer sources used by prefill. This uses - # common.cuh's native HIP path without enabling a second GGML HIP backend. - if(STRATA_GGML_DIR) - set(_strata_gg "${STRATA_GGML_DIR}/ggml") - else() - set(_strata_gg "${strata_llamacpp_SOURCE_DIR}/ggml") - endif() - set(_strata_mmq "${_strata_gg}/src/ggml-cuda") - if(NOT EXISTS "${_strata_mmq}/common.cuh" OR NOT EXISTS "${_strata_mmq}/quantize.cu") - message(FATAL_ERROR "The selected GGML source is missing the ggml-cuda MMQ sources: ${_strata_mmq}") - endif() - set(_strata_mmq_sources - src/prefill/moe_mmq.cu - src/prefill/ggml_cuda_host.cu - "${_strata_mmq}/quantize.cu") - foreach(_t q2_0 iq2_xxs iq2_xs iq2_s iq3_xxs iq3_s iq4_nl iq4_xs q8_0) - list(APPEND _strata_mmq_sources "${_strata_mmq}/template-instances/mmq-instance-${_t}.cu") - endforeach() - set_source_files_properties(${_strata_mmq_sources} PROPERTIES LANGUAGE HIP) - add_library(strata_mmq STATIC ${_strata_mmq_sources}) - target_include_directories(strata_mmq - PRIVATE "${_strata_gg}/include" "${_strata_gg}/src" "${_strata_mmq}" - PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include) - target_compile_definitions(strata_mmq PRIVATE GGML_USE_HIP=1) - target_link_libraries(strata_mmq PUBLIC ggml-base hip::host roc::hipblas) + target_link_libraries(strata_mmq PUBLIC ggml-base CUDA::cudart) target_link_libraries(strata_prefill PUBLIC strata_mmq) - target_compile_definitions(strata_prefill PUBLIC STRATA_PREFILL_MMQ=1) - - # Numerical diagnostic: synthetic layouts by default, optional first-expert IQ3_XXS/IQ4_NL check with a pack. - add_executable(hip_prefill_mmq_parity tests/hip/prefill_mmq_parity.cpp) - target_include_directories(hip_prefill_mmq_parity PRIVATE "${_strata_gg}/include") - target_link_libraries(hip_prefill_mmq_parity PRIVATE strata_prefill) - add_test(NAME hip_prefill_mmq_parity COMMAND hip_prefill_mmq_parity) - set_tests_properties(hip_prefill_mmq_parity PROPERTIES TIMEOUT 120) - list(APPEND _strata_hip_tests hip_prefill_mmq_parity) - endif() - if(STRATA_ENABLE_HIP) - # Optional real-model diagnostic: pass the GGUF shard containing the native PLE key. - add_executable(hip_ple_iq4 tests/hip/ple_iq4.cpp) - target_link_libraries(hip_ple_iq4 PRIVATE strata_kernels strata_hip_runtime ggml-base) - endif() - if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) add_executable(native_expert_parity src/kernels/native_expert_parity.cpp) - target_include_directories(native_expert_parity PRIVATE ${_strata_gpu_include_directories}) - target_link_libraries(native_expert_parity PRIVATE strata_kernels strata_kernels_cpu strata_artifact ${_strata_gpu_runtime_target}) + target_include_directories(native_expert_parity PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) + target_link_libraries(native_expert_parity PRIVATE strata_kernels strata_kernels_cpu strata_artifact CUDA::cudart) + # The grouped native expert kernels against the ones they replaced, kept in the test: bitwise (synthetic experts). + add_executable(expert_grouped_parity src/kernels/expert_grouped_parity.cu) + target_include_directories(expert_grouped_parity PRIVATE third_party/ggml src/kernels/cuda) + target_link_libraries(expert_grouped_parity PRIVATE strata_kernels strata_kernels_cpu CUDA::cudart) + add_test(NAME expert_grouped_parity COMMAND expert_grouped_parity) endif() endif() if(STRATA_ENABLE_CUDA AND STRATA_BUILD_TESTS) add_executable(hit_cpu_order_parity bench/micro/hit_cpu_order_parity.cu) - target_link_libraries(hit_cpu_order_parity PRIVATE strata_kernels strata_kernels_cpu ${_strata_gpu_runtime_target}) + target_link_libraries(hit_cpu_order_parity PRIVATE strata_kernels strata_kernels_cpu CUDA::cudart) endif() if(STRATA_BUILD_TESTS) # Plan v0.3 P6: one expert applied to several tokens, bitwise vs the single-token kernel (synthetic, CPU only). # plan v0.3 P5: the device dequantizers vs the artifact's CPU dequantizers on the real shards (not a ctest: needs # the 38 GB model; run it by hand or from tools/gpu_session.py) - if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) # #93: it needs the device dequantizers and a GPU runtime - add_executable(dequant_bf16_test src/kernels/dequant_bf16_test.cpp) - target_link_libraries(dequant_bf16_test PRIVATE strata_kernels strata_artifact ${_strata_gpu_runtime_target}) - endif() + add_executable(dequant_bf16_test src/kernels/dequant_bf16_test.cpp) + target_link_libraries(dequant_bf16_test PRIVATE strata_kernels strata_artifact CUDA::cudart) add_executable(expert_multi_test src/kernels/cpu/expert_multi_test.cpp) target_link_libraries(expert_multi_test PRIVATE strata_kernels_cpu) @@ -751,26 +581,11 @@ if(STRATA_BUILD_TESTS) add_test(NAME expert_parity COMMAND expert_parity --selftest WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) - add_executable(pool_test src/kernels/cpu/pool_test.cpp) - target_link_libraries(pool_test PRIVATE strata_kernels_cpu) add_test(NAME pool_test COMMAND pool_test --selftest WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) - # issue #29: tiny batches with sleeping workers, as on a big-VRAM card; a hang fails it - add_executable(pool_stress src/kernels/cpu/pool_stress.cpp) - target_link_libraries(pool_stress PRIVATE strata_kernels_cpu) - add_test(NAME pool_stress COMMAND pool_stress 20) - - if(UNIX AND NOT APPLE) - add_executable(direct_file_async_test tests/platform/direct_file_async.cpp src/platform/direct_file.cpp) - target_include_directories(direct_file_async_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include) - target_link_libraries(direct_file_async_test PRIVATE Threads::Threads) - add_test(NAME direct_file_async_test COMMAND direct_file_async_test) - endif() + add_test(NAME pool_stress COMMAND pool_test --stress) endif() if(STRATA_BUILD_TESTS) enable_testing() - # Model-level checks may be absent from source distributions. - if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/CMakeLists.txt") - add_subdirectory(tests) - endif() + add_subdirectory(tests) endif() diff --git a/README.md b/README.md index 2df0838cbf0..5224330c43e 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,11 @@ of a word): faster than you can read. - **Free and open source.** +> **This fork** (branch `custom`) is tuned for a PC with an RTX 3090 and an RTX 5070 Ti: the second card works as another +> expert tier, unsloth's K-quant files run too, and prompts and answers are faster. How it compares with llama.cpp and +> upstream Strata on that PC: [docs/COMPARISON.md](docs/COMPARISON.md). Its configs, and how to run them: +> [examples/](examples/README.md). + > **Jump to:** [How fast?](#how-fast-is-it) · [Which model?](#which-model-should-i-pick) · [Install](#install) · > [Using it](#using-it) · [Problems?](#something-went-wrong) · [How it works](#how-does-it-work) · > [All the details](docs/DETAILS.md) @@ -140,8 +145,9 @@ the same way - nothing big is downloaded again. - **Experimental speed projection (off by default):** an experimental control vector that setup can turn on; it changes how the model answers - read [what it does](docs/DETAILS.md#experimental-speed-projection-experimental-off-by-default) first. -**Good to know:** it answers one request at a time. The first message of a chat is read in full (about 1 minute per -30,000 tokens); after that it keeps the conversation and reads only what is new, so follow-ups start in seconds. +**Good to know:** it answers one request at a time. The next turn of a chat continues from what it has already read, +so only the new text is processed; a long document it reads for the first time takes about 1 minute per 30,000 tokens +before it starts writing. ## Something went wrong? diff --git a/data/draft_vocab.bin b/data/draft_vocab.bin index 217f2745362..c588b567f6e 100644 Binary files a/data/draft_vocab.bin and b/data/draft_vocab.bin differ diff --git a/docs/COMPARISON.md b/docs/COMPARISON.md new file mode 100644 index 00000000000..c55136e16b6 --- /dev/null +++ b/docs/COMPARISON.md @@ -0,0 +1,149 @@ +# This fork, upstream Strata and llama.cpp on one PC + +Qwen3.8-Flash-Next in two quantizations, IQ3_S and UD-Q4_K_XL, through three engines on an RTX 3090 + RTX 5070 Ti PC. +The llama.cpp and upstream columns were measured once (2026-09-30). This fork's columns are kept current: measured at +0bbd8f5 (2026-09-30). + +**In short**, with both GPUs: + +- **IQ3_S:** this fork reads prompts 4.2-5.0x as fast as llama.cpp and 2.3-2.4x as fast as upstream Strata, and + writes answers 4.2-7.9x as fast as llama.cpp and 1.5-2.0x as fast as upstream (the most on long contexts and on + answers that quote the prompt). On the 3090 alone, against upstream there: prompts 1.6x, answers 1.4-1.7x. +- **UD-Q4_K_XL:** prompts 5.0-6.2x and answers 4.8-8.8x as fast as llama.cpp; upstream cannot run the file. +- A follow-up turn after a 16K-token prompt starts in 0.2-0.3 s (llama.cpp 0.6-0.7 s; upstream 0.35 s on the 3090, + 1.1 s on both GPUs). +- The next-token distributions match llama.cpp's as closely as llama.cpp's own CPU and CUDA backends match each other + (IQ3_S; UD-Q4_K_XL somewhat further, from llama.cpp's 8-bit rounding of the file's Q8_0 hyper-connection weights). + +## The PC and the engines + +- **PC:** Ryzen 7 9700X (8 cores), 160 GB DDR5-5600 (2 x 32 + 2 x 48 GB), Windows 11. An RTX 3090 (24 GB, PCIe 4.0 + x4, nothing else on it) and an RTX 5070 Ti (16 GB, PCIe 5.0 x16, driving the display). +- **Models:** ISTA-DASLab's GSQ-RCO **IQ3_S** (shard 1 54.8 GB: experts 46.8 GiB, i-quants with Q2_0 and IQ4_NL down + projections) and unsloth's **UD-Q4_K_XL** (111 GB: experts 71.7 GiB, Q4_K / Q5_K with Q5_1 / Q8_0 downs). Each has + a 28.8 GB IQ4_NL n-gram table (the per-layer embeddings) besides. +- **llama.cpp** build 11116 (709fe755d, MSVC), its ggml-cuda.dll patched to pin the CPU experts' buffers (prompts + 1.8x faster than unpinned). Both GPUs with the layers split (the 5070 Ti first), the experts that do not fit in VRAM + in RAM: UD-Q4_K_XL by `--fit`, IQ3_S by fit's placement with its Q2_0 down projections kept on the GPUs (llama.cpp + has no x86 SIMD kernel for Q2_0: 30.5 instead of 27.2 t/s). 8-bit KV cache, 8 threads, 1024-token micro-batches, + the n-gram table in RAM. No speculative decoding: llama.cpp has no graph for this model's MTP layer. +- **Upstream Strata** a790805 (engine 0.1.27) with the config its setup writes for IQ3_S on this PC: MTP drafts, 8-bit + KV kept in RAM past 32K cells, the expert cache sized by the free VRAM. On the 3090, and on both GPUs as its setup + recommends here (a layer split with the newer 5070 Ti first: layers 0-25 there, 26-47 on the 3090). It cannot run + UD-Q4_K_XL: it has no K-quant expert kernels. +- **This fork** with its configs ([examples](../examples/README.md)): the 3090 runs the model and the 5070 Ti is a + second expert tier (12.2-12.4 GiB of experts; its share of each layer's experts, the next layer's likely experts + copied ahead, the prompt path's experts the 3090 lacks, and with IQ3_S part of the output head). MTP drafts and + prompt lookup, 8-bit KV in VRAM, 16K-token prompt chunks. And on the 3090 alone. + +## Speed through the API + +The same requests to each server's OpenAI API: greedy, one at a time, the prompt caches off except for the follow-up +turn. Prompts: 16K-128K tokens of llama.cpp's sources with a question (400 tokens of answer), a 28-token coding request +(up to 300), a coding task with thinking on (3,000), and a Python module to return with a class renamed (the edit, up +to 3,000 tokens that mostly quote the prompt). The follow-up turn sends the 16K prompt, its answer and a new question: +its prompt time is the processing of what the engine did not keep. Each engine writes its own text and the texts part +at the first near-tie, so the generation speeds compare different texts of one kind, one run each. The memory rows +come from a second start of each server (the model file then in the OS cache): the GPUs' memory taken, and the RAM +taken as the drop in the system's available memory. + +### IQ3_S + + +| | llama.cpp, both GPUs | upstream Strata, 3090 | upstream Strata, both GPUs | this fork, both GPUs | this fork, 3090 | +|---|---|---|---|---|---| +| Prompt, 16K tokens (t/s) | 581 | 1029 | 1060 | 2466 | 1605 | +| Prompt, 32K tokens (t/s) | 562 | 1120 | 1119 | 2666 | 1768 | +| Prompt, 64K tokens (t/s) | 551 | 1139 | 1148 | 2688 | 1791 | +| Prompt, 128K tokens (t/s) | 527 | 1111 | 1143 | 2650 | 1787 | +| Generation: code, 28-token prompt (t/s) | 31.0 | 74.1 | 84.8 | 167.7 | 113.9 | +| Generation: thinking, 3,000 tokens | 31.3 | 77.1 | 89.0 | 130.2 | 111.2 | +| Generation: edit, the prompt's file returned (2.4K tokens) | 31.1 | 99.4 | 117.2 | 234.8 | 171.3 | +| Generation after the 16K prompt | 29.8 | 75.7 | 85.4 | 151.2 | 108.1 | +| Generation after the 32K prompt | 26.9 | 80.6 | 94.2 | 150.2 | 122.1 | +| Generation after the 64K prompt | 22.0 | 73.3 | 85.3 | 139.6 | 113.9 | +| Generation after the 128K prompt | 16.4 | 67.6 | 76.8 | 129.2 | 110.2 | +| Follow-up turn after the 16K prompt: new tokens, prompt time | 34 in 0.62 s | 34 in 0.35 s | 440 in 1.12 s | 34 in 0.21 s | 34 in 0.34 s | +| Start to the API, the model file cached (s) | 21 | 16 | 21 | 21 | 20 | +| VRAM taken: 3090 / 5070 Ti (GiB) | 22.9 / 13.0 | 22.9 / 0.0 | 21.9 / 14.6 | 22.5 / 13.1 | 22.6 / 0.0 | +| RAM taken (GB) | 57 | 56 | 60 | 53 | 52 | + + +### UD-Q4_K_XL + +Upstream Strata cannot run this file: it has no kernels for K-quant experts. + + +| | llama.cpp, both GPUs | this fork, both GPUs | this fork, 3090 | +|---|---|---|---| +| Prompt, 16K tokens (t/s) | 467 | 2341 | 1099 | +| Prompt, 32K tokens (t/s) | 453 | 2645 | 1184 | +| Prompt, 64K tokens (t/s) | 444 | 2662 | 1209 | +| Prompt, 128K tokens (t/s) | 426 | 2627 | 1221 | +| Generation: code, 28-token prompt (t/s) | 23.5 | 121.3 | 76.8 | +| Generation: thinking, 3,000 tokens | 23.7 | 114.3 | 74.3 | +| Generation: edit, the prompt's file returned (2.4K tokens) | 23.7 | 160.0 | 113.4 | +| Generation after the 16K prompt | 22.6 | 126.7 | 67.6 | +| Generation after the 32K prompt | 21.7 | 138.4 | 75.7 | +| Generation after the 64K prompt | 18.6 | 119.7 | 66.4 | +| Generation after the 128K prompt | 14.5 | 127.5 | 69.2 | +| Follow-up turn after the 16K prompt: new tokens, prompt time | 34 in 0.69 s | 34 in 0.26 s | 34 in 0.48 s | +| Start to the API, the model file cached (s) | 68 | 41 | 36 | +| VRAM taken: 3090 / 5070 Ti (GiB) | 22.3 / 12.9 | 22.7 / 13.1 | 22.7 / 0.0 | +| RAM taken (GB) | 86 | 81 | 80 | + + +## Generation on fixed text (this fork) + +Each run emits the model's reference text (`--spec-follow`), so runs and builds compare on the same tokens and +experts: two runs, averaged, tokens a verify pass in parentheses. The code text's reference repeats its answer after +the end of the turn, which prompt lookup drafts and a server never generates. + + +| Text | IQ3_S t/s (tokens a pass) | UD-Q4_K_XL t/s (tokens a pass) | +|---|---|---| +| code (400 tokens) | 181.5 (4.02) | 114.6 (3.69) | +| thinking (1,200) | 137.8 (2.07) | 105.4 (1.71) | +| after the 16K prompt (400) | 145.2 (2.71) | 124.7 (2.65) | +| edit (2,397) | 243.7 (7.51) | 160.8 (7.51) | + + +## Agreement with llama.cpp + +Argmax agreement and mean KL divergence (nats) of the next-token distributions along llama.cpp's greedy text, which +this fork follows (`--spec-follow`, `--window-logits`), with `--adapt-every 0`. The two baselines differ from each +other only in rounding: llama.cpp's CPU backend against its CUDA backend, and this fork in one-token windows against +its four-token windows. Any change of rounding moves 2-4% of the thinking text's argmaxes, all near-ties. + + +| | Code, ~260 tokens | Thinking, 1,200 | After the 16K prompt, 300 | +|---|---|---|---| +| IQ3_S: this fork vs llama.cpp | 100.0%, 1.3e-04 | 98.0%, 2.1e-03 | 99.3%, 1.0e-03 | +| llama.cpp's CPU vs CUDA backend | 100.0%, 2.2e-04 | 96.9%, 3.2e-03 | 99.0%, 1.4e-03 | +| this fork, one-token vs four-token windows | 100.0%, 1.8e-04 | 98.3%, 1.9e-03 | 99.7%, 3.8e-04 | +| UD-Q4_K_XL: this fork vs llama.cpp | 99.3%, 5.3e-04 | 95.8%, 6.7e-03 | 97.7%, 3.5e-03 | +| llama.cpp's CPU vs CUDA backend | 99.6%, 1.3e-04 | 96.8%, 3.3e-03 | 100.0%, 1.6e-03 | +| this fork, one-token vs four-token windows | 99.6%, 1.0e-04 | 97.2%, 2.3e-03 | 98.7%, 5.4e-04 | + + +UD-Q4_K_XL's CPU baseline after the 16K prompt covers 181 positions: llama-server's greedy text itself differed between +two runs there, and the CPU backend was forced along the other one, which parts after 180 tokens. + +## The two files + +llama.cpp's perplexity over 28 chunks of 4,096 tokens of llama.cpp's own docs, C++ and Python sources (lower is +better): UD-Q4_K_XL 1.8849, IQ3_S 1.8946 (+0.5%), and ISTA-DASLab's IQ3_XXS, which IQ3_S replaced in this fork's +configs, 1.9215 (+1.9%). IQ3_S's experts take 46.8 GiB, UD-Q4_K_XL's 71.7: on this PC IQ3_S generates 16-58% +faster on fixed text, and reads prompts about as fast. + +## Features + +| | llama.cpp | Upstream Strata | This fork | +|---|---|---|---| +| Model files | any GGUF of the model | ISTA-DASLab's GSQ-RCO sizes (Q2_0 to IQ3_S), their Coder, Swift 1.5: i-quant and Q2_0 experts | those, and every quant of unsloth's (K-quant, Q5_1 and Q8_0 experts) | +| Speculative decoding | none: no graph for this model's MTP layer | the MTP layer's drafts, up to 3 a pass; prompt lookup (up to 5 tokens where it pays) | the MTP layer's drafts; prompt lookup (up to 7 tokens after a 16+-token match) | +| Several GPUs | a layer or row split | a layer split, each card holding its layers' experts | the second card as an expert tier: its share of each layer, the next layer's likely experts copied ahead, the prompt's experts the first card lacks, part of the output head | +| Experts in VRAM | a fixed placement | ranked by use, updated every 4 verify passes | ranked by use on both cards, updated every pass | +| KV cache | 8-bit, in VRAM | 8-bit, in RAM past 32K cells | 8-bit, in VRAM | +| Prompt cache | the live slot, states saved to RAM | checkpoints at turns and every 16K tokens | the live sequence, checkpoints at turns and chunks, replaced conversations kept in RAM with their K/V | +| Sampling, APIs | full; OpenAI and Anthropic | per request; OpenAI and Anthropic, a web app, MCP tools | upstream's | diff --git a/docs/DETAILS.md b/docs/DETAILS.md index c55d13e9532..5192c8df1fe 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -162,6 +162,39 @@ thought about for 1,524 tokens. Not a benchmark, but consistent with the claim. START-HERE.bat --setup --family swift --model IQ2_XS ``` +### unsloth's Unsloth Dynamic 3.0 files + +The engine also runs [unsloth's quantizations](https://huggingface.co/unsloth/Qwen3.8-Flash-Next-GGUF), which +START-HERE does not set up. All experts sit in RAM, pinned when they fit the driver's limit: + +| Model | Download | Experts in RAM | n-gram table | +| --- | ---: | ---: | --- | +| UD-IQ1_S / UD-IQ1_M | 73-75 GB | 40-42 GB | IQ4_NL, 28.8 GB | +| UD-Q2_K_XL / UD-IQ3_XXS | 79-82 GB | 46-49 GB | IQ4_NL | +| UD-Q3_K_XL / UD-IQ4_XS | 90-94 GB | 56-60 GB | IQ4_NL | +| UD-Q4_K_XL | 111 GB | 77 GB | IQ4_NL | +| UD-Q5_K_XL / UD-Q6_K_XL | 158-169 GB | 98-109 GB | Q8_0, 54.4 GB | +| Q8_0 | 188 GB | 128 GB | Q8_0 | + +Pack a model once (seconds: the tokenizer and the float tensors; the experts are read from the GGUF at start): + + python tools/iq_pack.py --gguf /-00001-of-0000N.gguf --out packs/ + +and start it with `--pack packs/ --native /-00001-of-0000N.gguf` plus the usual flags (`--spec`, +`--mtp`, `--prefill`, ...; this fork's config: [examples/ud-q4_k_xl.json](../examples/ud-q4_k_xl.json)). The +other shards, and the one holding the n-gram table, are found by name. +UD-Q4_K_XL's speeds on an RTX 3090 with 160 GB of RAM, with and without a second GPU: [COMPARISON.md](COMPARISON.md). + +### A second GPU + +With `--second-gpu N`, CUDA device N (the engine runs on `--main-gpu`, default 0) holds more experts and computes its share of +each layer while the CPU computes the rest, copies each next layer's likeliest missing experts to itself while the RAM +is idle, computes the prompt's experts the first GPU lacks, and holds the last rows of the output head (`--head-split`, +0.45 of them); it may be the card that drives the display. Build for both GPUs' architectures +(`-DCMAKE_CUDA_ARCHITECTURES="86;120"`). When the card has no free VRAM, the engine runs on the main GPU alone. An RTX +3090 with and without an RTX 5070 Ti as the second GPU, against upstream Strata and llama.cpp: +[COMPARISON.md](COMPARISON.md). + ## Before you start You need **only an NVIDIA driver** (version 580 or newer; update it with the NVIDIA App or from @@ -342,6 +375,15 @@ print(r.choices[0].message.content) `max_tokens` would run past the context is refused too (400); agents that always ask for their full output cap can instead get it shortened to the room left: add `"fit_max_tokens": true` to `strata-.json` (or pass `--fit-max-tokens` to `serve/server.py`). A prompt that leaves no room at all is still refused. +- **Prompt cache.** A request continues from the longest start of its prompt the engine still holds, so the next turn + of a chat or an agent processes only what is new (a 16K-token conversation: ~0.3 s instead of ~25 s before the first + token). It keeps checkpoints of the sequence state in RAM (at every turn's start and end, and every 16K tokens of a + long prompt: regenerated and edited answers resume from them) and moves a conversation that another request replaces + (a chat app's title request, a subagent) to RAM with its key/value cache, so switching back continues where it was. + Responses report the reused tokens (`usage.prompt_tokens_details.cached_tokens`; Anthropic: + `cache_read_input_tokens`). Engine options: `--prompt-cache N` checkpoints (default 16, ~110 MB each; 0 = off), + `--cache-ram MIB` for replaced conversations (default 16384), `--cache-every N` (default 16384). A changed system + prompt or reasoning effort changes the prompt from its first tokens, so nothing is reused. - **From other devices on your network.** The server listens on your PC only (`127.0.0.1`) unless you say otherwise: run setup with `START-HERE.bat --setup --host 0.0.0.0 --api-key some-long-secret` (or add `"host": "0.0.0.0"` and `"api_key": "..."` to `strata-.json`). The server window then prints this PC's addresses @@ -354,30 +396,17 @@ print(r.choices[0].message.content) add `"api_key": "some-long-secret"` to `strata-.json` (or set the `STRATA_API_KEY` environment variable); clients then send it as their API key. -**Conversation cache.** A request that continues a chat reads only the part after what the engine already holds: the -live session, or one of the checkpoints it keeps in RAM (up to 6, ~118 MB each, taken at the start of each new -assistant turn and every 16K prompt tokens). A checkpoint is used only when the prompt starts with exactly its tokens -and pictures. The oldest checkpoint - in practice the end of the system prompt, which every chat of the same client -shares - is kept for good while the rest rotates by least recent use, so a NEW chat that shares that prefix starts -reading after it instead of from token 0. A prompt read from the start is also checkpointed at the end of its system -prompt when that is 2,048 tokens or more (engine 0.1.20; PR #62 + #65), so that root exists for agent clients with long -system prompts and tool lists. Engine options: `--prompt-cache N` (0 = off), `--prompt-cache-every N`, -`--prompt-cache-root N` (0 = no system-prompt checkpoint), `--turn-token ID`. - -**Current limits (v1):** one request at a time, and one conversation's history in the KV cache at a time (switching -between two chats re-reads the part where they diverge; the shared prefix, such as the system prompt, is reused); images -only when set up with them (below); no video. **Temperature / top_p / top_k / min_p / -seed** are honored per request (OpenAI and Anthropic fields); with the default adaptive expert tier a sampled result -is not reproducible run to run - for seed-reproducible output add `--adapt-every 100000` (static residency) to the -engine arguments. The run config's optional `sampling` block sets the defaults for requests that leave the fields out -(`"sampling": {"temperature": 1.0, "top_p": 0.95, "top_k": 20}`); a request's own fields always win, and with no -block at all a request without sampling keys decodes greedy. The penalties (`presence_penalty`, `frequency_penalty`, -`repetition_penalty`, with `penalty_last_n` capping how many recent tokens they count over, default 64 when any -penalty is set) ride the same path; they count the tokens the request has consumed, so a repetition penalty -suppresses what the model itself just said, not the prompt alone. Since engine 0.1.19 they apply to every token -the speculative decoding checks at once, exactly as if it decoded one token at a time (before, only the first of -each batch got them). That makes requests with penalties 1-11% slower than in 0.1.18: the draft layer guesses -without penalties, so more of its guesses are now rejected. Requests without penalties are unchanged. `top_k` keeps at most 64 candidates: `0` ("off") or anything above 64 uses all 64. +**Current limits (v1):** one request at a time; requests with images are not cached; images only when set up with +them (below); no video. **Temperature / top_p / top_k / min_p / seed** are honored per request (OpenAI and Anthropic +fields); with the default adaptive expert tiers a sampled result is not reproducible run to run - for seed-reproducible +output add `--adapt-every 0` to the engine arguments. The run config's optional `sampling` block sets the defaults for +requests that leave the fields out (`"sampling": {"temperature": 1.0, "top_p": 0.95, "top_k": 20}`); a request's own +fields always win, and with no block at all a request without sampling keys decodes greedy. The penalties +(`presence_penalty`, `frequency_penalty`, `repetition_penalty`, with `penalty_last_n` capping how many recent tokens +they count over, default 64 when any penalty is set) ride the same path; they count the tokens the request has +consumed, so a repetition penalty suppresses what the model itself just said, not the prompt alone, and they apply to +every token the speculative decoding checks at once, exactly as if it decoded one token at a time. `top_k` keeps at +most 64 candidates: `0` ("off") or anything above 64 uses all 64. --- @@ -572,9 +601,8 @@ the document, +0.4% on the chat. Details: `bench/results/2026-09-27-esp/`. as the GPU works on the cached ones (AVX-512 / AVX2 kernels, ggml's for the i-quants). - **SSD:** the 28.8 GB n-gram table, read a few rows per token through the OS cache. - **Speculation:** the model's own MTP layer drafts up to 3 tokens; one pass over all 48 layers checks them. 2.4-3.2 - tokens per pass on average. When the reply repeats the context (code edits, quoted text), **prompt lookup** (engine - 0.1.7) drafts up to 5 tokens from the earlier copy, but only where its measured acceptance and cost say it pays: - code edits 6-11% faster, other text unchanged. The drafts are checked like the MTP's, so the output is the same. + tokens per pass on average. Where the answer repeats earlier text for 16+ tokens (a file returned with an edit), the + pass checks up to 7 tokens that followed it there instead. - **Prompts** are processed in 2,048-token chunks with the experts streamed to the GPU over PCIe. The full story, with measurements, bottlenecks and what comes next: **[docs/paper/Strata-Paper.pdf](paper/Strata-Paper.pdf)**. diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 00000000000..51cdfc09cdc --- /dev/null +++ b/examples/README.md @@ -0,0 +1,44 @@ +# This fork's configs + +The configs this fork runs with on a PC with an RTX 3090, an RTX 5070 Ti and 160 GB of RAM, for `serve/server.py` +(their speeds: [COMPARISON.md](../docs/COMPARISON.md)): + +- `iq3_s.json`: ISTA-DASLab's GSQ-RCO IQ3_S (~53 GB of RAM taken). +- `ud-q4_k_xl.json`: unsloth's UD-Q4_K_XL (~81 GB of RAM taken). + +Their paths are relative to the Strata folder, where the server starts: the engine in `build/`; the model files, the +packs and the MTP draft layer in setup's data folder, `../Strata-data` (unless setup was given another: `--data-dir`). +A pack is what `tools/iq_pack.py` makes from a model's GGUF files in seconds (~1.5 GB): the tokenizer with the chat +template, the model's float tensors in the forms the engine takes, and an index of the rest; the engine reads the +quantized tensors, the experts and most projections, from the GGUF files at start. + +1. Set up IQ3_S once: `START-HERE.bat --setup --model IQ3_S --no-start` downloads the model into + `../Strata-data/models/IQ3_S`, makes its pack in `../Strata-data/packs/iq3_s`, and makes the MTP draft layer in + `../Strata-data/mtp/rt` from the ~5 GB of MTP tensors of Qwen's original checkpoint. +2. Build the engine from this branch with code for both cards (the ready-made engine setup puts in `engine/` is + upstream's, without this fork's options), with MSVC, CUDA 13, CMake and Ninja (setup puts the last two in + `.venv\Scripts`): + + cmake -G Ninja -S . -B build -DCMAKE_BUILD_TYPE=Release -DSTRATA_ENABLE_CUDA=ON -DSTRATA_BUILD_TESTS=OFF ^ + "-DCMAKE_CUDA_ARCHITECTURES=86;120" + cmake --build build --target strata + + `86;120` is an RTX 30 card and an RTX 50 card. The engine needs the CUDA runtime DLLs (`cudart64_13.dll`, + `cublas64_13.dll`, `cublasLt64_13.dll`) beside it, on the PATH, or their folder in the config's `lib_dirs`. +3. For UD-Q4_K_XL, download unsloth's four `Qwen3.8-Flash-Next-UD-Q4_K_XL-0000N-of-00004.gguf` files into + `../Strata-data/models/UD-Q4_K_XL` and make its pack: + + .venv\Scripts\python.exe tools\iq_pack.py --out ../Strata-data/packs/ud-q4_k_xl ^ + --gguf ../Strata-data/models/UD-Q4_K_XL/Qwen3.8-Flash-Next-UD-Q4_K_XL-00001-of-00004.gguf + +4. Set the GPU numbers, as `nvidia-smi` numbers the cards (here the 5070 Ti is 0 and the 3090 is 1). `--main-gpu` + runs the model, and nothing else may use that card: a display or another program on it slows generation to a few + tokens per second. `--second-gpu` holds more experts and may drive the display. +5. Start a config with its launcher, `run-iq3_s.bat` (port 8081) or `run-ud-q4_k_xl.bat` (8082), or from the Strata + folder: + + .venv\Scripts\python.exe serve\server.py --engine strata --config examples\iq3_s.json --port 8081 + + (`START-HERE.bat` starts setup's own config.) + +On Linux the engine is `./build/strata` and Python `.venv/bin/python`. diff --git a/examples/iq3_s.json b/examples/iq3_s.json new file mode 100644 index 00000000000..b4cccab107b --- /dev/null +++ b/examples/iq3_s.json @@ -0,0 +1,44 @@ +{ + "exe": "./build/strata.exe", + "args": [ + "--pack", + "../Strata-data/packs/iq3_s", + "--native", + "../Strata-data/models/IQ3_S/Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-00001-of-00002.gguf", + "--expert-profile", + "./data/expert-profile.bin", + "--expert-cache", + "auto", + "--prefill", + "16384", + "--feed-max", + "80", + "--spec", + "4", + "--spec-min-p", + "0.5", + "--spec-lookup", + "16", + "--adapt-every", + "1", + "--mtp", + "../Strata-data/mtp/rt", + "--max-context", + "200192", + "--kv", + "int8", + "--pcie-mode", + "kernel", + "--pcie-frac", + "0", + "--second-gpu", + "0", + "--main-gpu", + "1" + ], + "cwd": ".", + "tokenizer": "../Strata-data/packs/iq3_s/tokenizer", + "model_name": "qwen3.8-flash-next-iq3_s", + "log": "./strata-iq3_s.log", + "lib_dirs": [] +} diff --git a/examples/run-iq3_s.bat b/examples/run-iq3_s.bat new file mode 100644 index 00000000000..22992850a63 --- /dev/null +++ b/examples/run-iq3_s.bat @@ -0,0 +1,10 @@ +@echo off +title Strata IQ3_S +rem Qwen3.8-Flash-Next, ISTA-DASLab's GSQ-RCO IQ3_S, on http://127.0.0.1:8081/v1 with iq3_s.json (README.md). +rem The config's --main-gpu runs the engine: another program on that card stalls its per-layer spin-waits. +rem Its --second-gpu holds more experts and may drive the display. +rem A PYTHONPATH set for another Python would mix that Python's packages into the venv's. +set "PYTHONPATH=" +cd /d "%~dp0.." +".venv\Scripts\python.exe" serve\server.py --engine strata --config examples\iq3_s.json --port 8081 +pause diff --git a/examples/run-ud-q4_k_xl.bat b/examples/run-ud-q4_k_xl.bat new file mode 100644 index 00000000000..48e1c2e6afa --- /dev/null +++ b/examples/run-ud-q4_k_xl.bat @@ -0,0 +1,10 @@ +@echo off +title Strata UD-Q4_K_XL +rem Qwen3.8-Flash-Next, unsloth's UD-Q4_K_XL, on http://127.0.0.1:8082/v1 with ud-q4_k_xl.json (README.md). +rem The config's --main-gpu runs the engine: another program on that card stalls its per-layer spin-waits. +rem Its --second-gpu holds more experts and may drive the display. +rem A PYTHONPATH set for another Python would mix that Python's packages into the venv's. +set "PYTHONPATH=" +cd /d "%~dp0.." +".venv\Scripts\python.exe" serve\server.py --engine strata --config examples\ud-q4_k_xl.json --port 8082 +pause diff --git a/examples/ud-q4_k_xl.json b/examples/ud-q4_k_xl.json new file mode 100644 index 00000000000..6a516293ea7 --- /dev/null +++ b/examples/ud-q4_k_xl.json @@ -0,0 +1,46 @@ +{ + "exe": "./build/strata.exe", + "args": [ + "--pack", + "../Strata-data/packs/ud-q4_k_xl", + "--native", + "../Strata-data/models/UD-Q4_K_XL/Qwen3.8-Flash-Next-UD-Q4_K_XL-00001-of-00004.gguf", + "--expert-profile", + "./data/expert-profile.bin", + "--expert-cache", + "auto", + "--prefill", + "16384", + "--feed-max", + "52", + "--spec", + "4", + "--spec-min-p", + "0.5", + "--spec-lookup", + "16", + "--adapt-every", + "1", + "--mtp", + "../Strata-data/mtp/rt", + "--max-context", + "200192", + "--kv", + "int8", + "--pcie-mode", + "kernel", + "--pcie-frac", + "0", + "--second-gpu", + "0", + "--main-gpu", + "1", + "--head-split", + "0" + ], + "cwd": ".", + "tokenizer": "../Strata-data/packs/ud-q4_k_xl/tokenizer", + "model_name": "qwen3.8-flash-next-ud-q4_k_xl", + "log": "./strata-ud-q4_k_xl.log", + "lib_dirs": [] +} diff --git a/include/strata/artifact/gguf_reader.hpp b/include/strata/artifact/gguf_reader.hpp index b7f35910022..bf0d91ba6cb 100644 --- a/include/strata/artifact/gguf_reader.hpp +++ b/include/strata/artifact/gguf_reader.hpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -30,6 +31,8 @@ #include #include +#include "strata/artifact/gguf_split.hpp" + #ifdef _WIN32 #define WIN32_LEAN_AND_MEAN #include @@ -339,6 +342,7 @@ class GgufFile { uint32_t version() const { return version_; } uint64_t data_start() const { return data_start_; } uint64_t file_size() const { return size_; } + uint64_t alignment() const { return alignment_; } const TensorInfo* find(const std::string& name) const { for (const auto& t : tensors_) @@ -356,6 +360,7 @@ class GgufFile { return it == meta_.end() ? nullptr : &it->second; } const uint8_t* tensor_data(const TensorInfo& t) const { return base_ + data_start_ + t.offset; } + const std::string& path() const { return path_; } private: void open() { @@ -457,6 +462,31 @@ class GgufFile { std::map meta_; }; +// The opened shards of one model, tensors looked up across all of them. Shard 0 holds the metadata. +class GgufModel { +public: + explicit GgufModel(const std::vector& paths) { + if (paths.empty()) throw std::runtime_error("GGUF: a model needs at least one shard"); + for (const auto& p : paths) shards_.push_back(std::make_unique(p)); + } + size_t size() const { return shards_.size(); } + const GgufFile& shard(size_t i) const { return *shards_[i]; } + /// The tensor named `name` (its shard index in `*shard`), or nullptr. Throws when two shards hold it. + const TensorInfo* find(const std::string& name, size_t* shard = nullptr) const { + const TensorInfo* found = nullptr; + for (size_t i = 0; i < shards_.size(); ++i) + if (const TensorInfo* t = shards_[i]->find(name)) { + if (found) throw std::runtime_error("GGUF: tensor " + name + " is in two shards"); + found = t; + if (shard) *shard = i; + } + return found; + } + +private: + std::vector> shards_; +}; + // ---- architecture guard (P1.S2). The engine is specialised to ONE model; anything else must be // refused with a precise error rather than silently mis-run. struct Qwen4ExpGuard { diff --git a/include/strata/artifact/gguf_split.hpp b/include/strata/artifact/gguf_split.hpp new file mode 100644 index 00000000000..5343393db47 --- /dev/null +++ b/include/strata/artifact/gguf_split.hpp @@ -0,0 +1,38 @@ +// include/strata/artifact/gguf_split.hpp - the shards of a split model. Separate from gguf_reader.hpp, which +// includes , so any translation unit can resolve a model's shards. +#pragma once + +#include +#include +#include +#include +#include + +namespace strata { + +// Split models: `-00001-of-0000N.gguf` ... `-0000N-of-0000N.gguf`, as llama.cpp's gguf-split +// names them. The shards in split order, from the first one's path; a file without that suffix is a model of +// one shard. Throws when a shard is missing. +inline std::vector gguf_split_paths(const std::string& first) { + static const char pat[] = "-00001-of-"; + const size_t at = first.rfind(pat); + const size_t digits = at == std::string::npos ? 0 : at + sizeof(pat) - 1; + if (at == std::string::npos || first.size() != digits + 10 || first.compare(digits + 5, 5, ".gguf") != 0) + return {first}; + int n = 0; + for (size_t i = digits; i < digits + 5; ++i) { + if (first[i] < '0' || first[i] > '9') return {first}; + n = n * 10 + (first[i] - '0'); + } + std::vector out; + for (int i = 1; i <= n; ++i) { + char no[8]; + std::snprintf(no, sizeof no, "%05d", i); + std::string p = first.substr(0, at + 1) + no + first.substr(at + 6); + if (!std::ifstream(p, std::ios::binary)) throw std::runtime_error("missing model shard " + p); + out.push_back(std::move(p)); + } + return out; +} + +} // namespace strata diff --git a/include/strata/core/adaptive_tier.hpp b/include/strata/core/adaptive_tier.hpp new file mode 100644 index 00000000000..f5c9c1c4ba5 --- /dev/null +++ b/include/strata/core/adaptive_tier.hpp @@ -0,0 +1,121 @@ +// include/strata/core/adaptive_tier.hpp - plan v0.3 P6: a VRAM expert tier follows the conversation. +// +// Every few rounds the missing experts routed most often (decayed counts, at least twice) move into the tier: into +// their layer's empty slots first, then into the slots of the layer's least-routed residents when they were routed +// clearly more often. The copies run on their own stream while the rounds go on; an evicted expert is a miss at once, +// a new one resident once its copy has landed (`apply_pending`). Both decode loops (--serve and the speculative loop) +// use it, for the first GPU's cache and for a second GPU's. +#pragma once + +#include "strata/core/expert_cache.hpp" +#include "strata/core/expert_source.hpp" + +#include + +#include +#include +#include +#include + +namespace strata::core { + +class AdaptiveTier { +public: + AdaptiveTier() = default; + ~AdaptiveTier(); + AdaptiveTier(const AdaptiveTier&) = delete; + AdaptiveTier& operator=(const AdaptiveTier&) = delete; + + /// `host_res` is the residency table (n_layers x n_expert, slot or kNotResident) the dispatch reads (the verify + /// windows take a snapshot of it; a device copy of it is its owner's to refresh); up to `max_moves` experts move + /// per call. `device`: the cache's GPU (its stream and copies), `main_device` current again afterwards. + bool init(ExpertCache& cache, ExpertSource& src, std::vector& host_res, int64_t n_layers, int64_t n_expert, + int max_moves, std::string& err, int device = -1, int main_device = 0); + bool on() const { return res_ != nullptr; } + /// An empty slot sized for `layer`'s experts. + void add_free(int64_t layer, int32_t slot) { free_[(size_t) layer].push_back(slot); } + int64_t free_slots() const; + /// A tier in front of this one: the experts it holds or is loading are not candidates here, and this tier's + /// copies of them are evicted first. + void set_upper(const AdaptiveTier* upper) { upper_ = upper; } + /// A tier behind this one: a candidate it holds counts half its use here (moving it only shifts work from that + /// tier's GPU to this one's). + void set_lower(const AdaptiveTier* lower) { lower_ = lower; } + /// An update waits while the previous one's moves are under way, instead of dropping the ones not staged yet (a + /// tier whose moves all go out before the next window). + void set_wait(bool on) { wait_ = on; } + + /// Ranks and queues this call's moves, then decays `usage` (n_layers x n_expert) unless `decay` is false (a tier + /// ranked before another on the same counts). The queued moves whose residents are still in place give way to + /// this call's (ranked again if still worth it); the experts still under way are not candidates. + bool adapt(std::vector& usage, std::string& err, bool decay = true); + + /// By default the copies go out in pieces (the first GPU's tier: its copies cross the link that the verify + /// window's own transfers need, and slowed them while they ran). `stage` evicts the next queued moves' residents + /// until `bytes` are staged and not sent: called before a window, whose residency snapshot then leaves their slots + /// alone. `pump_bytes` sends the staged moves' next `bytes`, a blob in as many calls as it takes. + void stage(uint64_t bytes); + bool pump_bytes(uint64_t bytes, std::string& err); + + /// Paced (the second GPU's tier, its copies between windows): whole moves, each resident evicted as its copy goes + /// out. `pump` submits them in order while they fit `budget` bytes, `sent` the bytes submitted, after `ev` (on + /// the cache's GPU: the last work that may read the slots they overwrite). + void set_paced(cudaEvent_t ev) { paced_ = true; after_ = ev; } + bool pump(uint64_t budget, uint64_t& sent, std::string& err); + + /// A running average of the copies' measured rate (bytes per ms), 0 until the first timed ones have landed. + double copy_rate() const { return rate_; } + /// Admits the moves whose copies have landed into `host_res`; `wait` sends the queued ones and + /// blocks until they have landed. + void apply_pending(bool wait); + + int64_t swaps = 0, fills = 0; ///< experts moved into an occupied / an empty slot + int64_t dropped = 0; ///< moves the next update dropped before their residents left + uint64_t sent_bytes = 0; ///< bytes copied + double ms = 0; ///< host time in `adapt` + +private: + struct Move { int32_t layer, in, out, slot; }; // out: the evicted expert, < 0 for an empty slot + ExpertCache* cache_ = nullptr; + ExpertSource* src_ = nullptr; + std::vector* res_ = nullptr; + int64_t n_layers_ = 0, n_expert_ = 0; + int max_moves_ = 0; + int dev_ = -1, main_ = 0; + bool paced_ = false; + cudaEvent_t after_ = nullptr; + const AdaptiveTier* upper_ = nullptr; + const AdaptiveTier* lower_ = nullptr; + bool wait_ = false; + std::vector blocked_; // per (layer, expert): not a candidate, rebuilt each call + std::vector upper_has_; // per (layer, expert), rebuilt each call + std::vector> free_; // per layer + // The queued moves in order and their (residency index, slot) once landed: [0, admitted_) admitted, up to sent_ + // sent (and off_ bytes of the next), up to staged_ with their residents evicted. + std::vector queued_; + std::vector> pending_; + size_t admitted_ = 0, sent_ = 0, staged_ = 0; + uint64_t off_ = 0; + cudaEvent_t t0_ = nullptr, t1_ = nullptr; // around the last timed copies, whose `timed_` bytes have not landed + uint64_t timed_ = 0; + double rate_ = 0; + std::string failed_; // a copy apply_pending could not submit + cudaStream_t stream_ = nullptr; + // the batches in flight, oldest first: a ring of events, each with the moves sent when it was recorded + static constexpr int kBatches = 64; + cudaEvent_t evs_[kBatches] = {}; + size_t batch_sent_[kBatches] = {}; + int first_ = 0, flying_ = 0; + // a call's copies, sent as one batch (copy_blobs) + std::vector cp_dst_; + std::vector cp_src_; + std::vector cp_bytes_; + uint64_t bytes_of(const Move& m) const; + void evict(const Move& m); + bool copy(const Move& m, uint64_t off, uint64_t n, std::string& err); + bool send(std::string& err); + bool time_begin(bool worth); + bool end_batch(bool timed, uint64_t bytes, std::string& err); +}; + +} // namespace strata::core diff --git a/include/strata/core/expert_cache.hpp b/include/strata/core/expert_cache.hpp index bf3b5831375..1675d84a50f 100644 --- a/include/strata/core/expert_cache.hpp +++ b/include/strata/core/expert_cache.hpp @@ -162,4 +162,9 @@ class ExpertCache { int64_t admitted_ = 0; }; +/// Copies `n` pinned host blobs into VRAM on `stream` as one submission (cudaMemcpyBatchAsync: the copies in any +/// order, the batch in the stream's). Under WDDM a copy submitted alone raises an interrupt when it completes, on the +/// CPU its GPU's interrupts go to (~10 us there, in the ISR and a DPC); a batch raises one. +bool copy_blobs(void* const* dst, const void* const* src, const size_t* bytes, size_t n, void* stream); + } // namespace strata::core diff --git a/include/strata/core/expert_source.hpp b/include/strata/core/expert_source.hpp index 0b357a1b6f9..63354985455 100644 --- a/include/strata/core/expert_source.hpp +++ b/include/strata/core/expert_source.hpp @@ -26,49 +26,12 @@ #include "strata/kernels/cpu/pool.hpp" #include -#include -#include #include #include namespace strata::core { -class RemoteExperts; - -namespace detail { - -/// Sentinel used by the pure complement planner for a blob that remains in the mmap fallback. -inline constexpr uint64_t kNoCacheComplement = ~uint64_t{0}; - -/// Required cgroup-v2 usage counters for the conservative cache-reclaim allowance. -struct CgroupMemoryStat { - uint64_t current = 0; - uint64_t inactive_file = 0; - uint64_t file_dirty = 0; - uint64_t file_writeback = 0; - bool valid = false; -}; - -/// Calculate additional bytes under a finite cgroup limit after reclaiming only clean inactive file cache. -/// Returns false when the required memory.stat counters were unavailable. -bool cgroup_available_bytes(uint64_t limit, const CgroupMemoryStat& stat, uint64_t& bytes); - -/// Build compact offsets for experts absent from both the primary GPU cache and an optional second GPU tier. -/// Kept CPU-only so selection and byte accounting can be tested without initializing a GPU. -bool make_cache_complement_plan( - int64_t n_layers, int64_t n_expert, const std::vector& layer_blob_bytes, - const std::vector>& primary_gpu_pairs, - const std::vector>& additional_gpu_pairs, - std::vector& offsets, uint64_t& bytes, std::string& err); - -/// Resolve one blob through the compact copy when present, otherwise preserve its exact mapped-file fallback. -const uint8_t* cache_complement_blob_or_fallback( - size_t index, const std::vector& offsets, const uint8_t* complement_host, - const uint8_t* mapped_fallback); - -} // namespace detail - - +class SecondGpu; /// Where one routed expert's bytes come from. /// @@ -79,7 +42,7 @@ class ExpertSource { public: virtual ~ExpertSource() = default; - /// The expert-layout blob for `(layer, expert)`, or nullptr if it cannot be produced. + /// The 1,382,400-byte blob for `(layer, expert)`, or nullptr if it cannot be produced. /// /// The pointer only has to stay valid until the next `blob()` call: with `h = 0` every expert is computed /// immediately and nothing is retained. A CACHING source must return pointers into the cache, not into a @@ -100,23 +63,22 @@ class ExpertSource { virtual const uint8_t* device_alias(int64_t layer, int64_t expert) const { (void) layer; (void) expert; return nullptr; } }; -/// Plan v0.3 P6: what the GPU computes in a verify window's layer, written by the pool (mapped host memory) right -/// after the ring and published before the CPU starts its own share. Groups of entries that share a blob; the -/// blob is a VRAM slot or a pinned host blob read over PCIe. +/// Plan v0.3 P6: the PCIe share of a verify window's layer - missed experts the main GPU reads over PCIe - written by +/// the pool (mapped host memory) right after the ring and published before the CPU starts its own share. (The +/// experts in its VRAM tier the window decides itself, from the residency the pool reads too.) Groups of entries +/// that share a blob: `ptr2[q]` is staging slot q (DMA) or the arena's device alias, its entries dst/tok +/// [start2[q], start2[q + 1]). counts[1] = entries, counts[2] = groups. struct GpuPlanSink { - int32_t* counts = nullptr; ///< [0] groups, [1] entries - int32_t* start = nullptr; ///< cap + 1 + int32_t* counts = nullptr; int32_t* dst = nullptr; ///< cap: the entry's row of parts (token * k + j) int32_t* tok = nullptr; ///< cap: the entry's token - unsigned long long* ptr = nullptr; ///< cap: the group's blob, device address - /// Plan v0.3 P6 (DMA): the PCIe share is a second list of groups - `ptr2[q]` is staging slot q, `start2` indexes - /// the same dst/tok entries - computed by the GPU after `fetch` has copied their blobs into staging with the - /// copy engine. counts[0] = VRAM groups, counts[1] = all entries, counts[2] = PCIe groups. - unsigned long long* ptr2 = nullptr; - int32_t* start2 = nullptr; + unsigned long long* ptr2 = nullptr; ///< cap + int32_t* start2 = nullptr; ///< cap + 1 unsigned long long staging = 0; int64_t staging_cap = 0; int64_t cap = 0; + /// Whether the window takes a PCIe share (`--pcie-frac` > 0): the pool publishes one every layer, else none. + bool pcie = false; void (*publish)(void* ctx) = nullptr; /// Starts the DMA copies of `n` host blobs (pinned) into staging slots 0..n-1 and signals the GPU when they land. void (*fetch)(void* ctx, const uint8_t* const* src, int n, size_t bytes) = nullptr; @@ -125,14 +87,31 @@ struct GpuPlanSink { /// kernel reads the mapped arena directly; 2 = a copy kernel stages it inside the graph. For 1 and 2 `ptr2` /// holds the arena's device alias. int pcie_mode = 0; + /// The GPU reads only the rows it did not compute, so the pool leaves the GPU's rows of `out` unwritten. + bool host_rows_only = false; + /// The second GPU's share of this token group (mapped; null without one): its entries, [n, pad x3, entries], and + /// the flag its graph raises to `ring` once their rows are in `out`, which the main GPU waits for. The pool + /// raises the flag itself when the second GPU takes no share, and sets `gpu2_ring` to `ring` once the flag is + /// raised or a share is submitted. + int32_t* gpu2_list = nullptr; + uint32_t* gpu2_flag = nullptr; + uint32_t ring = 0, gpu2_ring = 0; + /// This token group's activations as the window's router wrote them (mapped): ggml's Q8_K rows, `xk_stride` bytes + /// apart, when the layer takes them (`pool_takes_q8k`), else null (the pool quantizes the floats); and the q8_1 + /// rows (n_embd / 32 blocks each) the second GPU takes. + const uint8_t* xk = nullptr; + int64_t xk_stride = 0; + const uint8_t* x1 = nullptr; }; +/// Whether the pool takes a native layer's gate/up activation as ggml's Q8_K rows (`GpuPlanSink::xk`): its format's +/// vec_dot_type (the i-quants', the K-quants'), and not the Q2_0 kernels' own. +inline bool pool_takes_q8k(const strata::kernels::cpu::NativeFmt& f) { return f.gu_type != 42 && f.gu_act == 15; } + /// The adapter's own state. One per session, reused every layer so the token path allocates nothing (P2.T10). struct ExpertDispatch { strata::kernels::cpu::ExpertPool* pool = nullptr; ExpertSource* src = nullptr; - RemoteExperts* remote[3] = {}; ///< optional CUDA1..3 tiers for otherwise CPU-served rows - int remote_count = 0; int64_t n_expert = strata::kernels::cpu::NE; /// Counters, for the driver to report rather than for control flow. @@ -193,9 +172,9 @@ struct ExpertDispatch { /// so the decision is made exactly once, on this layer's ids, and neither side can re-decide it. std::vector is_hit; bool decided = false; - /// Plan v0.3 P4 token graph: the STATIC residency table (`n_layers x n_expert`, slot or -1), the host's copy - /// of what the device hit path reads. When set, the pool leaves a resident expert's row at zero (the GPU - /// computes it) without any `Launch` callback. + /// Plan v0.3 P4 token graph: the residency table (`n_layers x n_expert`, slot or -1), the host's copy of what the + /// device hit path reads (in verify windows `Verifier::residency`). When set, the pool leaves a resident + /// expert's row to the GPU without any `Launch` callback. const int32_t* host_res = nullptr; /// Plan v0.3 P4: split every expert by rows across the pool's threads (default on; A/B `--no-split-rows`). bool split_rows = true; @@ -242,8 +221,8 @@ struct ExpertDispatch { std::vector nact_multi; std::vector jobs_multi; std::vector job_of; - /// Plan v0.3 P6: the verify window's GPU plan (VRAM hits + the PCIe share of the misses); `pcie_num`/256 of - /// each layer's distinct missed experts (the last ones in routing order) are read by the GPU over PCIe. + /// Plan v0.3 P6: the verify window's PCIe share; `pcie_num`/256 of each layer's distinct missed experts (the last + /// ones in routing order) are read by the GPU over PCIe. GpuPlanSink* plan = nullptr; int pcie_num = 0; int64_t pcie_experts = 0; ///< distinct experts the GPU read over PCIe in verify windows @@ -253,6 +232,19 @@ struct ExpertDispatch { std::vector usage; int64_t multi_misses = 0; ///< distinct (layer, expert) pairs the CPU computed in verify windows int64_t multi_entries = 0; ///< routed (token, expert) entries the CPU served in verify windows + /// Routed (token, expert) entries per (layer, expert) in verify windows, when the caller sizes it: how many of + /// the entries a cache of the N most-routed experts would serve. + std::vector routed; + /// A second GPU's expert tier (verify windows) and its residency table (n_layers x n_expert, slot or -1): the + /// experts it holds or has prefetched and the first GPU's cache does not are computed there, beside the CPU pool. + SecondGpu* gpu2 = nullptr; + const int32_t* host_res2 = nullptr; + int64_t gpu2_entries = 0; ///< routed entries the second GPU served + int64_t gpu2_skipped = 0; ///< layers the CPU took the second GPU's share of (it was quicker) + bool gpu2_used = false; ///< this layer's decision + uint64_t gpu2_min_bytes = 0; ///< a layer's miss bytes from which the second GPU takes its share + uint64_t gpu2_prefetch_bytes = 0; ///< per layer, its likeliest experts no GPU holds, copied there ahead + std::string gpu2_err; ///< `fail` points here when the second GPU failed /// Set when `dispatch` could not produce an answer. The loop itself has no error channel, so this is /// where a source failure surfaces: the driver checks it after `session_loop` returns rather than the /// engine computing from a half-filled `parts`. @@ -273,9 +265,14 @@ void expert_pool_dispatch(void* user, const float* x_f, const int32_t* ids, cons /// Plan v0.3 P6: the pool for a verify window of `n_tok` tokens. `x_f` is (n_tok, n_embd), `ids` (n_tok, k) and /// `out` (n_tok * k, n_embd). Each distinct missed expert is computed once for all the tokens routed to it; -/// resident experts' rows are zeroed (the GPU adds them). Requires `host_res` (the token-graph residency). +/// resident experts' rows are the main GPU's (zeroed, or left unwritten with `GpuPlanSink::host_rows_only`), and the +/// second GPU writes its share's rows itself, after this returns (`GpuPlanSink::gpu2_flag`). void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k, float* out); +/// Copies `layer`'s likeliest experts that neither GPU holds to the second GPU's prefetch slots, where that layer's +/// dispatch finds them; `ids`/`w` (n_tok, k) are the layer's predicted routing (`Verifier::set_predict`). +void expert_prefetch_multi(ExpertDispatch& d, int64_t layer, const int32_t* ids, const float* w, int64_t n_tok, + int64_t k); /// **THE HITS, LAUNCHED AFTER THE MISSES ARE STAGED AND BEFORE `post[l]`.** Same shape as `PoolFn` and for the /// same reason: `session_loop` owns the ORDER and this owns the work, so the loop needs to know nothing about @@ -307,49 +304,26 @@ class FileExpertSource : public ExpertSource { FileExpertSource(const FileExpertSource&) = delete; FileExpertSource& operator=(const FileExpertSource&) = delete; - /// Maps `/experts.bin` and checks its size against the loaded expert layout. Canonical packs use - /// `n_layers * n_expert * BLOB`; native packs use their variable per-layer blob sizes and offsets. + /// Maps `/experts.bin` and checks its size against `n_layers * n_expert * BLOB`. /// /// The size check is not a formality: a short file would fault at the END of a long sequence, and an /// over-long one means the pack is not the one the geometry came from. Refuses with the two numbers. bool open(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, std::string& err); - /// Pin a compact host mirror of experts absent from a fully filled static GPU cache. The mmap remains open - /// as a fallback for later cache reloads. This is opt-in because the complement may still be a large allocation. - bool pin_cache_complement( - const ExpertCache& cache, std::string& err, bool pin = true, - const std::vector>& additional_gpu_pairs = {}); void close(); bool mapped() const { return base_ != nullptr; } int64_t blobs() const { return blobs_; } - uint64_t pinned_bytes() const { return complement_pinned_ ? complement_bytes_ : 0; } - uint64_t resident_bytes() const { return complement_bytes_; } - bool complement_pinned() const { return complement_pinned_; } const uint8_t* blob(int64_t layer, int64_t expert) override; - bool pinned(int64_t layer, int64_t expert) const override; - const uint8_t* device_alias(int64_t layer, int64_t expert) const override; /// Blobs touched, for the driver to report. With `h = 0` this is `48 * k` per token and the number is only /// interesting once Phase 3 makes it not so. int64_t reads() const override { return reads_; } private: - const uint8_t* mapped_blob(int64_t layer, int64_t expert) const; - static constexpr uint64_t kNoComplement = detail::kNoCacheComplement; const uint8_t* base_ = nullptr; int64_t blobs_ = 0; - int64_t n_layers_ = 0; int64_t n_expert_ = 0; - uint64_t mapped_bytes_ = 0; - std::vector layer_offsets_, layer_blob_bytes_; - void* complement_arena_ = nullptr; - const uint8_t* complement_host_ = nullptr; - const uint8_t* complement_device_ = nullptr; - uint64_t complement_bytes_ = 0; - std::vector complement_offsets_; - bool complement_pinned_ = false; - bool complement_ready_ = false; int64_t reads_ = 0; #if defined(_WIN32) void* file_ = nullptr; @@ -372,9 +346,12 @@ class FileExpertSource : public ExpertSource { // ordinary (anonymous, resident) memory; the engine reads `MapViewOfFile`. That is the whole difference, and it // is why this class exists. // -// It reads `experts.bin` into a `PinnedArena` once at startup, so the expert stream comes from anonymous memory -// the OS has no cheaper reason to evict. `PinnedArena` also tries `cudaHostRegister`, which the GPU needs for -// Phase 3's cache fills and the CPU/PCIe miss split - but registration is best-effort and reported, not assumed. +// It reads the experts once at startup into anonymous memory the OS has no cheaper reason to evict: one pinned +// block per layer (cudaHostAlloc, portable and mapped), which the GPUs need for cache fills, the prompt path and the +// PCIe share of the misses. Registering one reserved range instead (cudaHostRegister) is limited to ~42 GiB on +// Windows once a second GPU's context exists, and past that limit every CUDA allocation fails; allocated blocks reach +// ~75 GiB with one or two GPUs. The layers past the pinning budget share one locked, pageable range; the note says +// how many layers were pinned. class ArenaExpertSource : public ExpertSource { public: ArenaExpertSource() = default; @@ -384,13 +361,13 @@ class ArenaExpertSource : public ExpertSource { /// Allocates and loads `/experts.bin`. Prints nothing; the caller reports `note()` and the load /// rate, because those are the two numbers that say whether the arena is the one that was asked for. - bool open(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, int threads, std::string& err, - uint64_t max_pinned_bytes = 0); - /// Plan v0.3 P6: a native pack without experts.bin takes its experts from the model's shard 1. - void set_gguf(const std::string& shard1) { gguf_ = shard1; } + bool open(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, int threads, std::string& err); + /// Plan v0.3 P6: a native pack without experts.bin takes its experts from the model's GGUF shards + /// (strata::gguf_split_paths), found by tensor name. + void set_gguf(const std::vector& shards) { gguf_ = shards; } void close(); - bool mapped() const { return base_ != nullptr; } + bool mapped() const { return !layer_base_.empty(); } int64_t blobs() const { return blobs_; } const uint8_t* blob(int64_t layer, int64_t expert) override; int64_t reads() const { return reads_; } @@ -401,28 +378,19 @@ class ArenaExpertSource : public ExpertSource { /// only true if the engine says what it got. const std::string& note() const { return note_; } double load_gib_per_second() const { return gib_per_s_; } - // Loader fix: the load, split. `load_seconds()` is the wall clock of the load loop; the other two are - // sums over the reader threads (see LoadStats), so on their own they say how much of that wall was spent - // waiting for the disk and how much in memcpy + FNV-1a. - double load_seconds() const { return load_seconds_; } - double load_read_seconds() const { return load_read_s_; } - double load_copy_seconds() const { return load_copy_s_; } private: - void* arena_ = nullptr; ///< the PinnedArena, owned - std::vector dev_slice_; ///< device alias of each registered slice (or of the whole range) - uint64_t slice_bytes_ = 0; - const uint8_t* base_ = nullptr; + std::vector layer_base_; ///< each layer's experts: a pinned block, or a part of `tail_` + std::vector layer_dev_; ///< the device alias of each pinned layer (null for the others) + std::vector blocks_; ///< the pinned blocks, owned (layers [0, blocks_.size())) + void* tail_ = nullptr; ///< the layers past the pinning budget, locked resident, owned + uint64_t tail_bytes_ = 0; int64_t blobs_ = 0; int64_t n_expert_ = 0; int64_t reads_ = 0; std::string note_; double gib_per_s_ = 0.0; - double load_seconds_ = 0.0; - double load_read_s_ = 0.0; - double load_copy_s_ = 0.0; - uint64_t pinned_bytes_ = 0; - std::string gguf_; + std::vector gguf_; }; } // namespace strata::core diff --git a/include/strata/core/layer.hpp b/include/strata/core/layer.hpp index 1669256f340..b3f029c98a8 100644 --- a/include/strata/core/layer.hpp +++ b/include/strata/core/layer.hpp @@ -36,12 +36,9 @@ #include "strata/core/weights.hpp" #include "strata/kernels/gr.hpp" -#include "strata/kernels/kv_stream.hpp" #include "strata/kernels/ngram.hpp" #include "strata/kernels/ple.hpp" -#include -#include #include #include @@ -209,28 +206,10 @@ struct QsaState { int8_t* v_q = nullptr; uint16_t* k_scale = nullptr; uint16_t* v_scale = nullptr; - /// PR #21: Q4_0 KV with Walsh-Hadamard rotation (qsa_set_kv_q4, kv_q4.hpp): 144 B per cell and head - bool kv_q4 = false; - uint8_t* k_q4 = nullptr; - uint8_t* v_q4 = nullptr; - /// Hybrid K8V4: K in INT8 (unrotated - the scores stay exact), V in rotated Q4_0 (kv_q4.hpp): 816 B per cell. - /// Uses k_q/k_scale + v_q4. Mode 0 only (no KV streaming, no ring); under this setting the MTP drafter's - /// ring state stays plain INT8, so the block movers never see the hybrid layout. - bool kv_hybrid = false; - int32_t* page_table = nullptr; ///< (n_pages,) logical page -> physical page (-1: not resident, streamed) + int32_t* page_table = nullptr; ///< (n_pages,) logical page -> physical page int64_t n_pages = 0; int64_t max_cells = 0; - /// KV STREAMING (docs/kv-streaming-design.md, `kv_stream.hpp`). `kv_mode` 0: every page in VRAM, identity - /// table (n_slots == n_pages, no host copy). 1: streamed - the authoritative K/V in `host`, `n_slots` pages - /// in VRAM, `map` resolves the selection's blocks on device. 2: a ring (the MTP drafter's window) - `n_slots` - /// pages at `block % n_slots`, with the host copy for a resume. The pool pointers above point at the slots. - int kv_mode = 0; - int64_t n_slots = 0; - strata::kernels::KvHostPools host; - strata::kernels::KvStreamMap map; - int64_t idx_pooled_rows = 0; ///< rows of `idx_pooled` (a ring, which has no indexer, keeps 2) - float* idx_tail = nullptr; ///< (idx_block - 1, idx_dim): the raw tail of the block being filled float* idx_dead = nullptr; ///< (idx_dim,): the spare slot's key, CONSTANT for the sequence float* idx_pooled = nullptr; ///< (max_cells/idx_block + 2, idx_dim) @@ -263,44 +242,13 @@ struct QsaState { /// Plan v0.3 P7: the RoPE cos/sin table (max_cells x n_rot/2 x 2 floats, 64 MiB at 262K) is identical in every /// QSA layer. `with_rope = false` sizes a state that borrows it; `share_rope` points `st` at another state's table /// instead of building a copy (the session builds it once, in the first QSA layer). -uint64_t qsa_state_bytes(const ModelGeometry& g, int64_t max_cells, bool with_rope = true, int64_t ring_cells = 0); -/// KV streaming: keep `cells` cells of each QSA layer in VRAM and the rest in pinned host memory (0: all in VRAM, -/// the default). Set before sizing and initializing the session; a context that fits in `cells` is not streamed. -/// Also puts the MTP drafter's K/V in a ring of its window (`ring_cells` of qsa_state_bytes/init; -1 forces a fully -/// resident state). -void qsa_set_kv_resident(int64_t cells); -int64_t qsa_kv_resident(); -/// The fewest resident cells a streamed layer may have: one verify window's selections (8 queries x 2,051 cells -/// in whole blocks) must fit at once, with room to spare. -int64_t qsa_kv_resident_min(); -/// Pinned host bytes the streamed states hold (their host copies). -uint64_t qsa_kv_host_bytes(); +uint64_t qsa_state_bytes(const ModelGeometry& g, int64_t max_cells, bool with_rope = true); /// Plan v0.3 P7: store K/V as INT8 with FP16 scales per 64 values (half the VRAM of FP16). Set before sizing and /// initializing the session; default off until gate G-C accepts it. void qsa_set_kv_int8(bool enabled); bool qsa_kv_int8(); -/// PR #21: store K/V as Q4_0 after a Hadamard rotation (`--kv q4_0`): 576 B per cell, vs 1,056 in INT8. -void qsa_set_kv_q4(bool enabled); -bool qsa_kv_q4(); -/// Hybrid K8V4 (`--kv k8v4`): K in INT8, V in rotated Q4_0 - 816 B per cell. Not with --kv-resident. -void qsa_set_kv_hybrid(bool enabled); -bool qsa_kv_hybrid(); -/// The state's KV format for the block-moving functions of kv_stream.hpp (kKvF16 / kKvInt8 / kKvQ4). -inline int qsa_kv_format(const QsaState& st) { - // A hybrid K8V4 state is mode 0 only and never reaches the block movers; refuse rather than let it - // fall through to kKvF16 - a wrong layout silently applied is worse than a hard stop (PR review). - if (st.kv_hybrid) { - std::fprintf(stderr, "strata: qsa_kv_format: a hybrid K8V4 state must never reach the block movers\n"); - std::exit(1); // the kernels' own "unsupported geometry" convention (kv_q8.cu, qsa_decode_attn.cu) - } - return st.kv_q4 ? strata::kernels::kKvQ4 : st.kv_int8 ? strata::kernels::kKvInt8 : strata::kernels::kKvF16; -} uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, QsaState& st, - const QsaState* share_rope = nullptr, int64_t ring_cells = 0); -/// KV streaming: the pools a reader sees (the VRAM slots) and, when streamed, make the selection's blocks resident. -strata::kernels::QsaAttnPools qsa_attn_pools(const QsaState& st); -void qsa_kv_resolve(const QsaState& st, const ModelGeometry& g, const int32_t* ids, const int32_t* steps, int64_t n_q, - int64_t cap, void* stream); + const QsaState* share_rope = nullptr); /// Zeroes the pools AND the indexer, so a fresh sequence matches the reference's own `zeros()`. The KV pool /// matters even for cells that are never attended, because `kv_gather` reads whatever the selection names. void qsa_state_zero(const QsaState& st, const ModelGeometry& g, void* stream); @@ -433,6 +381,7 @@ struct BlockBuffers { /// because the fused read of one half reads the other half's injection while writing its own. float* inject2 = nullptr; ///< hc float* gr_rs = nullptr; ///< hc: the fused read's per-stream 1/rms + float* gr_scratch = nullptr; ///< `fused_gr_scratch_bytes()`: the fused read's scratch (session_zero) uint8_t* head_q8k = nullptr; ///< `q8k_bytes(n_embd)`: the head's Q8_K activation image strata::kernels::GrWorkspace gr; ///< `gr_read`'s scratch @@ -486,6 +435,9 @@ uint64_t block_buffers_init(const ModelGeometry& g, void* base, BlockBuffers& b) /// from the kernel's would make the prompt embed differently from the tokens generated after it. bool embed_row(const WeightTable& tables, const ModelGeometry& g, int64_t token, float* out_dev, void* stream, std::string& err); +/// embed_row for `n` token ids in DEVICE memory (checked by the caller) into out_dev [n, n_embd], one launch. +bool embed_rows(const WeightTable& tables, const ModelGeometry& g, const int32_t* tokens_dev, int64_t n, float* out_dev, + void* stream, std::string& err); /// `ref/model.py` L520: `fin = GR.gr_read(R, output_hc_*..., inject=None, rms_eps); logits = output @ fin`. /// diff --git a/include/strata/core/mtp.hpp b/include/strata/core/mtp.hpp index 6df843ba2f9..0480c668076 100644 --- a/include/strata/core/mtp.hpp +++ b/include/strata/core/mtp.hpp @@ -16,15 +16,20 @@ // * the attention is DENSE over every cell the layer has seen - identical to the model's sparse selection // below 2,051 cells - so speculative cells (draft steps, rejected window rows) never touch indexer state and // are simply overwritten when their positions are processed again; -// * all 512 routed experts live in VRAM (708 MB) and run through the grouped hit kernels. +// * all 512 routed experts live in VRAM and run through the grouped kernels: raw GGUF blocks in the formats +// `experts.txt` names, through the main model's native expert kernels, or else the Q2_0 blob layout (708 MB); +// * the K/V of a prompt's cells come from the batched prompt path (`prompt_kv`), in its GEMM arithmetic, as the +// model's own layers there. #pragma once #include "strata/core/layer.hpp" #include "strata/core/session.hpp" +#include "strata/kernels/iq_kernels.hpp" #include #include +#include #include #include @@ -32,6 +37,19 @@ namespace strata::core { class NativeHead; +/// What the batched prompt path needs for the draft layer's K/V over a prompt (strata/prefill/prefill.hpp): the two +/// input norms and projections, the attention hyper-connection read, the K and V projections and the K norm. +struct MtpPromptKv { + const float *norm_emb = nullptr, *norm_hidden = nullptr; ///< [n_embd], [hc * n_embd] + const void *fc_emb = nullptr, *fc_hidden = nullptr; ///< Q8_0 [n_embd, n_embd] + const float* hc_norm = nullptr; ///< [hc * n_embd] + const uint16_t *hc_down = nullptr, *hc_up = nullptr; ///< BF16 [hc_lr, hc * n_embd], [hc * n_embd, hc_lr] + const void *k_proj = nullptr, *v_proj = nullptr; ///< Q8_0 [n_head_kv * head_dim, n_embd] + const float* k_norm = nullptr; ///< [head_dim] + QsaState* kv = nullptr; ///< the layer's K/V cells + int64_t first_cell = 0; ///< cells before it the window never reaches again +}; + class MtpDrafter { public: MtpDrafter() = default; @@ -43,23 +61,18 @@ class MtpDrafter { /// Call before the VRAM expert tier is sized: this takes ~0.9 GB. bool load(const std::string& rt_dir, const ModelGeometry& g, SessionState& ss, int max_t, std::string& err, int64_t window = 32768); - /// The prompt's length: prefill() skips the cells the attention window can never reach again. + /// The prompt's length: the prompt's K/V skip the cells the attention window can never reach again. void set_prompt_len(int64_t n) { prompt_len_ = n; } - /// At most this many drafts per round (below max_t - 1): a window longer than the MTP's comes from elsewhere. - void set_max_drafts(int k) { max_drafts_ = k; } + /// The draft chain stops after `n` drafts (0: max_t - 1). + void set_max_drafts(int n) { max_drafts_ = n; } uint64_t vram_bytes() const { return vram_; } - /// The draft layer's K/V state (read-only: --serve's STRATA_STATE_HASH check hashes it) - const QsaState& kv_state() const { return st_; } - /// KV streaming: refill the ring of the drafter's window from its host copy for a sequence that continues at - /// `upto` (a conversation-cache resume). No-op unless the drafter's K/V is a ring. - void kv_restore(int64_t upto); + /// The layer's own K/V cells (the prompt cache stashes them with the model's). + QsaState& kv_state() { return st_; } + /// The weights and cells the batched prompt path computes the prompt's K/V with (after load and set_prompt_len). + MtpPromptKv prompt_kv(); /// The main model's embedding and head, and the verify window's final residuals (T rows, hc*n_embd each). bool bind(const WeightTable& wt, const NativeHead* head, const float* window_R, std::string& err); - /// Prompt cells [cell0, cell0 + n): residual rows `R_rows` (device, hc*n_embd each) and `next_tokens` (host, - /// the token at position cell+1). Runs in batches of up to max_t rows. - bool prefill(const float* R_rows, const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err); - /// One round: catch-up over T cells from `p` (rows = the window's final residuals, `tokens` = the window's /// argmaxes: row t pairs R_{p+t} with the token at p+t+1), then the draft chain from row `a` (the last /// accepted row) for T-1 drafts at cells p+a+1 ... `drafts` gets T-1 tokens. @@ -67,29 +80,29 @@ class MtpDrafter { float* probs = nullptr, float min_p = 0.0f, int* n_drafts = nullptr); /// The first round: one cell (`cell`) from `R_row` (device) and `token` -> T-1 drafts. + /// The draft head's rows gathered from the main head now (before the main GPU keeps only part of it); `bind` + /// then takes them as they are. + bool bind_head(const NativeHead* head, std::string& err); + bool has_draft_head() const { return dhead_ != nullptr; } bool draft_first(int T, const float* R_row, int32_t token, int64_t cell, int32_t* drafts, std::string& err, float* probs = nullptr, float min_p = 0.0f, int* n_drafts = nullptr); - double ms_draft = 0, ms_prefill = 0; - int64_t rounds = 0; + /// Called with each draft as soon as it is known (j = 0.., its probability under the draft layer), while the + /// chain goes on: the next window's PLE rows start there. + std::function on_draft; - /// E-9: the prompt path computes this layer's prompt K/V in batches (Prefill::draft_kv): its tensors, its K/V - /// state, the first cell a round can still read, its device, and a wait for its own stream. - const float* tensor_f32(const char* name) const { return f32(name); } - const uint16_t* tensor_bf16(const char* name) const { return bf16(name); } - const void* tensor_q8(const char* name) const { return q8(name); } - QsaState& kv_state_rw() { return st_; } - int64_t first_needed() const { return (window_ > 0 && prompt_len_ > 0) ? prompt_len_ - window_ - 64 : 0; } - int device() const { return device_; } - bool idle(std::string& err) { - if (cs_ && cudaStreamSynchronize(cs_) != cudaSuccess) { err = "mtp: its stream failed"; return false; } - return true; - } + double ms_draft = 0; + int64_t rounds = 0; private: bool record_forward(int T, int step_row0, cudaStream_t cs, std::string& err); - bool capture_prefill(int T, std::string& err); - bool capture_prefill_dev(int T, std::string& err); ///< E-4: without the mapped staging (inputs copied on device) + /// The layer's front for T rows at step rows [row0, +T): the embedding, the fc projections, the attention + /// hyper-connection read (R_, inj_, mixed_) and the K/V appended. + bool record_front(int T, int row0, cudaStream_t cs, std::string& err); + /// The rest for row 0 at step row `step_row`: the attention from the q projection, the MLP, the final mixer, + /// the head, the draft and its probability. Needs row 0's R_, inj_ and mixed_ (its q8_1 rows in xq_). + bool record_rest(int step_row, cudaStream_t cs, std::string& err); + void norm_rope(float* data, const float* gamma, int rows, int cols, const int32_t* p, cudaStream_t cs); bool capture_round(int T, std::string& err); bool capture_step(int j, std::string& err); cudaGraphExec_t step_exec_[9] = {}; @@ -102,22 +115,22 @@ class MtpDrafter { const WeightTable* wt_ = nullptr; const NativeHead* head_ = nullptr; const float* window_R_ = nullptr; - int max_t_ = 0; - int device_ = -1; ///< the device `load` ran on: the public calls switch to it (layer split) - int max_drafts_ = 1 << 30; + int max_t_ = 0, max_drafts_ = 0; int64_t n_vocab_ = 0; uint64_t vram_ = 0; cudaStream_t cs_ = nullptr; - cudaGraphExec_t prefill_exec_[9] = {}; - cudaGraphExec_t prefill_dev_exec_[9] = {}; - int32_t* pf_dev_ = nullptr; ///< E-4: a prompt's rows' token / step / position records, uploaded at once - int64_t pf_cap_ = 0; ///< its capacity in ints + cudaStream_t side_ = nullptr; // the shared expert's branch of the draft graphs + cudaEvent_t sh_fork_ = nullptr, sh_join_ = nullptr; + uint32_t* rseq_ = nullptr; // verify_router's sequence word and counter (device) + unsigned* rcount_ = nullptr; cudaGraphExec_t round_exec_[9] = {}; struct Tensor { std::string name, kind; int64_t rows = 0, cols = 0; uint64_t off = 0, bytes = 0; }; std::vector tensors_; uint8_t* dense_ = nullptr; uint8_t* experts_ = nullptr; + uint64_t blob_ = 0; // bytes per expert + strata::kernels::NativeExpertLayout nat_; // gu_type >= 0: GGUF blocks (experts.txt); else the Q2_0 blobs void* state_arena_ = nullptr; QsaState st_; void* arena_ = nullptr; @@ -136,17 +149,19 @@ class MtpDrafter { int64_t window_ = 0; // attention over the last window_ cells (0 = every cell) int64_t prompt_len_ = 0; float* probs_ = nullptr; + uint8_t* arg_scratch_ = nullptr; ///< argmax_rows' and row_top_prob_split's partials and counters + uint8_t* top_scratch_ = nullptr; // device int32_t *tok_ = nullptr, *step_ = nullptr, *pos_ = nullptr, *row_ = nullptr, *ident_ = nullptr; float *Rin_ = nullptr, *R_ = nullptr, *emb_ = nullptr, *en_ = nullptr, *e2_ = nullptr, *hn_ = nullptr, *h2_ = nullptr; float *mixed_ = nullptr, *inj_ = nullptr, *inj2_ = nullptr, *lo_ = nullptr, *rs_ = nullptr, *bo_ = nullptr; - float* xn_ = nullptr; + float* grs_ = nullptr; // the hyper-connection reads' scratch (zeroed with the arena) uint8_t* xq_ = nullptr; - float *qfull_ = nullptr, *qcur_ = nullptr, *kcur_ = nullptr, *vcur_ = nullptr, *attn_ = nullptr, *attn32_ = nullptr; + float *qfull_ = nullptr, *qcur_ = nullptr, *kcur_ = nullptr, *vcur_ = nullptr, *attn32_ = nullptr; float* attn_scratch_ = nullptr; float *logits_ = nullptr, *w_ = nullptr, *shared_ = nullptr, *parts_ = nullptr, *y_ = nullptr, *sample_ = nullptr; int32_t *ids_ = nullptr, *hit_slot_ = nullptr, *hit_dst_ = nullptr, *hit_count_ = nullptr, *out_ids_ = nullptr; - uint8_t* hit_xq_ = nullptr; + uint8_t *hit_xq_ = nullptr, *nat_xq_ = nullptr; unsigned long long* grp_ptr_ = nullptr; int32_t *grp_start_ = nullptr, *grp_counts_ = nullptr; float* hit_xs_ = nullptr; diff --git a/include/strata/core/native_head.hpp b/include/strata/core/native_head.hpp index 22b4c8db596..8a02843027f 100644 --- a/include/strata/core/native_head.hpp +++ b/include/strata/core/native_head.hpp @@ -3,12 +3,14 @@ #include #include #include +#include namespace strata::core { /// Opt-in native output projection. Uploads unchanged GGUF blocks once at startup and owns its /// activation scratch. The normal canonical head remains available for numerical A/B comparisons. -/// Plan v0.3 P6: any type the native MMVQ takes (Q5_K in the Q2_0 / IQ3_XXS files, IQ4_XS in IQ2_XS). +/// Plan v0.3 P6: any type the native MMVQ takes (Q5_K in the Q2_0 / IQ3_XXS files, IQ4_XS in IQ2_XS, Q8_0 in +/// UD-Q4_K_XL). `load` takes the model's shards (strata::gguf_split_paths). class NativeHead { public: NativeHead() = default; @@ -16,7 +18,7 @@ class NativeHead { NativeHead(const NativeHead&) = delete; NativeHead& operator=(const NativeHead&) = delete; - bool load(const std::string& gguf, int64_t n_in, int64_t n_out, std::string& err); + bool load(const std::vector& shards, int64_t n_in, int64_t n_out, std::string& err); bool run(const float* mixed, float* logits, void* stream, std::string& err) const; uint64_t weight_bytes() const { return bytes_; } bool loaded() const { return weights_ != nullptr; } @@ -25,6 +27,11 @@ class NativeHead { int type() const { return type_; } /// Bytes of one vocabulary row. size_t row_bytes() const { return n_out_ > 0 ? (size_t) (bytes_ / (uint64_t) n_out_) : 0; } + /// The rows held: all of the vocabulary, or the first ones after keep_rows. + int64_t rows() const { return n_out_; } + /// Keeps rows [0, rows) (the rest of the head runs on another GPU: split_head.hpp) and frees the others' VRAM. + /// `run` then refuses: it would give part of a token's logits. + bool keep_rows(int64_t rows, std::string& err); private: void* weights_ = nullptr; @@ -32,6 +39,7 @@ class NativeHead { uint64_t bytes_ = 0; int n_in_ = 0, n_out_ = 0; int type_ = -1; + bool part_ = false; }; /// Plan v0.3 P6: `token_embd.weight` in its GGUF form (the IQ model files), in mapped pinned host memory: a row @@ -43,7 +51,7 @@ class NativeEmbed { ~NativeEmbed(); NativeEmbed(const NativeEmbed&) = delete; NativeEmbed& operator=(const NativeEmbed&) = delete; - bool load(const std::string& gguf, int64_t n_embd, int64_t n_vocab, std::string& err); + bool load(const std::vector& shards, int64_t n_embd, int64_t n_vocab, std::string& err); /// Rows for device token ids. void gather_dev(const int32_t* tokens, int64_t n_tok, float* out, void* stream) const; /// One row for a host token id. diff --git a/include/strata/core/pinned.hpp b/include/strata/core/pinned.hpp index cc4ef1cb808..b4c094d76c6 100644 --- a/include/strata/core/pinned.hpp +++ b/include/strata/core/pinned.hpp @@ -39,8 +39,7 @@ struct PinnedArena { explicit PinnedArena(uint64_t bytes, uint64_t slice = 0); /// Plan v0.3 P6: slices of different sizes (one per layer of a native pack), given as their start offsets /// followed by the end of the last one. `slice_starts` holds the registered ones. - /// `max_pinned_bytes`: optional cap on CUDA registration. 0 preserves the normal unrestricted path. - PinnedArena(uint64_t bytes, const std::vector& bounds, uint64_t max_pinned_bytes = 0); + PinnedArena(uint64_t bytes, const std::vector& bounds); std::vector slice_starts; ~PinnedArena(); PinnedArena(const PinnedArena&) = delete; @@ -52,43 +51,22 @@ struct PinnedArena { struct LoadStats { double seconds = 0.0; - // Loader fix: the two halves of `seconds`, both SUMMED OVER THE READER THREADS rather than wall clock, so - // either can exceed `seconds` when the threads overlap. They exist to separate the read from the copy and - // the hash: an 8 MiB chunk that reaches the disk as ~2048 4 KiB reads spends its time in the first one. - double read_seconds = 0.0; // inside the read call only: no seek, no copy, no hash - double copy_seconds = 0.0; // memcpy + FNV-1a - // A short read is not a slow load, it is a WRONG one: the caller must refuse the pack rather than run - // on an arena whose tail was never written. `seconds` alone cannot say that (it is -1.0 on failure, but - // the caller's own check is on the REQUESTED byte count, which a refused read does not change). - bool ok = true; // false: the load failed, see `error` - std::string error; // why it failed, for the caller's message uint64_t bytes = 0; uint64_t layers = 0; std::vector layer_checksums; // one FNV-1a per layer double gib_per_second() const { return seconds > 0 ? (double) bytes / (1024.0 * 1024 * 1024) / seconds : 0.0; } - // Aggregate rate of the readers WHILE THEY WERE INSIDE THE READ CALL: `bytes` over the mean per-thread - // read time. `gib_per_second()` above divides the same bytes by the wall clock of the whole loop, which - // also contains the copy and the hash, so the two differ by how much of the loop was not I/O. - double read_gib_per_second(int threads) const { - if (read_seconds <= 0.0 || threads < 1) return 0.0; - return (double) bytes / (1024.0 * 1024 * 1024) / (read_seconds / (double) threads); - } }; // Load `layers` layers of the expert arena into `dst` with `threads` readers, `chunk` bytes at a time. // Each thread opens its OWN handle and seeks, which is the portable form of parallel pread: a shared handle // needs a lock around the seek and defeats the parallelism on Windows. -// -// The reads go through `fread()` on a per-thread `FILE*`, NOT through `std::ifstream`. MSVC's -// `basic_filebuf::xsgetn` splits every request larger than `_INTERNAL_BUFSIZ - 1` into 4095-byte `fread()` -// calls, so an 8 MiB `chunk` reached the disk as ~2048 4 KiB operations and the chunking did nothing at all -// (measured on a 42.9 GB pack: 1222 s, 4096 bytes per operation, 0.03 GiB/s). `fread()` hands a request -// larger than the stream buffer straight to `_read()`/`ReadFile()`, which is what `chunk` is for. The -// control flow, the copy, the per-layer FNV-1a and therefore the checksums are unchanged, so a pack loaded -// by the old and the new reader must produce identical `layer_checksums`. LoadStats load_experts(const std::string& path, uint8_t* dst, uint64_t blob_bytes, uint64_t blobs_per_layer, uint64_t layers, int threads, uint64_t chunk); /// Plan v0.3 P6: the same with one byte range per layer (`layer_off[L]`, `layer_bytes[L]`). +/// Layer L's `layer_bytes[L]` bytes at file offset `layer_off[L]` go to `layer_dst[L]`. +LoadStats load_experts_ranges(const std::string& path, const std::vector& layer_dst, + const std::vector& layer_off, const std::vector& layer_bytes, int threads, + uint64_t chunk); LoadStats load_experts_ranges(const std::string& path, uint8_t* dst, const std::vector& layer_off, const std::vector& layer_bytes, int threads, uint64_t chunk); diff --git a/include/strata/core/prompt_cache.hpp b/include/strata/core/prompt_cache.hpp new file mode 100644 index 00000000000..7d9ed19b455 --- /dev/null +++ b/include/strata/core/prompt_cache.hpp @@ -0,0 +1,97 @@ +// include/strata/core/prompt_cache.hpp - the --serve prompt cache. +// +// A request continues from the longest prefix of its prompt that the engine still holds, instead of processing +// the whole conversation again. Three things hold prefixes: +// +// * the LIVE sequence: the tokens whose cells the session state holds right now; +// * CHECKPOINTS of the live sequence (session_ckpt_*): the fixed-size state at one position. The K/V cells +// are indexed by position, so restoring a checkpoint at `pos` and overwriting the cells from `pos` on +// continues the sequence from there; +// * STASHED sequences: a sequence the live one replaced, kept in host memory with its K/V cells and its +// checkpoints, so a client that alternates between conversations (a chat and its title request, an agent +// and its subagents) does not start over on every switch. +// +// Checkpoint buffers are a shared pool of `max_ckpts`; the least recently used one goes when it is full. The +// stash's K/V copies share a byte budget in the same way. +#pragma once + +#include "strata/core/session.hpp" + +#include +#include +#include +#include + +namespace strata::core { + +class PromptCache { +public: + /// `kv`: every state whose cells a stashed sequence keeps (the QSA layers and the draft layer). + PromptCache(const ModelGeometry& g, SessionState& ss, std::vector kv, void* stream, int max_ckpts, + uint64_t stash_bytes); + ~PromptCache(); + PromptCache(const PromptCache&) = delete; + PromptCache& operator=(const PromptCache&) = delete; + + bool enabled() const { return max_ckpts_ > 0; } + uint64_t ckpt_bytes() const { return ckpt_bytes_; } + + /// Sets the session up for a prompt of `ids`: continues from the longest prefix (at most ids.size() - 1 + /// cells, the last token always goes through a verify window) that the live sequence, a checkpoint or a + /// stashed sequence holds, or zeroes the session. With `cacheable` false (an image request, whose cells the + /// ids do not describe) it zeroes the session and the live sequence is not kept. Returns the cells reused. + int64_t begin(const std::vector& ids, bool cacheable); + /// The cells [pos, pos + n) now hold `tokens`; pos <= live length, and anything after pos is gone. + void set(int64_t pos, const int64_t* tokens, int64_t n); + void set(int64_t pos, const int32_t* tokens, int64_t n); + /// Checkpoints the session, which is at `pos` == the live length. + void checkpoint(int64_t pos); + /// Forgets the live sequence (after an image request). + void forget(); + + /// One line for the log: checkpoints and stashed sequences. + std::string summary() const; + +private: + struct Ckpt { int64_t pos; void* host; uint64_t stamp; }; + struct Seq { + std::vector tokens; + std::vector ckpts; + std::unique_ptr kv; // cells [0, kv_cells) of every K/V state (stashed sequences only) + int64_t kv_cells = 0; + uint64_t kv_bytes = 0; + uint64_t stamp = 0; + }; + + /// One device region of the K/V cells, stored back to back in a stashed sequence's buffer. + struct Part { void* dev; uint64_t bytes; }; + + /// Cells of `seq` a prompt can continue from: the whole sequence when it is a prefix of the prompt (live + /// only: the state is there), else its best checkpoint within the common prefix (its position; -1: none). + int64_t reusable(const Seq& seq, const std::vector& ids, bool live) const; + static Ckpt* find(Seq& seq, int64_t pos); + void* take_buffer(); + void release(Seq& seq, int64_t keep_upto); // frees the checkpoints after `keep_upto` + /// Moves a copy of the live sequence to the stash: its K/V cells, a checkpoint of its end, and its + /// checkpoints after `keep_upto` (the ones before stay with the live sequence, which continues from there). + void stash_live(int64_t keep_upto); + std::vector kv_parts(int64_t cells) const; + bool kv_copy(const std::vector& parts, uint8_t* host, bool to_host); + bool restore(Seq& seq, int64_t pos); // K/V cells (stashed only) and the checkpoint at `pos` + + const ModelGeometry& g_; + SessionState& ss_; + std::vector kv_; + void* stream_; + int max_ckpts_; + uint64_t stash_budget_; + uint64_t ckpt_bytes_; + int allocated_ = 0; + uint64_t clock_ = 0; + Seq live_; + std::vector stash_; + std::vector free_; + std::vector> owned_; // every buffer, and whether it is pinned +}; + +} // namespace strata::core diff --git a/include/strata/core/second_gpu.hpp b/include/strata/core/second_gpu.hpp new file mode 100644 index 00000000000..7348bc2494d --- /dev/null +++ b/include/strata/core/second_gpu.hpp @@ -0,0 +1,99 @@ +// include/strata/core/second_gpu.hpp - a second GPU as another expert tier. +// +// It holds experts the first GPU's cache does not (the profile's next ranks, then the conversation's, through its +// own AdaptiveTier) and computes them for each layer of a verify window while the CPU pool computes the misses. +// It never spins on host memory: per layer the host launches captured graphs once the CPU pool's workers have +// started (copy the plan and the activations in, the grouped experts, their rows out), and the last one writes +// the rows into the window's own rows in mapped host memory and raises a flag there, which the first GPU waits for; +// so it keeps its latency on a GPU that also drives a display (a graph spinning there stalls whenever the desktop +// draws). The rows cross PCIe in 16-byte stores: written one float at a time they took ~0.2 ms a layer and slowed the +// CPU pool's RAM reads. One graph per format pair, group-count bucket and part: the grouped kernels' grids are sized +// for the bucket. Native packs only: the kernels are the first GPU's hit path. +// +// Prefetch: between a layer's CPU pool and the next layer's ring the RAM is idle (the first GPU's serial part, ~0.28 ms +// a layer), and this GPU's x16 link copies ~12 MiB in that time: 4 UD-Q4_K_XL or ~7 IQ3_XXS experts. The next layer's +// likeliest experts that neither GPU holds go into a few slots of their own, and that layer's share takes them from +// there: its own experts start at once, the prefetched ones once their copies have landed. +#pragma once + +#include "strata/core/expert_cache.hpp" + +#include + +#include +#include +#include + +namespace strata::core { + +class SecondGpu { +public: + SecondGpu() = default; + ~SecondGpu(); + SecondGpu(const SecondGpu&) = delete; + SecondGpu& operator=(const SecondGpu&) = delete; + + /// Buffers for windows of up to `max_t` tokens with `k` routed experts each, on `device`; the caller then opens + /// and fills `cache()` with that device current. `main_device` is made current again after every call. + bool init(int device, int main_device, int64_t n_embd, int64_t n_ff, int max_t, int64_t k, std::string& err); + ExpertCache& cache() { return cache_; } + int device() const { return dev_; } + + /// Room for `slots` prefetched experts (at most 16) of up to `blob_bytes`. After init(). + bool init_prefetch(int slots, uint64_t blob_bytes, std::string& err); + int prefetch_slots() const { return pre_max_; } + /// Starts copying `layer`'s experts ids[0..n) (n <= prefetch_slots(), `bytes` each from the pinned blobs `src`) into + /// the prefetch slots, once the last submitted layer is done with them. + bool prefetch(int64_t layer, const int32_t* ids, const uint8_t* const* src, int n, uint64_t bytes, std::string& err); + /// The prefetch slot holding (`layer`, `expert`), or -1; submit() takes prefetch slot p as slot -2 - p. + int prefetched(int64_t layer, int32_t expert) const { + if (layer != pre_layer_) return -1; + for (int p = 0; p < pre_n_; ++p) if (pre_ids_[p] == expert) return p; + return -1; + } + + /// One layer's share: group g is the expert in slot `slots[g]`, serving the routed entries + /// `entries[starts[g] .. starts[g+1])` (entry i = token i/k). `x` holds the window's n_tok activations as q8_1 + /// rows (n_embd / 32 blocks each). Entry i's row goes to row i of `out` (n_embd floats, mapped host memory), and + /// then `*flag` (mapped) is set to `ring`. Waits for the last share's graphs first (the plan and staging are + /// theirs until then). + bool submit(int64_t layer, const uint8_t* x, int n_tok, int64_t k, const int32_t* slots, const int32_t* starts, + const int32_t* entries, int n_groups, float* out, uint32_t* flag, uint32_t ring, std::string& err); + /// False, with the reason, when its work so far failed; does not wait. + bool healthy(std::string& err) const; + /// Recorded after the last submitted share: the last work that may read its cache and prefetch slots. + cudaEvent_t done_event() const { return ev_; } + + int64_t layers = 0, experts = 0, entries_done = 0; + int64_t prefetch_copied = 0, prefetch_used = 0; ///< experts copied ahead / of those, computed here + +private: + bool graph_for(int gu_type, int d_type, int groups, int part, bool prep, cudaGraphExec_t& exec, std::string& err); + + int dev_ = -1, main_ = 0, max_t_ = 0; + int64_t n_embd_ = 0, n_ff_ = 0, cap_ = 0; + ExpertCache cache_; + cudaStream_t s_ = nullptr; + cudaEvent_t ev_ = nullptr; + // pinned staging (portable, mapped): the activations (q8_1 rows) and the plan + uint8_t *h_x_ = nullptr, *m_x_ = nullptr; + uint8_t *h_plan_ = nullptr, *m_plan_ = nullptr; + uint8_t *d_xq_ = nullptr, *d_plan_ = nullptr, *d_scratch_ = nullptr; + float* d_rows_ = nullptr; // the rows, then copied out in full transactions + unsigned* d_count_ = nullptr; // the rows-out kernel's block counter + // part 0: the experts in its cache, launched at once; part 1: the prefetched ones, after their copies. `prep`: the + // graph copies the plan and the activations in (the first launch of a layer) + struct Graph { int gu, d, groups, part; bool prep; cudaGraphExec_t exec; }; + std::vector graphs_; + // prefetch slots, their copy stream, and the event of the last copies + static constexpr int kPrefetchMax = 16; + uint8_t* d_pre_ = nullptr; + uint64_t pre_cap_ = 0; + int pre_max_ = 0, pre_n_ = 0; + int64_t pre_layer_ = -1; + int32_t pre_ids_[kPrefetchMax] = {}; + cudaStream_t pre_s_ = nullptr; + cudaEvent_t pre_ev_ = nullptr; +}; + +} // namespace strata::core diff --git a/include/strata/core/session.hpp b/include/strata/core/session.hpp index e9c39ced286..26c1a4b72ec 100644 --- a/include/strata/core/session.hpp +++ b/include/strata/core/session.hpp @@ -83,6 +83,22 @@ uint64_t session_init(const ModelGeometry& g, int64_t max_cells, int64_t k, void /// from the reference's own `zeros()`. void session_zero(SessionState& s, const ModelGeometry& g, const float* R_init, void* stream); +// ================================ prompt-cache checkpoints ================================ +// +// The K/V caches are indexed by position, so the cells before a position stay valid when a later sequence +// overwrites the ones after it. What a position does NOT recover is the state that folds the whole prefix +// into a fixed size: every GDN layer's recurrence and conv history, each QSA indexer's raw block tail (plus +// its spare key, the pooled row of the open block and the block position), and the PLE's conv history and +// token window. A checkpoint is exactly that state at one position; restoring it and overwriting the cells +// from that position on continues the sequence as if it had been processed up to there. + +/// Bytes of one checkpoint (host memory). +uint64_t session_ckpt_bytes(const ModelGeometry& g); +/// Copies the sequence state at `pos` (positions [0, pos) processed) into `host`; synchronous. +bool session_ckpt_save(const SessionState& s, const ModelGeometry& g, int64_t pos, void* host, void* stream); +/// The inverse: the session continues at `pos`. The K/V cells [0, pos) must still hold the same prefix. +bool session_ckpt_load(SessionState& s, const ModelGeometry& g, int64_t pos, const void* host, void* stream); + /// One token: layers 0..47 in order, each a `block_layer`, and the residual is updated in place. /// /// `parts` is (k, n_embd) DEVICE memory, filled by the caller - by the CPU expert pool and the VRAM-resident @@ -348,8 +364,9 @@ struct SessionLoopScratch { long long pinned_core = -1; ///< the affinity to restore, or -1 if the host was never pinned bool pinned = false; - /// Allocates the buffers and pins the host thread. Call ONCE, at session setup. - bool init(size_t parts_bytes_in, std::string& err); + /// Allocates the buffers and pins the host thread to logical processor `host` (CorePlacement::host(); -1: the + /// last physical core's first). Call ONCE, at session setup. + bool init(size_t parts_bytes_in, std::string& err, int host = -1); /// Frees everything and restores the affinity. Safe to call twice, or on a default-constructed object. void free(); }; diff --git a/include/strata/core/split_head.hpp b/include/strata/core/split_head.hpp new file mode 100644 index 00000000000..3d45522c2b2 --- /dev/null +++ b/include/strata/core/split_head.hpp @@ -0,0 +1,65 @@ +// include/strata/core/split_head.hpp - the verify window's output head across two GPUs. +// +// The head streams the whole output matrix after the last layer (IQ3_S's Q6_K 521 MB, UD-Q4_K_XL's Q8_0 675 MB): the +// main GPU keeps rows [0, split) and a second GPU holds rows [split, n_vocab), so each streams part of it. The main +// GPU's window graph writes the head's input rows (the final hyper-connection mix, T x n_in floats) to mapped memory +// and raises a count; the host then launches this GPU's part (the same quantization, projection and per-row argmax +// kernels, so every logit is the one-GPU head's), which writes each token's pick and logit to mapped memory and raises +// a count of its own. The host keeps the larger logit, on equality the main GPU's (the lower index). It never waits +// on host memory on this GPU (a display may draw there). With `logits`, a window's rows are copied to mapped memory +// too, for sampled requests and --window-logits. +#pragma once + +#include + +#include +#include +#include + +namespace strata::core { + +class SplitHead { +public: + SplitHead() = default; + ~SplitHead(); + SplitHead(const SplitHead&) = delete; + SplitHead& operator=(const SplitHead&) = delete; + + /// Rows [split, n_vocab) of `output.weight` in the model's shards on `device`, for windows of up to `max_t` tokens. + /// `main_device` is current again after every call. + bool init(int device, int main_device, const std::vector& shards, int64_t n_in, int64_t n_vocab, + int64_t split, int max_t, std::string& err); + bool on() const { return dev_ >= 0; } + int64_t split() const { return split_; } + uint64_t bytes() const { return bytes_; } + /// Mapped: the head's input rows the main GPU writes (max_t x n_in floats) and the count it raises after them. + float* input() const { return h_in_; } + uint32_t* handoff() const { return h_hand_; } + /// Launches this GPU's part for a window of T tokens (after the handoff count has risen); `logits` copies its + /// rows to logits() too. + bool submit(int T, bool logits, std::string& err); + /// Until this GPU's count reaches `n` (the submits so far). + bool wait(uint32_t n, std::string& err); + /// Per token: the pick's row within this part, its logit; T x (n_vocab - split) logits after a `logits` submit. + const int32_t* idx() const { return h_idx_; } + const float* val() const { return h_val_; } + const float* logits() const { return h_logits_; } + +private: + bool capture(int T, bool logits, std::string& err); + + int dev_ = -1, main_ = 0, max_t_ = 0, type_ = -1; + int64_t n_in_ = 0, split_ = 0, rows_ = 0; + uint64_t bytes_ = 0; + cudaStream_t s_ = nullptr; + void* w_ = nullptr; // the rows' GGUF blocks + float *mixed_ = nullptr, *out_ = nullptr; // device: the input, the logits + uint8_t *xq_ = nullptr, *xil_ = nullptr, *scratch_ = nullptr; + float* h_in_ = nullptr; // mapped (portable) + uint32_t *h_hand_ = nullptr, *h_done_ = nullptr; + int32_t* h_idx_ = nullptr; + float *h_val_ = nullptr, *h_logits_ = nullptr; + cudaGraphExec_t exec_[9][2] = {}; +}; + +} // namespace strata::core diff --git a/include/strata/core/verify.hpp b/include/strata/core/verify.hpp index bc0051db14a..c8e42118c52 100644 --- a/include/strata/core/verify.hpp +++ b/include/strata/core/verify.hpp @@ -17,15 +17,22 @@ // and the PLE history is set to its snapshot after token n_keep-1. K/V cells past the accepted prefix are simply // overwritten when those positions are processed again, before any query can read them. // +// EXPERTS. The main GPU computes the routed experts its VRAM tier holds, decided in the graph from a snapshot of the +// residency table that the pool reads too (`residency`), so neither waits for the other; its shared expert runs on a +// branch of its own, forked before the router, and the wait for the CPU's rows on another. The pool computes the rest +// on the CPU and the second GPU, and with a PCIe share publishes the missed experts the GPU reads over PCIe. With the +// native router and a native pack the router kernel writes their input in the forms they take: ggml's Q8_K rows for +// the CPU (`pool_takes_q8k` layers) and q8_1 rows for the GPUs, so neither quantizes it. +// // Requires the default native decode configuration (native projections, fused GR, fused GDN, fast attention and -// selection, native indexer) and a profile-filled VRAM expert tier with its residency table on the device. +// selection, native indexer, QSA norms, RoPE, combine, shared expert gate and PLE block) and a profile-filled VRAM +// expert tier. #pragma once -#include - #include "strata/core/expert_source.hpp" #include "strata/core/layer.hpp" #include "strata/core/session.hpp" +#include "strata/core/split_head.hpp" #include "strata/kernels/sampler.hpp" #include @@ -38,15 +45,26 @@ namespace strata::core { class NativeHead; -/// The CPU pool for a window: x_f (n_tok, n_embd), ids (n_tok, k) -> out (n_tok * k, n_embd), hit rows zeroed. +/// The CPU pool for a window: x_f (n_tok, n_embd), ids (n_tok, k) -> out (n_tok * k, n_embd): the CPU's rows, the +/// main GPU's zeroed unless `GpuPlanSink::host_rows_only`; a second GPU writes its own (`set_gpu2`). using PoolMultiFn = void (*)(void* user, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k, float* out, int64_t layer); +/// False, with the reason, when work the window waits for failed (the second GPU's). +using WatchFn = bool (*)(void* user, std::string& err); +/// Layer `layer`'s likely experts, its router applied to an estimate of its FFN input (`Verifier::set_predict`): ids +/// and weights (n_tok, k), host memory. +using PredictFn = void (*)(void* user, int64_t layer, const int32_t* ids, const float* w, int64_t n_tok, int64_t k); +/// The window's last CPU rows are in. +using TailFn = void (*)(void* user); +/// Layer `layer`'s CPU rows are in (`Verifier::set_rows_in`). +using RowsFn = void (*)(void* user, int64_t layer); +/// The main GPU's VRAM expert tier. struct VerifyHits { - const int32_t* d_res = nullptr; ///< device [n_layers * n_expert] slot or -1 - const uint8_t* cache_base = nullptr; ///< slot 0 of the VRAM expert arena - const uint64_t* slot_off = nullptr; ///< E-6: host per-slot offsets when slots differ in size (null: slot * blob) - int64_t n_slots = 0; ///< E-6: how many (for the device copy) + const int32_t* res = nullptr; ///< host [n_layers * n_expert]: slot or -1, the live residency table + const uint8_t* cache_base = nullptr; ///< slot 0 of the VRAM expert arena + const uint64_t* slot_off = nullptr; ///< [slots]: each slot's byte offset from `cache_base`; null: slot * blob + int64_t slots = 0; int64_t blob = 0; }; @@ -57,65 +75,58 @@ class Verifier { Verifier(const Verifier&) = delete; Verifier& operator=(const Verifier&) = delete; - /// The watchdog's view of the window in flight (issue #31): the layer, the GPU's sequence, the flags. - void diag(std::FILE* f) const; - /// `max_t` <= kVerifyMaxT. `head` may be null (the canonical head is then run per token). bool init(const WeightTable& wt, const ModelGeometry& g, SessionState& ss, const VerifyHits& hits, const NativeHead* head, int max_t, std::string& err); /// One window: `tokens[0..T)` at positions pos0.., the pool served per layer; `out[t]` = argmax after token t. - /// The PLE rows are gathered here from `ss.ple_prev` and the tokens. Captures the T-token graph on first use. + /// The PLE rows follow `ss.ple_prev` and the tokens: read while layer 0 runs (the graph waits for them before + /// layer 1), unless `ple_ahead` already started them. Captures the T-token graph on first use. bool run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, void* user, int32_t* out, std::string& err); - /// The sampling the verify window's head applies (temperature / top_p / top_k / seed). Set per - /// request; greedy by default. The sampling itself runs OUTSIDE the captured graph - its - /// parameters would otherwise be baked forever - so this can change between requests freely. - void set_sampling(const strata::kernels::SamplerParams& sp) { - sampling_ = sp; // row t of a window at pos0 draws Philox(seed, pos0 + t): see run() - if (next_) next_->set_sampling(sp); - } - - /// The penalty histories for `sampling_.penalty_last_n`: ONE ROW PER WINDOW ROW, T rows of `history_len` - /// int32 slots at that stride (`strata::kernels::penalty_rows` builds them), most recent token LAST, unused - /// front slots -1 (the kernel reads only the tail window). Row t follows the window's drafts 1..t - staging - /// row 0 alone (before 0.1.19) left the drafted rows with unwritten histories. Null disables the penalties - /// entirely - the neutral run's sampling call is byte-for-byte what it was. The engine re-uploads the rows - /// before every window; the buffer must hold kVerifyMaxT rows and stay alive across the request. - void set_history(const int32_t* history, int history_len) { - hist_d_ = history; - hist_len_ = history_len; - if (next_) next_->set_history(history, history_len); - } - /// Off: `run` skips the request's head sampling and `out` is the recorded greedy pick. For windows whose - /// picks are discarded - a prompt read through windows commits every token - so they cost no sampler launch - /// or sync and never read a history staged for another position. - void set_head_sampling(bool on) { head_sampling_ = on; if (next_) next_->set_head_sampling(on); } - - /// LAYER SPLIT (multi-GPU): this verifier runs layers [layer_begin, layer_end) of every window. A stage that - /// does not start at layer 0 takes its residual from `handoff_in` instead of embedding the tokens; a stage that - /// does not end at the last layer writes its residual to `handoff_out` and has no head. The hand-off holds, - /// per token, the residual R (hc x n_embd), the last layer's pending write bo (n_embd) and inject (hc): the next - /// stage folds that write into its first read exactly as the unsplit window does, so the split is bit-exact. - /// Both pointers must be device-visible (mapped pinned memory, portable when the stages are on two devices). - /// Set before `init`. Default: the whole model, no hand-off. - void set_stage(int64_t layer_begin, int64_t layer_end, const float* handoff_in, float* handoff_out) { - lb_ = layer_begin; le_ = layer_end; hand_in_ = handoff_in; hand_out_ = handoff_out; - } - /// The next stage: `run` and `commit` continue into it (its pool calls get `next_user`); sampling settings - /// and `final_R` are the last stage's. - void set_next(Verifier* next, void* next_user) { next_ = next; next_user_ = next_user; } - /// floats per token in a hand-off buffer - static int64_t handoff_floats(const ModelGeometry& g) { return (int64_t) g.hc * g.n_embd + g.n_embd + g.hc; } - - /// Keep the first `n_keep` (1..T) tokens of the last window; advances `ss.ple_prev` by them. - bool commit(int n_keep, std::string& err); + /// Starts reading the PLE rows of the next window's token t (0, 1, ... in order, after `commit`): `run` takes + /// them if its token t, and the ones before it, are these. For the drafts, while the draft layer runs. + void ple_ahead(int t, int32_t token); + + /// The sampling the head applies (temperature / top_p / top_k / min_p / seed), per request; greedy by default. + /// It runs after the captured graph (a captured sampler would bake its parameters in), so it can change between + /// requests freely: row t of a window at pos0 draws Philox(seed, pos0 + t), tied to the position it samples, not + /// to how the text was cut into windows. A rejected row's draw is discarded. + void set_sampling(const strata::kernels::SamplerParams& sp) { sampling_ = sp; } + /// The penalty histories for `penalty_last_n`: one row per window row, T rows of `history_len` int32 slots + /// (`strata::kernels::penalty_rows`), device memory staged before every window; null: no penalties. + void set_history(const int32_t* history, int history_len) { hist_d_ = history; hist_len_ = history_len; } + /// Off: `run` skips the head sampling and `out` is the greedy pick (windows whose picks are discarded: a prompt + /// read through windows commits every token). + void set_head_sampling(bool on) { head_sampling_ = on; } + + /// Keep the first `n_keep` (1..T) tokens of the last window; advances `ss.ple_prev` by them. With `wait` false + /// the commit graph is only launched (the MTP draft, on its own stream, reads nothing it writes): `wait_commit` + /// before anything outside the next `run` reads the sequence state. + bool commit(int n_keep, std::string& err, bool wait = true); + bool wait_commit(std::string& err); + + /// The residency the window decides the main GPU's hits by: `VerifyHits::res` as `init` and the last `run` found + /// it. The pool reads this one (`ExpertDispatch::host_res`), so it leaves exactly the GPU's rows to the GPU. + const int32_t* residency() const { return h_res_; } /// Token t's residual after the last layer, (hc, n_embd) on the device, valid until the next `run`. const float* final_R(int t) const; - const float* final_R_all() const { return next_ ? next_->final_R_all() : R_; } + const float* final_R_all() const { return R_; } + /// With `set_host_rows(true)` (before the first `run`), the same rows in host memory once `run` returns: the + /// graph writes them there itself, so reading them waits for no copy engine. + void set_host_rows(bool on) { host_rows_ = on; } + const float* final_R_host() const { return h_rows_; } + /// The head's logits of the last window's first `T` tokens, T * n_vocab floats, copied to host memory `out`. + bool copy_logits(int T, float* out, std::string& err) const; + /// The head's rows [split, n_vocab) on another GPU (`head` then holds rows [0, split)): the window hands the head's + /// input over, and `run` keeps each token's larger pick. Set before `init`. + void set_split_head(SplitHead* sh) { shead_ = sh; } + /// Whole logits rows after every window (copy_logits, --window-logits): with a split head the second GPU's part + /// then copies its rows back. Set before `init`. + void set_logits_wanted(bool on) { logits_wanted_ = on; } - /// The GPU plan the pool writes each layer (VRAM hits + the PCIe share of the misses); give it to the - /// dispatch (`ExpertDispatch::plan`) before the first `run`. + /// Where the pool publishes each layer's PCIe share of the misses; give it to the dispatch + /// (`ExpertDispatch::plan`) before the first `run`. GpuPlanSink* plan_sink() { return &sink_; } /// Plan v0.3 P6: split the window into two token groups and pipeline the CPU experts of one with the GPU work /// of the other (default on). Set before the first `run`. @@ -125,43 +136,57 @@ class Verifier { /// arena directly, 2 = a copy kernel stages it inside the graph (no API calls on the pool's thread; best when /// the CPU is RAM-bound, Q2_0). Set before the first `run`. void set_pcie_mode(int mode) { sink_.pcie_mode = mode; } - /// the pool never plans a PCIe share (--pcie-frac 0): the window skips that path. Before the first run. + /// Whether the pool may give the GPU a PCIe share of the misses (`--pcie-frac` > 0); without one the window + /// has no PCIe stage. Set before the first `run`. + void set_pcie_share(bool on) { pcie_share_ = on; sink_.pcie = on; } + /// Diagnostics: GPU timestamps between the stages of every layer (`--window-profile`; each costs a kernel + /// launch, ~2 us). Set before the first `run`; `print_profile` writes the per-stage table to stdout. + void set_profile(bool on) { profile_ = on; } + void print_profile() const; + /// Each layer's router also runs on an estimate of its FFN input from the previous layer's state, on a side branch + /// of the graph: the layer's FFN read of that state, its gates their running average over the windows, and the + /// router's logits corrected by a running average of their error. Once the CPU's rows of a layer are in, `fn` + /// gets the next layer's prediction (the RAM is idle until the next ring), waited for when it comes later + /// (`predict_late`). Without `learn` the averages stay at their start (gates of one half, no correction), so the + /// prediction depends on the window alone. Windows of one token group. Set before the first `run`. + void set_predict(PredictFn fn, void* user, bool learn) { predict_ = fn; predict_user_ = user; learn_ = learn; } + /// `fn` runs once the window's last CPU rows are in: the RAM stays idle until the next window's first pool (the + /// head, the commit, the drafts). The main GPU may still run the window then: `window_done` follows its graph. + /// Set before the first `run`. + void set_tail(TailFn fn, void* user) { tail_ = fn; tail_user_ = user; } + /// `fn` runs once each layer's CPU rows are in and the next layer's prediction has gone out, but the last's (the + /// tail's): the main GPU's link then has nothing to carry until the next layer's router writes its rows to the + /// host. Set before the first `run`. + void set_rows_in(RowsFn fn, void* user) { rows_in_ = fn; rows_user_ = user; } + cudaEvent_t window_done() const { return done_; } + /// The pool gives a second GPU a share of each layer (`GpuPlanSink::gpu2_flag`): it writes its rows into the + /// pool's rows and raises its token group's flag, and a branch of the window, forked at the ring, takes them + /// into VRAM while the CPU works. Set before the first `run`. + void set_gpu2(bool on) { gpu2_ = on; } + /// `fn` is asked every ~2 ms while the host waits for a ring: on failure the window's waits are released and `run` + /// returns the reason. Set before the first `run`. + void set_watch(WatchFn fn, void* user) { watch_ = fn; watch_user_ = user; } - double ms_wait = 0, ms_pool = 0, ms_host = 0, ms_commit = 0; - int64_t windows = 0; - /// STRATA_VERIFY_PROFILE=1 - GPU stage times of the windows since the last call (ms per - /// window), as one line; empty when off. - std::string profile_report(); + double ms_wait = 0, ms_pool = 0, ms_host = 0, ms_commit = 0, ms_predict = 0; + int64_t windows = 0, predict_late = 0; + /// Layers whose second-GPU rows came after the CPU's, and the host time from the CPU's rows to theirs. + int64_t late2 = 0; + double ms_late2 = 0; private: bool capture(int T, std::string& err); + bool capture_commit(std::string& err); + bool record_window(int T, cudaStream_t cs, std::string& err); + strata::kernels::SamplerParams sampling_ = [] { strata::kernels::SamplerParams s; s.greedy = true; s.temperature = 0.0f; return s; - }(); ///< greedy by default; per-request via set_sampling - const int32_t* hist_d_ = nullptr; ///< penalty-history row (set_history); null = no penalties apply + }(); + const int32_t* hist_d_ = nullptr; int hist_len_ = 0; - bool head_sampling_ = true; ///< set_head_sampling - int device_ = -1; ///< the device `init` ran on: run/commit switch to it (layer split) - bool device_plan_ = false; ///< E-6: resident-only layers planned on the device (STRATA_VERIFY_DEVICE_PLAN) - uint32_t* skip_ = nullptr; ///< E-6: per group, the ring whose plan the device built (0: the host's) - unsigned long long* slot_off_d_ = nullptr; ///< E-6: the slot offsets on the device - int64_t lb_ = 0, le_ = -1; ///< set_stage: the layers this verifier runs (-1: to the last) - const float* hand_in_ = nullptr; - float* hand_out_ = nullptr; - Verifier* next_ = nullptr; - void* next_user_ = nullptr; - bool ple_stage() const { return lb_ <= 1 && 1 < le_; } ///< holds layer 1, where the PLE block runs - bool capture_commit(std::string& err); - bool record_window(int T, cudaStream_t cs, std::string& err); - static constexpr int kProfPer = 32; // stamps per layer - bool prof_on_ = false; - unsigned long long* prof_ = nullptr; // device: n_layers * kProfPer + 4 stamps - std::vector prof_h_; - double prof_sum_[2][kProfPer] = {}; // [GDN / QSA layers][stage] - int64_t prof_windows_ = 0; + bool head_sampling_ = true; const WeightTable* wt_ = nullptr; const ModelGeometry* g_ = nullptr; @@ -176,6 +201,7 @@ class Verifier { cudaStream_t cs_ = nullptr; cudaGraphExec_t exec_[9] = {}; cudaGraphExec_t commit_exec_ = nullptr; + bool commit_pending_ = false; // mapped staging (host pointer, device alias) int32_t* h_tok_ = nullptr; int32_t* m_tok_ = nullptr; // T @@ -184,43 +210,116 @@ class Verifier { int32_t* h_commit_ = nullptr; int32_t* m_commit_ = nullptr; // [n_keep, n_keep-1, pos_0 .. pos_{T-1}] float* h_ple_ = nullptr; float* m_ple_ = nullptr; // T * n_embd int32_t* h_out_ = nullptr; int32_t* m_out_ = nullptr; // T argmax ids + float* h_val_ = nullptr; float* m_val_ = nullptr; // their logits (a split head's merge) + SplitHead* shead_ = nullptr; + bool logits_wanted_ = false; + int64_t head_rows_ = 0; // the head's rows here: n_vocab, or the split + float* head_full_ = nullptr; // T whole logits rows (a split head, when wanted) + uint32_t hand_n_ = 0, done_n_ = 0; // the split head's counts so far float* h_x_ = nullptr; float* m_x_ = nullptr; // doorbell payload: T * n_embd int32_t* h_ids_ = nullptr; int32_t* m_ids_ = nullptr; // T * k float* h_w_ = nullptr; float* m_w_ = nullptr; // T * k uint32_t* h_seq_ = nullptr; uint32_t* m_seq_ = nullptr; uint32_t* h_flag_ = nullptr; uint32_t* m_flag_ = nullptr; - uint32_t* h_flagA_ = nullptr; uint32_t* m_flagA_ = nullptr; // the GPU plan is in place + uint32_t* h_flagA_ = nullptr; uint32_t* m_flagA_ = nullptr; // the PCIe share's plan is in place uint32_t* h_flagB_ = nullptr; uint32_t* m_flagB_ = nullptr; // the PCIe share's DMA copies have landed + uint32_t* h_pleflag_ = nullptr; uint32_t* m_pleflag_ = nullptr; // the PLE rows are in h_ple_ + bool host_rows_ = false; + float* h_rows_ = nullptr; float* m_rows_ = nullptr; // T * hc * n_embd: the final rows (set_host_rows) + int32_t ahead_tok_[8] = {}; // the tokens ple_ahead read for cudaStream_t copy_ = nullptr; // the copy engine's stream (DMA of missed experts) struct FlagSet { uint32_t* flag; uint32_t value; }; FlagSet flag_sets_[2 * 64 * 2] = {}; // host-function arguments, one per (layer, group) static void fetch_dma(void* ctx, const uint8_t* const* src, int n, size_t bytes); static void raise_flag(uint32_t* flag, uint32_t value); - int32_t* h_plan_ = nullptr; int32_t* m_plan_ = nullptr; // counts | start | dst | tok | ptr (as int32 pairs) - int64_t plan_i32_ = 0; // int32 words in the plan block + int32_t* h_plan_ = nullptr; int32_t* m_plan_ = nullptr; // the PCIe share's plan block per token group + int64_t plan_i32_ = 0, plan_ptr_off_ = 0, plan_stride_ = 0; // a plan block's int32 words, its ptr field, the + // device blocks' stride (see `init`) + int32_t* h_res_ = nullptr; int32_t* m_res_ = nullptr; // the residency snapshot (`residency`) + int64_t res_words_ = 0; GpuPlanSink sink_; uint32_t cur_layer_ = 0; static void publish_plan(void* ctx); void set_plan_slot(int grp); bool split_ = false; // opt-in (--spec-split): exact but slower, see the overlap study + bool pcie_share_ = true; int groups_[9] = {}; float* h_ymiss_ = nullptr; float* m_ymiss_ = nullptr; // T * k * n_embd + // the FFN input as the router kernel writes it (quant_router_): Q8_K rows `xk_stride_` bytes apart for the layers + // the pool takes them (xk_layer_), q8_1 rows for the second GPU + bool quant_router_ = false; + std::vector xk_layer_; + int64_t xk_stride_ = 0; + uint8_t* h_xk_ = nullptr; uint8_t* m_xk_ = nullptr; // T rows + uint8_t* h_xq1_ = nullptr; uint8_t* m_xq1_ = nullptr; // T rows of n_embd / 32 blocks + // the second GPU's share (set_gpu2), per token group: the flag it raises (64 bytes apart) and its entries, + // [n, pad x3, entries] (mapped; `list2_` their device copies, which the combine reads) + bool gpu2_ = false; + uint32_t* h_flag2_ = nullptr; uint32_t* m_flag2_ = nullptr; + int32_t* h_list2_ = nullptr; int32_t* m_list2_ = nullptr; + int32_t* list2_ = nullptr; + int64_t list2_stride_ = 0; + WatchFn watch_ = nullptr; + void* watch_user_ = nullptr; + // the next layer's predicted experts: two halves by layer parity, so a prediction stays put until the host has + // taken it; the counter holds the predictions published in this window + PredictFn predict_ = nullptr; + void* predict_user_ = nullptr; + bool learn_ = false; + TailFn tail_ = nullptr; + void* tail_user_ = nullptr; + RowsFn rows_in_ = nullptr; + void* rows_user_ = nullptr; + int32_t* h_pids_ = nullptr; int32_t* m_pids_ = nullptr; // 2 x T * k + float* h_pw_ = nullptr; float* m_pw_ = nullptr; // 2 x T * k + uint32_t* h_pseq_ = nullptr; uint32_t* m_pseq_ = nullptr; + // short branches beside the main stream: on side_ the GDN gates or the QSA indexer (the mixer) and the experts' + // activations (beside the router), then the prediction; on shs_ the GDN gate projection or the QSA queries (the + // mixer), then the shared expert + cudaStream_t side_ = nullptr; + cudaEvent_t join_ = nullptr; // the prediction's + cudaEvent_t bfork_ = nullptr, bjoin_ = nullptr; // side_'s others + cudaEvent_t pfork_ = nullptr, pjoin_ = nullptr; // shs_'s projections + cudaStream_t shs_ = nullptr; // the shared expert's branch (and the snapshot's copy) + cudaEvent_t shfork_ = nullptr, shjoin_ = nullptr, res_ready_ = nullptr; + cudaStream_t g2s_ = nullptr; // the second GPU's rows, per token group + cudaEvent_t g2fork_ = nullptr, g2join_[2] = {}; + cudaStream_t cps_ = nullptr; // the wait for the CPU's rows, per token group + cudaEvent_t cpfork_ = nullptr, cpjoin_[2] = {}; + cudaStream_t kst_ = nullptr; // the QSA mixer's keys + cudaEvent_t kjoin_ = nullptr, iproj_ = nullptr; // and the end of its indexer's projections + cudaStream_t sgs_ = nullptr; // the shared expert's scalar gate + cudaEvent_t sgjoin_ = nullptr; + cudaEvent_t efork_ = nullptr, rdone_ = nullptr; // the FFN read's end, the router's + cudaEvent_t done_ = nullptr; // after the last window's graph // device void* arena_ = nullptr; int32_t *tok_ = nullptr, *step_ = nullptr, *pos_ = nullptr, *commit_ = nullptr; float *ple_ = nullptr, *emb_ = nullptr, *R_ = nullptr, *mixed_ = nullptr, *bo_ = nullptr; - float *inj_ = nullptr, *inj2_ = nullptr, *lo_ = nullptr, *rs_ = nullptr, *xn_ = nullptr; + float *inj_ = nullptr, *inj2_ = nullptr, *lo_ = nullptr, *rs_ = nullptr; + float* grs_ = nullptr; // the hyper-connection reads' scratch (zeroed with the arena) uint8_t* xq_ = nullptr; // T columns of q8_1 + uint8_t* xil_ = nullptr; // their interleaved copy (2+ tokens) float *qkv_L_ = nullptr, *h_L_ = nullptr, *gate_L_ = nullptr, *beta_L_ = nullptr; // per GDN layer float *z_ = nullptr, *y_ = nullptr, *y_dummy_ = nullptr; float *qfull_ = nullptr, *qcur_ = nullptr, *kcur_ = nullptr, *vcur_ = nullptr, *idx_raw_L_ = nullptr; - float *qidx_ = nullptr, *scores_ = nullptr, *attn_ = nullptr, *attn32_ = nullptr, *attn_scratch_ = nullptr; + float *qidx_ = nullptr, *scores_ = nullptr, *attn32_ = nullptr, *attn_scratch_ = nullptr; float* tail_snap_ = nullptr; // per QSA layer int32_t* sel_ = nullptr; - float *logits_ = nullptr, *w_ = nullptr, *shared_ = nullptr, *parts_ = nullptr, *hit_out_ = nullptr; + float *logits_ = nullptr, *w_ = nullptr, *shared_ = nullptr, *hit_out_ = nullptr; int32_t *ids_ = nullptr, *hit_slot_ = nullptr, *hit_dst_ = nullptr, *hit_count_ = nullptr; - int32_t* plan_ = nullptr; // device copy of the plan block + unsigned* ring_count_ = nullptr; // verify_router's block counters: the layer's, the prediction's + float* plogits_ = nullptr; // the predictions' logits, by layer parity + float* pw_ = nullptr; // the next layer's router (prediction) + int32_t* pids_ = nullptr; + float* pmixed_ = nullptr; // its input: the next layer's FFN read, estimated + float* gate_ema_ = nullptr; // [layer][hc][n_embd]: the FFN reads' gates, averaged + float* pbias_ = nullptr; // [layer][n_expert]: the predictions' logit bias + int32_t* plan_ = nullptr; // device copies of the PCIe share's plan blocks + int32_t* dplan_ = nullptr; // the main GPU's hit plans, decided on the device + int32_t* res_ = nullptr; // the residency snapshot, copied as a window starts + unsigned long long* slot_ptr_ = nullptr; // each VRAM slot's address uint8_t* staging_ = nullptr; // VRAM slots for the PCIe share of the misses static constexpr int64_t kStagingBlobs = 16; uint8_t* hit_xq_ = nullptr; @@ -228,10 +327,27 @@ class Verifier { float* hit_xs_ = nullptr; void* hit_scratch_ = nullptr; float *head_mixed_ = nullptr, *head_inj_ = nullptr, *head_logits_ = nullptr; - uint16_t* sh_bf16_ = nullptr; + uint8_t* arg_scratch_ = nullptr; ///< argmax_rows' partials and counters + int32_t* one_ = nullptr; ///< device {1}: the n_keep of a one-token window, which commits itself float *sh_gate_ = nullptr, *sh_up_ = nullptr, *sh_g_ = nullptr; float* hist_snap_ = nullptr; // T * NG_HIST * NG_HC_DIM + // the PLE block's projections (keys T * NG_HC_DIM, values T * n_embd) and the rest of the block's rows + float *ple_key_ = nullptr, *ple_val_ = nullptr, *ple_nkey_ = nullptr, *ple_norm_ = nullptr, *ple_gated_ = nullptr; + float* ple_gate_ = nullptr; int64_t cap_ = 0, max_blocks_ = 0, attn_scratch_floats_ = 0; + + // --window-profile: stamps [window size][layer, then one row for the window][stage]: a layer's main-stream stages, + // then the ends of its shared expert's branch and of the second GPU's rows' branch + static constexpr int kLayerStamps = 13, kStamps = 15; + bool profile_ = false; + unsigned long long* stamps_ = nullptr; + std::vector h_stamps_; + std::vector prof_ns_; // [mixer kind (GDN, QSA)][stage], summed + double prof_window_ns_[4] = {}; // inputs + embeddings, head, whole window, stamp gap + int64_t prof_windows_ = 0, prof_layers_[2] = {}; + static constexpr int kWaitBuckets = 7; // the wait for the CPU's rows, by bucket (us) + static constexpr double kWaitEdges[kWaitBuckets - 1] = {6, 8, 12, 20, 50, 100}; + int64_t prof_wait_hist_[2][kWaitBuckets] = {}; }; } // namespace strata::core diff --git a/include/strata/kernels/bf16_gemv.hpp b/include/strata/kernels/bf16_gemv.hpp index 906a8f0b150..822202ea955 100644 --- a/include/strata/kernels/bf16_gemv.hpp +++ b/include/strata/kernels/bf16_gemv.hpp @@ -44,9 +44,10 @@ void bf16_gemv_split(const uint16_t* x, const uint16_t* w, float* y, int64_t n_i /// No allocations or synchronization, including when stream is null (the CUDA default stream). void bf16_gemv_fp32_mmvf(const float* x, const uint16_t* w, float* y, int64_t n_in, int64_t n_out, void* stream); -/// bf16_gemv_fp32_mmvf for n_tok (1..8) activation rows x[t * ldx], outputs y[t * ldy + j]; one launch, -/// the weight read once, every output bit-identical to its own single-row call. -void bf16_gemv_fp32_mmvf_multi(const float* x, int64_t ldx, const uint16_t* w, float* y, int64_t ldy, - int64_t n_in, int64_t n_out, int n_tok, void* stream); +/// `bf16_gemv_fp32_mmvf` for n_tok (1..8) activation rows x (n_tok, n_in) -> y (n_tok, n_out), each weight read once; +/// every row bitwise what the single-row call gives. With `w2` the same launch also gives y2 (n_tok, n_out2) from a +/// second matrix of the same input. +void bf16_gemv_fp32_mmvf_multi(const float* x, const uint16_t* w, float* y, int64_t n_in, int64_t n_out, int n_tok, + void* stream, const uint16_t* w2 = nullptr, float* y2 = nullptr, int64_t n_out2 = 0); } // namespace strata::kernels diff --git a/include/strata/kernels/cpu/expert_layout.hpp b/include/strata/kernels/cpu/expert_layout.hpp index 126549ae473..b96a5c9f1be 100644 --- a/include/strata/kernels/cpu/expert_layout.hpp +++ b/include/strata/kernels/cpu/expert_layout.hpp @@ -20,12 +20,6 @@ struct ExpertLayout { int64_t n_layers = 0, n_expert = NE; std::vector fmt; ///< per layer (native packs) std::vector offset, bytes; ///< per layer: where its 512 blobs start, bytes per blob - /// Plan v0.3 P6: per layer, the absolute offsets of the gate / up / down tensors in the model's shard 1, so - /// the arena can be filled from the GGUF itself when the pack has no experts.bin (3 x n_layers, 0 = unknown). - std::vector gguf_off; - /// Per layer, the GGUF file (a name beside the --native shard) that holds its experts when the model's - /// shards split the layers (Swift's GGUFs: layers 13-47 in shard 2). Empty = the --native shard itself. - std::vector gguf_file; uint64_t max_blob = BLOB; uint64_t total = 0; ///< experts.bin size diff --git a/include/strata/kernels/cpu/pool.hpp b/include/strata/kernels/cpu/pool.hpp index dce9c2bfd8b..7fbf49c8cbe 100644 --- a/include/strata/kernels/cpu/pool.hpp +++ b/include/strata/kernels/cpu/pool.hpp @@ -11,32 +11,28 @@ // one batch. A ring would add a wrap-around to get wrong and buy nothing. What is kept from the phase is // the part that matters: `head`/`done` are single fetch_add counters, one claim per worker, no lock. // -// THE COMPLETION PROTOCOL, because this is where a pool usually goes wrong. `run()` waits for `done == n` -// AND for every worker to PARK. Waiting only for `done` is not enough: a worker can still be inside the -// drain loop after its last `done` increment, and the host resetting `head` underneath it would let that -// worker claim a job from the NEXT batch before the next batch has been published. +// THE PHASE PROTOCOL, because this is where a pool usually goes wrong. One atomic word holds the phase's +// epoch, whether it is open, and how many workers are parked. The host writes a phase's jobs while the +// previous phase is closed, then opens the next epoch; a worker joins with a compare-exchange on the whole +// word, so it can only join the phase it saw open. The host closes a phase once `done == n` AND every +// worker that joined has parked again. Waiting only for `done` is not enough: a worker can still be inside +// the drain loop after its last `done` increment, and the host resetting `head` underneath it would let it +// claim a job from the NEXT batch. Closing is what makes a late worker harmless: one that notices a phase +// only after the host finished it (it was asleep, or descheduled) finds it closed and waits for the next. // -// **AND `parked` IS NOT ENOUGH EITHER (issue #29).** A worker that went to sleep (after `kSpinBeforeSleep`) is -// still counted as parked when it wakes, so for a moment after it has seen a new epoch the host believes it is -// idle. On a card with most experts in VRAM the workers sleep in the middle of a request, tiny batches finish -// before a sleeper is awake, and that moment comes round constantly: a late worker could claim from the NEXT -// batch while the host was still writing it, run a job twice, or add to `done` after the host had reset it - -// and `done != n` then never ended, with the GPU waiting on the pool forever. So every claim carries its -// batch: `head` is one 64-bit word `epoch | njobs | index`, a claim is a CAS that only succeeds for the epoch -// the worker woke for, and a late worker's claim simply fails. A claim that succeeds belongs to the current -// batch, whose description the host cannot change until that job's `done` has landed. +// IDLE WORKERS SLEEP. A parked worker spins for `kIdleSpin`, then blocks on the word (`std::atomic::wait`: +// WaitOnAddress, a futex) until the host opens a phase; the host wakes them only when one sleeps. Generation +// opens phases every few hundred microseconds, and a round's longest gap (head, drafts, the first layer) is +// a few ms, so they sleep only between requests and while a prompt runs on the GPUs. #pragma once #include "strata/kernels/cpu/expert.hpp" #include "strata/kernels/cpu/native_expert.hpp" #include -#include -#include #include -#include #include -#include +#include #include #include @@ -63,19 +59,52 @@ struct ExpertJobMulti { const void* nact[MAXT] = {}; }; -/// One logical processor per PHYSICAL core, so a worker is never scheduled onto an SMT sibling of another -/// worker. On the 6-core/12-thread machine this project measures on, `hardware_concurrency()/2` workers on -/// logical processors 0..5 would put every worker on a sibling pair and halve the useful bandwidth - which is -/// exactly the kind of error that shows up as "the CPU path is slower than the model says" with no clue why. -/// -/// `skip_first` drops the first core, which P2.S3 reserves for the host loop. -std::vector physical_cores(bool skip_first); +/// Each PHYSICAL core's logical processors (SMT siblings together), in the OS's order. The pool runs one thread +/// per physical core, so a worker is never scheduled onto an SMT sibling of another worker. On the 6-core/12-thread +/// machine this project measures on, `hardware_concurrency()/2` workers on logical processors 0..5 would put every +/// worker on a sibling pair and halve the useful bandwidth - which is exactly the kind of error that shows up as +/// "the CPU path is slower than the model says" with no clue why. +std::vector> physical_cores(); + +class ExpertPool; + +/// Where the host loop and the pool's workers run - one logical processor each, one thread per physical core - and +/// how they move off the processors that take interrupts. A GPU's copies and launches interrupt the CPU when they +/// complete (~10-40 us each in the ISR and a DPC, thousands a second during generation), on processors Windows picks +/// and moves as the load shifts (the power plan's interrupt steering; here the first few). So `tick`, called on the +/// host thread between windows, reads each processor's interrupt and DPC time once a second, and a thread whose +/// processor spent over 5% there two seconds in a row moves to a quiet processor: a free core's, else its own core's +/// other one. Without the counters (not Windows) the threads stay where they started. +class CorePlacement { +public: + /// The start: the host on the last physical core, the workers on the others; with `spare_first` (a second GPU, + /// whose work raises most of the interrupts) and more than four cores, the first core stays free: where Windows + /// sends interrupts first, and a place to move to. + explicit CorePlacement(bool spare_first); + int host() const { return host_; } + /// The workers' logical processors, by worker (the pool pins worker i to `workers()[i]`). + const std::vector& workers() const { return workers_; } + /// At most once a second: samples each processor's interrupt and DPC time and moves the threads whose processors + /// took them (`pool`'s workers; the calling thread for the host's). Returns what moved, or an empty string. + std::string tick(ExpertPool& pool); + int64_t moves = 0; + +private: + int quiet_processor(int from, const std::vector& load) const; + std::vector> cores_; + std::vector core_of_; // logical processor -> physical core, -1 when unknown + int host_ = -1; + std::vector workers_; + std::vector busy_; // each processor's interrupt and DPC time at the last sample (100 ns units) + std::vector hot_; // the samples in a row a processor spent over 5% there + std::chrono::steady_clock::time_point at_{}; +}; /// **THE RESERVATION IS A FICTION UNLESS THE HOST IS ACTUALLY PUT THERE.** /// -/// `physical_cores(true)` keeps the workers off the first physical core so that the host loop can spin on -/// `cudaEventQuery` without stealing a worker's cycles. Nothing in the pool can enforce the other half of -/// that, so this is it: the host loop calls this on entry and restores on exit. +/// `CorePlacement` keeps the workers off the host's core so that the host loop can spin on `cudaEventQuery` +/// without stealing a worker's cycles. Nothing in the pool can enforce the other half of that, so this is it: +/// the host loop calls this on entry and restores on exit. /// /// MEASURED, and this is why it exists: the pool runs at **36.32 GB/s on 5 workers with nothing else running** /// - exactly 5/6 of L9's 44.14 on 6 - and at **26.9 GB/s inside the host loop**, where the unpinned spinning @@ -85,15 +114,19 @@ std::vector physical_cores(bool skip_first); long long pin_current_thread(int core); void restore_thread_affinity(long long previous); +#ifdef _MSC_VER +#pragma warning(push) +#pragma warning(disable : 4324) // the alignas(64) members pad the class on purpose (one cache line each) +#endif class ExpertPool { public: - /// `n_workers <= 0` means "every physical core except the first". Workers are pinned to physical cores - /// (minus core 0 by default) and each owns one `ExpertScratch`, so nothing in the token path allocates. + /// Worker i is pinned to `cores[i]` (none of `cores`: `CorePlacement(false).workers()`); `n_workers <= 0` means + /// one on each. Each owns one `ExpertScratch`, so nothing in the token path allocates. /// /// **`host_works` PUTS THE HOST THREAD INTO THE DRAIN (R2.2's FIRST HALF).** /// - /// The pool reserves core 0 for the host loop so the doorbell spin cannot steal a worker's cycles - but - /// during `run()` the host does not spin, it waits, so core 0 is idle for the whole drain. Measured on the + /// The pool keeps a core for the host loop so the doorbell spin cannot steal a worker's cycles - but + /// during `run()` the host does not spin, it waits, so that core is idle for the whole drain. Measured on the /// 6-core machine this project targets: the engine's pool drains at **33.7 GB/s** (663.6 MB of expert /// blobs in 19.71 ms/token) where the same kernel on 5 workers should reach 5/6 x 44.14 = 36.8 and the /// machine measures 44.14 GB/s on all six. So the sixth core is being paid for and not used. @@ -101,19 +134,19 @@ class ExpertPool { /// With `host_works`, `run()` claims jobs itself instead of spinning on `done_`, and the pool is six /// threads on six cores. `false` is the A/B arm and exists so the change is measurable rather than /// asserted - the counter it moves is `pool phases ... drain`, which is host-side and needs no profiler. - explicit ExpertPool(int n_workers = 0, bool pin = true, bool host_works = true); - /// The watchdog's view of the pool (issue #31): the batch, the counters, every thread's state. - void diag(std::FILE* f) const; + explicit ExpertPool(int n_workers = 0, bool pin = true, bool host_works = true, std::vector cores = {}); ~ExpertPool(); ExpertPool(const ExpertPool&) = delete; ExpertPool& operator=(const ExpertPool&) = delete; int workers() const { return n_; } + /// Moves worker `i` to logical processor `core` (CorePlacement::tick). + void pin_worker(int i, int core); /// Whether the host thread also drains. Reported at startup, because "the engine adapts to the machine it /// is on" is only true if the engine says which adaptation it took. bool host_works() const { return host_works_; } - /// Publish `n` jobs, then block until every one has been claimed AND every worker has parked. + /// Publish `n` jobs, then block until every one is done AND the phase is closed. /// `jobs` must outlive the call (it does, and the workers never touch it afterwards). void run(ExpertJob* jobs, int n); @@ -126,56 +159,37 @@ class ExpertPool { /// Plan v0.3 P6: `run_split` for multi-token jobs (at most `kMaxSplitMulti`); the rows of each expert are /// read once for all of its tokens. void run_split_multi(ExpertJobMulti* jobs, int n); - /// Plan v0.3 P6: the same for a native pack's layer (ggml-cpu arithmetic, `nact` activations). - void run_split_multi_native(const NativeFmt& f, ExpertJobMulti* jobs, int n); + /// Plan v0.3 P6: the same for a native pack's layer (ggml-cpu arithmetic, `nact` activations). `first`, when + /// given, runs once on the host: after the workers have started on the first phase, before it joins them (or + /// alone when there are no jobs). + void run_split_multi_native(const NativeFmt& f, ExpertJobMulti* jobs, int n, void (*first)(void*) = nullptr, + void* ctx = nullptr); static constexpr int kMaxSplitMulti = 96; /// run_split_multi's phases, accumulated ms: gate/up rows, the intermediate quantization, down rows. double ms_multi_gu = 0, ms_multi_q = 0, ms_multi_down = 0; int64_t multi_bytes = 0; - /// Total `_mm_pause` iterations spent waiting, over all workers, is no longer counted - see the note on the - /// atomics below. It was a LOCKED read-modify-write in the park loop, so measuring the contention added to - /// it. - long long pauses() const { return 0; } + /// How long a parked worker spins before it sleeps. + static constexpr std::chrono::milliseconds kIdleSpin{50}; + /// How many times a worker has gone to sleep (each worker counts once per idle stretch). + uint32_t sleeps() const { return sleeps_.load(std::memory_order_relaxed); } - /// **WHERE `run()` SPENDS ITS TIME, in milliseconds accumulated over its lifetime.** Three phases per - /// layer - wait for every worker to be parked, wait for the drain, wait for them to re-park - and until now - /// all three were reported as one number. Without the split there is no way to tell a pool that is slow at - /// the WORK from one that is slow at the SYNCHRONISATION, and those need opposite fixes: the first is a - /// kernel problem and the second is a barrier problem. + /// **WHERE `run()` SPENDS ITS TIME, in milliseconds accumulated over its lifetime**: the drain, and closing + /// the phase (the joined workers parking again). Without the split there is no way to tell a pool that is + /// slow at the WORK from one that is slow at the SYNCHRONISATION, and those need opposite fixes. /// /// Only the host thread touches these, in `run()`, so they need no atomics. - void phase_ms(double& wait_park, double& drain, double& repark) const { - wait_park = ms_wait_park_; + void phase_ms(double& drain, double& close) const { drain = ms_drain_; - repark = ms_repark_; + close = ms_close_; } - /// **A PARKED WORKER SPINS FOR THIS LONG, THEN SLEEPS.** The park is a `_mm_pause` spin because a layer's - /// batches are microseconds apart and a wake-up from the OS costs more than that. But a spin that never ends - /// keeps every worker's core at 100% while the engine waits for a request - issue #4, "CPU 50% even when - /// doing nothing" (7 of the 5700X's 16 threads). Between requests the workers block on `sleep_cv_`. NOT only - /// between requests: with most experts in VRAM (a 24 GB card) many layers have no CPU work, so the workers - /// also sleep mid-request - which the claim protocol above must survive (issue #29). - /// `STRATA_POOL_SPIN_US` overrides it (a test knob: a short spin makes the workers sleep constantly). - static constexpr std::chrono::milliseconds kSpinBeforeSleep{20}; - /// A pool wait that sees no completion for this long is a bug; the engine stops with a message instead of - /// spinning forever, and the server starts it again (issue #29). - static constexpr std::chrono::seconds kStall{60}; - private: void worker(int id); - void drain(int id, ExpertScratch& scratch, uint32_t epoch); - void run_phase(int mode, int n_tasks); - /// Claim the next job of batch `epoch`, or -1 (that batch is exhausted, or it is not the current one). - int claim(uint32_t epoch); - /// Publish the batch whose description the caller has just written: reset `done`, then `head`, then the epoch. - uint32_t begin_batch(int n); - /// The host's waits, bounded by `kStall`. - void wait_parked(const char* what); - void wait_done(int n); - /// Bump `epoch_`, and wake the workers that went to sleep. Every publish goes through here. - void publish(); + void drain(int id, ExpertScratch& scratch); + void run_phase(int mode, int n_tasks, void (*first)(void*) = nullptr, void* ctx = nullptr); + void open_phase(); + void close_phase(); int n_ = 0; bool host_works_ = true; @@ -184,40 +198,22 @@ class ExpertPool { /// The host's own scratch when `host_works_`. A separate object rather than a share of `scratch_[i]`, /// because a worker may own any index and the two must not be able to collide. ExpertScratch host_scratch_; - // `run()`'s three phases, accumulated. Host-thread only; see `phase_ms`. - double ms_wait_park_ = 0.0; + // `run()`'s phases, accumulated. Host-thread only; see `phase_ms`. double ms_drain_ = 0.0; - double ms_repark_ = 0.0; - // ---- EACH ATOMIC GETS ITS OWN CACHE LINE, AND THE SPIN COUNTER IS GONE. (Review finding C3.) - // - // These were six adjacent atomics, which put `head_`, `done_`, `parked_` and `epoch_` on ONE cache line - - // the four that workers and the host actually contend on, invalidating each other on every access. - // - // Worse, the parked spin did `pauses_.fetch_add(1)` on EVERY iteration: a LOCKED read-modify-write, five - // workers against one line, at roughly one iteration per `_mm_pause`. So the line the host must WRITE to - // publish work (`epoch_`) and READ to confirm the workers are parked (`parked_`) was being hammered by the - // very threads waiting for it. That is contention the pool imposes on itself. + double ms_close_ = 0.0; + // ---- EACH ATOMIC GETS ITS OWN CACHE LINE, AND THE PARK LOOP WRITES NOTHING SHARED. (Review finding C3.) // - // The counter was diagnostic only - `pauses()` was read in one place, to print a number nothing branched on - // - so it is deleted rather than amortised. `alignas(64)` then stops the remaining four sharing. - // issue #31 diagnostics: each worker's state (kParked, kSleeping, kBetween, or the job it runs) and the host's - // (kIdle, kWaitParked, kWaitDone, or its job), printed by the serve watchdog through `diag` - static constexpr int32_t kParked = -1, kSleeping = -2, kBetween = -3, kIdle = -10, kWaitParked = -11, - kWaitDone = -12; - std::unique_ptr[]> wstate_; - std::atomic hstate_{kIdle}; - std::atomic hstate_ms_{0}; - alignas(64) std::atomic head_{0}; // epoch << 32 | njobs << 16 | next index (issue #29) + // Adjacent atomics that workers and the host contend on invalidate each other on every access, and a + // diagnostic counter the park loop once incremented on every iteration hammered the very line the host + // writes to publish work. `alignas(64)` keeps them apart; the park loop only reads `state_` (a worker + // writes it once to join a phase and once to park again, and `sleepers_` only when it goes to sleep). + alignas(64) std::atomic head_{0}; alignas(64) std::atomic done_{0}; - alignas(64) std::atomic parked_{0}; - alignas(64) std::atomic epoch_{0}; - alignas(64) std::atomic stop_{false}; - // The sleep after `kSpinBeforeSleep`. `sleepers_` is how `publish` knows whether anyone needs waking, so the - // token path pays one uncontended load per publish and never takes the mutex while the workers spin. + /// epoch << 32 | kClosed (bit 31) | parked workers + alignas(64) std::atomic state_{0}; alignas(64) std::atomic sleepers_{0}; - std::mutex sleep_mu_; - std::condition_variable sleep_cv_; - std::chrono::microseconds spin_before_sleep_{kSpinBeforeSleep}; + std::atomic sleeps_{0}; + alignas(64) std::atomic stop_{false}; std::vector threads_; std::vector scratch_; // one per worker: no allocation, no false sharing of the hot data // run_split state: mode 0 = whole experts, 1 = gate/up row parts, 2 = down row parts @@ -240,5 +236,8 @@ class ExpertPool { const NativeFmt* nfmt_ = nullptr; std::vector split_multi_; }; +#ifdef _MSC_VER +#pragma warning(pop) +#endif } // namespace strata::kernels::cpu diff --git a/include/strata/kernels/elementwise.hpp b/include/strata/kernels/elementwise.hpp index 0136176c47c..6db5abd9dfa 100644 --- a/include/strata/kernels/elementwise.hpp +++ b/include/strata/kernels/elementwise.hpp @@ -116,5 +116,13 @@ void doorbell_publish(const float* x, const int32_t* ids, const float* weights, /// Plan v0.3 P3: copy `n` int32 from mapped pinned host memory into device memory with a kernel (the QSA /// per-token step and positions), instead of a host-to-device memcpy node in the middle of a layer. void copy_i32_from_mapped(int32_t* dst, const int32_t* src, int64_t n, void* stream); +/// Up to 8 copies from mapped host memory in one launch: `words` 4-byte words each (16 bytes a load where both +/// ends are 16-byte aligned and the count a multiple of 4). Graph-capturable. +struct MappedCopy { + void* dst = nullptr; + const void* src = nullptr; + int64_t words = 0; +}; +void copy_from_mapped_multi(const MappedCopy* copies, int n, void* stream); } // namespace strata::kernels diff --git a/include/strata/kernels/fused_gr.hpp b/include/strata/kernels/fused_gr.hpp index ba4e7c9d9e4..abe486a616a 100644 --- a/include/strata/kernels/fused_gr.hpp +++ b/include/strata/kernels/fused_gr.hpp @@ -1,28 +1,32 @@ -// include/strata/kernels/fused_gr.hpp - plan v0.3 P3: the hyper-connection read in TWO kernels, with the -// previous half's write folded in. +// include/strata/kernels/fused_gr.hpp - the hyper-connection read in two kernels, with the previous half's write +// folded in, for up to 8 tokens that share the weights. // -// The native path spends six kernels per `gr_read` (norm, down MMVF, silu, up MMVF, gate+mean, inject MMVF) and -// one per `gr_write`, 96 + 96 times per token, and its up projection (10240 rows of 320) runs one 160-thread -// block per row: 20.6 us for 6.5 MB. Here: +// down : R' = R + bo_prev * 2 sigmoid(inj_prev / hc) (only when `apply`, computed on the fly) +// rs[c] = rsqrt(mean(R'[c]^2) + eps) +// lo[k] = silu((w_down[k] . xn) / hc), xn = R' * w_norm * rs k < hc_lr +// inject[c] = w_inject[c] . xn when w_inject is given +// up : R <- R' in place (when `apply`) +// mixed[d] = mean_c xn[c,d] * sigmoid(w_up[c*n_embd + d] . lo) +// gate_ema[c,d] += (mean over the tokens of sigmoid(w_up[c*n_embd + d] . lo) - gate_ema[c,d]) / 4 +// when `gate_ema` is given +// est[d] = mean_c R'[c,d] * rs[c] * est_norm[c,d] * est_gates[c,d] when `est_norm` is given: the mixed of +// another read of R' (its norm weights) with the gates `est_gates` // -// fused_gr_down : R' = R + bo_prev * 2 sigmoid(inj_prev / hc) (only when `apply`, computed on the fly) -// rs[c] = rsqrt(mean(R'[c]^2) + eps), xn = R' * w_norm * rs -// lo[k] = silu((w_down[k] . xn) / hc) k < hc_lr -// inject[c] = w_inject[c] . xn when w_inject is given -// fused_gr_up : R <- R' in place for this block's columns (when `apply`) -// mixed[d] = mean_c xn[c,d] * sigmoid(w_up[c*n_embd + d] . lo) -// -// FP32 activations and BF16 weights, like the native MMVF contract; the summation order differs from it (G-C -// judges the result). Geometry is the artifact's: n_embd 2560, hc 4, hc_lr 320. `inj_prev` and `inject_out` -// must be different buffers (every block reads the former while one block writes the latter). +// The down projection is split by input slices: 80 blocks each take a quarter of the 324 rows over a twentieth of +// the input, so each block reads 1/20 of every token's R' instead of all of it, and the stream's rs multiplies the +// slice sums (a slice lies in one stream). The last block of each quarter adds the partials in a fixed order. +// FP32 activations and BF16 weights; every token's outputs are the same whatever its place in the window and the +// window's size. Geometry is the artifact's: n_embd 2560, hc 4, hc_lr 320. `inj_prev` and `inject_out` must be +// different buffers (every block reads the former while one block writes the latter). #pragma once +#include #include namespace strata::kernels { struct FusedGrArgs { - const float* R = nullptr; ///< (hc, n_embd), read by `down`; `up` updates it in place when apply + const float* R = nullptr; ///< (hc, n_embd); `up` updates it in place when apply float* R_out = nullptr; ///< == R for the in-place update bool apply = false; ///< fold the previous half's gr_write const float* bo_prev = nullptr; ///< that half's block output, n_embd @@ -36,17 +40,20 @@ struct FusedGrArgs { float* rs = nullptr; ///< workspace, hc floats float* inject_out = nullptr; ///< hc floats (when w_inject) float* mixed = nullptr; ///< n_embd + float* gate_ema = nullptr; ///< (hc, n_embd): the gates' running average over the reads, or null + const float* est_norm = nullptr; ///< another read's w_norm, or null + const float* est_gates = nullptr; ///< (hc, n_embd): the gates for its estimate + float* est = nullptr; ///< n_embd: the estimate }; bool fused_gr_supported(int64_t n_embd, int64_t hc, int64_t hc_lr); -void fused_gr_read(const FusedGrArgs& a, void* stream); -/// Plan v0.3 P6: the same read for up to 8 tokens that share the weights (a verify window): the weights are read -/// once for all of them. `a[t]` is token t's arguments (its own R, pending write, lo, rs, inject, mixed; the four -/// weight pointers and eps must be the same for every t); `xn_scratch` is n_tok * hc * n_embd floats. Every -/// token's outputs are bitwise `fused_gr_read(a[t])`. +/// The read of `n_tok` tokens (1-8): `a[t]` is token t's arguments (its own R, pending write, lo, rs, inject, +/// mixed, est; the four weight pointers, gate_ema, est_norm, est_gates and eps must be the same for every t). +/// `scratch` is `fused_gr_scratch_bytes()`, zeroed once when allocated (the kernels leave it zeroed); reads that share +/// it must run in stream order. constexpr int kFusedGrMaxT = 8; -void fused_gr_read_multi(const FusedGrArgs* a, int n_tok, float* xn_scratch, void* stream, - unsigned long long* stamp_buf = nullptr, int stamp_i0 = 0); +size_t fused_gr_scratch_bytes(); +void fused_gr_read_multi(const FusedGrArgs* a, int n_tok, float* scratch, void* stream); } // namespace strata::kernels diff --git a/include/strata/kernels/iq_kernels.hpp b/include/strata/kernels/iq_kernels.hpp index 020d78d522f..721790f834f 100644 --- a/include/strata/kernels/iq_kernels.hpp +++ b/include/strata/kernels/iq_kernels.hpp @@ -1,5 +1,6 @@ -// include/strata/kernels/iq_kernels.hpp - the i-quant formats (IQ1_M, IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ3_S, -// IQ4_NL) and Q2_0 on the GPU for the IQ2_XS / IQ3_XXS model files. +// include/strata/kernels/iq_kernels.hpp - the GGUF expert and embedding formats on the GPU: the i-quants (IQ1_M, +// IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ3_S, IQ4_NL) and Q2_0 of the IQ2_XS / IQ3_XXS model files, and Q4_K, Q5_K, +// Q5_1 and Q8_0 of the UD-Q4_K_XL file. // // The block layouts, codebook grids and dot products are llama.cpp's (ggml-common.h, ggml-cuda/vecdotq.cuh, // ggml-cuda/dequantize.cuh; MIT, see third_party/ggml/LICENSE and VERSION.txt), so a weight means exactly what it @@ -21,6 +22,11 @@ void quantize_q8_1_rows(const float* x, int64_t n_rows, int64_t n_cols, void* y, /// y[c][r] = W[r] . x[c] for `ncols` columns of q8_1 activations (x stride n_in/32 blocks per column). void iq_mmvq(int ggml_type, const void* w, const void* x_q8_1, float* y, int n_in, int n_out, int ncols, void* stream); +/// iq_mmvq for 2..8 columns from their interleaved copy (native_quantize_q8_1_il), or one from its q8_1 blocks, +/// bitwise, a warp taking `rows` rows (1, 2 or 4): IQ3_S. +bool iq_mmvq_il_supported(int ggml_type) noexcept; +void iq_mmvq_il(int ggml_type, const void* w, const void* x_q8_1, const void* x_il, float* y, int n_in, int n_out, + int ncols, void* stream, int rows); /// Dequantize `n` contiguous values (n a multiple of 256) to fp16 / fp32. void iq_dequant_f16(int ggml_type, const void* src, int64_t n, uint16_t* dst, void* stream); @@ -33,6 +39,10 @@ void iq_embed_rows(int ggml_type, const void* table, size_t row_bytes, const int /// uses: row 2r = gate row r, row 2r+1 = up row r. void iq_dequant_gu_f16(int ggml_type, const void* gate, const void* up, int64_t n_ff, int64_t n_embd, uint16_t* dst, void* stream); +/// iq_dequant_gu_f16 and iq_dequant_f16 (the down matrix) of n experts in one launch: expert i's blob at blobs[i] +/// (a DEVICE array), its gate/up into gu + i * 2 * n_ff * n_embd and its down into dn + i * n_embd * n_ff. +void iq_dequant_experts_f16(int gu_type, int d_type, const uint8_t* const* blobs, int n, size_t up_off, size_t down_off, + int64_t n_ff, int64_t n_embd, uint16_t* gu, uint16_t* dn, void* stream); /// The layout of one native expert blob: [gate rows | up rows | down rows], raw GGUF blocks. struct NativeExpertLayout { @@ -43,6 +53,9 @@ struct NativeExpertLayout { size_t bytes = 0; // the whole blob }; NativeExpertLayout native_expert_layout(int gu_type, int d_type, int64_t n_embd, int64_t n_ff); +/// Whether the GPU has the dot products (native_expert_grouped) and the dequantizers (the prompt path) for a layer's +/// gate/up and down formats at this geometry. +bool native_expert_supported(int gu_type, int d_type, int64_t n_embd, int64_t n_ff) noexcept; /// Bytes of scratch `native_expert_grouped` needs for `cap_entries` entries. size_t native_expert_scratch_bytes(int64_t cap_entries, int64_t n_ff); @@ -53,5 +66,13 @@ size_t native_expert_scratch_bytes(int64_t cap_entries, int64_t n_ff); void native_expert_grouped(const NativeExpertLayout& L, const unsigned long long* grp_ptr, const int32_t* grp_start, const int32_t* n_groups, const int32_t* ent_dst, const int32_t* ent_tok, int64_t cap_groups, int64_t cap_entries, const void* x_q8_1, void* scratch, float* out, void* stream); +/// Entries [grp_start[0], grp_start[*n_groups]): entry e's row ent_dst[e] of `rows` (n_embd floats each, +/// n_embd % 4 == 0) to the same row of `*out`, at most `max_rows` rows, in 16-byte stores: rows bound for mapped host +/// memory cross PCIe in full transactions (the grouped kernels' one-float stores there cost ~0.2 ms a layer). Then, +/// unless `*flag` is null, `**flag = *ring` once every row is visible to the host. Counts and pointers are read on +/// the device; `counter` (device, zero) is left zero. +void native_expert_rows_out(const float* rows, const int32_t* n_groups, const int32_t* grp_start, + const int32_t* ent_dst, int64_t n_embd, int64_t max_rows, float* const* out, + uint32_t* const* flag, const int32_t* ring, unsigned* counter, void* stream); } // namespace strata::kernels diff --git a/include/strata/kernels/kv_q8.hpp b/include/strata/kernels/kv_q8.hpp index a535cc32d8f..e94d23abe9f 100644 --- a/include/strata/kernels/kv_q8.hpp +++ b/include/strata/kernels/kv_q8.hpp @@ -12,7 +12,6 @@ // the FP16 pools: codes `[page][kv_head][page_size][head_dim]`, scales `[page][kv_head][page_size][head_dim/64]`. #pragma once -#include "strata/kernels/kv_stream.hpp" #include "strata/kernels/qsa.hpp" #include @@ -26,11 +25,14 @@ inline uint64_t kv_q8_bytes_per_cell(const QsaShapes& s) { return (uint64_t) s.n_head_kv * s.head_dim * 2 + (uint64_t) s.n_head_kv * (s.head_dim / KV_Q8_GROUP) * 2 * 2; } -/// Append the cell at step[kStepPos] (graph-capturable: position and page come from device memory). With a host -/// copy (KV streaming, `kv_stream.hpp`) the cell is written there too, and to VRAM only if its block is resident. +/// Append the cell at step[kStepPos] (graph-capturable: position and page come from device memory). void kv_append_q8_step(int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, const int32_t* page_table, - const int32_t* step, const float* kcur, const float* vcur, const QsaShapes& s, void* stream, - const KvHostPools* host = nullptr); + const int32_t* step, const float* kcur, const float* vcur, const QsaShapes& s, void* stream); +/// `kv_append_q8_step` for n_tok cells in one launch: token t's step record at step + t * step_stride, its K and V at +/// kcur/vcur + t * n_head_kv * head_dim. Bitwise the per-token appends. +void kv_append_q8_steps(int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, const int32_t* page_table, + const int32_t* step, int step_stride, const float* kcur, const float* vcur, int n_tok, + const QsaShapes& s, void* stream); /// Gather step[kStepWidth] cells named by `ids` into FP16 scratch `[id][kv_head][head_dim]`; the grid is sized by /// `max_ids` (capacity), the kernel reads the real count from `step`. diff --git a/include/strata/kernels/native_gr_norm.hpp b/include/strata/kernels/native_gr_norm.hpp index b8bb4ed1b6c..dc5b3ab9cc3 100644 --- a/include/strata/kernels/native_gr_norm.hpp +++ b/include/strata/kernels/native_gr_norm.hpp @@ -14,5 +14,9 @@ namespace strata::kernels { // default stream, preserving gr_read's existing API; this function never syncs. void native_gr_rms_norm_weighted(const float* input, const float* gamma, float* output, int n_cols, int n_rows, float epsilon, void* stream); +// The same with gamma holding gamma_rows rows: row r is scaled by gamma row r % gamma_rows (several tokens' +// rows against one token's gamma). +void native_gr_rms_norm_weighted_repeat(const float* input, const float* gamma, float* output, + int n_cols, int n_rows, int gamma_rows, float epsilon, void* stream); } // namespace strata::kernels diff --git a/include/strata/kernels/native_mmvq.hpp b/include/strata/kernels/native_mmvq.hpp index 9f69dd940cf..959e1d05672 100644 --- a/include/strata/kernels/native_mmvq.hpp +++ b/include/strata/kernels/native_mmvq.hpp @@ -109,4 +109,20 @@ std::size_t native_mmvq_weight_bytes(int ggml_type, int n_in, int n_out); void native_mmvq(int ggml_type, const void* weights, const void* x_q8_1, float* y, int n_in, int n_out, int ncols, void* stream); +// Plan v0.3 P6: 2..8 columns also INTERLEAVED, so one 16-byte load reads a position of every column: CP = 2 (2 +// columns), 4 (3-4) or 8 (5-8), zero-padded; nb = n_in / 32 blocks; int32 bm[nb][8][CP] (block-major), int32 +// pm[8][nb][CP] (position-major), float d[nb][CP] (each block's scale as the dots read it, __low2float of ds). +constexpr int native_q8_1_il_cp(int ncols) { return ncols <= 2 ? 2 : ncols <= 4 ? 4 : 8; } +std::size_t native_q8_1_il_bytes(int n_in, int ncols); +// native_quantize_q8_1 that also writes the interleaved copy to x_il (native_q8_1_il_bytes); 2 <= ncols <= 8. +void native_quantize_q8_1_il(const float* x, void* x_q8_1, void* x_il, int n_in, int ncols, void* stream); +// native_mmvq for 1 column (x_il unused) or 2..8 quantized by native_quantize_q8_1_il, bitwise native_mmvq's output: +// kernels in which a warp takes 1, 2 or 4 rows, reading 2+ columns from the interleaved copy, where they are the faster +// on the RTX 3090 (a table by format, columns and rows: IQ4_XS, Q4_K, Q5_K, Q6_K and Q8_0 matrices of 2+ columns and +// 2048+ rows, IQ3_S ones); native_mmvq otherwise. +void native_mmvq_il(int ggml_type, const void* weights, const void* x_q8_1, const void* x_il, float* y, int n_in, + int n_out, int ncols, void* stream); +// Tests and benchmarks: every native_mmvq_il call takes `rows` a warp (0: the table), and with `one` 1 column too. +void native_mmvq_il_tune(int rows, bool one); + } // namespace strata::kernels diff --git a/include/strata/kernels/native_moe.hpp b/include/strata/kernels/native_moe.hpp index 34760efe64a..cb2dc21895c 100644 --- a/include/strata/kernels/native_moe.hpp +++ b/include/strata/kernels/native_moe.hpp @@ -15,7 +15,17 @@ bool native_moe_combine_enabled(); // Requires a nonnull ordered stream and disjoint output. No allocation or sync. void native_moe_combine(const float* parts, const float* weights, const float* shared, float* output, int64_t n_embd, int64_t k, void* stream); -/// n_tok rows (parts [n,k,N], weights [n,k], shared/output [n,N]) in one launch, each as the single call. -void native_moe_combine_multi(const float* parts, const float* weights, const float* shared, float* output, - int64_t n_embd, int64_t k, int n_tok, void* stream); + +// The verify window's combine for n_tok tokens, bitwise `native_moe_combine` per token over the rows the window +// used to assemble: row t*k + j comes from `gpu_rows` (row-indexed) as the host's zeroed row plus the hit when an +// entry of dst[0..*count) or dst2[0..*count2) names it, as it is when an entry of dst3[0..*count3) does, and +// otherwise from `host_rows` (mapped memory: only these rows cross PCIe). `dst2` and `dst3` may be null. +// weights (n_tok, k), shared and output (n_tok, n_embd). With `shared_gate` (n_tok logits of the shared expert's +// scalar gate, shared_expert_multi) the shared rows are its down projection's, and each is gated first, bitwise +// shared_expert_gate_rows. Graph-capturable. +void native_moe_gather_combine(const float* gpu_rows, const float* host_rows, const int32_t* dst, + const int32_t* count, const int32_t* dst2, const int32_t* count2, const int32_t* dst3, + const int32_t* count3, const float* weights, const float* shared, + const float* shared_gate, float* output, int64_t n_embd, int64_t k, int n_tok, + void* stream); } diff --git a/include/strata/kernels/native_ple_postops.hpp b/include/strata/kernels/native_ple_postops.hpp index 97c49903469..2fe69748206 100644 --- a/include/strata/kernels/native_ple_postops.hpp +++ b/include/strata/kernels/native_ple_postops.hpp @@ -27,12 +27,21 @@ void native_ple_postops(const float* projected_key, const float* hidden, const PleWeights& weights, const NativePlePostopsBuffers& buffers, void* stream); -/// The prompt path: T tokens at once, element for element the arithmetic of T calls of native_ple_postops with -/// the history advanced after each. `key` (T x 10240) holds the projected keys and is normalized in place; -/// `hidden` (T x 10240) becomes the result; `value` T x 2560; `query_norm` (T x 10240, the normalized hidden, then -/// the normalized gated values), `gated` (T x 10240) and `gate` (T x 4) are scratch; `history` (9 rows per -/// channel) is advanced past the T tokens. -void native_ple_postops_batch(float* key, float* hidden, const float* value, float* history, const PleWeights& w, - float* query_norm, float* gated, float* gate, int T, void* stream); +// n_tok consecutive tokens (the prompt path, a verify window's group), row-major with the token slow: projected_key, +// hidden, key, query, gated, normalized and result [n_tok, 10240], value [n_tok, 2560], gate [n_tok, 4]. Bitwise +// n_tok calls of native_ple_postops, each followed by ple_history_advance: history is read, then left holding the +// last nine normalized rows; `snaps` (optional, [n_tok, 9 * 10240]) receives the history after each token. result +// may equal hidden and query may equal normalized; nothing else may overlap. +struct NativePleTokensBuffers { + float* key; + float* query; + float* gate; + float* gated; + float* normalized; + float* result; +}; +void native_ple_postops_tokens(const float* projected_key, const float* hidden, const float* value, + float* history, const PleWeights& weights, const NativePleTokensBuffers& buffers, + int n_tok, void* stream, float* snaps = nullptr); } // namespace strata::kernels diff --git a/include/strata/kernels/native_qsa.hpp b/include/strata/kernels/native_qsa.hpp index c56eb4d29c3..961eed7eeee 100644 --- a/include/strata/kernels/native_qsa.hpp +++ b/include/strata/kernels/native_qsa.hpp @@ -21,6 +21,13 @@ bool native_qsa_enabled(); void native_qsa_rms_norm_weighted(const float* input, const float* gamma, float* output, int n_cols, int n_rows, float epsilon, void* stream); +// native_qsa_rms_norm_weighted, then native_rope_apply_tokens (64 rotated columns) on its output, in one launch and +// bitwise both: n_rows rows of n_cols (128 or 256) read at input + r * in_stride, written contiguously to output; +// row r's position is positions[(r / heads) * pos_stride + r % heads]. output == input needs in_stride == n_cols. +void native_qsa_norm_rope_tokens(const float* input, int in_stride, const float* gamma, float* output, int n_cols, + int n_rows, float epsilon, int n_rot, float freq_base, const int* positions, int heads, + int pos_stride, void* stream); + // attn/output [head_dim,n_head]; q_full [2*head_dim,n_head], each row is // [query channels, gate channels]. Output is attn*sigmoid(second-half gate). // Exact output==attn is supported; all other spans must be disjoint. diff --git a/include/strata/kernels/native_qsa_indexer.hpp b/include/strata/kernels/native_qsa_indexer.hpp index 560b12b5bec..b5cbd104f29 100644 --- a/include/strata/kernels/native_qsa_indexer.hpp +++ b/include/strata/kernels/native_qsa_indexer.hpp @@ -38,5 +38,12 @@ void native_qsa_indexer_append(const float* raw, const int32_t* relative_pos_dev int32_t pos_base, const float* gamma, float epsilon, const QsaIndexerBuffers& buffers, const QsaShapes& shapes, int64_t max_cells, float freq_base, void* stream); +// n_tok appends in one launch, in order: token t's raw key at raw + t * 128, its position at +// relative_pos_device[t * pos_stride]. Bitwise n_tok calls of native_qsa_indexer_append. With `tail_snap` +// (3 * 128 floats) the key tail as it was before them is copied there first. +void native_qsa_indexer_append_multi(const float* raw, const int32_t* relative_pos_device, int pos_stride, + int n_tok, int32_t pos_base, const float* gamma, float epsilon, + const QsaIndexerBuffers& buffers, const QsaShapes& shapes, + int64_t max_cells, float freq_base, void* stream, float* tail_snap = nullptr); } // namespace strata::kernels diff --git a/include/strata/kernels/native_rope.hpp b/include/strata/kernels/native_rope.hpp index 24bffe61e67..c4ab78c6605 100644 --- a/include/strata/kernels/native_rope.hpp +++ b/include/strata/kernels/native_rope.hpp @@ -11,4 +11,9 @@ bool native_rope_enabled(); // Explicit stream required. No allocation or synchronization. void native_rope_apply(const float* x, float* out, int rows, int head_dim, int n_rot, float freq_base, const int* positions, void* stream); +// The same for several tokens' heads in one launch: row r is head r % heads of token r / heads, and takes +// positions[(r / heads) * pos_stride + r % heads] (a verify window's per-token position vectors). Each row is +// bitwise what native_rope_apply gives it. +void native_rope_apply_tokens(const float* x, float* out, int rows, int head_dim, int n_rot, float freq_base, + const int* positions, int heads, int pos_stride, void* stream); } diff --git a/include/strata/kernels/native_router.hpp b/include/strata/kernels/native_router.hpp index 0511b0b0103..c7d02db91b3 100644 --- a/include/strata/kernels/native_router.hpp +++ b/include/strata/kernels/native_router.hpp @@ -14,6 +14,39 @@ bool native_router_enabled(); // four-byte aligned and outputs disjoint from each other and the input. // Requires a nonnull ordered CUDA stream. No allocation or synchronization. void native_router_top10(const float* logits, int32_t* ids, float* weights, void* stream); -/// n_tok rows at once (logits [n,512], ids/weights [n,10]); each row exactly as the single call. +// `native_router_top10` for n_tok tokens in one launch: logits (n_tok, 512) -> ids, weights (n_tok, 10). void native_router_top10_multi(const float* logits, int32_t* ids, float* weights, int n_tok, void* stream); + +// Plan v0.3 P6: a verify window layer's router in one kernel: bf16_gemv_fp32_mmvf_multi's logits, +// native_router_top10_multi's ids and weights and doorbell_publish's ring, bitwise, and verify_hit_plan's plan. The +// forms of x are written while the logits run (with `xq1`, its q8_1 rows bitwise quantize_q8_1_rows' and its Q8_K +// rows ggml's quantize_row_q8_K's); the ids and weights (and their mapped copies) once every token's logits are in; +// then the sequence number; then the plan. +struct VerifyRouterArgs { + const float* x = nullptr; // (n_tok, n_embd) + const uint16_t* w = nullptr; // the router: bf16 (n_expert, n_embd) + float* logits = nullptr; // (n_tok, n_expert) scratch + int32_t* ids = nullptr; // (n_tok, 10) + float* weights = nullptr; // (n_tok, 10) + float* x_out = nullptr; // mapped copy of x, or null + uint8_t* xq1 = nullptr; // x as q8_1 rows (n_embd / 32 blocks each), or null + uint8_t* xq1_out = nullptr; // mapped copy of them, or null (needs xq1) + uint8_t* xk_out = nullptr; // mapped: x as Q8_K rows, `xk_stride` bytes apart (a multiple + int xk_stride = 0; // of 16), or null (needs xq1) + int32_t* ids_out = nullptr; // mapped copies of ids and weights, or null + float* w_out = nullptr; + uint32_t* seq = nullptr; // mapped: set to `ring` once those copies are visible + uint32_t ring = 0; // (doorbell_publish's increment: the rings so far, fixed at capture) + const int32_t* res = nullptr; // verify_hit_plan's arguments; plan null: no plan + const unsigned long long* slot_ptr = nullptr; + int32_t* plan = nullptr; + int cap = 0, ptr_off = 0; + unsigned* counter = nullptr; // device, zero before the first launch (each launch leaves it so) + int n_tok = 0, n_embd = 0, n_expert = 0; // 1..8 tokens, 512 experts, n_embd a multiple of 512 + const float* bias = nullptr; // (n_expert) added to the logits, or null +}; +void verify_router(const VerifyRouterArgs& a, void* stream); +// A prediction's bias: bias[e] = fmaf(1/8 / n_tok, sum over the n_tok tokens of (logits[t][e] - predicted[t][e]), +// bias[e]), where predicted holds the logits a verify_router with this bias computed for the same tokens (512 experts). +void verify_router_bias_update(const float* logits, const float* predicted, int n_tok, float* bias, void* stream); } diff --git a/include/strata/kernels/ngram.hpp b/include/strata/kernels/ngram.hpp index d15491dedc3..ef6fc180176 100644 --- a/include/strata/kernels/ngram.hpp +++ b/include/strata/kernels/ngram.hpp @@ -5,9 +5,9 @@ // * `ngram_rows` - the HOST hash. Sixteen row indices per token from the last three token ids, computed // with 64-bit multiply/xor and no tensor op at all. `qwen4exp.cpp` L1092-1124, where the source says // outright that it is host-side because "ggml has no int64 and no xor". -// * `PleTable` - the row gather. `per_layer_token_embd.weight` is IQ4_NL and is NOT IN THE PACK: it is -// 51.2e9 elements (28.8 GB) in the ORIGINAL second GGUF shard, and it is the only tensor this engine -// reads from the GGUF rather than from the canonical pack. +// * `PleTable` - the row gather. `per_layer_token_embd.weight` is IQ4_NL (Q8_0 in the unsloth UD-Q5_K_XL and +// larger files) and is NOT IN THE PACK: it is 51.2e9 elements (28.8 GB; 54.4 GB as Q8_0) in the model's +// GGUF, and the engine reads it from there. // // THE HASH IS THE ONE PLACE IN THIS MODEL WHERE A WRONG BIT IS SILENTLY PLAUSIBLE. Every one of its // properties has a rival reading that produces a valid index in range: @@ -24,6 +24,7 @@ #include #include +#include namespace strata::kernels { @@ -45,6 +46,7 @@ inline constexpr float NG_RMS_EPS = 1e-6f; // 5 blocks of 32 at 18 bytes = 90 bytes. The head-slowest flatten then makes 16 rows exactly n_embd = 2560. inline constexpr uint64_t PLE_TABLE_ROWS = 320001536ull; inline constexpr int PLE_ROW_BYTES = (PLE_HEAD_DIM / 32) * 18; // 90 +inline constexpr int PLE_ROW_BYTES_Q8_0 = (PLE_HEAD_DIM / 32) * 34; // 170, the Q8_0 tables /// The artifact's own hash constants, transcribed from `docs/gguf-dump-shard1.txt`: /// @@ -95,15 +97,19 @@ int iq4nl_code(int code); /// so byte j carries elements j and j+16 - NOT 2j and 2j+1, which is the natural reading and which produces /// a perfectly plausible embedding of the wrong 160 values. void iq4nl_dequant_row(const uint8_t* row, float* out160); +/// One row of a table of GGML type `type` (IQ4_NL = 20 or Q8_0 = 8) -> 160 floats. +void ple_dequant_row(int type, const uint8_t* row, float* out160); /// How the table's rows are read (plan v0.3 P2). `Direct` is the default: unbuffered 4 KiB reads from the SSD, /// so the table never occupies RAM or the OS file cache. `Mmap` is the earlier memory-mapped path, kept as the -/// A/B arm; it returns the same bytes. +/// A/B arm; it returns the same bytes. `Ram`: Direct until `start_ram_load` has read the whole table into RAM +/// (unbuffered, so the file cache holds no second copy), then from RAM: a 16K-token prompt chunk's 262K random +/// row reads take ~0.6 s from the SSD. /// /// NEVER KEEP THE SHARD MAPPED WHILE READING IT DIRECT: a live section on the same file serializes the unbuffered /// reads (311 -> 1,575 us per token, bench/results/2026-09-23-p2-ssd-direct). Direct mode drops its own mapping /// after the header parse; nothing else in the process may hold one. -enum class PleIo { Direct, Mmap }; +enum class PleIo { Direct, Mmap, Ram }; struct PleIoOptions { PleIo mode = PleIo::Direct; @@ -112,6 +118,9 @@ struct PleIoOptions { bool io_thread = true; ///< reads submitted by a worker thread, not the caller }; +/// The shard among `shards` that holds `per_layer_token_embd.weight`, or "" (headers only). +std::string ple_table_shard(const std::vector& shards); + /// The PLE table. Held by pointer-to-impl so this header does not drag `` into every /// translation unit that wants the hash. class PleTable { @@ -133,9 +142,22 @@ class PleTable { /// while a single-token `issue` is pending. The mapped mode gathers row by row. bool gather_batch(const uint32_t* rows, size_t n_tokens, float* out, std::string& err); + /// Tokens read ahead, each in a slot of its own beside issue/collect: `ahead` starts reading a token's 16 rows + /// (a slot still holding other rows waits for them first; outside Direct mode the rows are only recorded), + /// `ahead_holds` tells whether a slot holds exactly these rows, `ahead_collect` waits and gives the 2560 floats. + static constexpr int kAheadSlots = 8; + void ahead(int slot, const uint32_t* rows16); + bool ahead_holds(int slot, const uint32_t* rows16) const; + bool ahead_collect(int slot, float* out2560, std::string& err); + /// Fault injection (Direct mode): every row read completes no earlier than `us` after issue. void set_injected_delay_us(double us); + /// Ram mode: reads the table into (pageable) RAM on a thread of its own; the rows come from the SSD until it is + /// done. Reports on stderr when done. + bool start_ram_load(std::string& err); + bool ram_ready() const; + /// Maps the ORIGINAL second GGUF shard read-only. The tensor's data does NOT start at file offset 0: /// the manifest's `shard2_tensor.offset` is relative to the GGUF's DATA SECTION, and the header before it /// is 192 bytes. Reading at 0 would silently decode the header plus 192 bytes of shifted rows - still diff --git a/include/strata/kernels/ple.hpp b/include/strata/kernels/ple.hpp index fabad838a3e..09c2da99e1c 100644 --- a/include/strata/kernels/ple.hpp +++ b/include/strata/kernels/ple.hpp @@ -33,6 +33,7 @@ namespace strata::kernels { /// Opt in to pinned CUDA BF16/F32 MMVF for the PLE value projection only (default false). /// Configure before session capture; captured graphs retain their selected projection kernels. void ple_set_native_bf16(bool enabled); +bool ple_native_bf16(); /// Opt in to pinned CUDA postprojection norms, gate, convolution and residual /// arithmetic (default false). Set before capture; existing graphs keep their diff --git a/include/strata/kernels/qsa_decode_attn.hpp b/include/strata/kernels/qsa_decode_attn.hpp index 45c0f301220..6637f1cb84e 100644 --- a/include/strata/kernels/qsa_decode_attn.hpp +++ b/include/strata/kernels/qsa_decode_attn.hpp @@ -1,17 +1,15 @@ -// include/strata/kernels/qsa_decode_attn.hpp - plan v0.3 P3/P7: split-K decode attention over the selected cells. +// include/strata/kernels/qsa_decode_attn.hpp - the QSA attention over each query's selected cells. // -// The first attention kernel runs one block per query head (24 blocks on a 48-SM part), reads each key row with one -// thread (uncoalesced) and walks every selected cell serially for the values; at a 4K context (2,051 selected -// cells) that costs ~0.3-0.5 ms per QSA layer. This one reads the KV POOLS directly through the page table and the -// selection ids (no gather copy), in chunks of CHUNK cells per block, and serves all `n_head / n_head_kv` query -// heads that share a KV head from one read of the chunk: +// The kernels read the KV POOLS directly through the page table and the selection ids (no gather copy) and serve +// all `n_head / n_head_kv` query heads that share a KV head from one read of each cell: a block per (KV head, query) +// walks the query's cells in tiles of 64 with an online softmax. INT8 pools on the tensor cores: q and the +// probabilities as 24-bit fixed point (three 8-bit limbs) against the int8 codes, the integer sums exact, ~2.5e-6 of +// a head's largest output from FP64 on real prompts (FP32 kernels ~3e-6). FP16 pools in FP32. // -// grid (ceil(cap / CHUNK), n_head_kv): s = q.k * scale for 12 heads x CHUNK cells, chunk max m and sum l, -// acc = p . V -> partials -// grid (n_head): merge the chunks with the usual log-sum-exp rescale -// -// FP16 pools or INT8 pools (codes + fp16 scale per 64 values, `kv_q8.hpp`). The grid is sized by the capacity; -// the real count comes from `step[kStepWidth]` as for every other capturable QSA kernel. +// The prompt path's queries fill the GPU by their number. A verify window's 1-4 queries would leave most of it idle, +// so the decode splits each query's cells among up to 20 blocks (whole tiles, all queries' blocks in one wave) whose +// partial sums a merge adds with the usual log-sum-exp rescale. Grids are sized by the capacity; the real count comes +// from `step[kStepWidth]` as for every other capturable QSA kernel. #pragma once #include "strata/kernels/qsa.hpp" @@ -32,15 +30,23 @@ struct QsaAttnPools { const int32_t* page_table = nullptr; }; -/// Scratch floats for `cap` selected cells: partial accumulators, maxima and sums. +/// Scratch floats a query with `cap` selected cells needs: its splits' partial sums, maxima and sums. uint64_t qsa_decode_attn_scratch_floats(int64_t cap, const QsaShapes& s); void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* step, int64_t cap, const QsaShapes& s, float* scratch, float* attn, void* stream); /// Plan v0.3 P5: `n_q` queries at once, each with its own selection: q [n_q, n_head, 256], ids [n_q, cap], steps -/// [n_q, kStepCount], attn [n_q, n_head, 256]; scratch is `n_q` times the single-query size. +/// [n_q, kStepCount], attn [n_q, n_head, 256]; scratch is `n_q` times the single-query size. With `gate` (the +/// query projection's rows [n_q, n_head, 2 * 256]: a head's queries, then its gate logits), attn is the gated output, +/// bitwise native_qsa_gate_apply's. void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, - int64_t cap, const QsaShapes& s, float* scratch, float* attn, int64_t n_q, void* stream); + int64_t cap, const QsaShapes& s, float* scratch, float* attn, int64_t n_q, void* stream, + const float* gate = nullptr); + +/// The prompt path's form of qsa_decode_attn_batch (no scratch): a block per (KV head, query) walks all of the +/// query's cells, so no partial sums go through memory. +void qsa_prefill_attn(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, int64_t cap, + const QsaShapes& s, float* attn, int64_t n_q, void* stream); } // namespace strata::kernels diff --git a/include/strata/kernels/qsa_select.hpp b/include/strata/kernels/qsa_select.hpp index e31d3f3bca4..4f37d0a4729 100644 --- a/include/strata/kernels/qsa_select.hpp +++ b/include/strata/kernels/qsa_select.hpp @@ -5,12 +5,17 @@ // single-block kernel that makes 32 bit-serial passes over every CELL (topk_kernel): ~0.1-0.2 ms per query at 32K, // once per QSA layer and token - 1-2 ms of a decode token and most of a 32K prompt's time. Here: // -// qsa_block_scores : one warp per (query, block), FP32; relu per indexer head, summed; block n_bid (the -// incomplete tail) scores the `dead` key and gets +1e9 when it has cells - which is exactly -// what the per-token path reads there, because pooled[n_bid] holds `dead` at that time (in -// a chunk it may already hold a block completed later, so it is not read). -// qsa_block_topk : one block per query; a 4-pass radix select over the query's n_bid + 1 blocks, each -// weighted by its cell count, then the cells emitted in ascending order with ties to the +// qsa_block_scores : FP32; relu per indexer head, summed; block n_bid (the incomplete tail) scores the `dead` key +// and gets +1e9 when it has cells - which is exactly what the per-token path reads there, +// because pooled[n_bid] holds `dead` at that time (in a chunk it may already hold a block +// completed later, so it is not read). Each score sums in the order one warp per +// (query, block) did. Decode windows (< 12 queries): an 8-lane group per key block reads the +// key once for all queries. Batches: a per-thread GEMM form, a lane's 2 queries x 4 key +// blocks over all heads, the tree on a stack of partial sums, the queries read by the whole +// warp at once; one persistent block an SM with an equal share of the (key tile, query chunk) +// units. +// qsa_block_topk : one 1024-thread block per query; a 4-pass radix select over the query's n_bid + 1 blocks, +// each weighted by its cell count, then the cells emitted in ascending order with ties to the // lowest index - the same selection as topk_kernel, over a quarter of the elements. // // Queries carry their own step record (pos, n_kv, n_bid, width) as everywhere else in QSA. @@ -23,21 +28,17 @@ namespace strata::kernels { /// scores [nq, max_blocks]; q_idx [nq, idx_n_head, idx_dim] (normed and rotated); steps [nq, kStepCount]. +/// The grid covers the first `grid_blocks` blocks and goes round for the rest: a captured graph passes a fixed size +/// (qsa_score_grid_blocks), the prompt path the batch's largest n_bid + 1. void qsa_block_scores(const float* pooled, const float* dead, const float* q_idx, const int32_t* steps, int64_t nq, - int64_t max_blocks, const QsaShapes& s, float* scores, void* stream, - int64_t active_blocks = -1); ///< perf-review C-1: > 0 launches only this many blocks (the - ///< batch's largest n_bid + 1; not for a captured graph) + int64_t max_blocks, const QsaShapes& s, float* scores, void* stream, int64_t grid_blocks); -/// The same scores on tensor cores (3xTF32, FP32-level accuracy but another summation order: not bitwise; the tail -/// block n_bid is the warp kernel's arithmetic). For the prompt path; false (nothing launched) on another geometry. -bool qsa_block_scores_tc(const float* pooled, const float* dead, const float* q_idx, const int32_t* steps, int64_t nq, - int64_t max_blocks, const QsaShapes& s, float* scores, void* stream, int64_t active_blocks); +/// The block-score grid of the decode paths, whatever the context: 164 thread blocks of 64 key blocks, one wave on +/// the RTX 3090 (two per SM); longer contexts go round. +constexpr int64_t qsa_score_grid_blocks = 10496; /// ids [nq, cap] (cells, ascending); `cap` >= the largest selection width. void qsa_block_topk(const float* scores, const int32_t* steps, int64_t nq, int64_t max_blocks, int64_t cap, const QsaShapes& s, int32_t* ids, void* stream); -/// The original kernel (keys read from memory on every radix pass), for tests: the same ids. -void qsa_block_topk_ref(const float* scores, const int32_t* steps, int64_t nq, int64_t max_blocks, int64_t cap, - const QsaShapes& s, int32_t* ids, void* stream); } // namespace strata::kernels diff --git a/include/strata/kernels/shared_expert.hpp b/include/strata/kernels/shared_expert.hpp index 8909ed9e607..e0c990b8d51 100644 --- a/include/strata/kernels/shared_expert.hpp +++ b/include/strata/kernels/shared_expert.hpp @@ -24,6 +24,7 @@ namespace strata::kernels { /// Configure before session capture; captured graphs retain their selected kernels. /// In this mode shared_expert requires its optional unrounded x_f32 input. void shared_expert_set_native_bf16(bool enabled); +bool shared_expert_native_bf16(); /// Optional native GGUF projections. Each supported type with nonnull data /// replaces only that canonical projection; absent or unsupported entries fall @@ -37,6 +38,7 @@ struct NativeSharedWeights { const void* up_data = nullptr; const void* down_data = nullptr; void* q8_1 = nullptr; + void* q8_1_il = nullptr; // shared_expert_multi: the interleaved copy of 2+ tokens' inputs (native_q8_1_il_bytes) }; /// Bytes of caller-owned scratch `shared_expert` needs. **THE KERNEL USED TO `cudaMalloc` FOUR BUFFERS ON @@ -68,13 +70,20 @@ void shared_expert(const uint8_t* x_q8_0, const uint8_t* x_q8k, const uint16_t* int64_t n_embd, int64_t n_ff, int tpr, void* stream, const float* x_f32 = nullptr, const NativeSharedWeights* native = nullptr); -/// Plan v0.3 P6: the shared expert for `n_tok` <= 8 tokens (a verify window) with all three projections native: -/// multi-column MMVQ, so the weights are read once; every token is bitwise `shared_expert` on that token. -/// `x` (n_tok, n_embd) f32, `x_bf16` the same rounded (only read when the native BF16 gate is off), `gate`/`up` -/// (n_tok, n_ff) scratch, `g` n_tok floats, `out` (n_tok, n_embd). `nw.q8_1` must hold n_tok columns of n_embd. -void shared_expert_multi(int n_tok, const float* x, const uint16_t* x_bf16, const NativeSharedWeights& nw, - const uint16_t* gate_inp_bf16, float* gate, float* up, float* g, float* out, int64_t n_embd, - int64_t n_ff, void* stream); +/// The native scalar gate applied: `g` (n_tok logits) becomes sigmoid(g), and each token's row of `out` (n_tok, +/// n_embd) is multiplied by it. +void shared_expert_gate_rows(float* out, float* g, int64_t n_embd, int n_tok, void* stream); + +/// Plan v0.3 P6: the shared expert for `n_tok` <= 8 tokens (a verify window) with all three projections and the gate +/// native: multi-column MMVQ, so the weights are read once. `x` (n_tok, n_embd) f32, `gate`/`up` (n_tok, n_ff) +/// scratch; `out` (n_tok, n_embd) receives the down projection's rows and `g` the scalar gate's n_tok logits, +/// computed on `gate_stream` (may be `stream`), so that the chain ends with the down projection: +/// `shared_expert_gate_rows` on them, or `native_moe_gather_combine` applying them, gives every token bitwise +/// `shared_expert` on that token. `nw.q8_1` must hold n_tok columns of n_embd; with `nw.q8_1_il` 2+ tokens' gate and +/// up projections run through `native_mmvq_il` (IQ3_S's weights decoded once for all tokens). +void shared_expert_multi(int n_tok, const float* x, const NativeSharedWeights& nw, const uint16_t* gate_inp_bf16, + float* gate, float* up, float* g, float* out, int64_t n_embd, int64_t n_ff, void* stream, + void* gate_stream); /// The MoE block's final combination, `ref/moe.py::moe` L156: /// diff --git a/include/strata/kernels/verify_kernels.hpp b/include/strata/kernels/verify_kernels.hpp index 50d951a6966..b837a103d4f 100644 --- a/include/strata/kernels/verify_kernels.hpp +++ b/include/strata/kernels/verify_kernels.hpp @@ -21,9 +21,10 @@ namespace strata::kernels { inline constexpr int kVerifyMaxT = 8; /// For token t of T: conv over [history(3) | qkv_0 .. qkv_t] -> SiLU -> L2 norm of the q/k heads -> h[t]. -/// `history` is NOT written. Bitwise `fused_gdn_conv_l2` per token. +/// `history` is NOT written, but with `commit` (one token): it then keeps the token, as gdn_conv_commit. Bitwise +/// `fused_gdn_conv_l2` per token. Up to 65535 tokens (the prompt path runs a whole chunk through it). void gdn_conv_l2_multi(const float* history, const float* qkv, const float* conv_w, float* h, int channels, - int qk_heads, float eps, int n_tok, void* stream, int t_begin = 0); + int qk_heads, float eps, int n_tok, void* stream, int t_begin = 0, bool commit = false); /// history <- the last 3 entries of [history | qkv_0 .. qkv_{n-1}], n = *n_keep (0 leaves it as it was). void gdn_conv_commit(float* history, const float* qkv, int channels, const int32_t* n_keep, void* stream); /// alpha/beta for T columns of x (T, n_embd): gate (T, h_v), beta (T, h_v). Bitwise `fused_gdn_ab` per column. @@ -38,24 +39,8 @@ void gdn_step_norm_multi(float* state, const float* h, int conv_channels, const /// Spin until *flag >= value (a mapped host flag). The value is fixed at capture, so several rings can be /// outstanding at once (the split verify window keeps two). void wait_flag_ge(const uint32_t* flag, uint32_t value, void* stream); -/// the GPU's %globaltimer (ns) into buf[i] (a one-thread kernel: the verify window's stage profiler). -void gpu_stamp(unsigned long long* buf, int i, void* stream); - -// ---- perf-review E-6: a layer whose routed experts are all in VRAM needs nothing from the host -/// One group's plan, built on the device when every routed expert of its n*k entries is resident: the host pool's -/// layout (counts | start | dst | tok | pad | ptr | ptr2 | start2, `capx` entries) and order (distinct experts in -/// routing order, their entries ascending), no PCIe groups. *skip = ring when it did, else 0. -void resident_plan(const int32_t* ids, int n_entries, int k, const int32_t* res_layer, int n_expert, - const uint8_t* cache_base, const unsigned long long* slot_off, long long blob, int32_t* plan, - long long capx, uint32_t* skip, uint32_t ring, void* stream); -/// wait_flag_ge that also returns when *skip == value (device memory). -void wait_flag_ge_or(const uint32_t* flag, uint32_t value, const uint32_t* skip, void* stream); -/// copy_i32_from_mapped unless *skip == value. -void copy_i32_from_mapped_unless(int32_t* dst, const int32_t* src, long long n, const uint32_t* skip, uint32_t value, - void* stream); -/// copy_from_mapped, or zeros when *skip == value (n a multiple of 4, 16-byte aligned). -void copy_or_zero_from_mapped(float* dst, const float* src, long long n, const uint32_t* skip, uint32_t value, - void* stream); +/// *dst = the GPU's global nanosecond timer when the stream reaches this point (`--window-profile`). +void gpu_stamp(unsigned long long* dst, void* stream); /// Rows of the S2/S4/S8 embedding for T token ids read from DEVICE memory; out (T, n). Bitwise `embedding_gather`. void embedding_gather_dev(const uint8_t* codes, const float* scales, const float* offsets, const int32_t* tokens, @@ -70,8 +55,20 @@ void copy_indexed(float* dst, const float* src, int64_t stride, const int32_t* i /// into dst + k * blob_bytes with coalesced 16-byte loads - the PCIe share of a layer's missed experts, staged /// into VRAM before the grouped expert kernel reads them. Launched for a capacity of `cap` blobs. void fetch_blobs(const unsigned long long* src, const int32_t* n, uint8_t* dst, int64_t blob_bytes, int cap, void* stream); +/// A list in mapped host memory, [n, pad x3, rows...]: each listed row of `src` (mapped, n_embd floats, a multiple of +/// 4) to the same row of `dst`, and the list to `list_dev` (device, the same layout). At most `cap` rows; rows +/// outside [0, cap) are not copied. The second GPU's rows of a verify window's layer, taken while the CPU works. +void fetch_listed_rows(const int32_t* list, const float* src, float* dst, int32_t* list_dev, int cap, int64_t n_embd, + void* stream); /// ptr[k] = base + k * blob_bytes for k < *n (the staged copies `fetch_blobs` made). void rebase_ptrs(unsigned long long* ptr, const int32_t* n, uint8_t* base, int64_t blob_bytes, void* stream); +/// The main GPU's share of a verify window's layer, decided on the device: the entries of `ids` (n = tokens * k, +/// n <= 128) whose expert its VRAM tier holds (`res[e]` = slot, else < 0; `slot_ptr[slot]` = the slot's address). One +/// group per distinct resident expert, in the order of their first entries, its entries in routing order; written in +/// the plan block's layout (`Verifier`): counts [groups, entries, 0, 0] | start (cap + 1) | dst (cap) | tok (cap) | +/// ptr (cap u64, at int32 offset `ptr_off`). +void verify_hit_plan(const int32_t* ids, int n, int k, const int32_t* res, int n_expert, + const unsigned long long* slot_ptr, int32_t* plan, int64_t cap, int64_t ptr_off, void* stream); // ---- the MTP draft layer (src/core/mtp.cpp) /// R[t][c][:] = h[t][c][:] + e[t][:] (the embedding branch added to every stream). @@ -83,12 +80,33 @@ void ident_hits(const int32_t* ids, int n, int32_t* slot, int32_t* dst, int32_t* void mtp_select(const float* R_src, int64_t R_stride, const int32_t* ids, const int32_t* row_dev, float* R_dst, int32_t* tok_dst, int32_t* out, int j, void* stream, const float* probs = nullptr, float* out_p = nullptr); -/// dst row i = src row ids[i] (row_bytes each, multiple of 16), for n rows. +/// Row *row_dev of a, b and c (a_n, b_n, c_n floats a row) copied to their row 0 (the draft layer's rest runs on +/// row 0). Graph-capturable: the row is read on the device. +void copy_row_to_first(const int32_t* row_dev, float* a, int64_t a_n, float* b, int64_t b_n, float* c, int64_t c_n, + void* stream); +/// dst row i = src row ids[i] (row_bytes each, a multiple of 4), for n rows. void gather_rows(const uint8_t* src, int64_t row_bytes, const int32_t* ids, int64_t n, uint8_t* dst, void* stream); /// ids[t] = table[ids[t]] for n entries (a subset index back to a token id). void map_ids(int32_t* ids, const int32_t* table, int n, void* stream); /// probs[t] = softmax(logits[t])[ids[t]] for n_rows rows of n_vocab (the probability of each row's argmax). void row_top_prob(const float* logits, int n_rows, int n_vocab, const int32_t* ids, float* probs, void* stream); +/// out[t] = the argmax of row t (n floats a row; ties to the lowest index, NaN never picked, 0 when no value is above +/// -inf): bitwise sample_tokens' greedy pick. Up to 128 blocks a row scan slices of it, the row's last block merges +/// their picks. `scratch`: argmax_rows_scratch_bytes(n_rows), zero before the first launch (each launch leaves it +/// so). Graph-capturable. +/// With `out_val` each pick's logit too (-inf when no value is above it). +uint64_t argmax_rows_scratch_bytes(int n_rows); +void argmax_rows(const float* logits, int n_rows, int n, void* scratch, int32_t* out, void* stream, + float* out_val = nullptr); +/// *flag += 1 once the stream's earlier work is visible to the host (mapped memory). Graph-capturable. +void mapped_bump(uint32_t* flag, void* stream); +/// dst row t (n floats) = a row t (na floats) then b row t (n - na floats), for `rows` rows. +void join_rows(float* dst, int64_t n, const float* a, int64_t na, const float* b, int rows, void* stream); +/// row_top_prob over 8 blocks a row, each computing 4 of its 32 warps' sums, the row's last block adding the 32 in +/// order: bitwise row_top_prob. `scratch`: row_top_prob_scratch_bytes(n_rows), zero before the first launch. +uint64_t row_top_prob_scratch_bytes(int n_rows); +void row_top_prob_split(const float* logits, int n_rows, int n_vocab, const int32_t* ids, float* probs, void* scratch, + void* stream); /// Dense-attention step records for `n` cells: [cell, cell+1, (cell+1)/4, cell+1] from cells[i] (device memory). void dense_steps(const int32_t* cells, int n, int32_t* steps, void* stream); /// A sliding attention window: for `n` step records (kStepCount ints each, n_kv at [1]) the selection becomes the diff --git a/include/strata/ngram/ple_reader.hpp b/include/strata/ngram/ple_reader.hpp index 0552ff1e4c5..4a0ee2ddf17 100644 --- a/include/strata/ngram/ple_reader.hpp +++ b/include/strata/ngram/ple_reader.hpp @@ -1,6 +1,7 @@ // include/strata/ngram/ple_reader.hpp - plan v0.3 P2: the n-gram table read straight from the SSD. // -// The table is 320,001,536 rows of 90 bytes (26.8 GiB) and is never held in RAM: every row comes from an +// The table is 320,001,536 rows of 90 bytes (IQ4_NL, 26.8 GiB; 170 bytes as Q8_0) and is never held in RAM: every +// row comes from an // unbuffered 4 KiB read (platform::DirectFile). A token needs 16 rows on 16 different pages, and all 16 depend // on the token itself, so the only time to hide them is the embedding plus layer 0. Hence the split API: // @@ -11,7 +12,7 @@ // `issue` also serves prefill: pass all 16 x N rows of a chunk; pages are deduplicated, sorted by offset and // kept at most `max_inflight` deep, so a chunk's reads can run while the previous chunk computes. // -// THE ROW CACHE IS NOT THE TABLE. It keeps rows this process has already fetched (90 bytes each, bounded, +// THE ROW CACHE IS NOT THE TABLE. It keeps rows this process has already fetched (one row each, bounded, // clock eviction). Measured on the frozen corpus (bench/results/2026-09-23-ngram-io): 1M rows (~95 MB) // would serve up to ~82% of reads; within one long prompt 20-34% of rows recur. Capacity 0 disables it. #pragma once @@ -24,7 +25,7 @@ namespace strata::ngram { -inline constexpr uint32_t ROW_BYTES = 90; +inline constexpr uint32_t ROW_BYTES = 90; ///< an IQ4_NL row, the default inline constexpr uint32_t PAGE = 4096; struct ReaderStats { @@ -52,17 +53,17 @@ class PleReader { PleReader(const PleReader&) = delete; PleReader& operator=(const PleReader&) = delete; - /// `table_offset` is the byte offset of row 0 in the file and `n_rows` the row count; both come from a - /// validated GGUF parse (PleTable::open checks the table exactly fills the file from there). + /// `table_offset` is the byte offset of row 0 in the file, `n_rows` the row count and `row_bytes` a row's + /// size; all come from a validated GGUF parse (PleTable::open checks the table exactly fills its span). /// `io_thread` (default): a worker thread submits and reaps reads, so `issue` costs the caller no ReadFile /// calls. false: the caller's thread does it (A/B arm). bool open(const std::string& path, uint64_t table_offset, uint64_t n_rows, uint32_t max_inflight, - uint64_t cache_rows, std::string& err, bool io_thread = true); + uint64_t cache_rows, std::string& err, bool io_thread = true, uint32_t row_bytes = ROW_BYTES); void close(); bool is_open() const; - /// Start fetching `n` rows; row i's 90 raw bytes land at `out_raw + 90 * i`. `out_raw` must stay valid - /// until `collect` returns. Out-of-range rows produce 90 zero bytes (the mmap path's behaviour). + /// Start fetching `n` rows; row i's raw bytes land at `out_raw + row_bytes * i`. `out_raw` must stay valid + /// until `collect` returns. Out-of-range rows produce zero bytes (the mmap path's behaviour). Ticket issue(const uint32_t* rows, size_t n, uint8_t* out_raw); /// Block until every row of the ticket is in `out_raw`. Returns false on an I/O error (message in `err`). diff --git a/include/strata/prefill/experts.hpp b/include/strata/prefill/experts.hpp new file mode 100644 index 00000000000..68263da0eb7 --- /dev/null +++ b/include/strata/prefill/experts.hpp @@ -0,0 +1,95 @@ +// include/strata/prefill/experts.hpp - the prompt path's routed experts on one GPU. +// +// A prompt chunk uses nearly every expert of every layer. Its experts go through in batches of consecutive ids: per +// batch one staging copy for each run of adjacent arena blobs among the experts the GPU's cache does not hold, one +// dequantization launch, one gather of the batch's rows, the two GEMMs of each expert with a SwiGLU between them, and +// one ordered add of every row, times its routing weight, into its token's sum. A token's experts add up in ascending +// id order, whatever the batches. A chunk of a few hundred tokens or more uses nearly every expert, so the blobs a +// layer will stream can be copied ahead (`prefetch`), while the prompt path's dense work runs and before the routing +// is known; experts that do not fit there stream batch by batch. +// +// Local: on the prompt path's own GPU and stream, reading the chunk's MoE input and adding into its MoE output there. +// Remote: on a second GPU. The first GPU (the 3090) sits on a PCIe 4.0 x4 link, over which an expert its cache lacks +// streams at 6.4 GB/s (~0.65 GB (IQ3_XXS) to ~1.4 GB (UD-Q4_K_XL) a layer); the second GPU sits on x16. The first GPU +// copies the chunk's MoE input (FP16) to pinned host memory; the second copies it in, computes the experts the first +// one's cache lacks - from its own cache when it holds them - and returns their sums (FP16: they stay below ~20) the +// same way, while the first computes the experts its cache holds. Neither GPU reads the other's memory: Windows gives GeForce cards no +// peer access. +#pragma once + +#include "strata/core/expert_cache.hpp" +#include "strata/core/expert_source.hpp" + +#include + +#include +#include +#include +#include + +namespace strata::prefill { + +class ExpertRunner { +public: + ExpertRunner(); + ~ExpertRunner(); + ExpertRunner(const ExpertRunner&) = delete; + ExpertRunner& operator=(const ExpertRunner&) = delete; + + /// For chunks of up to `max_chunk` tokens on `device`. `stream`: the prompt path's stream on that device (local), + /// or null (remote: its own streams; `main_device` is made current again after every call). `cache`/`res`: the + /// device's expert cache and its residency (n_layers x n_expert slots, -1 when not held), or null. `streams`: + /// whether experts its cache does not hold come its way (staging buffers). + bool init(int device, int main_device, void* stream, core::ExpertSource* src, const core::ExpertCache* cache, + const int32_t* res, int64_t n_expert, int64_t max_chunk, bool streams, std::string& err); + /// Chunks of at least this many tokens prefetch: they use nearly every expert, and their dense steps last long + /// enough to hide the copies (a 521-token chunk prefetching was 2.2 s slower on IQ3_XXS: routing is skewed, so + /// many of the experts no cache holds go unused in a short chunk). + static constexpr int64_t kPrefetchMin = 4096; + + /// Device bytes of the buffers for chunks of `chunk` tokens. + static uint64_t bytes_needed(int64_t chunk, int64_t n_expert, bool remote, bool streams); + /// The same with its prefetch area, for the residency `skip` (see prefetch_need) and its own as they are now. + uint64_t bytes_for(int64_t chunk, const int32_t* skip) const; + /// The buffers for chunks of up to `chunk` tokens, carved from `region` (`bytes` long, e.g. lent cache slots), or + /// with a null region allocated for `max_chunk` tokens (once). Before the first layer of every prompt that + /// uses a region. `area`: device memory for the prefetched blobs; null: what the region has left. + bool bind(void* region, uint64_t bytes, int64_t chunk, std::string& err, void* area = nullptr, + uint64_t area_bytes = 0); + /// Bytes of prefetch area that hold a layer's streamed experts: the largest layer's experts that neither its + /// cache nor `skip` (another GPU's residency, n_layers x n_expert, or null) holds, and a margin. + uint64_t prefetch_need(const int32_t* skip) const; + /// Before a layer's routing: copies the blobs of `experts` (ascending ids) its cache does not hold into the + /// prefetch area, as many as fit, on its copy stream; run_layer of that layer reads them there. + bool prefetch(int64_t layer, const std::vector& experts, std::string& err); + + /// Remote: pinned [max_chunk, n_embd]: the MoE input goes in here, the sums come out here (both FP16). + uint16_t* host_input() const; + uint16_t* host_sum() const; + /// Remote, before run_layer of the same layer: rows [t0, t0 + rows) of the input, copied in once `ready` (an + /// event of the first GPU, recorded once host_input holds them) has passed; a chunk's rows in pieces, as they + /// reach host_input. + void stage_input(int64_t t0, int64_t rows, cudaEvent_t ready); + + /// Layer `layer` of a chunk of T tokens: expert experts[j] serves rows [off[j], off[j + 1]), row r being token + /// src[r] with routing weight w[r]; each row times its weight adds to its token's sum. Local: reads `input` + /// [T, n_embd] (FP16) and adds into `sum` [T, n_embd], queued on the prompt path's stream. Remote: `input` and + /// `sum` are null, the input came through stage_input, and the sums reach host_sum in pieces of `piece` rows, + /// piece_done(k) recorded as piece k lands. Returns once the work is queued. + bool run_layer(int64_t layer, int64_t T, const std::vector& experts, const std::vector& off, + const std::vector& src, const std::vector& w, int64_t piece, + const uint16_t* input, float* sum, std::string& err); + cudaEvent_t piece_done(int64_t k) const; + + int64_t experts_resident = 0; ///< expert-layer groups served from its cache + int64_t experts_streamed = 0; ///< expert blobs copied from the arena + int64_t experts_prefetched = 0; ///< ...of which ahead of the routing + double ms_host = 0; ///< host time queuing its work + double ms_gpu = 0; ///< remote: its GPU time from the input's arrival (all layers but the last) + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace strata::prefill diff --git a/include/strata/prefill/gemm.hpp b/include/strata/prefill/gemm.hpp index c9ddc6f4dc3..70054e8af0f 100644 --- a/include/strata/prefill/gemm.hpp +++ b/include/strata/prefill/gemm.hpp @@ -6,7 +6,6 @@ // rounded to BF16, which is also what llama.cpp's batched CUDA path does. Tensor-core GEMM through cuBLAS. #pragma once -#include #include #include @@ -19,11 +18,12 @@ class Gemm { Gemm(const Gemm&) = delete; Gemm& operator=(const Gemm&) = delete; - /// `scratch_elems`: BF16 elements of the dequantization scratch (the largest weight dequantized at once). - bool init(void* stream, int64_t scratch_elems, std::string& err); - /// The same with caller-owned device buffers (the prompt path borrowing expert-cache slots). - bool init_external(void* stream, uint16_t* scratch, int64_t scratch_elems, void* workspace, size_t ws_bytes, - std::string& err); + /// The cuBLAS handle on `stream`, and with `side_streams` the side streams f16_grouped() spreads over; + /// set_buffers() before the first product. + bool init(void* stream, std::string& err, bool side_streams = false); + /// Caller-owned device buffers (e.g. lent expert-cache slots): the dequantization scratch (`scratch_elems` FP16 + /// elements: the largest weight dequantized at once; may be null for no native() calls) and the cuBLAS workspace. + void set_buffers(uint16_t* scratch, int64_t scratch_elems, void* workspace, size_t ws_bytes); /// Y[T, N] (fp32, row stride ldy) = X[T, K] (bf16, row-major) . W[N, K]^T (bf16, row-major). `beta` = 1 adds. void bf16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy = 0, @@ -33,27 +33,33 @@ class Gemm { void f16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy = 0, float beta = 0.0f); + /// Y_i[rows[i], N] = X_i[rows[i], K] . W_i[N, K]^T (FP16 in, FP32 out) for n problems of one shape but their + /// own row counts, in one cuBLAS grouped call; the pointer arrays are DEVICE memory, `rows` host memory. Where + /// cuBLAS does not group these types it runs f16() per problem, from the host copies of the pointers, spread + /// over the side streams when the problems average 128 rows or more (a GEMM of a few hundred rows leaves most + /// of the GPU idle; below that the launches, not the GPU, set the pace). + void f16_grouped(const uint16_t* const* X, const uint16_t* const* W, float* const* Y, const uint16_t* const* X_host, + const uint16_t* const* W_host, float* const* Y_host, const int* rows, int n, int64_t N, int64_t K); + /// W given as native GGUF blocks of `ggml_type`, dequantized to FP16 in the scratch, X in FP16. void native(const uint16_t* X, int ggml_type, const void* W_blocks, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy = 0, float beta = 0.0f); - /// Caller-owned buffers only: the scratch and workspace moved (the prompt path laid its buffers out again). - void rebind(uint16_t* scratch, int64_t scratch_elems, void* workspace, size_t ws_bytes); - uint16_t* scratch() const { return scratch_; } int64_t scratch_elems() const { return scratch_elems_; } void* stream() const { return stream_; } private: + static constexpr int AUX = 3; // side streams for independent products void* handle_ = nullptr; void* stream_ = nullptr; + void* aux_handle_[AUX] = {}; + void* aux_stream_[AUX] = {}; + void* fork_ = nullptr; + void* join_[AUX] = {}; uint16_t* scratch_ = nullptr; int64_t scratch_elems_ = 0; - void* workspace_ = nullptr; - bool external_ = false; - void* hipblaslt_state_ = nullptr; + bool grouped_ = true; // cuBLAS groups FP16 -> FP32 GEMMs (cleared at the first refusal) }; - - } // namespace strata::prefill diff --git a/include/strata/prefill/kernels.hpp b/include/strata/prefill/kernels.hpp index f4256945220..1f4ea38de73 100644 --- a/include/strata/prefill/kernels.hpp +++ b/include/strata/prefill/kernels.hpp @@ -5,8 +5,6 @@ // the kernel walks the chunk's tokens in order inside one launch. #pragma once -#include "strata/kernels/kv_stream.hpp" - #include namespace strata::prefill { @@ -14,22 +12,22 @@ namespace strata::prefill { // ---- hyper-connection (n_embd 2560, hc 4, hc_lr 320) /// xn[t, c*2560 + d] = R[t,c,d] * rsqrt(mean_d R[t,c,:]^2 + eps) * w_norm[c*2560 + d]; also its BF16 image. void gr_norm(const float* R, const float* w_norm, float eps, float* xn, uint16_t* xn16, int64_t T, void* stream); -/// F-1: gr_norm without its FP32 output: the row scales rs[t*4 + c] and the BF16 image; gr_mix_r then reads R. +/// gr_norm without the FP32 rows: the row scales rs[t*4 + c] and the BF16 image; gr_mix_r reads R instead. void gr_norm_rs(const float* R, const float* w_norm, float eps, float* rs, uint16_t* xn16, int64_t T, void* stream); -/// gr_mix with xn recomputed from R, rs and w_norm exactly as gr_norm computes it (the same bits). +/// gr_mix with xn recomputed from R, rs and w_norm as gr_norm computes it (the same bits). void gr_mix_r(const float* R, const float* rs, const float* w_norm, const float* gated, float* mixed, uint16_t* mixed16, int64_t T, void* stream, uint16_t* mixed_h = nullptr); -/// F-2: gr_write, then gr_norm_rs of the next half (its norm weights) over the rows just written - the same bits as -/// the two calls, without reading R back. -void gr_write_norm_rs(float* R, const float* bo, const float* inj, int64_t inj_ld, const float* w_norm_next, float eps, - float* rs, uint16_t* xn16, int64_t T, void* stream); +/// gr_write, then gr_norm_rs of the written rows with the next read's norm weights, R read once. +void gr_write_norm_rs(float* R, const float* bo, const float* inj, int64_t inj_ld, const float* w_norm, float eps, + float* rs, uint16_t* xn16, int64_t T, void* stream, const uint16_t* partial = nullptr); /// lo16[t, k] = bf16(silu(lo[t, k] / hc)) void gr_silu(const float* lo, uint16_t* lo16, int64_t T, void* stream); -/// mixed[t, d] = mean_c xn[t, c, d] * sigmoid(gated[t, c, d]); FP32, BF16 and FP16 (either image may be null). +/// mixed[t, d] = mean_c xn[t, c, d] * sigmoid(gated[t, c, d]); FP32, BF16 and FP16 (any of them may be null). void gr_mix(const float* xn, const float* gated, float* mixed, uint16_t* mixed16, int64_t T, void* stream, uint16_t* mixed_h = nullptr); -/// R[t, c, d] += bo[t, d] * 2 sigmoid(inj[t, c] / hc) (inj has row stride inj_ld) -void gr_write(float* R, const float* bo, const float* inj, int64_t inj_ld, int64_t T, void* stream); +/// R[t, c, d] += (bo[t, d] (+ partial[t, d], FP16)) * 2 sigmoid(inj[t, c] / hc) (inj has row stride inj_ld) +void gr_write(float* R, const float* bo, const float* inj, int64_t inj_ld, int64_t T, void* stream, + const uint16_t* partial = nullptr); /// R[t, c, :] = e[t, :] for all four streams (the embedding broadcast). void gr_broadcast(const float* e, float* R, int64_t T, void* stream); @@ -37,30 +35,38 @@ void gr_broadcast(const float* e, float* R, int64_t T, void* stream); /// gate[t,h] = softplus(ab[t,h] + dt[h]) * ssm_a[h]; beta[t,h] = sigmoid(ab[t, 48 + h]) (ab: [T, 96]) void gdn_gates(const float* ab, const float* dt, const float* ssm_a, float* gate, float* beta, int64_t T, void* stream); /// The 4-tap causal conv + SiLU over the chunk (history [C][3] in, updated to the chunk's last three inputs), then -/// the L2 norm of the q and k heads of every token. h: [T, C]. +/// the L2 norm of the q and k heads of every token: the verify window's gdn_conv_l2_multi over all T tokens at once. +/// h: [T, C]. void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream); -/// The recurrence over the chunk, block per value head, state in registers; y[t] = rmsnorm(o) * gamma * sigmoid(z) -/// (FP32 and FP16 bits: the out projection is quantized). -void gdn_recurrence(float* state, const float* h, const float* gate, const float* beta, const float* z, - const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, void* stream); +/// The recurrence over the chunk (state in registers): y[t] = o (before the norm). With gdn_out_norm, bitwise +/// gdn_step_norm_multi over the tokens. +void gdn_scan(float* state, const float* h, const float* gate, const float* beta, float* y, int64_t T, void* stream); +/// In place: y[t] = rmsnorm(o) * gamma * sigmoid(z) per head, and its FP16 bits (the out projection is quantized). +void gdn_out_norm(float* y, const float* z, const float* gamma, float eps, uint16_t* y16, int64_t T, void* stream); // ---- MoE /// softmax over 512, top-10 (ties to the lower id), weights renormalised over the ten (the native router). -void route(const float* logits, int32_t* ids, float* weights, int64_t T, int64_t n_expert, void* stream); +void route(const float* logits, int32_t* ids, float* weights, int64_t T, void* stream); /// Expert blob (Strata pack layout, Q2_0) -> BF16 matrices: gate/up interleaved [1280, 2560], down [2560, 640]. void blob_dequant(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, void* stream); /// h16[n, r] = fp16(silu(gu[n, 2r]) * gu[n, 2r + 1]) (the interleaved expert gate/up) void swiglu_interleaved(const float* gu, uint16_t* h16, int64_t n, void* stream); /// h16[n, r] = fp16(silu(g[n, r]) * u[n, r]) (the shared expert, gate and up separate, width 640) void swiglu_pair(const float* g, const float* u, uint16_t* h16, int64_t n, void* stream); -/// dst[i] = src[i] for n int32s, as a kernel: either side may be mapped host memory, and the copy never waits -/// behind the copy engine's queue (the prompt path's grouping tables, while the expert stream fills it). -void copy_i32(int32_t* dst, const int32_t* src, int64_t n, void* stream); /// Gather rows: dst16[i, :] = x16[src[i], :] (n rows of `width` BF16). void gather_rows16(const uint16_t* x16, const int32_t* src, uint16_t* dst16, int64_t n, int64_t width, void* stream); -/// bo[t, :] = shared[t, :] * sigmoid(sg[t]) + sum_k w[t, k] * D[slot[t, k], :] -void moe_combine(const float* D, const int32_t* slot, const float* w, const float* shared, const float* sg, float* bo, - int64_t T, void* stream); +/// The same from FP16 rows into FP32 ones (MMQ's quantizer reads floats). +void gather_rows16_f32(const uint16_t* x16, const int32_t* src, float* dst, int64_t n, int64_t width, void* stream); +/// y = FP16 bits of x, saturated at +-65504 (the second GPU's expert sums on their way back: below ~20 measured) +void sums_to_f16(const float* x, uint16_t* y, int64_t n, void* stream); +/// In place: x[t, :] *= sigmoid(sg[t]) (the shared expert's scalar gate) +void moe_shared_gate(float* x, const float* sg, int64_t T, void* stream); +/// sum[src[r], :] += w[r] * rows[r, :] for n rows (one expert's rows: no token twice). +void moe_scatter_add(float* sum, const float* rows, const float* w, const int32_t* src, int64_t n, void* stream); +/// The same for several experts' rows (row r at rows[r - r0]), where a token may recur: token tok[b] adds its rows +/// list[start[b] .. start[b + 1]) in that order, so it matches moe_scatter_add per expert in the list's order. +void moe_gather_add(float* sum, const float* rows, int64_t r0, const float* w, const int32_t* tok, const int32_t* start, + const int32_t* list, int64_t n_tok, void* stream); // ---- QSA helpers /// In place: x[r, :] = x[r, :] * rsqrt(mean x^2 + eps) * w over rows of `cols` (row stride `ld`). @@ -77,14 +83,11 @@ void gate_attn(const float* attn, const float* q_full, uint16_t* out16, int64_t /// K, V: [T, 2, 256]. void kv_append(const float* K, const float* V, int64_t T, int64_t pos0, const int32_t* page_table, int64_t page_size, uint16_t* k_pool, uint16_t* v_pool, int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, - void* stream, const strata::kernels::KvHostPools* host = nullptr, - const strata::kernels::KvHostPools* stage = nullptr); + void* stream); /// fp32 -> fp16 bits and fp32 -> bf16, n elements (the two activation images of the prompt GEMMs). void to_f16(const float* x, uint16_t* y, int64_t n, void* stream); void to_bf16(const float* x, uint16_t* y, int64_t n, void* stream); -/// y = fp32(fp16(x)): what an FP16 store of x would read back (the FP16 indexer-key experiment) -void round_f16(const float* x, float* y, int64_t n, void* stream); /// Expert blob -> FP16 (Q2_0 values are exact in FP16). void blob_dequant_f16(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, void* stream); diff --git a/include/strata/prefill/moe_mmq.hpp b/include/strata/prefill/moe_mmq.hpp index 4efa1baf857..48e882fda15 100644 --- a/include/strata/prefill/moe_mmq.hpp +++ b/include/strata/prefill/moe_mmq.hpp @@ -12,7 +12,8 @@ namespace strata::prefill::mmq { /// This build has the MMQ path (the ggml sources were available to the build). bool built(); -/// MMQ covers this ggml type (the i-quants and Q2_0 the packs use; IQ1_M is not covered). +/// MMQ covers this ggml type (the i-quants and Q2_0 of ISTA-DASLab's files, the K-quants, Q5_1 and Q8_0 of +/// unsloth's; IQ1_M is not covered). bool supported(int ggml_type); /// Bytes of one expert's gate+up ([2*n_ff, n_embd]) or down ([n_embd, n_ff]) weights in `ggml_type`. size_t matrix_bytes(int ggml_type, int64_t rows, int64_t cols); diff --git a/include/strata/prefill/prefill.hpp b/include/strata/prefill/prefill.hpp index b792ff043b5..4879f44a099 100644 --- a/include/strata/prefill/prefill.hpp +++ b/include/strata/prefill/prefill.hpp @@ -1,11 +1,12 @@ // include/strata/prefill/prefill.hpp - plan v0.3 P5: batched prompt processing. // -// The prompt's positions [pos0, pos0 + n) are processed in chunks of `chunk` tokens through all 48 layers, leaving -// the session state (GDN recurrence and conv state, QSA KV pools and indexer, PLE history) where the token path -// would have left it; the decode loop then continues with the next token. Per layer: the projections are -// tensor-core GEMMs (quantized weights dequantized to FP16 on the fly, BF16 weights as they are), the recurrences -// walk the chunk inside one kernel, and the routed experts are grouped by expert: resident ones are read from the -// VRAM tier, the others streamed from the host arena through a pinned ring on a copy stream. +// The prompt's positions [pos0, pos0 + n) are processed in chunks through all 48 layers, leaving the session state +// (GDN recurrence and conv state, QSA KV pools and indexer, PLE history) where the token path would have left it; the +// decode loop then continues with the next token. Per layer: the projections are tensor-core GEMMs (quantized +// weights dequantized to FP16 on the fly, BF16 weights as they are) and the recurrences walk the tokens inside one +// kernel, in sub-chunks of up to 2048 tokens whose buffers are reused; the routed experts then run over the whole +// chunk (strata/prefill/experts.hpp), so an expert the cache does not hold is streamed once per chunk. Long chunks +// cost little VRAM besides the residual stream (~55 KB a token). // // Requires the native weights (`--native`): every quantized projection must carry its GGUF blocks. #pragma once @@ -13,8 +14,10 @@ #include "strata/core/expert_cache.hpp" #include "strata/core/expert_source.hpp" #include "strata/core/layer.hpp" +#include "strata/core/mtp.hpp" #include "strata/core/session.hpp" #include "strata/core/weights.hpp" +#include "strata/prefill/experts.hpp" #include #include @@ -23,96 +26,79 @@ namespace strata::prefill { +/// The sections `Prefill::profile` times, in the order a layer runs them. +enum PrefillSection { + kPsPle, kPsHcRead, kPsGdn, kPsGdnConv, kPsGdnScan, kPsGdnNorm, kPsGdnOut, kPsQsaProj, kPsQsaIndexer, kPsQsaScores, kPsQsaAttn, kPsQsaOut, kPsHcFfn, kPsRouter, + kPsExperts, kPsCombine, kPsDraft, kPsCount +}; +const char* prefill_section_name(int section); + struct PrefillStats { int64_t tokens = 0; int64_t chunks = 0; double ms_total = 0; - double ms_experts_host = 0; ///< host time staging non-resident experts - int64_t experts_streamed = 0; ///< expert blobs copied host -> device - int64_t experts_dma = 0; ///< ...of which straight from the pinned arena (no CPU copy) - int64_t experts_resident = 0; ///< expert-layer groups served from the VRAM tier + double ms_experts_host = 0; ///< host time queuing this GPU's experts + int64_t experts_streamed = 0; ///< expert blobs copied host -> this GPU + int64_t experts_prefetched = 0; ///< ...of which ahead of the routing + int64_t experts_resident = 0; ///< expert-layer groups served from its VRAM tier double ms_ple = 0; + double ms_section[kPsCount] = {}; ///< GPU ms per section, with `Prefill::profile` }; -} // namespace strata::prefill -namespace strata::core { class MtpDrafter; } -namespace strata::prefill { - class Prefill { public: - /// E-9: the draft layer's K/V for prompt cells [cell0, cell0 + n) from their final residual rows `R_rows` - /// (device) and `next_tokens` (host: the token at cell+1), in batches through this path's GEMMs and its idle - /// scratch - call it from on_chunk. false with `err` empty: not applicable here (a ring or hybrid K/V, another - /// device, too little scratch; STRATA_MTP_BATCH=0), the caller runs the drafter's own pass. Not bit-identical to - /// that pass (FP16 GEMMs instead of Q8_1 activations): the drafts may differ, never the target's tokens' logits. - bool draft_kv(core::MtpDrafter& mtp, const float* R_rows, const int32_t* next_tokens, int64_t n, int64_t cell0, - std::string& err); Prefill(); ~Prefill(); Prefill(const Prefill&) = delete; Prefill& operator=(const Prefill&) = delete; - /// `host_res`: the static residency table (n_layers x n_expert, slot or -1) or null; `cache` its slots. - /// `borrow`/`borrow_bytes`: device memory to carve every buffer from (the top slots of the expert cache, - /// lent for the prompt and refilled after it); null = allocate normally. + /// For chunks of up to `max_chunk` tokens. `host_res`: the static residency table (n_layers x n_expert, slot + /// or -1) or null; `cache` its slots. `offload`: a runner on a second GPU (initialized, bound by the caller) + /// that computes the experts this GPU's cache does not hold; null: they stream here. bool init(const core::WeightTable& wt, const core::ModelGeometry& g, core::SessionState& ss, - core::ExpertSource* src, const core::ExpertCache* cache, const int32_t* host_res, int64_t chunk, - void* stream, std::string& err, void* borrow = nullptr, uint64_t borrow_bytes = 0); - - /// With borrowed buffers: lay them out again for chunks of `chunk` tokens (at most `init`'s) in `borrow` - a - /// request lends only the slots its prompt needs. The stream must be idle (between prompts). - bool relayout(int64_t chunk, void* borrow, uint64_t borrow_bytes, std::string& err); - int64_t chunk() const; - - /// The share of the streamed experts' bytes DMA-able straight from pinned RAM (1 = all). Sizes the streamed - /// ring (a big one only pays when the copy engine, not the host copies, is the limit); set before bytes_needed. - static void set_pinned_share(double share); - static double pinned_share(); - - /// Device bytes `init` needs for a chunk of `chunk` tokens (what a borrowed region must hold). - static uint64_t bytes_needed(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk); - - /// Positions [pos0, pos0 + n) holding `tokens`; `ss.ple_prev` must be the two tokens before pos0 (oldest - /// first, -1 for none) and is advanced to the last two of these. + core::ExpertSource* src, const core::ExpertCache* cache, const int32_t* host_res, int64_t max_chunk, + void* stream, ExpertRunner* offload, std::string& err); + + /// Device bytes of the buffers for chunks of `chunk` tokens. + static uint64_t bytes_needed(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk, + bool offload); + /// The same with the area its experts are prefetched into (without a second GPU), for the residency as it is now. + uint64_t bytes_for(int64_t chunk) const; + /// The buffers for chunks of up to `chunk` tokens, carved from `region` (`bytes` long: lent expert-cache + /// slots; what is left over takes prefetched experts), or with a null region allocated for the longest chunk + /// (once). Before the first chunk of every prompt that uses a region. + bool bind(void* region, uint64_t bytes, int64_t chunk, std::string& err); + + /// Positions [pos0, pos0 + n) holding `tokens`, in chunks of the bound length; `ss.ple_prev` must be the two + /// tokens before pos0 (oldest first, -1 for none) and is advanced to the last two of these. With a `draft` + /// layer tokens[n], the token after the last position, is read too. bool run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err); + /// Positions the last `run` completed: n, or fewer after `should_stop` (the session state holds them). + int64_t processed() const { return processed_; } const PrefillStats& stats() const { return stats_; } - /// Plan v0.3 P6: called after every chunk with the chunk's final multi-stream residual rows (device, - /// T x hc*n_embd, valid until the next chunk) and the chunk's first position; the MTP draft layer builds its - /// K/V from them. The prefill stream is synchronized before the call. - std::function on_chunk; - - /// Layer split: called by every stage when it has read a chunk, with the position reached, while its own state - /// is still at that chunk's end (its stream synchronized; the last stage calls it just before `on_chunk`). An - /// earlier stage is a chunk or more ahead of the last one by the time `on_chunk` runs, so this is where a - /// mid-prompt checkpoint takes each stage's part. Runs on that stage's thread, with its device current. - std::function on_stage_chunk; + /// The MTP draft layer (or null): after every chunk, its K/V over the chunk's cells in the chunk's batched + /// arithmetic (cell t pairs the chunk's final residual at t with the token at t + 1). + core::MtpDrafter* draft = nullptr; - /// Checked before every chunk: true stops the prompt early (`run` returns false with err "cancelled"). + /// Checked before every chunk: true stops the prompt there (`run` returns false with err "cancelled", and the + /// state and `ss.ple_prev` hold the chunks before it, `processed()` positions). std::function should_stop; /// The vision path: HOST rows (n_embd floats) indexed by absolute position, read in place of the token /// embedding where non-null (an image's <|image_pad|> cells). Null (default): every position embeds its token. const float* const* embd_rows = nullptr; - /// LAYER SPLIT (multi-GPU): this prompt path runs layers [layer_begin, layer_end) (-1: to the last) on the - /// device `init` runs on. A stage that does not start at layer 0 reads each chunk's residual rows from the - /// previous stage instead of embedding the tokens; a stage that does not end at the last layer copies its rows - /// to pinned host buffers (two, allocated by `init`) and runs `next` on them - on a thread, so the next stage - /// reads chunk c while this one reads chunk c + 1. `on_chunk` belongs on the last stage. Set before `init`. - void set_stage(int64_t layer_begin, int64_t layer_end, Prefill* next) { - stage_lb_ = layer_begin; stage_le_ = layer_end; next_ = next; - } + /// GPU time per section into `stats().ms_section` (events between the sections; the host waits for the last + /// one after every chunk). + bool profile = false; private: - int64_t stage_lb_ = 0, stage_le_ = -1; - Prefill* next_ = nullptr; - const float* hand_in_ = nullptr; ///< the previous stage's rows of the chunk being read (host, pinned) - bool carve(std::size_t T, void* alloc); // the device buffers of a chunk (prefill.cpp's Alloc) struct Impl; std::unique_ptr impl_; PrefillStats stats_; + int64_t processed_ = 0; }; } // namespace strata::prefill diff --git a/include/strata/spec/suffix_drafter.hpp b/include/strata/spec/suffix_drafter.hpp index 731bbef1b3b..c033832c2f9 100644 --- a/include/strata/spec/suffix_drafter.hpp +++ b/include/strata/spec/suffix_drafter.hpp @@ -4,10 +4,15 @@ // output quotes its input (edits, refactors, repeated code) the continuation is usually exact, and a verify pass // accepts many tokens at once. It needs no model and costs microseconds. // -// Index: every trigram of the history (prompt + accepted output) maps to its WAYS most recent end positions in a -// fixed-size open-addressing table, so memory is bounded (~20 bytes per history token) and appends are O(1). -// A proposal checks those candidates, extends each match backwards up to `max_match`, and takes the longest -// (most recent on ties). Matches shorter than `min_match` propose nothing. +// Index: an open-addressing table maps every trigram of the history (prompt + accepted output) to its most recent +// end position, and each position links to its trigram's previous one, so appends are O(1) (the table is rebuilt, +// larger if the history needs it, when half full). A proposal checks the trigram's WAYS most recent earlier +// occurrences, extends each match backwards up to `max_match`, and takes the longest (most recent on ties). +// Matches shorter than `min_match` propose nothing. +// +// `assign` replaces the history by another text that may share a prefix with it (a conversation's next request): +// the positions past the common prefix are undone newest first, which leaves the index as appending the prefix +// alone would, and the rest is appended. Proposals therefore depend on the history alone. #pragma once #include @@ -18,7 +23,7 @@ namespace strata::spec { class SuffixDrafter { public: - static constexpr int WAYS = 4; + static constexpr int WAYS = 16; // earlier occurrences a proposal checks (in code a trigram recurs often) explicit SuffixDrafter(int min_match = 3, int max_match = 32, size_t capacity_tokens = 1u << 19); @@ -26,6 +31,8 @@ class SuffixDrafter { /// Add tokens to the history (the prompt, then every accepted token). void append(const int32_t* tokens, size_t n); void append(int32_t token) { append(&token, 1); } + /// Make the history `tokens` (see above). + void assign(const int32_t* tokens, size_t n); /// Write up to `max_k` proposed next tokens to `out`; returns how many (0 = no match of at least min_match). int propose(int max_k, int32_t* out); @@ -36,17 +43,20 @@ class SuffixDrafter { private: struct Slot { - uint64_t key = 0; // trigram hash + 1 (0 = empty) - uint32_t pos[WAYS] = {}; // end positions, most recent first - uint8_t n = 0; + uint64_t key = 0; // trigram hash (0 = empty) + int32_t last = -1; // the trigram's most recent end position (-1: none left after `assign`) }; Slot* find_slot(uint64_t key, bool insert); uint64_t key_at(size_t end) const; + void link(size_t end); // index hist_[end], the newest token + void rebuild(); int min_match_, max_match_; std::vector hist_; + std::vector prev_; // per end position: its trigram's previous end position (-1: none) std::vector table_; size_t mask_ = 0; + size_t used_ = 0; // occupied slots, those of undone positions included int last_match_ = 0; }; diff --git a/serve/server.py b/serve/server.py index cb2617650a3..421008f9e35 100644 --- a/serve/server.py +++ b/serve/server.py @@ -1768,7 +1768,9 @@ def main() -> int: import webbrowser webbrowser.open(f"http://{'127.0.0.1' if a.host in ('0.0.0.0', '') else a.host}:{a.port}/") try: - threading.Event().wait() + # a sleep, not threading.Event().wait(): on Windows Ctrl+C interrupts only the former + while True: + time.sleep(1) except KeyboardInterrupt: httpd.shutdown() if hasattr(engine, "close"): diff --git a/setup.py b/setup.py index 0c30e244181..7185b535f62 100644 --- a/setup.py +++ b/setup.py @@ -1968,7 +1968,7 @@ def main() -> int: fail("the model has no per_layer_token_embd tensor (is this a Qwen3.8-Flash-Next GGUF?)") args = ["--pack", str(pack), "--native", str(shards[0]), "--ple-gguf", str(ple), "--expert-profile", str(ROOT / "data" / fam.get("profile", "expert-profile.bin")), "--expert-cache", "auto", - "--prefill", "auto", "--spec", "4", "--spec-min-p", "0.5", "--mtp", str(rt), + "--prefill", "auto", "--spec", "4", "--spec-min-p", "0.5", "--spec-lookup", "16", "--mtp", str(rt), "--max-context", str(ctx)] if ctx > 8192: args += ["--kv", kv] diff --git a/src/artifact/dequant.cpp b/src/artifact/dequant.cpp index ff058c0ac71..8419ad6a54e 100644 --- a/src/artifact/dequant.cpp +++ b/src/artifact/dequant.cpp @@ -38,7 +38,6 @@ int main(int argc, char** argv) { // implementation on identical bytes. A single implementation being self-consistent // proves little - two disagreeing is what finds bugs (round 67). double sum = 0, sumabs = 0; - const uint8_t* base = g.tensor_data(t); for (uint64_t b = 0; b < nblocks && b < 64; ++b) { float out[64]; strata::dequantize_q2_0(base + b * 18, out); diff --git a/src/core/adaptive_tier.cpp b/src/core/adaptive_tier.cpp new file mode 100644 index 00000000000..a298afecf3f --- /dev/null +++ b/src/core/adaptive_tier.cpp @@ -0,0 +1,272 @@ +// src/core/adaptive_tier.cpp - see include/strata/core/adaptive_tier.hpp. +#include "strata/core/adaptive_tier.hpp" + +#include "strata/kernels/cpu/expert_layout.hpp" + +#include +#include + +namespace strata::core { +namespace { +struct OnDevice { // the tier's GPU current for one call (when it has one), the main one again afterwards + int main; + bool set; + OnDevice(int dev, int main_dev) : main(main_dev), set(dev >= 0) { if (set) cudaSetDevice(dev); } + ~OnDevice() { if (set) cudaSetDevice(main); } +}; +} // namespace + +AdaptiveTier::~AdaptiveTier() { + OnDevice on(dev_, main_); + if (stream_) cudaStreamSynchronize(stream_); + for (cudaEvent_t e : evs_) if (e) cudaEventDestroy(e); + cudaEvent_t evs[] = {t0_, t1_}; + for (cudaEvent_t e : evs) if (e) cudaEventDestroy(e); + if (stream_) cudaStreamDestroy(stream_); +} + +bool AdaptiveTier::init(ExpertCache& cache, ExpertSource& src, std::vector& host_res, int64_t n_layers, + int64_t n_expert, int max_moves, std::string& err, int device, int main_device) { + cache_ = &cache; + src_ = &src; + res_ = &host_res; + n_layers_ = n_layers; + n_expert_ = n_expert; + max_moves_ = max_moves; + dev_ = device; + main_ = main_device; + free_.assign((size_t) n_layers, {}); + OnDevice on(dev_, main_); + bool ok = cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking) == cudaSuccess && + cudaEventCreate(&t0_) == cudaSuccess && cudaEventCreate(&t1_) == cudaSuccess; + for (cudaEvent_t& e : evs_) ok = ok && cudaEventCreateWithFlags(&e, cudaEventDisableTiming) == cudaSuccess; + if (!ok) { + res_ = nullptr; // off + err = "adaptive tier: cannot create the refill stream"; + } + return ok; +} + +int64_t AdaptiveTier::free_slots() const { + int64_t n = 0; + for (const auto& f : free_) n += (int64_t) f.size(); + return n; +} + +uint64_t AdaptiveTier::bytes_of(const Move& m) const { + return strata::kernels::cpu::expert_layout().blob_bytes(m.layer); +} + +void AdaptiveTier::evict(const Move& m) { + if (m.out >= 0) (*res_)[(size_t) (m.layer * n_expert_ + m.out)] = kNotResident; // a miss from now on +} + +bool AdaptiveTier::copy(const Move& m, uint64_t off, uint64_t n, std::string& err) { + const uint8_t* b = src_->blob(m.layer, m.in); + if (b == nullptr) { + err = "adaptive tier: a refill copy failed"; + return false; + } + cp_dst_.push_back(cache_->device_slot(m.slot) + off); + cp_src_.push_back(b + off); + cp_bytes_.push_back((size_t) n); + sent_bytes += n; + return true; +} + +bool AdaptiveTier::send(std::string& err) { + const bool ok = copy_blobs(cp_dst_.data(), cp_src_.data(), cp_bytes_.data(), cp_dst_.size(), stream_); + cp_dst_.clear(); + cp_src_.clear(); + cp_bytes_.clear(); + if (!ok) err = "adaptive tier: a refill copy failed"; + return ok; +} + +// Takes the last timed copies' rate once they have landed; then starts timing these when `worth` (long enough for +// their start and end not to dominate) and none are being timed. +bool AdaptiveTier::time_begin(bool worth) { + float took = 0; + if (timed_ > 0 && cudaEventQuery(t1_) == cudaSuccess && cudaEventElapsedTime(&took, t0_, t1_) == cudaSuccess && + took > 0) { + const double r = (double) timed_ / (double) took; + rate_ = rate_ > 0 ? 0.8 * rate_ + 0.2 * r : r; + timed_ = 0; + } + return worth && timed_ == 0 && cudaEventRecord(t0_, stream_) == cudaSuccess; +} + +// After a batch of copies: its timing, and its event in the ring, with the moves sent so far (a full ring's newest +// event takes this batch too). +bool AdaptiveTier::end_batch(bool timed, uint64_t bytes, std::string& err) { + if (timed) { + if (cudaEventRecord(t1_, stream_) != cudaSuccess) { err = "adaptive tier: a refill copy failed"; return false; } + timed_ = bytes; + } + if (flying_ < kBatches) ++flying_; + const int i = (first_ + flying_ - 1) % kBatches; + batch_sent_[i] = sent_; + if (cudaEventRecord(evs_[i], stream_) != cudaSuccess) { err = "adaptive tier: a refill copy failed"; return false; } + return true; +} + +void AdaptiveTier::stage(uint64_t bytes) { + uint64_t have = 0; + for (size_t i = sent_; i < staged_; ++i) have += bytes_of(queued_[i]); + have -= off_; + for (; staged_ < queued_.size() && have < bytes; ++staged_) { + evict(queued_[staged_]); + have += bytes_of(queued_[staged_]); + } +} + +bool AdaptiveTier::pump_bytes(uint64_t bytes, std::string& err) { + if (sent_ >= staged_ || bytes < (64u << 10)) return true; + OnDevice on(dev_, main_); + const bool timed = time_begin(bytes >= (256u << 10)); + uint64_t done = 0; + while (sent_ < staged_ && done < bytes) { + const Move& m = queued_[sent_]; + const uint64_t blob = bytes_of(m); + uint64_t n = std::min(bytes - done, blob - off_); + if (off_ + n < blob) n &= ~(uint64_t) 4095; // whole pages, but the blob's end + if (n == 0) break; + if (!copy(m, off_, n, err)) return false; + done += n; + off_ += n; + if (off_ == blob) { + off_ = 0; + ++sent_; + } + } + return done == 0 || (send(err) && end_batch(timed, done, err)); +} + +bool AdaptiveTier::pump(uint64_t budget, uint64_t& sent, std::string& err) { + sent = 0; + size_t n = 0; + for (; sent_ + n < queued_.size(); ++n) { + const uint64_t b = bytes_of(queued_[sent_ + n]); + if (sent + b > budget) break; + sent += b; + } + if (n == 0) return true; + OnDevice on(dev_, main_); + // timed from when the work that may read its slots is done + if (after_ != nullptr && cudaStreamWaitEvent(stream_, after_, 0) != cudaSuccess) { + err = "adaptive tier: a refill copy failed"; + return false; + } + const bool timed = time_begin(n >= 4); // a batch of a few copies: its start and end dominate + for (size_t i = 0; i < n; ++i, ++sent_) { + evict(queued_[sent_]); + if (!copy(queued_[sent_], 0, bytes_of(queued_[sent_]), err)) return false; + } + staged_ = sent_; + return send(err) && end_batch(timed, sent, err); +} + +bool AdaptiveTier::adapt(std::vector& usage, std::string& err, bool decay) { + const auto t0 = std::chrono::steady_clock::now(); + if (!failed_.empty()) { err = failed_; return false; } + apply_pending(false); + if (wait_ && admitted_ < queued_.size()) return true; + // the moves whose residents are still in place give way to this call's (ranked again if still worth it) + for (size_t i = staged_; i < queued_.size(); ++i) { + const Move& m = queued_[i]; + if (m.out >= 0) { + --swaps; + } else { + free_[(size_t) m.layer].push_back(m.slot); + --fills; + } + } + dropped += (int64_t) (queued_.size() - staged_); + queued_.resize(staged_); + pending_.resize(staged_); + queued_.erase(queued_.begin(), queued_.begin() + (ptrdiff_t) admitted_); + pending_.erase(pending_.begin(), pending_.begin() + (ptrdiff_t) admitted_); + for (int b = 0; b < flying_; ++b) batch_sent_[(first_ + b) % kBatches] -= admitted_; + sent_ -= admitted_; + staged_ -= admitted_; + admitted_ = 0; + blocked_.assign((size_t) (n_layers_ * n_expert_), 0); // the experts still under way + for (const auto& pr : pending_) blocked_[(size_t) pr.first] = 1; + if (upper_ != nullptr) { + upper_has_.assign((size_t) (n_layers_ * n_expert_), 0); + for (size_t i = 0; i < upper_has_.size(); ++i) upper_has_[i] = (*upper_->res_)[i] >= 0; + for (size_t i = 0; i < upper_->staged_; ++i) upper_has_[(size_t) upper_->pending_[i].first] = 1; + } + struct Ranked { float gain; int32_t layer, in, out, slot; }; // out < 0: an empty slot + std::vector moves; + std::vector> cand, vict; + std::vector& res = *res_; + for (int64_t l = 0; l < n_layers_; ++l) { + cand.clear(); + vict.clear(); + const float* u = usage.data() + l * n_expert_; + const int32_t* r = res.data() + l * n_expert_; + const uint8_t* bl = blocked_.data() + l * n_expert_; + const uint8_t* up = upper_ != nullptr ? upper_has_.data() + l * n_expert_ : nullptr; + const int32_t* lo = lower_ != nullptr ? lower_->res_->data() + l * n_expert_ : nullptr; + for (int32_t e = 0; e < (int32_t) n_expert_; ++e) { + const float ue = lo != nullptr && lo[e] >= 0 ? 0.5f * u[e] : u[e]; + if (r[e] < 0) { if (ue >= 2.0f && !bl[e] && (up == nullptr || !up[e])) cand.emplace_back(ue, e); } + else vict.emplace_back(up != nullptr && up[e] ? -1.0f : u[e], e); // the upper tier's copy goes first + } + if (cand.empty()) continue; + std::sort(cand.begin(), cand.end(), [](auto& a, auto& b) { return a.first > b.first; }); + const std::vector& fr = free_[(size_t) l]; + size_t c = 0; + for (; c < cand.size() && c < fr.size(); ++c) + moves.push_back({cand[c].first, (int32_t) l, cand[c].second, -1, fr[fr.size() - 1 - c]}); + const size_t nv = std::min(cand.size() - c, vict.size()); + std::partial_sort(vict.begin(), vict.begin() + (ptrdiff_t) nv, vict.end(), + [](auto& a, auto& b) { return a.first < b.first; }); + for (size_t i = 0; i < nv; ++i, ++c) { + if (cand[c].first < vict[i].first + 1.5f) break; + moves.push_back({cand[c].first - vict[i].first, (int32_t) l, cand[c].second, vict[i].second, + r[vict[i].second]}); + } + } + std::sort(moves.begin(), moves.end(), [](const Ranked& a, const Ranked& b) { return a.gain > b.gain; }); + if ((int) moves.size() > max_moves_) moves.resize((size_t) max_moves_); + // a move counts once queued; its resident leaves when it is staged (or sent, paced), its expert comes once landed + for (const Ranked& m : moves) { + if (m.out >= 0) { + ++swaps; + } else { + std::vector& fr = free_[(size_t) m.layer]; + fr.erase(std::find(fr.begin(), fr.end(), m.slot)); + ++fills; + } + queued_.push_back({m.layer, m.in, m.out, m.slot}); + pending_.emplace_back((int32_t) (m.layer * n_expert_ + m.in), m.slot); + } + if (decay) + for (float& v : usage) v *= 0.7f; + ms += std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + return true; +} + +void AdaptiveTier::apply_pending(bool wait) { + if (wait && admitted_ < queued_.size()) { // send the rest and wait for it + std::string err; + uint64_t sent = 0; + if (!paced_) stage(UINT64_MAX); + if (!(paced_ ? pump(UINT64_MAX, sent, err) : pump_bytes(UINT64_MAX, err))) { + failed_ = err; + return; + } + OnDevice on(dev_, main_); + cudaStreamSynchronize(stream_); + } + size_t landed = admitted_; // the batches landed, oldest first + for (; flying_ > 0 && cudaEventQuery(evs_[first_]) == cudaSuccess; first_ = (first_ + 1) % kBatches, --flying_) + landed = std::max(landed, batch_sent_[first_]); + if (landed == admitted_) return; + std::vector& res = *res_; + for (; admitted_ < landed; ++admitted_) res[(size_t) pending_[admitted_].first] = pending_[admitted_].second; +} + +} // namespace strata::core diff --git a/src/core/expert_cache.cpp b/src/core/expert_cache.cpp index 6eb092949ad..f2bf53e2e92 100644 --- a/src/core/expert_cache.cpp +++ b/src/core/expert_cache.cpp @@ -278,6 +278,15 @@ bool ExpertCache::fill_slot(int32_t slot, const uint8_t* host_blob, void* stream return true; } +bool copy_blobs(void* const* dst, const void* const* src, const size_t* bytes, size_t n, void* stream) { + if (n == 0) return true; + cudaMemcpyAttributes attr{}; + attr.srcAccessOrder = cudaMemcpySrcAccessOrderStream; + attr.flags = cudaMemcpyFlagPreferOverlapWithCompute; // the copy engines, beside the kernels + size_t first = 0; + return cudaMemcpyBatchAsync(dst, src, bytes, n, &attr, &first, 1, (cudaStream_t) stream) == cudaSuccess; +} + bool ExpertCache::fill_slot_blocking(int32_t slot, const uint8_t* host_blob, std::string& err, int64_t bytes) { const size_t n = (size_t) (bytes > 0 && bytes <= blob_ ? bytes : blob_); uint8_t* dst = device_slot(slot); diff --git a/src/core/expert_source.cpp b/src/core/expert_source.cpp index 6e007928ebd..2a3ca69aa92 100644 --- a/src/core/expert_source.cpp +++ b/src/core/expert_source.cpp @@ -1,9 +1,11 @@ // src/core/expert_source.cpp - the adapter. See the header for the three clauses of the contract. #include "strata/core/expert_source.hpp" -#include "strata/core/remote_experts.hpp" +#include "strata/core/second_gpu.hpp" +#include "strata/artifact/gguf_reader.hpp" #include "strata/kernels/cpu/expert_layout.hpp" #include "strata/core/pinned.hpp" +#include "strata/platform/memory.hpp" #include "strata/kernels/elementwise.hpp" #include "strata/kernels/quantize_act.hpp" #include "strata/kernels/s2_expert_grouped.hpp" @@ -17,18 +19,11 @@ #include #include #include -#include -#include -#include #include -#include #include #if defined(_WIN32) #define WIN32_LEAN_AND_MEAN -#ifndef NOMINMAX -#define NOMINMAX // std::numeric_limits::max() below -#endif #include #else #include @@ -39,213 +34,6 @@ namespace strata::core { -namespace detail { - -bool cgroup_available_bytes(uint64_t limit, const CgroupMemoryStat& stat, uint64_t& bytes) { - bytes = 0; - if (!stat.valid) return false; - - // memory.stat's inactive_file can race memory.current, so bound it to charged usage first. - uint64_t reclaimable = std::min(stat.inactive_file, stat.current); - reclaimable = stat.file_dirty >= reclaimable ? 0 : reclaimable - stat.file_dirty; - reclaimable = stat.file_writeback >= reclaimable ? 0 : reclaimable - stat.file_writeback; - - // Reclaiming clean file pages reduces usage; saturating subtraction also handles a transient over-limit read. - const uint64_t usage_after_reclaim = stat.current - reclaimable; - bytes = usage_after_reclaim < limit ? limit - usage_after_reclaim : 0; - return true; -} - -bool make_cache_complement_plan( - int64_t n_layers, int64_t n_expert, const std::vector& layer_blob_bytes, - const std::vector>& primary_gpu_pairs, - const std::vector>& additional_gpu_pairs, - std::vector& offsets, uint64_t& bytes, std::string& err) { - offsets.clear(); - bytes = 0; - err.clear(); - if (n_layers <= 0 || n_expert <= 0 || layer_blob_bytes.size() != (size_t) n_layers) { - err = "FileExpertSource: invalid geometry for the cache complement plan"; - return false; - } - if ((uint64_t) n_layers > (uint64_t) std::numeric_limits::max() / (uint64_t) n_expert) { - err = "FileExpertSource: cache complement index table is too large"; - return false; - } - const size_t count = (size_t) n_layers * (size_t) n_expert; - std::vector omitted(count, 0); - auto mark_pairs = [&](const std::vector>& pairs, uint8_t bit, - const char* label) -> bool { - for (const auto& pair : pairs) { - if (pair.first < 0 || pair.second < 0 || pair.first >= n_layers || pair.second >= n_expert) { - err = std::string("FileExpertSource: ") + label + " pair is outside the expert geometry"; - return false; - } - const size_t index = (size_t) pair.first * (size_t) n_expert + (size_t) pair.second; - if ((omitted[index] & bit) != 0) { - err = std::string("FileExpertSource: duplicate ") + label + " pair in the cache complement plan"; - return false; - } - if (bit == 2 && (omitted[index] & 1) != 0) { - err = "FileExpertSource: the primary and additional GPU expert tiers overlap"; - return false; - } - omitted[index] |= bit; - } - return true; - }; - if (!mark_pairs(primary_gpu_pairs, 1, "primary GPU") || - !mark_pairs(additional_gpu_pairs, 2, "additional GPU")) return false; - for (uint64_t blob_bytes : layer_blob_bytes) { - if (blob_bytes == 0) { - err = "FileExpertSource: cache complement layer has zero-sized expert blobs"; - return false; - } - } - - offsets.assign(count, kNoCacheComplement); - for (int64_t layer = 0; layer < n_layers; ++layer) { - const uint64_t blob_bytes = layer_blob_bytes[(size_t) layer]; - for (int64_t expert = 0; expert < n_expert; ++expert) { - const size_t index = (size_t) layer * (size_t) n_expert + (size_t) expert; - if (omitted[index] != 0) continue; - if (bytes > std::numeric_limits::max() - blob_bytes) { - offsets.clear(); - bytes = 0; - err = "FileExpertSource: cache complement size overflows"; - return false; - } - offsets[index] = bytes; - bytes += blob_bytes; - } - } - if (bytes > (uint64_t) std::numeric_limits::max()) { - offsets.clear(); - bytes = 0; - err = "FileExpertSource: cache complement exceeds the host address space"; - return false; - } - return true; -} - -const uint8_t* cache_complement_blob_or_fallback( - size_t index, const std::vector& offsets, const uint8_t* complement_host, - const uint8_t* mapped_fallback) { - if (complement_host != nullptr && index < offsets.size() && offsets[index] != kNoCacheComplement) - return complement_host + (size_t) offsets[index]; - return mapped_fallback; -} - -} // namespace detail - -namespace { -constexpr uint64_t kPinnedMemoryHeadroom = 8ull << 30; - -#if defined(__linux__) -bool read_cgroup_memory_stat(const std::filesystem::path& path, uint64_t current, - detail::CgroupMemoryStat& stat) { - std::ifstream input(path / "memory.stat"); - if (!input) return false; - - bool inactive_file = false, file_dirty = false, file_writeback = false; - std::string line; - while (std::getline(input, line)) { - std::istringstream fields(line); - std::string key; - uint64_t value = 0; - if (!(fields >> key >> value)) return false; - fields >> std::ws; - if (!fields.eof()) return false; - - if (key == "inactive_file") { - if (inactive_file) return false; - inactive_file = true; - stat.inactive_file = value; - } else if (key == "file_dirty") { - if (file_dirty) return false; - file_dirty = true; - stat.file_dirty = value; - } else if (key == "file_writeback") { - if (file_writeback) return false; - file_writeback = true; - stat.file_writeback = value; - } - } - if (!input.eof() || !inactive_file || !file_dirty || !file_writeback) return false; - stat.current = current; - stat.valid = true; - return true; -} -#endif - -bool available_memory_bytes(uint64_t& bytes) { -#if defined(_WIN32) - MEMORYSTATUSEX status{}; - status.dwLength = sizeof(status); - if (!GlobalMemoryStatusEx(&status)) return false; - bytes = (uint64_t) status.ullAvailPhys; - return bytes > 0; -#elif defined(__linux__) - // MemAvailable includes reclaimable page cache, unlike _SC_AVPHYS_PAGES. - std::ifstream info("/proc/meminfo"); - std::string line; - bytes = 0; - while (std::getline(info, line)) { - std::istringstream fields(line); - std::string key, unit; - uint64_t value = 0; - if (fields >> key >> value >> unit && key == "MemAvailable:" && unit == "kB" && - value <= std::numeric_limits::max() / 1024) bytes = value * 1024; - } - if (bytes == 0) return false; - // Account for the tightest cgroup-v2 ancestor limit when its normal mount is visible. - // This is a point-in-time guard, not a reservation against concurrent allocations. - std::ifstream groups("/proc/self/cgroup"); - if (!groups) return false; - bool resolved_v2 = false; - while (std::getline(groups, line)) { - if (line.rfind("0::/", 0) != 0) continue; - const std::filesystem::path root("/sys/fs/cgroup"); - auto path = (root / line.substr(4)).lexically_normal(); - if (path.string().rfind(root.string(), 0) != 0 || !std::filesystem::is_directory(path)) return false; - resolved_v2 = true; - while (path.string().rfind(root.string(), 0) == 0) { - std::ifstream limit_file(path / "memory.max"), current_file(path / "memory.current"); - std::string limit; - uint64_t current = 0; - const bool readable = bool(limit_file >> limit) && bool(current_file >> current); - // The host's root cgroup has no memory.max; ordinary child groups must expose their limits. - if (!readable && !(path == root && !std::filesystem::exists(path / "memory.max") && - std::filesystem::exists(path / "cgroup.controllers"))) return false; - if (readable && limit != "max") { - try { - size_t consumed = 0; - const uint64_t cap = std::stoull(limit, &consumed); - if (consumed != limit.size()) return false; - detail::CgroupMemoryStat stat; - if (!read_cgroup_memory_stat(path, current, stat)) return false; - uint64_t cgroup_available = 0; - if (!detail::cgroup_available_bytes(cap, stat, cgroup_available)) return false; - bytes = std::min(bytes, cgroup_available); - } catch (...) { return false; } - } - if (path == root) break; - path = path.parent_path(); - } - } - return resolved_v2; -#else - const long pages = sysconf(_SC_AVPHYS_PAGES); - const long page_bytes = sysconf(_SC_PAGESIZE); - if (pages <= 0 || page_bytes <= 0 || - (uint64_t) pages > std::numeric_limits::max() / (uint64_t) page_bytes) return false; - bytes = (uint64_t) pages * (uint64_t) page_bytes; - return bytes > 0; -#endif -} - -} // namespace - // ================================ THE FILE-BACKED SOURCE ================================ FileExpertSource::~FileExpertSource() { close(); } @@ -253,74 +41,9 @@ FileExpertSource::~FileExpertSource() { close(); } bool FileExpertSource::open(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, std::string& err) { close(); if (n_layers <= 0 || n_expert <= 0) { err = "FileExpertSource: the geometry is empty"; return false; } - const auto& layout = strata::kernels::cpu::expert_layout(); - if (layout.n_layers != n_layers || layout.n_expert != n_expert) { - err = "FileExpertSource: the requested geometry does not match the loaded expert layout"; - return false; - } - if ((uint64_t) n_layers > (uint64_t) std::numeric_limits::max() / (uint64_t) n_expert) { - err = "FileExpertSource: the expert count overflows"; - return false; - } - const uint64_t blob_count = (uint64_t) n_layers * (uint64_t) n_expert; - if (blob_count > (uint64_t) std::numeric_limits::max() || - (uint64_t) n_layers > (uint64_t) std::numeric_limits::max()) { - err = "FileExpertSource: the expert count overflows"; - return false; - } - - std::vector layer_offsets((size_t) n_layers), layer_blob_bytes((size_t) n_layers); - const uint64_t want = layout.total; - if (want == 0 || want > (uint64_t) std::numeric_limits::max()) { - err = "FileExpertSource: the loaded expert layout has an invalid size"; - return false; - } - if (!layout.native) { - if (blob_count > std::numeric_limits::max() / (uint64_t) strata::kernels::cpu::BLOB) { - err = "FileExpertSource: the canonical expert size overflows"; - return false; - } - const uint64_t canonical_size = blob_count * (uint64_t) strata::kernels::cpu::BLOB; - if (want != canonical_size) { - err = "FileExpertSource: the canonical expert layout has an inconsistent size"; - return false; - } - const uint64_t bytes = (uint64_t) strata::kernels::cpu::BLOB; - const uint64_t layer_bytes = (uint64_t) n_expert * bytes; - for (int64_t layer = 0; layer < n_layers; ++layer) { - layer_offsets[(size_t) layer] = (uint64_t) layer * layer_bytes; - layer_blob_bytes[(size_t) layer] = bytes; - } - } else { - if (layout.offset.size() != (size_t) n_layers || layout.bytes.size() != (size_t) n_layers || - layout.fmt.size() != (size_t) n_layers) { - err = "FileExpertSource: the native expert layout is incomplete"; - return false; - } - uint64_t at = 0; - for (int64_t layer = 0; layer < n_layers; ++layer) { - const size_t i = (size_t) layer; - const uint64_t bytes = (uint64_t) layout.fmt[i].bytes; - if (layout.offset[i] != at || bytes == 0 || layout.bytes[i] != bytes || - bytes > std::numeric_limits::max() / (uint64_t) n_expert) { - err = "FileExpertSource: the native expert layout is invalid at layer " + std::to_string(layer); - return false; - } - const uint64_t layer_bytes = bytes * (uint64_t) n_expert; - if (at > want || layer_bytes > want - at) { - err = "FileExpertSource: the native expert layout exceeds its declared size at layer " + - std::to_string(layer); - return false; - } - layer_offsets[i] = layout.offset[i]; - layer_blob_bytes[i] = bytes; - at += layer_bytes; - } - if (at != want) { - err = "FileExpertSource: the native expert layout has an inconsistent size"; - return false; - } - } + n_expert_ = n_expert; + blobs_ = n_layers * n_expert; + const uint64_t want = (uint64_t) blobs_ * (uint64_t) strata::kernels::cpu::BLOB; const std::string path = pack_dir + "/experts.bin"; #if defined(_WIN32) @@ -355,9 +78,10 @@ bool FileExpertSource::open(const std::string& pack_dir, int64_t n_layers, int64 if ((uint64_t) sz.QuadPart != want) { char buf[400]; std::snprintf(buf, sizeof buf, - "FileExpertSource: %s is %llu B but the loaded expert layout requires %llu B - this is not " - "the pack this geometry came from", - path.c_str(), (unsigned long long) sz.QuadPart, (unsigned long long) want); + "FileExpertSource: %s is %llu B but %lld layers x %lld experts x %d B is %llu B - this " + "is not the pack this geometry came from", + path.c_str(), (unsigned long long) sz.QuadPart, (long long) n_layers, + (long long) n_expert, (int) strata::kernels::cpu::BLOB, (unsigned long long) want); CloseHandle(f); err = buf; return false; @@ -383,13 +107,13 @@ bool FileExpertSource::open(const std::string& pack_dir, int64_t n_layers, int64 if (fd < 0) { err = "FileExpertSource: cannot open " + path; return false; } struct stat st{}; if (fstat(fd, &st) != 0) { ::close(fd); err = "FileExpertSource: cannot stat " + path; return false; } - if (st.st_size < 0 || (uint64_t) st.st_size != want) { + if ((uint64_t) st.st_size != want) { char buf[400]; std::snprintf(buf, sizeof buf, - "FileExpertSource: %s is %llu B but the loaded expert layout requires %llu B - this is not " - "the pack this geometry came from", - path.c_str(), (unsigned long long) (st.st_size < 0 ? 0 : st.st_size), - (unsigned long long) want); + "FileExpertSource: %s is %llu B but %lld layers x %lld experts x %d B is %llu B - this " + "is not the pack this geometry came from", + path.c_str(), (unsigned long long) st.st_size, (long long) n_layers, (long long) n_expert, + (int) strata::kernels::cpu::BLOB, (unsigned long long) want); ::close(fd); err = buf; return false; @@ -399,27 +123,10 @@ bool FileExpertSource::open(const std::string& pack_dir, int64_t n_layers, int64 fd_ = fd; base_ = (const uint8_t*) view; #endif - blobs_ = (int64_t) blob_count; - n_layers_ = n_layers; - n_expert_ = n_expert; - mapped_bytes_ = want; - layer_offsets_ = std::move(layer_offsets); - layer_blob_bytes_ = std::move(layer_blob_bytes); return true; } void FileExpertSource::close() { - if (complement_arena_ != nullptr) { - if (complement_pinned_) (void) cudaFreeHost(complement_arena_); - else std::free(complement_arena_); - } - complement_arena_ = nullptr; - complement_host_ = nullptr; - complement_device_ = nullptr; - complement_bytes_ = 0; - complement_offsets_.clear(); - complement_pinned_ = false; - complement_ready_ = false; #if defined(_WIN32) if (base_ != nullptr) UnmapViewOfFile((LPCVOID) base_); if (mapping_ != nullptr) CloseHandle((HANDLE) mapping_); @@ -427,225 +134,27 @@ void FileExpertSource::close() { mapping_ = nullptr; file_ = nullptr; #else - if (base_ != nullptr) munmap((void*) base_, (size_t) mapped_bytes_); + if (base_ != nullptr) munmap((void*) base_, (size_t) blobs_ * (size_t) strata::kernels::cpu::BLOB); if (fd_ >= 0) ::close(fd_); fd_ = -1; #endif base_ = nullptr; blobs_ = 0; - n_layers_ = 0; - n_expert_ = 0; - mapped_bytes_ = 0; - layer_offsets_.clear(); - layer_blob_bytes_.clear(); reads_ = 0; } -const uint8_t* FileExpertSource::mapped_blob(int64_t layer, int64_t expert) const { - if (base_ == nullptr || layer < 0 || expert < 0 || layer >= n_layers_ || expert >= n_expert_) return nullptr; - const size_t i = (size_t) layer; - if (i >= layer_offsets_.size() || i >= layer_blob_bytes_.size()) return nullptr; - const uint64_t blob_bytes = layer_blob_bytes_[i]; - if (blob_bytes == 0 || (uint64_t) expert > std::numeric_limits::max() / blob_bytes) return nullptr; - const uint64_t expert_offset = (uint64_t) expert * blob_bytes; - const uint64_t layer_offset = layer_offsets_[i]; - if (layer_offset > mapped_bytes_ || expert_offset > mapped_bytes_ - layer_offset) return nullptr; - const uint64_t offset = layer_offset + expert_offset; - if (blob_bytes > mapped_bytes_ - offset) return nullptr; - return base_ + (size_t) offset; -} - -bool FileExpertSource::pin_cache_complement( - const ExpertCache& cache, std::string& err, bool pin, - const std::vector>& additional_gpu_pairs) { - err.clear(); - if (base_ == nullptr) { err = "FileExpertSource: open the mapped experts before pinning a complement"; return false; } - if (complement_ready_) { err = "FileExpertSource: the cache complement is already pinned"; return false; } - if (!cache.valid()) { err = "FileExpertSource: the GPU expert cache is not open"; return false; } - if (cache.fills() != cache.resident()) { - err = "FileExpertSource: the GPU expert cache is not fully filled"; - return false; - } - const cudaError_t sync = cudaDeviceSynchronize(); - if (sync != cudaSuccess) { - err = std::string("FileExpertSource: GPU expert cache is not ready: ") + cudaGetErrorString(sync); - (void) cudaGetLastError(); - return false; - } - - std::vector> primary_gpu_pairs; - primary_gpu_pairs.reserve((size_t) cache.resident()); - for (int64_t layer = 0; layer < n_layers_; ++layer) { - for (int64_t expert = 0; expert < n_expert_; ++expert) { - if (cache.slot_of(layer, expert) != kNotResident) - primary_gpu_pairs.emplace_back((int32_t) layer, (int32_t) expert); - } - } - std::vector offsets; - uint64_t bytes = 0; - if (!detail::make_cache_complement_plan(n_layers_, n_expert_, layer_blob_bytes_, primary_gpu_pairs, - additional_gpu_pairs, offsets, bytes, err)) return false; - - if (bytes > 0) { - uint64_t physical = 0; - if (!available_memory_bytes(physical)) { - err = "FileExpertSource: cannot determine available RAM for the resident-memory safety check"; - return false; - } - if (physical <= kPinnedMemoryHeadroom || bytes > physical - kPinnedMemoryHeadroom) { - char message[256]; - std::snprintf(message, sizeof message, - "FileExpertSource: resident complement %.2f GiB exceeds available RAM minus the 8 GiB safety headroom", - (double) bytes / 1073741824.0); - err = message; - return false; - } - } - - void* arena = nullptr; - const uint8_t* host = nullptr; - const uint8_t* device = nullptr; - if (bytes > 0) { - std::fprintf(stderr, "FileExpertSource: allocating %.2f GiB %s cache complement\n", - (double) bytes / 1073741824.0, pin ? "mapped pinned" : "pageable resident"); - std::fflush(stderr); - if (pin) { - const cudaError_t allocated = cudaHostAlloc(&arena, (size_t) bytes, - cudaHostAllocMapped | cudaHostAllocPortable); - if (allocated != cudaSuccess) { - err = std::string("FileExpertSource: mapped pinned complement allocation failed: ") + - cudaGetErrorString(allocated); - std::fprintf(stderr, "%s\n", err.c_str()); - std::fflush(stderr); - (void) cudaGetLastError(); - return false; - } - host = (const uint8_t*) arena; - void* alias = nullptr; - const cudaError_t aliased = cudaHostGetDevicePointer(&alias, arena, 0); - if (aliased != cudaSuccess || alias == nullptr) { - err = std::string("FileExpertSource: mapped pinned complement device alias failed: ") + - cudaGetErrorString(aliased); - std::fprintf(stderr, "%s\n", err.c_str()); - std::fflush(stderr); - (void) cudaGetLastError(); - (void) cudaFreeHost(arena); - return false; - } - device = (const uint8_t*) alias; - } else { - arena = std::malloc((size_t) bytes); - if (arena == nullptr) { - err = "FileExpertSource: pageable resident complement allocation failed"; - return false; - } - host = (const uint8_t*) arena; - } - } - - uint64_t copied = 0; - for (int64_t layer = 0; layer < n_layers_; ++layer) { - const uint64_t blob_bytes = layer_blob_bytes_[(size_t) layer]; - for (int64_t expert = 0; expert < n_expert_; ++expert) { - const size_t index = (size_t) layer * (size_t) n_expert_ + (size_t) expert; - const uint64_t offset = offsets[index]; - if (offset == kNoComplement) continue; - const uint8_t* source = mapped_blob(layer, expert); - if (source == nullptr || offset > bytes || blob_bytes > bytes - offset) { - err = "FileExpertSource: invalid blob bounds while building the cache complement"; - if (arena != nullptr) { - if (pin) (void) cudaFreeHost(arena); - else std::free(arena); - } - return false; - } - std::memcpy((uint8_t*) host + (size_t) offset, source, (size_t) blob_bytes); - copied += blob_bytes; - } -#if !defined(_WIN32) - const uint64_t layer_offset = layer_offsets_[(size_t) layer]; - const uint64_t layer_bytes = blob_bytes * (uint64_t) n_expert_; - const uint64_t layer_end = layer_offset + layer_bytes; - const long page_bytes = sysconf(_SC_PAGESIZE); - if (page_bytes <= 0) { - err = "FileExpertSource: cannot determine page size for mapped-page release"; - if (arena != nullptr) { - if (pin) (void) cudaFreeHost(arena); - else std::free(arena); - } - return false; - } - const uint64_t page = (uint64_t) page_bytes; - const uint64_t advice_start = layer_offset - layer_offset % page; - const uint64_t end_remainder = layer_end % page; - const uint64_t extra = end_remainder == 0 ? 0 : page - end_remainder; - const uint64_t advice_end = extra > mapped_bytes_ - layer_end ? mapped_bytes_ : layer_end + extra; - if (advice_end > advice_start && - madvise((void*) (base_ + (size_t) advice_start), (size_t) (advice_end - advice_start), MADV_DONTNEED) != 0) { - err = "FileExpertSource: madvise could not release mapped expert layer " + std::to_string(layer); - if (arena != nullptr) { - if (pin) (void) cudaFreeHost(arena); - else std::free(arena); - } - return false; - } - const int advice = posix_fadvise(fd_, (off_t) layer_offset, (off_t) layer_bytes, POSIX_FADV_DONTNEED); - if (advice != 0) { - err = "FileExpertSource: posix_fadvise could not release expert layer " + std::to_string(layer); - if (arena != nullptr) { - if (pin) (void) cudaFreeHost(arena); - else std::free(arena); - } - return false; - } -#endif - if ((layer + 1) % 8 == 0 || layer + 1 == n_layers_) { - std::fprintf(stderr, "FileExpertSource: copied cache complement through layer %lld/%lld (%.2f GiB)\n", - (long long) (layer + 1), (long long) n_layers_, (double) copied / 1073741824.0); - std::fflush(stderr); - } - } - - complement_arena_ = arena; - complement_host_ = host; - complement_device_ = device; - complement_bytes_ = bytes; - complement_offsets_ = std::move(offsets); - complement_pinned_ = pin && bytes > 0; - complement_ready_ = true; - std::fprintf(stderr, "FileExpertSource: %s cache complement ready: resident %.2f GiB, pinned %.2f GiB\n", - complement_pinned_ ? "mapped pinned" : "pageable resident", - (double) resident_bytes() / 1073741824.0, (double) pinned_bytes() / 1073741824.0); - if (!additional_gpu_pairs.empty()) { - std::fprintf(stderr, "FileExpertSource: %zu verified additional-GPU experts remain on the mmap fallback\n", - additional_gpu_pairs.size()); - } - std::fflush(stderr); - return true; -} - const uint8_t* FileExpertSource::blob(int64_t layer, int64_t expert) { - if (base_ == nullptr || layer < 0 || expert < 0 || layer >= n_layers_ || expert >= n_expert_) return nullptr; - const size_t index = (size_t) layer * (size_t) n_expert_ + (size_t) expert; - const uint8_t* mapped_fallback = mapped_blob(layer, expert); - const uint8_t* result = complement_ready_ - ? detail::cache_complement_blob_or_fallback(index, complement_offsets_, complement_host_, mapped_fallback) - : mapped_fallback; - if (result != nullptr) ++reads_; - return result; -} - -bool FileExpertSource::pinned(int64_t layer, int64_t expert) const { - if (!complement_ready_ || !complement_pinned_ || complement_host_ == nullptr || layer < 0 || expert < 0 || - layer >= n_layers_ || expert >= n_expert_) return false; - const size_t index = (size_t) layer * (size_t) n_expert_ + (size_t) expert; - return index < complement_offsets_.size() && complement_offsets_[index] != kNoComplement; -} - -const uint8_t* FileExpertSource::device_alias(int64_t layer, int64_t expert) const { - if (!pinned(layer, expert) || complement_device_ == nullptr) return nullptr; - const size_t index = (size_t) layer * (size_t) n_expert_ + (size_t) expert; - return complement_device_ + (size_t) complement_offsets_[index]; + if (base_ == nullptr) return nullptr; + if (layer < 0 || expert < 0) return nullptr; + // **`expert >= n_expert_` IS CHECKED SEPARATELY, AND THE FLAT INDEX ALONE DOES NOT CATCH IT.** The blob + // index is `layer * n_expert + expert`, so `blob(0, 512)` has flat index 512 - which is in range, and is + // `blob(1, 0)`. A router id one past the end of a layer would then read the NEXT LAYER's first expert: + // finite, correctly sized, and wrong. Layer and expert are separate axes and are validated as such. + if (expert >= n_expert_) return nullptr; + const int64_t i = layer * n_expert_ + expert; + if (i >= blobs_) return nullptr; + ++reads_; + return base_ + (size_t) i * strata::kernels::cpu::BLOB; } // ================================ THE ADAPTER ================================ @@ -693,22 +202,6 @@ void expert_pool_dispatch(void* user, const float* x_f, const int32_t* ids, cons const bool use_hits = graph_hits || (d.hits_ready() && d.decided); int64_t njobs = 0; - if (d.remote_count > 0) { - int32_t kind[32]; - if (k > 32) { - d.failed = true; d.fail = "remote experts: routing width exceeds 32"; return; - } - for (int64_t i = 0; i < k; ++i) - kind[i] = use_hits && ids[i] >= 0 && ids[i] < d.n_expert && (graph_hits - ? d.host_res[(size_t) d.layers * (size_t) d.n_expert + (size_t) ids[i]] >= 0 - : d.is_hit[(size_t) i] != 0) ? 0 : -1; - static thread_local std::string remote_error; - for (int r = 0; r < d.remote_count; ++r) - if (!d.remote[r]->begin(d.layers, x_f, ids, 1, k, kind, d.host_res, remote_error)) { - d.failed = true; d.fail = remote_error.c_str(); d.fail_layer = d.layers; return; - } - } - for (int64_t i = 0; i < k; ++i) { const int64_t e = ids[i]; if (e < 0 || e >= d.n_expert) { @@ -741,13 +234,6 @@ void expert_pool_dispatch(void* user, const float* x_f, const int32_t* ids, cons } if (graph_hits) ++d.cache_refused; // token graph: a miss (nothing is admitted during a token) - bool remote_owns = false; - for (int r = 0; r < d.remote_count; ++r) remote_owns |= d.remote[r]->owns(i); - if (remote_owns) { - std::memset(out + (size_t) i * (size_t) n_embd, 0, (size_t) n_embd * sizeof(float)); - continue; - } - // `njobs` indexes the JOB ARRAY and `i` indexes the OUTPUT - they are the same only when nothing is a // hit, and using one for the other is how a hit's row would get two experts summed into it. ExpertJob& j = d.jobs[(size_t) njobs++]; @@ -761,13 +247,6 @@ void expert_pool_dispatch(void* user, const float* x_f, const int32_t* ids, cons // Plan v0.3 P4: rows of every expert across all threads (bitwise the same as `run`). if (d.split_rows) d.pool->run_split(d.jobs.data(), (int) njobs); else d.pool->run(d.jobs.data(), (int) njobs); - if (d.remote_count > 0) { - static thread_local std::string remote_error; - for (int r = 0; r < d.remote_count; ++r) - if (!d.remote[r]->finish(out, remote_error)) { - d.failed = true; d.fail = remote_error.c_str(); d.fail_layer = d.layers; return; - } - } ++d.layers; d.experts += k; } @@ -799,113 +278,149 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32 if (!d.usage.empty()) for (int64_t i = 0; i < n_tok * k; ++i) if (ids[i] >= 0 && ids[i] < d.n_expert) d.usage[(size_t) d.layers * (size_t) d.n_expert + (size_t) ids[i]] += 1.0f; - // ---- plan v0.3 P6: the GPU's share, decided and published FIRST so the GPU starts while the CPU works. - // Distinct experts in routing order; resident ones and the last pcie_num/256 of the missed ones go to the GPU. + if (!d.routed.empty()) + for (int64_t i = 0; i < n_tok * k; ++i) + if (ids[i] >= 0 && ids[i] < d.n_expert) ++d.routed[(size_t) d.layers * (size_t) d.n_expert + (size_t) ids[i]]; + // ---- plan v0.3 P6: who computes each routed entry. Distinct experts in routing order: the ones the main GPU's + // VRAM tier holds are its own (the window decides them itself, from the same residency), the second GPU's when it + // takes this layer, the last pcie_num/256 of the rest the PCIe share (published FIRST, so the GPU reads them while + // the CPU works), the others the CPU's. const int64_t n = n_tok * k; - int32_t kind[128]; // per entry: -1 CPU, 0 VRAM, 1 PCIe - if (d.plan != nullptr && n <= 128 && n <= d.plan->cap) { - int64_t distinct[128], first_of[128]; - int nd = 0, nmiss = 0; - for (int64_t i = 0; i < n; ++i) { - first_of[i] = i; - for (int64_t j = 0; j < i; ++j) - if (ids[j] == ids[i]) { first_of[i] = first_of[j]; break; } - if (first_of[i] == i) { - distinct[nd++] = i; - const int32_t e = ids[i]; - if (e >= 0 && e < d.n_expert && d.host_res[(size_t) d.layers * (size_t) d.n_expert + (size_t) e] < 0) ++nmiss; - } - } - const bool pcie_ok = d.pcie_num > 0 && d.src->device_alias(d.layers, 0) != nullptr; - const int m = pcie_ok ? (nmiss * d.pcie_num) >> 8 : 0; - int miss_rank = 0, groups = 0, entries = 0, fetches = 0; - GpuPlanSink& P = *d.plan; - const uint8_t* dma_src[64]; - int64_t pcie_i0[64]; + if (n > 128) { + d.failed = true; + d.fail = "a verify window routes more than 128 entries"; + d.fail_layer = d.layers; + return; + } + const int32_t* res = d.host_res != nullptr ? d.host_res + (size_t) d.layers * (size_t) d.n_expert : nullptr; + auto in_vram = [&](int32_t e) { return res != nullptr && res[e] >= 0; }; // e in range + int32_t kind[128]; // per entry: -1 CPU, 0 VRAM, 1 PCIe, 2 the second GPU + // the second GPU's share: group g = the expert in its slot g2_slot[g], entries g2_ent[g2_start[g] ..) + int n2g = 0, n2e = 0; + int32_t g2_slot[128], g2_start[129], g2_ent[128]; + auto slot2 = [&](int32_t e) -> int32_t { // its slot there, -2 - p for prefetch slot p, or -1 + if (d.gpu2 == nullptr) return -1; + const int32_t s = d.host_res2[(size_t) d.layers * (size_t) d.n_expert + (size_t) e]; + if (s >= 0) return s; + const int p = d.gpu2->prefetched(d.layers, e); + return p >= 0 ? -2 - p : -1; + }; + auto on_gpu2 = [&](int32_t e) { return slot2(e) != -1; }; + int64_t distinct[128], first_of[128]; + int nd = 0; + for (int64_t i = 0; i < n; ++i) { + first_of[i] = i; + for (int64_t j = 0; j < i; ++j) + if (ids[j] == ids[i]) { first_of[i] = first_of[j]; break; } + if (first_of[i] == i) distinct[nd++] = i; + } + // The second GPU takes this layer's share only when the layer's misses are big enough that the CPU would need + // longer for all of them than its round trip (~100 us: ~4 MB of experts at the pool's ~45 GB/s); otherwise it + // would only add its latency. Its rows reach the main GPU through the sink. + GpuPlanSink* const S2 = d.gpu2 != nullptr && d.plan != nullptr && d.plan->gpu2_flag != nullptr ? d.plan : nullptr; + d.gpu2_used = false; + if (S2 != nullptr) { + uint64_t miss_bytes = 0, gpu2_bytes = 0; for (int q = 0; q < nd; ++q) { - const int64_t i0 = distinct[q]; - const int32_t e = ids[i0]; - int kd = -1; - unsigned long long ptr = 0; - if (e >= 0 && e < d.n_expert) { - const int32_t slot = d.host_res[(size_t) d.layers * (size_t) d.n_expert + (size_t) e]; - if (slot >= 0) { - kd = 0; - ptr = (unsigned long long) (d.cache_base + (d.cache_slot_off ? (size_t) d.cache_slot_off[slot] - : (size_t) slot * (size_t) d.cache_blob)); - } else { - if (miss_rank >= nmiss - m && fetches < P.staging_cap && fetches < 64) { - const uint8_t* src = d.src->blob(d.layers, e); - if (src != nullptr && d.src->pinned(d.layers, e)) { - kd = 1; - dma_src[fetches] = src; - pcie_i0[fetches] = i0; - ++fetches; - } + const int32_t e = ids[distinct[q]]; + if (e < 0 || e >= d.n_expert || in_vram(e)) continue; + miss_bytes += lay.blob_bytes(d.layers); + if (on_gpu2(e)) gpu2_bytes += lay.blob_bytes(d.layers); + } + d.gpu2_used = gpu2_bytes > 0 && miss_bytes >= d.gpu2_min_bytes; + if (gpu2_bytes > 0 && !d.gpu2_used) ++d.gpu2_skipped; + } + auto on_gpu2_now = [&](int32_t e) { return d.gpu2_used && on_gpu2(e); }; + GpuPlanSink* P = d.plan != nullptr && d.plan->pcie && n <= d.plan->cap ? d.plan : nullptr; + int nmiss = 0; + for (int q = 0; q < nd; ++q) { + const int32_t e = ids[distinct[q]]; + if (e >= 0 && e < d.n_expert && !in_vram(e) && !on_gpu2_now(e)) ++nmiss; + } + const bool pcie_ok = P != nullptr && d.pcie_num > 0 && d.src->device_alias(d.layers, 0) != nullptr; + const int m = pcie_ok ? (nmiss * d.pcie_num) >> 8 : 0; + int miss_rank = 0, fetches = 0; + const uint8_t* dma_src[64]; + int64_t pcie_i0[64]; + for (int q = 0; q < nd; ++q) { + const int64_t i0 = distinct[q]; + const int32_t e = ids[i0]; + int kd = -1; + if (e >= 0 && e < d.n_expert) { + if (in_vram(e)) { + kd = 0; + } else if (on_gpu2_now(e)) { + kd = 2; + g2_slot[n2g] = slot2(e); + g2_start[n2g++] = n2e; + for (int64_t i = i0; i < n; ++i) + if (first_of[i] == i0) g2_ent[n2e++] = (int32_t) i; + } else { + if (pcie_ok && miss_rank >= nmiss - m && fetches < P->staging_cap && fetches < 64) { + const uint8_t* src = d.src->blob(d.layers, e); + if (src != nullptr && d.src->pinned(d.layers, e)) { + kd = 1; + dma_src[fetches] = src; + pcie_i0[fetches] = i0; + ++fetches; } - ++miss_rank; } + ++miss_rank; } - for (int64_t i = i0; i < n; ++i) - if (first_of[i] == i0) kind[i] = kd; - if (kd != 0) continue; // the VRAM groups first; the PCIe groups below - P.ptr[groups] = ptr; - P.start[groups] = entries; - for (int64_t i = i0; i < n; ++i) - if (first_of[i] == i0) { - P.dst[entries] = (int32_t) i; - P.tok[entries] = (int32_t) (i / k); - ++entries; - } - ++groups; } - P.start[groups] = entries; + for (int64_t i = i0; i < n; ++i) + if (first_of[i] == i0) kind[i] = kd; + } + g2_start[n2g] = n2e; + if (S2 != nullptr) { // the main GPU takes the second GPU's rows once their flag rises: at once without a share + S2->gpu2_list[0] = n2e; + for (int i = 0; i < n2e; ++i) S2->gpu2_list[4 + i] = g2_ent[i]; + if (n2g == 0) { + std::atomic_thread_fence(std::memory_order_seq_cst); + *(volatile uint32_t*) S2->gpu2_flag = S2->ring; + S2->gpu2_ring = S2->ring; + } + } + if (P != nullptr) { // the PCIe groups: staging slot q, their entries in routing order const uint64_t bb = lay.blob_bytes(d.layers); - for (int q = 0; q < fetches; ++q) { // the PCIe groups: staging slot q, entries after the VRAM ones + int entries = 0; + for (int q = 0; q < fetches; ++q) { const int64_t i0 = pcie_i0[q]; - P.ptr2[q] = P.pcie_mode != 0 ? (unsigned long long) d.src->device_alias(d.layers, ids[i0]) - : P.staging + (unsigned long long) q * (unsigned long long) bb; - P.start2[q] = entries; + P->ptr2[q] = P->pcie_mode != 0 ? (unsigned long long) d.src->device_alias(d.layers, ids[i0]) + : P->staging + (unsigned long long) q * (unsigned long long) bb; + P->start2[q] = entries; for (int64_t i = i0; i < n; ++i) if (first_of[i] == i0) { - P.dst[entries] = (int32_t) i; - P.tok[entries] = (int32_t) (i / k); + P->dst[entries] = (int32_t) i; + P->tok[entries] = (int32_t) (i / k); ++entries; } ++d.pcie_experts; } - P.start2[fetches] = entries; - P.counts[0] = groups; - P.counts[1] = entries; - P.counts[2] = fetches; + P->start2[fetches] = entries; + P->counts[0] = 0; + P->counts[1] = entries; + P->counts[2] = fetches; std::atomic_thread_fence(std::memory_order_seq_cst); pt("publish", fetches); - if (P.publish) P.publish(P.ctx); + if (P->publish) P->publish(P->ctx); pt("fetch", fetches); - if (P.fetch) P.fetch(P.ctx, dma_src, P.pcie_mode != 0 ? 0 : fetches, (size_t) bb); // the copy engine, beside the CPU's work - } else { - for (int64_t i = 0; i < n; ++i) { - const int32_t e = ids[i]; - kind[i] = (e >= 0 && e < d.n_expert && d.host_res != nullptr && - d.host_res[(size_t) d.layers * (size_t) d.n_expert + (size_t) e] >= 0) ? 0 : -1; - } - } - if (d.remote_count > 0) { - static thread_local std::string remote_error; - for (int r = 0; r < d.remote_count; ++r) { - if (!d.remote[r]->begin(d.layers, x_f, ids, n_tok, k, kind, d.host_res, remote_error)) { - d.failed = true; d.fail = remote_error.c_str(); d.fail_layer = d.layers; return; - } - for (int64_t i = 0; i < n; ++i) if (d.remote[r]->owns(i)) kind[i] = 2; - } + if (P->fetch) P->fetch(P->ctx, dma_src, P->pcie_mode != 0 ? 0 : fetches, (size_t) bb); // the copy engine, beside the CPU's work } + const bool skip_gpu_rows = d.plan != nullptr && d.plan->host_rows_only; const auto c1 = std::chrono::steady_clock::now(); + // the gate/up activations: the window's router wrote them as Q8_K rows when the layer takes those + const uint8_t* xk = native && d.plan != nullptr ? d.plan->xk : nullptr; + auto nact = [&](int64_t t) -> const void* { + return xk != nullptr ? xk + (size_t) t * (size_t) d.plan->xk_stride + : d.nact_multi.data() + (size_t) t * kNativeActBytes; + }; if (native && lay.fmt[(size_t) d.layers].gu_type == 42) // a native Q2_0 pack: the Q2_0 kernels' activations for (int64_t t = 0; t < n_tok; ++t) act_quant_any(x_f + (size_t) t * H, H, d.act_multi[(size_t) t]); - else if (native) + else if (native && xk == nullptr) for (int64_t t = 0; t < n_tok; ++t) native_quant_act(lay.fmt[(size_t) d.layers], x_f + (size_t) t * H, d.nact_multi.data() + (size_t) t * kNativeActBytes); - else + else if (!native) for (int64_t t = 0; t < n_tok; ++t) act_quant_q8_1(x_f + (size_t) t * H, H, d.act_multi[(size_t) t]); const auto c2 = std::chrono::steady_clock::now(); int njobs = 0; @@ -921,9 +436,10 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32 d.fail_expert = e; return; } - if (kind[i] >= 0) { // CUDA0, PCIe, or a remote result staged into this row below + if (kind[i] >= 0) { // a GPU computes this entry (a VRAM hit, a PCIe read, the second GPU) if (kind[i] == 0) ++d.cache_hits; - std::memset(row, 0, (size_t) H * sizeof(float)); + if (kind[i] == 2) ++d.gpu2_entries; // its row is the second GPU's to write + else if (!skip_gpu_rows) std::memset(row, 0, (size_t) H * sizeof(float)); continue; } ++d.cache_refused; @@ -945,21 +461,42 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32 } ExpertJobMulti& jb = d.jobs_multi[(size_t) jo]; jb.act[jb.nt] = &d.act_multi[(size_t) t]; - jb.nact[jb.nt] = native ? d.nact_multi.data() + (size_t) t * kNativeActBytes : nullptr; + jb.nact[jb.nt] = native ? nact(t) : nullptr; jb.out[jb.nt] = row; ++jb.nt; ++d.multi_entries; } const auto c3 = std::chrono::steady_clock::now(); + // the second GPU's share goes out once the CPU's workers have started on theirs: its launch takes the host + // tens of microseconds. It takes the q8_1 rows the window's router wrote, and writes its rows into `out` itself. + struct Submit2 { + ExpertDispatch* d; + int n_tok, n_groups; + int64_t k; + const int32_t *slots, *starts, *entries; + float* out; + bool ok; + } s2{&d, (int) n_tok, n2g, k, g2_slot, g2_start, g2_ent, out, true}; + void (*submit2)(void*) = [](void* p) { + Submit2& s = *(Submit2*) p; + GpuPlanSink& S = *s.d->plan; + s.ok = s.d->gpu2->submit(s.d->layers, S.x1, s.n_tok, s.k, s.slots, s.starts, s.entries, s.n_groups, s.out, + S.gpu2_flag, S.ring, s.d->gpu2_err); + if (s.ok) S.gpu2_ring = S.ring; + }; + if (n2g <= 0) submit2 = nullptr; pt("run", njobs); - if (native) d.pool->run_split_multi_native(lay.fmt[(size_t) d.layers], d.jobs_multi.data(), njobs); - else d.pool->run_split_multi(d.jobs_multi.data(), njobs); - if (d.remote_count > 0) { - static thread_local std::string remote_error; - for (int r = 0; r < d.remote_count; ++r) - if (!d.remote[r]->finish(out, remote_error)) { - d.failed = true; d.fail = remote_error.c_str(); d.fail_layer = d.layers; return; - } + if (native) { + d.pool->run_split_multi_native(lay.fmt[(size_t) d.layers], d.jobs_multi.data(), njobs, submit2, &s2); + } else { + if (submit2 != nullptr) submit2(&s2); + d.pool->run_split_multi(d.jobs_multi.data(), njobs); + } + if (!s2.ok) { + d.failed = true; + d.fail = d.gpu2_err.c_str(); + d.fail_layer = d.layers; + return; } const auto c4 = std::chrono::steady_clock::now(); pt("ran"); @@ -977,6 +514,40 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32 d.experts += n_tok * k; } +void expert_prefetch_multi(ExpertDispatch& d, int64_t layer, const int32_t* ids, const float* w, int64_t n_tok, + int64_t k) { + if (d.failed || d.gpu2 == nullptr || d.gpu2->prefetch_slots() <= 0) return; + // the predicted experts neither GPU holds, ranked by their summed routing weight over the window's tokens + int32_t cand[128]; + float score[128]; + int nc = 0; + for (int64_t i = 0; i < n_tok * k && i < 128; ++i) { + const int32_t e = ids[i]; + if (e < 0 || e >= d.n_expert) continue; + const size_t at = (size_t) layer * (size_t) d.n_expert + (size_t) e; + if (d.host_res[at] >= 0 || d.host_res2[at] >= 0 || !d.src->pinned(layer, e)) continue; + int c = 0; + while (c < nc && cand[c] != e) ++c; + if (c == nc) { cand[nc] = e; score[nc++] = 0.0f; } + score[c] += w[i]; + } + const int m = (int) (std::min)((uint64_t) (std::min)(nc, d.gpu2->prefetch_slots()), + d.gpu2_prefetch_bytes / strata::kernels::cpu::expert_layout().blob_bytes(layer)); + const uint8_t* src[128]; + for (int q = 0; q < m; ++q) { + int best = q; + for (int c = q + 1; c < nc; ++c) if (score[c] > score[best]) best = c; + std::swap(cand[q], cand[best]); + std::swap(score[q], score[best]); + src[q] = d.src->blob(layer, cand[q]); + } + if (!d.gpu2->prefetch(layer, cand, src, m, strata::kernels::cpu::expert_layout().blob_bytes(layer), d.gpu2_err)) { + d.failed = true; + d.fail = d.gpu2_err.c_str(); + d.fail_layer = layer; + } +} + void expert_hit_run(void* user, void* stream, HitPhase phase, const int32_t* ids, int64_t k) { ExpertDispatch& d = *(ExpertDispatch*) user; if (d.failed) return; @@ -1086,44 +657,72 @@ void expert_hit_run(void* user, void* stream, HitPhase phase, const int32_t* ids // ================================ THE RESIDENT ARENA (R2.1) ================================ -// Plan v0.3 P6: the arena from the model's shard 1. Each layer's gate, up and down tensors hold the 512 experts -// one after another; they are read in chunks and each expert's slice lands at its place in the blob +namespace { +// Plan v0.3 P6: the arena straight from the model's GGUF shards. Each layer's gate, up and down tensors hold the +// 512 experts one after another; they are read in chunks and each expert's slice lands at its place in the blob // [gate rows | up rows | down rows] - the layout tools/iq_pack.py would have written to experts.bin. -LoadStats load_experts_gguf(const std::string& gguf, uint8_t* dst, const strata::kernels::cpu::ExpertLayout& lay, +struct GgufSpan { + size_t shard = 0; + uint64_t offset = 0; ///< in the shard's file +}; + +/// Where each layer's gate, up and down tensors are (3 per layer), found by name in whichever shard holds them +/// and checked against the pack's formats, so a pack and a model that do not belong together are refused. +bool locate_experts_gguf(const std::vector& shards, const strata::kernels::cpu::ExpertLayout& lay, + std::vector& out, std::string& err) { + static const char* roles[3] = {"gate", "up", "down"}; + out.assign((size_t) (3 * lay.n_layers), GgufSpan{}); + try { + const strata::GgufModel model(shards); + for (int64_t l = 0; l < lay.n_layers; ++l) { + const auto& fm = lay.fmt[(size_t) l]; + const uint64_t per[3] = {fm.up_off, fm.up_off, lay.bytes[(size_t) l] - fm.down_off}; + for (int r = 0; r < 3; ++r) { + const std::string name = "blk." + std::to_string(l) + ".ffn_" + roles[r] + "_exps.weight"; + size_t s = 0; + const strata::TensorInfo* t = model.find(name, &s); + int be = 0, bb = 0; + if (t == nullptr || (int) t->type != (r < 2 ? fm.gu_type : fm.d_type) || + !strata::block_geometry(t->type, be, bb) || + t->elements() / (uint64_t) be * (uint64_t) bb != per[r] * (uint64_t) lay.n_expert) { + err = "ArenaExpertSource: " + name + " is missing from the model or differs from the pack's " + "native_experts.txt"; + return false; + } + out[(size_t) (3 * l + r)] = GgufSpan{s, model.shard(s).data_start() + t->offset}; + } + } + return true; + } catch (const std::exception& e) { + err = std::string("ArenaExpertSource: ") + e.what(); + return false; + } +} + +LoadStats load_experts_gguf(const std::vector& shards, const std::vector& spans, + const std::vector& layer_base, const strata::kernels::cpu::ExpertLayout& lay, int threads) { LoadStats st; st.layers = (uint64_t) lay.n_layers; const auto t0 = std::chrono::steady_clock::now(); std::atomic next{0}; std::atomic bad{false}; - // a layer's experts may sit in another shard of the model (native_experts.txt v3): a name beside `gguf` - const size_t cut = gguf.find_last_of("/\\"); - const std::string dir = cut == std::string::npos ? std::string() : gguf.substr(0, cut + 1); - auto file_of = [&](int64_t l) -> std::string { - if (lay.gguf_file.empty() || lay.gguf_file[(size_t) l].empty()) return gguf; - return dir + lay.gguf_file[(size_t) l]; - }; auto worker = [&]() { - std::ifstream f; - std::string open_name; + std::vector files(shards.size()); std::vector buf; for (;;) { const int64_t l = next.fetch_add(1); if (l >= lay.n_layers || bad) break; - const std::string name = file_of(l); - if (name != open_name) { - f.close(); - f.clear(); - f.open(name, std::ios::binary); - if (!f) { bad = true; return; } - open_name = name; - } const auto& fm = lay.fmt[(size_t) l]; const uint64_t blob = lay.bytes[(size_t) l]; const uint64_t per[3] = {fm.up_off, fm.up_off, blob - fm.down_off}; const uint64_t at[3] = {0, fm.up_off, fm.down_off}; for (int r = 0; r < 3; ++r) { - const uint64_t src = lay.gguf_off[(size_t) (3 * l + r)]; + const GgufSpan& sp = spans[(size_t) (3 * l + r)]; + std::ifstream& f = files[sp.shard]; + if (!f.is_open()) f.open(shards[sp.shard], std::ios::binary); + if (!f) { bad = true; return; } + const uint64_t src = sp.offset; const uint64_t total = per[r] * (uint64_t) lay.n_expert; const uint64_t chunk = per[r] * 16; // 16 experts per read buf.resize((size_t) chunk); @@ -1134,7 +733,7 @@ LoadStats load_experts_gguf(const std::string& gguf, uint8_t* dst, const strata: if ((uint64_t) f.gcount() != n) { bad = true; return; } for (uint64_t k = 0; k < n / per[r]; ++k) { const uint64_t e = done / per[r] + k; - std::memcpy(dst + lay.blob_offset(l, (int64_t) e) + at[r], buf.data() + k * per[r], (size_t) per[r]); + std::memcpy(layer_base[(size_t) l] + e * blob + at[r], buf.data() + k * per[r], (size_t) per[r]); } } } @@ -1146,19 +745,36 @@ LoadStats load_experts_gguf(const std::string& gguf, uint8_t* dst, const strata: for (auto& t : pool) t.join(); if (bad) { st.seconds = -1.0; - st.ok = false; - st.error = "short read or unreadable shard while reading the experts from the GGUF"; return st; } st.bytes = lay.total; st.seconds = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); return st; } +// Pageable memory for the arena's layers past the pinning budget. +void* reserve_pageable(uint64_t bytes) { +#if defined(_WIN32) + return VirtualAlloc(nullptr, (SIZE_T) bytes, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE); +#else + void* p = mmap(nullptr, (size_t) bytes, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + return p == MAP_FAILED ? nullptr : p; +#endif +} + +void release_pageable(void* p, uint64_t bytes) { +#if defined(_WIN32) + (void) bytes; + VirtualFree(p, 0, MEM_RELEASE); +#else + munmap(p, (size_t) bytes); +#endif +} +} // namespace ArenaExpertSource::~ArenaExpertSource() { close(); } bool ArenaExpertSource::open(const std::string& pack_dir, int64_t n_layers, int64_t n_expert, int threads, - std::string& err, uint64_t max_pinned_bytes) { + std::string& err) { close(); const std::string path = pack_dir + "/experts.bin"; // plan v0.3 P6: the layout (canonical, or a native pack's per-layer blobs) was loaded by the driver @@ -1171,10 +787,13 @@ bool ArenaExpertSource::open(const std::string& pack_dir, int64_t n_layers, int6 const uint64_t want = lay.total; // plan v0.3 P6: no experts.bin in a native pack -> the experts come straight from the GGUF - const bool from_gguf = !std::ifstream(path, std::ios::binary) && lay.native && !lay.gguf_off.empty() && !gguf_.empty(); + const bool from_gguf = !std::ifstream(path, std::ios::binary) && lay.native && !gguf_.empty(); // SIZE CHECK BEFORE THE ALLOCATION, not after. A wrong pack should name the two numbers rather than spend // 34 GB and a minute of loading first. - if (!from_gguf) { + std::vector spans; + if (from_gguf) { + if (!locate_experts_gguf(gguf_, lay, spans, err)) return false; + } else { std::ifstream f(path, std::ios::binary | std::ios::ate); if (!f) { err = "ArenaExpertSource: cannot open " + path; return false; } const uint64_t got = (uint64_t) f.tellg(); @@ -1190,96 +809,91 @@ bool ArenaExpertSource::open(const std::string& pack_dir, int64_t n_layers, int6 } } - // one layer per registration slice, so no expert straddles two registrations. The arena is one blob - // longer than the file: a copy of a whole VRAM slot (the largest blob) may then start at any expert. - std::vector bounds, loff, lbytes; + // one pinned block per layer, each followed by room for a whole slot-sized copy (the largest blob) from its last + // expert; the layers past the pinning budget in one pageable range, locked resident + const uint64_t maxb = lay.max_blob; + std::vector loff, lbytes, room; for (int64_t l = 0; l < n_layers; ++l) { - bounds.push_back(lay.layer_offset(l)); loff.push_back(lay.layer_offset(l)); lbytes.push_back(lay.blob_bytes(l) * (uint64_t) n_expert); + room.push_back(lbytes.back() + (maxb - lay.blob_bytes(l))); } - bounds.push_back(want); - PinnedArena* a = new PinnedArena(want + (uint64_t) blob, bounds, max_pinned_bytes); - if (!a->valid()) { - delete a; - err = "ArenaExpertSource: the arena could not be reserved (" + std::to_string(want) + " B)"; - return false; - } - const LoadStats st = from_gguf ? load_experts_gguf(gguf_, a->data(), lay, threads) - : load_experts_ranges(path, a->data(), loff, lbytes, threads, /*chunk=*/8u << 20); - if (!st.ok) { - delete a; - err = "ArenaExpertSource: the expert load was refused: " + (st.error.empty() ? std::string("unknown") : st.error); - return false; + uint64_t pinned_total = 0; + std::string tail_note; + for (int64_t l = 0; l < n_layers; ++l) { + void* p = nullptr; + if (cudaHostAlloc(&p, (size_t) room[(size_t) l], cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess) { + (void) cudaGetLastError(); + for (int64_t m = l; m < n_layers; ++m) tail_bytes_ += room[(size_t) m]; + tail_ = reserve_pageable(tail_bytes_); + if (tail_ == nullptr) { + close(); + err = "ArenaExpertSource: the arena could not be reserved (" + std::to_string(tail_bytes_) + " B)"; + return false; + } + tail_note = strata::platform::lock_resident(tail_, tail_bytes_).note; + uint8_t* at = (uint8_t*) tail_; + for (int64_t m = l; m < n_layers; ++m) { + layer_base_.push_back(at); + layer_dev_.push_back(nullptr); + at += room[(size_t) m]; + } + break; + } + blocks_.push_back((uint8_t*) p); + layer_base_.push_back((uint8_t*) p); + void* d = nullptr; + layer_dev_.push_back(cudaHostGetDevicePointer(&d, p, 0) == cudaSuccess ? (const uint8_t*) d : nullptr); + (void) cudaGetLastError(); + pinned_total += room[(size_t) l]; } + const LoadStats st = from_gguf ? load_experts_gguf(gguf_, spans, layer_base_, lay, threads) + : load_experts_ranges(path, layer_base_, loff, lbytes, threads, /*chunk=*/8u << 20); if (st.bytes != want) { - delete a; + close(); err = "ArenaExpertSource: the load read " + std::to_string(st.bytes) + " B of " + std::to_string(want); return false; } - arena_ = a; - base_ = a->data(); - pinned_bytes_ = a->registered_bytes; - // plan v0.3 P6: device aliases of the mapped registration, for the PCIe share of the misses - dev_slice_.clear(); - slice_bytes_ = a->slice_bytes; - if (a->registered_bytes > 0) { - std::vector starts = a->slice_bytes > 0 ? a->slice_starts : std::vector{0}; - for (uint64_t off : starts) { - void* d = nullptr; - if (cudaHostGetDevicePointer(&d, (void*) (base_ + off), 0) != cudaSuccess) { - (void) cudaGetLastError(); - dev_slice_.clear(); - break; - } - dev_slice_.push_back((const uint8_t*) d); - } - } + note_ = std::to_string(blocks_.size()) + " of " + std::to_string(n_layers) + " layers pinned (" + + std::to_string(pinned_total >> 30) + " GiB, a cudaHostAlloc block each)"; + if (tail_ != nullptr) note_ += "; the others pageable: " + tail_note; blobs_ = n_layers * n_expert; n_expert_ = n_expert; reads_ = 0; - note_ = a->note; gib_per_s_ = st.gib_per_second(); - load_seconds_ = st.seconds; - load_read_s_ = st.read_seconds; - load_copy_s_ = st.copy_seconds; return true; } void ArenaExpertSource::close() { - if (arena_ != nullptr) { - delete (PinnedArena*) arena_; - arena_ = nullptr; - } - base_ = nullptr; + for (uint8_t* p : blocks_) cudaFreeHost(p); + blocks_.clear(); + if (tail_ != nullptr) { + strata::platform::unlock_resident(tail_, tail_bytes_); + release_pageable(tail_, tail_bytes_); + tail_ = nullptr; + tail_bytes_ = 0; + } + layer_base_.clear(); + layer_dev_.clear(); blobs_ = 0; n_expert_ = 0; } bool ArenaExpertSource::pinned(int64_t layer, int64_t expert) const { - if (base_ == nullptr || layer < 0 || expert < 0 || expert >= n_expert_) return false; - const auto& lay = strata::kernels::cpu::expert_layout(); - return lay.blob_offset(layer, expert) + lay.blob_bytes(layer) <= pinned_bytes_; + return layer >= 0 && layer < (int64_t) blocks_.size() && expert >= 0 && expert < n_expert_; } const uint8_t* ArenaExpertSource::device_alias(int64_t layer, int64_t expert) const { - if (dev_slice_.empty() || !pinned(layer, expert)) return nullptr; - const auto& lay = strata::kernels::cpu::expert_layout(); - if (slice_bytes_ == 0) return dev_slice_[0] + lay.blob_offset(layer, expert); - // one registration slice per layer - if ((size_t) layer >= dev_slice_.size()) return nullptr; - return dev_slice_[(size_t) layer] + (uint64_t) expert * lay.blob_bytes(layer); + if (!pinned(layer, expert) || layer_dev_[(size_t) layer] == nullptr) return nullptr; + return layer_dev_[(size_t) layer] + (uint64_t) expert * strata::kernels::cpu::expert_layout().blob_bytes(layer); } const uint8_t* ArenaExpertSource::blob(int64_t layer, int64_t expert) { - if (base_ == nullptr) return nullptr; - if (layer < 0 || expert < 0 || expert >= n_expert_) return nullptr; - const int64_t idx = layer * n_expert_ + expert; - if (idx < 0 || idx >= blobs_) return nullptr; + if (layer < 0 || layer >= (int64_t) layer_base_.size() || expert < 0 || expert >= n_expert_) return nullptr; ++reads_; // Pointer arithmetic into resident memory. No fault, no copy, no mapping - which is the entire point of // this class over `FileExpertSource`. - return base_ + strata::kernels::cpu::expert_layout().blob_offset(layer, expert); + return layer_base_[(size_t) layer] + (uint64_t) expert * strata::kernels::cpu::expert_layout().blob_bytes(layer); } } // namespace strata::core diff --git a/src/core/layer.cpp b/src/core/layer.cpp index 5c6c06ef112..eee96a65997 100644 --- a/src/core/layer.cpp +++ b/src/core/layer.cpp @@ -7,7 +7,6 @@ #include "strata/kernels/elementwise.hpp" #include "strata/kernels/gdn.hpp" #include "strata/kernels/qsa.hpp" -#include "strata/kernels/kv_q4.hpp" #include "strata/kernels/kv_q8.hpp" #include "strata/kernels/quantize_act.hpp" #include "strata/kernels/rope.hpp" @@ -25,12 +24,10 @@ #include "strata/kernels/native_rope.hpp" #include "strata/kernels/native_flash_attn.hpp" #include "strata/kernels/fused_gr.hpp" -#include "strata/kernels/cvec.hpp" #include "strata/kernels/qsa_decode_attn.hpp" #include "strata/kernels/fused_gdn.hpp" #include "strata/kernels/qsa_select.hpp" -#include -#include +#include "strata/kernels/verify_kernels.hpp" #include #include #include @@ -172,13 +169,9 @@ if (!w.wants_q8k()) { s_gemv_q8_0_split(x80, p.codes, p.scales, p.offset, // namespace void layer_set_native_bf16(bool enabled) { native_bf16_projections = enabled; } void layer_set_native_flash_attn_short(bool enabled) { native_flash_attn_short = enabled; } -namespace { bool g_kv_int8 = false, g_kv_q4 = false, g_kv_hybrid = false; } +namespace { bool g_kv_int8 = false; } void qsa_set_kv_int8(bool enabled) { g_kv_int8 = enabled; } bool qsa_kv_int8() { return g_kv_int8; } -void qsa_set_kv_q4(bool enabled) { g_kv_q4 = enabled; } -bool qsa_kv_q4() { return g_kv_q4; } -void qsa_set_kv_hybrid(bool enabled) { g_kv_hybrid = enabled; } -bool qsa_kv_hybrid() { return g_kv_hybrid; } uint64_t gdn_buffers_bytes(const ModelGeometry& g) { const int64_t C = g.ssm_conv_channels; const int64_t V = g.ssm_value_dim; const uint64_t parts[] = { q8k_bytes(g.n_embd), // x_q8k (uint64_t) (g.n_embd / 32) * 34, @@ -364,8 +357,11 @@ if (!native_bf16_projections) f32_to_bf16_bulk(x, b.x_bf16, g.n_embd, stream); // LEDGER L41 -> L42. project_bf16(x, b.x_bf16, (const uint16_t*) w_router->data, b.logits, g.n_embd, g.n_expert, true, stream); // ---- routing: softmax over ALL experts, stable descending argsort with ties by index, gather, renormalise -// the native fused router is canonical-512x10 only; anything else takes the generic top-k kernel -if (native_router_enabled() && g.n_expert == 512 && k == 10) { +if (native_router_enabled()) { + if (g.n_expert != 512 || k != 10) { + err = v.name("router") + ": native router requires 512 experts and k=10"; + return false; + } try { native_router_top10(b.logits, b.ids, b.weights, stream); } catch (const std::exception& error) { err = v.name("router") + ": " + error.what(); return false; } } else router_top10(b.logits, 1, (int) g.n_expert, (int) k, b.ids, b.weights, stream); @@ -451,6 +447,10 @@ bool layer_verify_compatible(std::string& why) { else if (!g_fast_attn || native_flash_attn_short) why = "the split-K decode attention is off"; else if (!g_fast_select) why = "the block top-k selection is off"; else if (!strata::kernels::native_qsa_indexer_enabled()) why = "the native QSA indexer is off"; + else if (!strata::kernels::native_qsa_enabled() || !strata::kernels::native_rope_enabled()) + why = "the native QSA norms or RoPE are off"; + else if (!strata::kernels::native_moe_combine_enabled()) why = "the native MoE combine is off"; + else if (!strata::kernels::shared_expert_native_bf16()) why = "the shared expert's native gate is off"; else return true; return false; } @@ -480,249 +480,39 @@ struct Cursor { uint8_t* p; uint64_t used = 0; template // dispatch too and the two must not be able to drift. } // namespace -namespace { -// KV streaming (docs/kv-streaming-design.md): 0 keeps every cell in VRAM. -int64_t g_kv_resident = 0; -uint64_t g_kv_host_bytes = 0; - -/// How one state holds its K/V: `mode` as in QsaState::kv_mode, `slots` VRAM pages of `pages` logical ones. -struct KvPlan { - int mode = 0; - int64_t pages = 0, slots = 0, pooled_rows = 0; -}; -KvPlan kv_plan(const QsaShapes& s, int64_t max_cells, int64_t ring_cells) { - KvPlan p; - p.pages = (max_cells + s.page_size - 1) / s.page_size; - p.slots = p.pages; - p.pooled_rows = max_cells / s.idx_block + 2; - if (g_kv_resident <= 0 || ring_cells < 0) return p; // ring_cells < 0: always fully resident - // A/B only: STRATA_KV_RING_OFF keeps the drafter fully resident, STRATA_KV_MAIN_OFF the main layers - static const bool ring_off = std::getenv("STRATA_KV_RING_OFF") != nullptr; - static const bool main_off = std::getenv("STRATA_KV_MAIN_OFF") != nullptr; - if (ring_cells > 0 ? ring_off : main_off) return p; - if (ring_cells > 0) { - const int64_t r = (ring_cells + s.page_size - 1) / s.page_size; - if (r < p.pages) { p.mode = 2; p.slots = r; p.pooled_rows = 2; } // the drafter has no indexer - } else { - const int64_t r = (std::max(g_kv_resident, qsa_kv_resident_min()) + s.page_size - 1) / s.page_size; - if (r < p.pages) { p.mode = 1; p.slots = r; } - } - return p; -} -uint64_t kv_pool_bytes(const QsaShapes& s, int64_t pages, bool hybrid, bool int8) { - if (hybrid) { // K8V4: the INT8 K half (codes + scales) plus the Q4_0 V half (kv_q4.hpp's rotation) - const uint64_t rows = (uint64_t) pages * s.page_size * s.n_head_kv; - return rows * (uint64_t) s.head_dim + rows * (uint64_t) (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2 + - rows * strata::kernels::kv_q4_bytes_per_head((int) s.head_dim) + 64; - } - if (g_kv_q4) return (uint64_t) pages * s.page_size * strata::kernels::kv_q4_bytes_per_cell(s) + 64; - return int8 ? (uint64_t) pages * s.page_size * strata::kernels::kv_q8_bytes_per_cell(s) + 64 - : (uint64_t) pages * s.n_head_kv * s.page_size * s.head_dim * 2 * 2; -} -} // namespace - -void qsa_set_kv_resident(int64_t cells) { g_kv_resident = cells > 0 ? cells : 0; } -int64_t qsa_kv_resident() { return g_kv_resident; } -int64_t qsa_kv_resident_min() { return 20480; } -uint64_t qsa_kv_host_bytes() { return g_kv_host_bytes; } - -uint64_t qsa_state_bytes(const ModelGeometry& g, int64_t max_cells, bool with_rope, int64_t ring_cells) { - const QsaShapes s = qsa_shapes(g); - const KvPlan p = kv_plan(s, max_cells, ring_cells); - uint64_t n = 0; - n += kv_pool_bytes(s, p.slots, g_kv_hybrid && ring_cells <= 0, - g_kv_int8 || g_kv_hybrid) + 4 * 16; // K/V pools (the VRAM slots) - n += (uint64_t) p.pages * 4; // page_table - if (p.mode == 1) n += strata::kernels::kv_stream_map_bytes(p.slots) + 6 * 16; // the residency map - n += (uint64_t) (s.idx_block - 1) * s.idx_dim * 4; // tail - n += (uint64_t) s.idx_dim * 4; // dead - n += (uint64_t) p.pooled_rows * s.idx_dim * 4; // pooled - n += 16; // block_pos - if (with_rope) n += (uint64_t) max_cells * (s.n_rot / 2) * 4 * 2; // cos + sin tables - n += strata::kernels::qsa_step_bytes() + 16; // counts and aligned attention status - n += (uint64_t) s.n_head * 4; // pos_dev - return align_up16(n) + 256; -} - -uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, QsaState& st, - const QsaState* share_rope, int64_t ring_cells) { - const QsaShapes s = qsa_shapes(g); - const KvPlan p = kv_plan(s, max_cells, ring_cells); - const int64_t pages = p.pages; - Cursor c{(uint8_t*) base}; - st.kv_int8 = g_kv_int8 && !g_kv_q4; - st.kv_q4 = g_kv_q4; - // Hybrid K8V4, main layers only (the drafter's state is created with the globals toggled to INT8 - - // mtp.cpp). Streamed mode is refused outright; generate.cpp validates it too, this is the backstop. - if (g_kv_hybrid && ring_cells <= 0) { - if (p.mode == 1) { - std::fprintf(stderr, "strata: hybrid K8V4 KV does not support --kv-resident streaming\n"); - return 0; - } - st.kv_hybrid = true; - st.kv_int8 = false; - st.kv_q4 = false; - } - st.kv_mode = p.mode; - st.n_slots = p.slots; - const uint64_t rows = (uint64_t) p.slots * s.n_head_kv * s.page_size; // VRAM rows: the slots - const uint64_t q4_row = strata::kernels::kv_q4_bytes_per_head((int) s.head_dim); - if (st.kv_hybrid) { - st.k_q = c.take(rows * s.head_dim); - st.k_scale = c.take(rows * (s.head_dim / strata::kernels::KV_Q8_GROUP)); - st.v_q4 = c.take(rows * q4_row); - } else if (st.kv_q4) { - st.k_q4 = c.take(rows * q4_row); - st.v_q4 = c.take(rows * q4_row); - } else if (st.kv_int8) { // st, not the global: the drafter's ring is INT8 under --kv k8v4 too - st.k_q = c.take(rows * s.head_dim); - st.v_q = c.take(rows * s.head_dim); - st.k_scale = c.take(rows * (s.head_dim / strata::kernels::KV_Q8_GROUP)); - st.v_scale = c.take(rows * (s.head_dim / strata::kernels::KV_Q8_GROUP)); - } else { - st.k_pool = c.take(rows * s.head_dim); - st.v_pool = c.take(rows * s.head_dim); - } - st.page_table = c.take((uint64_t) pages); - st.n_pages = pages; - st.max_cells = max_cells; - st.map = strata::kernels::KvStreamMap{}; - st.map.page_table = st.page_table; - st.map.n_blocks = pages; - st.map.n_slots = p.slots; - if (p.mode == 1) { - st.map.slot_block = c.take((uint64_t) p.slots); - st.map.slot_stamp = c.take((uint64_t) p.slots); - st.map.slot_ref = c.take((uint64_t) p.slots); - st.map.miss_block = c.take((uint64_t) p.slots); - st.map.miss_slot = c.take((uint64_t) p.slots); - st.map.ctl = c.take((uint64_t) strata::kernels::kKvCtlInts); - } - st.idx_tail = c.take((uint64_t) (s.idx_block - 1) * s.idx_dim); - st.idx_dead = c.take((uint64_t) s.idx_dim); - st.idx_pooled_rows = p.pooled_rows; - st.idx_pooled = c.take((uint64_t) p.pooled_rows * s.idx_dim); - st.idx_block_pos = c.take(1); - if (share_rope != nullptr) { - st.cos_tab = share_rope->cos_tab; - st.sin_tab = share_rope->sin_tab; - } else { - st.cos_tab = c.take((uint64_t) max_cells * (s.n_rot / 2)); - st.sin_tab = c.take((uint64_t) max_cells * (s.n_rot / 2)); - } - st.step = c.take((uint64_t) strata::kernels::kStepCount); - st.attention_status = c.take(1); - st.pos_dev = c.take((uint64_t) s.n_head); - // the pinned staging the uploads copy FROM - see the note on `host_step` in the header - if (cudaHostAlloc((void**) &st.host_step, strata::kernels::qsa_step_bytes() + sizeof(int32_t), - cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess || - cudaHostAlloc((void**) &st.host_pos, (size_t) s.n_head * 4, cudaHostAllocMapped | cudaHostAllocPortable) != - cudaSuccess) { - return 0; // the caller sees a zero byte count; a half-built state is worse than none - } - st.host_step[strata::kernels::kStepCount] = 0; - // KV streaming: the authoritative K/V of every cell, pinned and device-mapped, in the identity layout - st.host = strata::kernels::KvHostPools{}; - if (p.mode != 0) { - const uint64_t hrows = (uint64_t) pages * s.n_head_kv * s.page_size; - const uint64_t bytes = (uint64_t) pages * strata::kernels::kv_block_bytes(s, qsa_kv_format(st)) + 4 * 256; - uint8_t* h = nullptr; - uint8_t* d = nullptr; - if (cudaHostAlloc((void**) &h, bytes, cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess || - cudaHostGetDevicePointer((void**) &d, h, 0) != cudaSuccess) { - // under WSL the NVIDIA driver pins only ~1 GiB in all, which is less than 128K of 8-bit KV needs - if (p.mode == 1) std::fprintf(stderr, "strata: KV streaming: cannot pin %.2f GiB of RAM for a layer's KV copy " - "(%.2f GiB pinned so far) - lower the context, or run without --kv-resident (under " - "WSL the driver pins only about 1 GiB in all)\n", (double) bytes / 1073741824.0, - (double) g_kv_host_bytes / 1073741824.0); - return 0; - } - g_kv_host_bytes += bytes; - Cursor hc{d}; - if (st.kv_q4) { - st.host.k_q4 = hc.take(hrows * q4_row); - st.host.v_q4 = hc.take(hrows * q4_row); - } else if (st.kv_int8) { - st.host.k_q = hc.take(hrows * s.head_dim); - st.host.v_q = hc.take(hrows * s.head_dim); - st.host.k_scale = hc.take(hrows * (s.head_dim / strata::kernels::KV_Q8_GROUP)); - st.host.v_scale = hc.take(hrows * (s.head_dim / strata::kernels::KV_Q8_GROUP)); - } else { - st.host.k_pool = hc.take(hrows * s.head_dim); - st.host.v_pool = hc.take(hrows * s.head_dim); - } - } - // THE ROPE TABLE IS BUILT ON THE HOST IN FLOAT64 and uploaded once, because the reference computes its - // frequencies in float64 and reproducing that on device means double-precision `pow`/`cos` that need not - // agree with the host's libm. A table shorter than the sequence would have the rotation read past it. - if (share_rope == nullptr) { - std::vector hc((size_t) max_cells * (s.n_rot / 2)), hs((size_t) max_cells * (s.n_rot / 2)); - strata::kernels::build_rope_table((int) s.n_rot, strata::kernels::qsa_freq_base(), (int) max_cells, - hc.data(), hs.data()); - cudaMemcpy(st.cos_tab, hc.data(), hc.size() * 4, cudaMemcpyHostToDevice); - cudaMemcpy(st.sin_tab, hs.data(), hs.size() * 4, cudaMemcpyHostToDevice); - } - // the page table starts as the IDENTITY, which is the simplest legal mapping and what a caller that does - // not page at all wants; a streamed state starts with nothing resident, a ring at `block % n_slots`. - if (p.mode == 0) { - std::vector tab((size_t) pages); - for (int64_t i = 0; i < pages; ++i) tab[(size_t) i] = (int32_t) i; - cudaMemcpy(st.page_table, tab.data(), tab.size() * 4, cudaMemcpyHostToDevice); - } else if (p.mode == 1) { - strata::kernels::kv_stream_reset(st.map, nullptr); - } else { - strata::kernels::kv_ring_table(st.page_table, pages, p.slots, nullptr); - } - cudaDeviceSynchronize(); - return c.used; -} - -void qsa_state_zero(const QsaState& st, const ModelGeometry& g, void* stream) { - const QsaShapes s = qsa_shapes(g); - cudaStream_t cs = (cudaStream_t) stream; - const size_t rows = (size_t) st.n_slots * s.n_head_kv * s.page_size; - if (st.kv_hybrid) { - cudaMemsetAsync(st.k_q, 0, rows * s.head_dim, cs); - cudaMemsetAsync(st.k_scale, 0, rows * (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2, cs); - cudaMemsetAsync(st.v_q4, 0, rows * strata::kernels::kv_q4_bytes_per_head((int) s.head_dim), cs); - } else if (st.kv_q4) { - cudaMemsetAsync(st.k_q4, 0, rows * strata::kernels::kv_q4_bytes_per_head((int) s.head_dim), cs); - cudaMemsetAsync(st.v_q4, 0, rows * strata::kernels::kv_q4_bytes_per_head((int) s.head_dim), cs); - } else if (st.kv_int8) { - cudaMemsetAsync(st.k_q, 0, rows * s.head_dim, cs); - cudaMemsetAsync(st.v_q, 0, rows * s.head_dim, cs); - cudaMemsetAsync(st.k_scale, 0, rows * (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2, cs); - cudaMemsetAsync(st.v_scale, 0, rows * (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2, cs); - } else { - cudaMemsetAsync(st.k_pool, 0, rows * s.head_dim * 2, cs); - cudaMemsetAsync(st.v_pool, 0, rows * s.head_dim * 2, cs); - } - // A streamed state starts over with nothing resident. Its host copy is not cleared (GBs over PCIe per new - // conversation): no reader names a cell before this sequence has written it, and a block copied in whole - // carries the unwritten cells past the end, which nothing reads. - if (st.kv_mode == 1) strata::kernels::kv_stream_reset(st.map, stream); - cudaMemsetAsync(st.idx_tail, 0, (size_t) (s.idx_block - 1) * s.idx_dim * 4, cs); - cudaMemsetAsync(st.idx_dead, 0, (size_t) s.idx_dim * 4, cs); - cudaMemsetAsync(st.idx_pooled, 0, (size_t) st.idx_pooled_rows * s.idx_dim * 4, cs); -} - -strata::kernels::QsaAttnPools qsa_attn_pools(const QsaState& st) { - strata::kernels::QsaAttnPools pools; - pools.page_table = st.page_table; - if (st.kv_hybrid) { pools.k_q = st.k_q; pools.k_scale = st.k_scale; pools.v_q4 = st.v_q4; } - else if (st.kv_q4) { pools.k_q4 = st.k_q4; pools.v_q4 = st.v_q4; } - else if (st.kv_int8) { pools.k_q = st.k_q; pools.v_q = st.v_q; pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; } - else { pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; } - return pools; -} - -void qsa_kv_resolve(const QsaState& st, const ModelGeometry& g, const int32_t* ids, const int32_t* steps, int64_t n_q, - int64_t cap, void* stream) { - if (st.kv_mode != 1) return; - strata::kernels::kv_stream_resolve(st.map, qsa_attn_pools(st), st.host, qsa_kv_format(st), ids, steps, n_q, cap, - qsa_shapes(g), stream); +uint64_t qsa_state_bytes(const ModelGeometry& g, int64_t max_cells, bool with_rope) { const QsaShapes s = qsa_shapes(g); const int64_t pages = (max_cells + s.page_size - 1) / s.page_size; uint64_t n = 0; n += g_kv_int8 ? (uint64_t) pages * s.page_size * strata::kernels::kv_q8_bytes_per_cell(s) + 64 : (uint64_t) pages * s.n_head_kv * s.page_size * s.head_dim * 2 * 2; +// k_pool + v_pool, fp16 +n += (uint64_t) pages * 4; +// page_table +n += (uint64_t) (s.idx_block - 1) * s.idx_dim * 4; +// tail +n += (uint64_t) s.idx_dim * 4; +// dead +n += (uint64_t) ((max_cells / s.idx_block) + 2) * s.idx_dim * 4; +// pooled +n += 16; +// block_pos +if (with_rope) n += (uint64_t) max_cells * (s.n_rot / 2) * 4 * 2; +// cos + sin tables +n += strata::kernels::qsa_step_bytes() + 16; // counts and aligned attention status +// the per-token counts +n += (uint64_t) s.n_head * 4; +// pos_dev +return align_up16(n) + 256;} +uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, QsaState& st, const QsaState* share_rope) { const QsaShapes s = qsa_shapes(g); const int64_t pages = (max_cells + s.page_size - 1) / s.page_size; Cursor c{(uint8_t*) base}; st.kv_int8 = g_kv_int8; if (g_kv_int8) { const uint64_t cells = (uint64_t) pages * s.n_head_kv * s.page_size; st.k_q = c.take(cells * s.head_dim); st.v_q = c.take(cells * s.head_dim); st.k_scale = c.take(cells * (s.head_dim / strata::kernels::KV_Q8_GROUP)); st.v_scale = c.take(cells * (s.head_dim / strata::kernels::KV_Q8_GROUP)); } else { st.k_pool = c.take((uint64_t) pages * s.n_head_kv * s.page_size * s.head_dim); st.v_pool = c.take((uint64_t) pages * s.n_head_kv * s.page_size * s.head_dim); } st.page_table = c.take((uint64_t) pages); st.n_pages = pages; st.max_cells = max_cells; st.idx_tail = c.take((uint64_t) (s.idx_block - 1) * s.idx_dim); st.idx_dead = c.take((uint64_t) s.idx_dim); st.idx_pooled = c.take((uint64_t) ((max_cells / s.idx_block) + 2) * s.idx_dim); st.idx_block_pos = c.take(1); if (share_rope != nullptr) { st.cos_tab = share_rope->cos_tab; st.sin_tab = share_rope->sin_tab; } else { st.cos_tab = c.take((uint64_t) max_cells * (s.n_rot / 2)); st.sin_tab = c.take((uint64_t) max_cells * (s.n_rot / 2)); } st.step = c.take((uint64_t) strata::kernels::kStepCount); st.attention_status = c.take(1); st.pos_dev = c.take((uint64_t) s.n_head); +// the pinned staging the uploads copy FROM - see the note on `host_step` in the header +if (cudaHostAlloc((void**) &st.host_step, strata::kernels::qsa_step_bytes() + sizeof(int32_t), cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess || cudaHostAlloc((void**) &st.host_pos, (size_t) s.n_head * 4, cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess) { return 0; +// the caller sees a zero byte count; a half-built state is worse than none } - +st.host_step[strata::kernels::kStepCount] = 0; +// THE ROPE TABLE IS BUILT ON THE HOST IN FLOAT64 and uploaded once, because the reference computes its +// frequencies in float64 and reproducing that on device means double-precision `pow`/`cos` that need not +// agree with the host's libm. A table shorter than the sequence would have the rotation read past it. +if (share_rope == nullptr) { std::vector hc((size_t) max_cells * (s.n_rot / 2)), hs((size_t) max_cells * (s.n_rot / 2)); strata::kernels::build_rope_table((int) s.n_rot, strata::kernels::qsa_freq_base(), (int) max_cells, hc.data(), hs.data()); cudaMemcpy(st.cos_tab, hc.data(), hc.size() * 4, cudaMemcpyHostToDevice); cudaMemcpy(st.sin_tab, hs.data(), hs.size() * 4, cudaMemcpyHostToDevice); } +// the page table starts as the IDENTITY, which is the simplest legal mapping and what a caller that does +// not page at all wants. A real allocator re-points it; nothing outside `kv_append`/`kv_gather` may care. +std::vector tab((size_t) pages); for (int64_t i = 0; i < pages; ++i) tab[(size_t) i] = (int32_t) i; cudaMemcpy(st.page_table, tab.data(), tab.size() * 4, cudaMemcpyHostToDevice); return c.used;} +void qsa_state_zero(const QsaState& st, const ModelGeometry& g, void* stream) { const QsaShapes s = qsa_shapes(g); cudaStream_t cs = (cudaStream_t) stream; if (st.kv_int8) { const size_t cells = (size_t) st.n_pages * s.n_head_kv * s.page_size; cudaMemsetAsync(st.k_q, 0, cells * s.head_dim, cs); cudaMemsetAsync(st.v_q, 0, cells * s.head_dim, cs); cudaMemsetAsync(st.k_scale, 0, cells * (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2, cs); cudaMemsetAsync(st.v_scale, 0, cells * (s.head_dim / strata::kernels::KV_Q8_GROUP) * 2, cs); } else { cudaMemsetAsync(st.k_pool, 0, (size_t) st.n_pages * s.n_head_kv * s.page_size * s.head_dim * 2, cs); cudaMemsetAsync(st.v_pool, 0, (size_t) st.n_pages * s.n_head_kv * s.page_size * s.head_dim * 2, cs); } cudaMemsetAsync(st.idx_tail, 0, (size_t) (s.idx_block - 1) * s.idx_dim * 4, cs); cudaMemsetAsync(st.idx_dead, 0, (size_t) s.idx_dim * 4, cs); cudaMemsetAsync(st.idx_pooled, 0, (size_t) ((st.max_cells / s.idx_block) + 2) * s.idx_dim * 4, cs);} uint64_t qsa_buffers_bytes(const ModelGeometry& g, int64_t max_cells) { const QsaShapes s = qsa_shapes(g); const int64_t cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s); uint64_t n = 0; n += (uint64_t) q8k_bytes(g.n_embd); n += (uint64_t) (g.n_embd / 32) * 34; // block_q8_0 n += (uint64_t) g.n_embd * 2; n += (uint64_t) g.n_head * 2 * g.head_dim * 4; n += (uint64_t) g.n_head * g.head_dim * 4; n += (uint64_t) g.n_head_kv * g.head_dim * 4 * 2; n += (uint64_t) g.idx_key_dim * 4; n += (uint64_t) g.idx_q_heads * g.idx_key_dim * 4; n += (uint64_t) max_cells * 4; n += (uint64_t) cap * 4; n += (uint64_t) cap * g.n_head_kv * g.head_dim * 2 * 2; n += (uint64_t) g.n_head * g.head_dim * 4; @@ -838,7 +628,19 @@ void stage_timing_report(int64_t n_layers) { // needs nothing from the layer it is called for beyond its index. static void dump_slot(float* dump, const ModelGeometry& g, int64_t layer, const float* src, uint64_t off, uint64_t n, void* stream); -bool qsa_layer(const WeightTable& tables, const ModelGeometry& g, int64_t layer, int64_t pos, int32_t pos_base, const QsaState& st, const QsaBuffers& b, const float* x, float* out, void* stream, std::string& err, float* dump) { using namespace strata::kernels; const QsaShapes s = qsa_shapes(g); const LayerView v(tables, layer); const int64_t n_kv = pos + 1; /* P7 audit: RoPE reads cos/sin row pos_base + pos, and the table holds max_cells rows. */ if ((int64_t) pos_base + pos >= st.max_cells || pos_base < 0) { err = "qsa_layer: position " + std::to_string((long long) pos_base + pos) + " is outside the RoPE table (" + std::to_string((long long) st.max_cells) + " rows)"; return false; } const int64_t n_bid = n_kv / s.idx_block; const int64_t width = qsa_selection_width(n_kv, s); const int64_t cap = qsa_selection_width(kTopkMaxCells, s); +bool qsa_layer(const WeightTable& tables, const ModelGeometry& g, int64_t layer, int64_t pos, int32_t pos_base, + const QsaState& st, const QsaBuffers& b, const float* x, float* out, void* stream, + std::string& err, float* dump) { + using namespace strata::kernels; + const QsaShapes s = qsa_shapes(g); + const LayerView v(tables, layer); + /* P7 audit: RoPE reads cos/sin row pos_base + pos, and the table holds max_cells rows. */ + if ((int64_t) pos_base + pos >= st.max_cells || pos_base < 0) { + err = "qsa_layer: position " + std::to_string((long long) pos_base + pos) + " is outside the RoPE table (" + + std::to_string((long long) st.max_cells) + " rows)"; + return false; + } + const int64_t cap = qsa_selection_width(kTopkMaxCells, s); const auto normalize_rotate = [&](float* data, const WeightRef* norm, int rows, int cols) { try { if (native_qsa_enabled()) native_qsa_rms_norm_weighted(data, (const float*) norm->data, data, cols, rows, RMS_EPS, stream); @@ -891,18 +693,7 @@ if (!gemv_quantized(*w_attnk, p_k, f_k, b.x_q8_0, b.x_q8k, b.kcur, g.n_embd, g.n // EVERY ENTRY POINT FROM HERE ON IS THE CAPTURABLE ONE: the per-token counts come from `st.step` and every // launch is sized from a capacity in `st`/`b`, so this sequence can be captured and replayed. The // host-scalar wrappers would be correct here today and silently wrong in a graph. -if (st.kv_hybrid) { - // K8V4: only V is rotated (kv_q4.hpp's H); the scores pair unrotated q with unrotated INT8 K, and the - // output - a mix of rotated values - is rotated back after attention. Each append/gather call folds the - // unused half's lanes onto the used pool (a bit-identical duplicate write), so no kernel variants exist. - strata::kernels::fwht256_inplace_cuda(b.vcur, g.n_head_kv, stream); - kv_append_q8_step(st.k_q, st.k_q, st.k_scale, st.k_scale, st.page_table, st.step, b.kcur, b.kcur, s, stream, nullptr); // mode 0: no host mirror - strata::kernels::kv_append_q4_step(st.v_q4, st.v_q4, st.page_table, st.step, b.vcur, b.vcur, s, stream, nullptr); -} else if (st.kv_q4) { - strata::kernels::fwht256_inplace_cuda(b.kcur, g.n_head_kv, stream); // Q4_0: rotated K and V (kv_q4.hpp) - strata::kernels::fwht256_inplace_cuda(b.vcur, g.n_head_kv, stream); - strata::kernels::kv_append_q4_step(st.k_q4, st.v_q4, st.page_table, st.step, b.kcur, b.vcur, s, stream, &st.host); -} else if (st.kv_int8) kv_append_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, st.step, b.kcur, b.vcur, s, stream, &st.host); else kv_append_step(st.k_pool, st.v_pool, st.page_table, st.step, b.kcur, b.vcur, s, stream, &st.host); { const uint64_t nvk = (uint64_t) g.n_head_kv * g.head_dim; const uint64_t base = (uint64_t) 2 * g.n_embd + 2 * g.hc + (uint64_t) g.n_head * g.head_dim + 2 * nvk + 8; if (!st.kv_int8 && !st.kv_hybrid) dump_slot(dump, g, layer, (const float*) st.k_pool, base, nvk / 2, stream); if (!st.kv_int8 && !st.kv_hybrid) dump_slot(dump, g, layer, (const float*) st.v_pool, base + nvk / 2, nvk / 2, stream); dump_slot(dump, g, layer, b.vcur, base + nvk, nvk, stream); dump_slot(dump, g, layer, b.kcur, base + 2 * nvk, nvk, stream); } { const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; if (native_qsa_indexer_enabled()) { +if (st.kv_int8) kv_append_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, st.step, b.kcur, b.vcur, s, stream); else kv_append_step(st.k_pool, st.v_pool, st.page_table, st.step, b.kcur, b.vcur, s, stream); { const uint64_t nvk = (uint64_t) g.n_head_kv * g.head_dim; const uint64_t base = (uint64_t) 2 * g.n_embd + 2 * g.hc + (uint64_t) g.n_head * g.head_dim + 2 * nvk + 8; if (!st.kv_int8) dump_slot(dump, g, layer, (const float*) st.k_pool, base, nvk / 2, stream); if (!st.kv_int8) dump_slot(dump, g, layer, (const float*) st.v_pool, base + nvk / 2, nvk / 2, stream); dump_slot(dump, g, layer, b.vcur, base + nvk, nvk, stream); dump_slot(dump, g, layer, b.kcur, base + 2 * nvk, nvk, stream); } { const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; if (native_qsa_indexer_enabled()) { try { native_qsa_indexer_append(b.idx_raw, st.step + kStepPos, pos_base, (const float*) w_ikn->data, RMS_EPS, ib, s, st.max_cells, (float) qsa_freq_base(), stream); @@ -929,26 +720,20 @@ if (!normalize_rotate(b.q_idx, w_iqn, (int) g.idx_q_heads, (int) g.idx_key_dim)) int64_t max_blocks = (st.max_cells / s.idx_block) + 2; if (g_fast_select) { // Plan v0.3 P7: block-level FP32 scores and a radix selection over blocks (qsa_select.hpp). - qsa_block_scores(st.idx_pooled, st.idx_dead, b.q_idx, st.step, 1, max_blocks, s, b.cell_scores, stream); + qsa_block_scores(st.idx_pooled, st.idx_dead, b.q_idx, st.step, 1, max_blocks, s, b.cell_scores, stream, + qsa_score_grid_blocks); qsa_block_topk(b.cell_scores, st.step, 1, max_blocks, cap, s, b.ids, stream); } else { qsa_index_step(st.idx_pooled, b.q_idx, nullptr, s, st.step, max_blocks, b.cell_scores, stream); topk_512_step(b.cell_scores, s, cap, st.step, b.ids, stream); } - // KV streaming: every block the selection names is made resident before anything reads it - qsa_kv_resolve(st, g, b.ids, st.step, 1, cap, stream); - if (st.kv_q4) strata::kernels::fwht256_inplace_cuda(b.qcur, g.n_head, stream); // = if (g_fast_attn && !native_flash_attn_short && dump == nullptr) { - const strata::kernels::QsaAttnPools pools = qsa_attn_pools(st); + strata::kernels::QsaAttnPools pools; + pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; pools.k_q = st.k_q; pools.v_q = st.v_q; + pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; pools.page_table = st.page_table; + if (!st.kv_int8) { pools.k_q = nullptr; pools.v_q = nullptr; } strata::kernels::qsa_decode_attn_step(b.qcur, pools, b.ids, st.step, cap, s, b.attn_scratch, b.attn, stream); } else { - if (st.kv_hybrid) { - kv_gather_q8_step(st.k_q, st.k_q, st.k_scale, st.k_scale, st.page_table, b.ids, st.step, cap, s, - b.k_scratch, b.k_scratch, stream); - strata::kernels::kv_gather_q4_step(st.v_q4, st.v_q4, st.page_table, b.ids, st.step, cap, s, - b.v_scratch, b.v_scratch, stream); - } - else if (st.kv_q4) strata::kernels::kv_gather_q4_step(st.k_q4, st.v_q4, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); - else if (st.kv_int8) kv_gather_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); else kv_gather_step(st.k_pool, st.v_pool, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); if (native_flash_attn_short) { + if (st.kv_int8) kv_gather_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); else kv_gather_step(st.k_pool, st.v_pool, st.page_table, b.ids, st.step, cap, s, b.k_scratch, b.v_scratch, stream); if (native_flash_attn_short) { if (st.max_cells < 1 || st.max_cells > 256 || !st.attention_status || !st.host_step) { err = v.name("native_flash_attn") + ": short adapter requires context <=256 and persistent status storage"; return false; @@ -963,7 +748,6 @@ int64_t max_blocks = (st.max_cells / s.idx_block) + 2; } } else qsa_attend_step(b.qcur, b.k_scratch, b.v_scratch, st.step, cap, s, b.attn, nullptr, stream); } - if (st.kv_q4 || st.kv_hybrid) strata::kernels::fwht256_inplace_cuda(b.attn, g.n_head, stream); // the output back: H is self-inverse dump_slot(dump, g, layer, b.attn, (uint64_t) 2 * g.n_embd + 2 * g.hc, (uint64_t) g.n_head * g.head_dim, stream); { const uint64_t vs = (uint64_t) 2 * g.n_embd + 2 * g.hc + (uint64_t) g.n_head * g.head_dim + (uint64_t) 2 * g.n_head_kv * g.head_dim; dump_slot(dump, g, layer, (const float*) b.v_scratch, vs, (uint64_t) g.n_head_kv * g.head_dim / 2, stream); dump_slot(dump, g, layer, (const float*) b.ids, vs + (uint64_t) g.n_head_kv * g.head_dim / 2, 4, stream); dump_slot(dump, g, layer, (const float*) st.step, vs + (uint64_t) g.n_head_kv * g.head_dim / 2 + 4, 4, stream); } // ---- 9. Gate the attention output before its projection. Native CUDA keeps F32 // sigmoid/multiply arithmetic and uses Q8_1 for the native quantized projection. @@ -992,34 +776,25 @@ void doorbell_reset(const Doorbell& db) { if (db.h_flag) *(volatile uint32_t*) db.h_flag = 0; } // ================================ THE TWO ENDS OF A TOKEN ================================ -bool embed_row(const WeightTable& tables, const ModelGeometry& g, int64_t token, float* out_dev, - void* stream, std::string& err) { - if (const NativeEmbed* ne = native_embed()) { // plan v0.3 P6: the GGUF-form table (IQ model files) - if (token < 0 || out_dev == nullptr) { err = "embed_row: invalid token or output"; return false; } - ne->gather_one(token, out_dev, stream); - return true; - } +namespace { +// The pack's S2/S4/S8 embedding table, checked against its planes; null with `err` when it cannot be read. +const WeightRef* s_embedding(const WeightTable& tables, const ModelGeometry& g, std::string& err) { const WeightRef* w = tables.find(EMBEDDING_NAME); - if (w == nullptr) { err = "token_embd.weight is missing"; return false; } + if (w == nullptr) { err = "token_embd.weight is missing"; return nullptr; } if (w->code_bits != 2 && w->code_bits != 4 && w->code_bits != 8) { err = "embed_row: token_embd.weight is not an S2/S4/S8 tensor"; - return false; + return nullptr; } if (w->codebook_iq4nl) { err = "embed_row: an IQ4NL embedding is not supported"; - return false; + return nullptr; } - if (w->data == nullptr || out_dev == nullptr || w->ne0 <= 0 || w->ne1 <= 0 || + if (w->data == nullptr || w->ne0 <= 0 || w->ne1 <= 0 || w->ne0 != g.n_embd || w->group_elems <= 0 || w->ne0 % w->group_elems != 0 || w->ne0 % (8 / w->code_bits) != 0) { err = "embed_row: invalid embedding pointers, dimensions or group size"; - return false; - } - if (token < 0 || token >= w->ne1) { - err = "embed_row: token " + std::to_string(token) + " is outside 0.." + std::to_string(w->ne1 - 1); - return false; + return nullptr; } - const uint64_t row_codes = (uint64_t) (w->ne0 / (8 / w->code_bits)); const uint64_t row_groups = (uint64_t) (w->ne0 / w->group_elems); const uint64_t rows = (uint64_t) w->ne1; @@ -1034,9 +809,48 @@ bool embed_row(const WeightTable& tables, const ModelGeometry& g, int64_t token, w->has_offset != (w->offset_bytes != 0) || (w->has_offset && w->offset_bytes != w->scales_bytes)) { err = "embed_row: embedding planes do not match the loaded tensor shape"; - return false; + return nullptr; } + return w; +} +} // namespace +bool embed_rows(const WeightTable& tables, const ModelGeometry& g, const int32_t* tokens_dev, int64_t n, float* out_dev, + void* stream, std::string& err) { + if (n <= 0) return true; + if (tokens_dev == nullptr || out_dev == nullptr) { err = "embed_rows: invalid tokens or output"; return false; } + if (const NativeEmbed* ne = native_embed()) { + ne->gather_dev(tokens_dev, n, out_dev, stream); + return true; + } + const WeightRef* w = s_embedding(tables, g, err); + if (w == nullptr) return false; + const auto* codes = static_cast(w->data); + const auto* scales = reinterpret_cast(codes + w->codes_bytes); + const auto* offsets = w->has_offset ? reinterpret_cast(codes + w->codes_bytes + w->scales_bytes) + : nullptr; + strata::kernels::embedding_gather_dev(codes, scales, offsets, tokens_dev, (int) n, w->ne0, w->code_bits, w->code_bias, + w->group_elems, (uint64_t) (w->ne0 / (8 / w->code_bits)), + (uint64_t) (w->ne0 / w->group_elems), out_dev, stream); + return true; +} + +bool embed_row(const WeightTable& tables, const ModelGeometry& g, int64_t token, float* out_dev, + void* stream, std::string& err) { + if (const NativeEmbed* ne = native_embed()) { // plan v0.3 P6: the GGUF-form table (IQ model files) + if (token < 0 || out_dev == nullptr) { err = "embed_row: invalid token or output"; return false; } + ne->gather_one(token, out_dev, stream); + return true; + } + const WeightRef* w = s_embedding(tables, g, err); + if (w == nullptr) return false; + if (out_dev == nullptr) { err = "embed_row: invalid output"; return false; } + if (token < 0 || token >= w->ne1) { + err = "embed_row: token " + std::to_string(token) + " is outside 0.." + std::to_string(w->ne1 - 1); + return false; + } + const uint64_t row_codes = (uint64_t) (w->ne0 / (8 / w->code_bits)); + const uint64_t row_groups = (uint64_t) (w->ne0 / w->group_elems); const auto* codes = static_cast(w->data); const auto* scales = reinterpret_cast(codes + w->codes_bytes); const auto* offsets = w->has_offset ? reinterpret_cast(codes + w->codes_bytes + w->scales_bytes) @@ -1098,8 +912,10 @@ n += (uint64_t) g.hc * 4 * 2 + 32; // inject2, gr_rs n += q8k_bytes(g.n_embd); // head_q8k -n += strata::kernels::gr_workspace_bytes(s); return align_up16(n) + 256;} -uint64_t block_buffers_init(const ModelGeometry& g, void* base, BlockBuffers& b) { const strata::kernels::GrShapes s{g.n_embd, g.hc, g.hc_lr}; Cursor c{(uint8_t*) base}; b.R = c.take((uint64_t) g.hc * g.n_embd); b.mixed = c.take((uint64_t) g.n_embd); b.block_out = c.take((uint64_t) g.n_embd); b.inject = c.take((uint64_t) g.hc); b.inject2 = c.take((uint64_t) g.hc); b.gr_rs = c.take((uint64_t) g.hc); b.head_q8k = c.take_bytes(q8k_bytes(g.n_embd)); uint8_t* grw = c.take_bytes(strata::kernels::gr_workspace_bytes(s)); strata::kernels::gr_workspace_init(s, grw, b.gr); return c.used;} +n += strata::kernels::gr_workspace_bytes(s); +// gr_scratch +n += align_up16(strata::kernels::fused_gr_scratch_bytes()); return align_up16(n) + 256;} +uint64_t block_buffers_init(const ModelGeometry& g, void* base, BlockBuffers& b) { const strata::kernels::GrShapes s{g.n_embd, g.hc, g.hc_lr}; Cursor c{(uint8_t*) base}; b.R = c.take((uint64_t) g.hc * g.n_embd); b.mixed = c.take((uint64_t) g.n_embd); b.block_out = c.take((uint64_t) g.n_embd); b.inject = c.take((uint64_t) g.hc); b.inject2 = c.take((uint64_t) g.hc); b.gr_rs = c.take((uint64_t) g.hc); b.gr_scratch = (float*) c.take_bytes(strata::kernels::fused_gr_scratch_bytes()); b.head_q8k = c.take_bytes(q8k_bytes(g.n_embd)); uint8_t* grw = c.take_bytes(strata::kernels::gr_workspace_bytes(s)); strata::kernels::gr_workspace_init(s, grw, b.gr); return c.used;} // ---- THE HALF-LEVEL C1 ORACLE. `dump` is a HOST buffer of `2 * n_embd + 2 * hc + n_head * head_dim` floats // per layer; this enqueues one slice of it as a device-to-host copy. **IT RUNS INSIDE THE CAPTURE**, which is // the point: the copy becomes a node of that layer's graph and replays with it, so the layer functions stay one @@ -1149,15 +965,13 @@ bool block_layer_pre(const WeightTable& tables, const ModelGeometry& g, int64_t // the device half: `ple_block` reading `emb_dev`, and the history shift. const bool fused = g_fused_gr && stage_prefix == 0 && half == 0 && strata::kernels::fused_gr_supported(g.n_embd, g.hc, g.hc_lr); - // layer-1's FFN write has not been applied to R yet - unless a control vector follows it, which needs R - bool pending_ffn = fused && layer > 0 && !strata::kernels::cvec().covers(layer - 1); + bool pending_ffn = fused && layer > 0; // layer-1's FFN write has not been applied to R yet if (ple != nullptr && ple->ready() && layer == 1 && (half == 0 || half == 1)) { if (pending_ffn) { const strata::kernels::GrShapes gs0{g.n_embd, g.hc, g.hc_lr}; gr_write(bb.R, bb.block_out, bb.inject2, gs0, bb.R, stream); pending_ffn = false; } - const int64_t hcd = strata::kernels::NG_HC_DIM; strata::kernels::PleOut po; // Exports must not alias the block's internal workspace. The previous diagnostic views used a // different layout inside that workspace: exporting gated values overwrote normalized values @@ -1212,7 +1026,7 @@ st_begin(layer, 0, stream); fa.w_norm = (const float*) w_norm[0]->data; fa.w_down = (const uint16_t*) w_down[0]->data; fa.w_up = (const uint16_t*) w_up[0]->data; fa.w_inject = (const uint16_t*) w_inject[0]->data; fa.eps = RMS_EPS; fa.lo = bb.gr.lo; fa.rs = bb.gr_rs; fa.inject_out = bb.inject; fa.mixed = bb.mixed; - strata::kernels::fused_gr_read(fa, stream); + strata::kernels::fused_gr_read_multi(&fa, 1, bb.gr_scratch, stream); } else { gr_read(R, (const float*) w_norm[0]->data, (const uint16_t*) w_down[0]->data, (const uint16_t*) w_up[0]->data, (const uint16_t*) w_inject[0]->data, RMS_EPS, gs, bb.gr, bb.mixed, bb.inject, stream); } @@ -1235,7 +1049,7 @@ st_begin(layer, 3, stream); fa.w_norm = (const float*) w_norm[1]->data; fa.w_down = (const uint16_t*) w_down[1]->data; fa.w_up = (const uint16_t*) w_up[1]->data; fa.w_inject = (const uint16_t*) w_inject[1]->data; fa.eps = RMS_EPS; fa.lo = bb.gr.lo; fa.rs = bb.gr_rs; fa.inject_out = bb.inject2; fa.mixed = bb.mixed; - strata::kernels::fused_gr_read(fa, stream); + strata::kernels::fused_gr_read_multi(&fa, 1, bb.gr_scratch, stream); } else { gr_read(R, (const float*) w_norm[1]->data, (const uint16_t*) w_down[1]->data, (const uint16_t*) w_up[1]->data, (const uint16_t*) w_inject[1]->data, RMS_EPS, gs, bb.gr, bb.mixed, bb.inject, stream); } @@ -1285,23 +1099,10 @@ st_begin(layer, 5, stream); if (g_shared_early ? !moe_combine_parts(g, layer, k, mb, parts, bb.block_out, stream, err) : !moe_finish(tables, g, layer, k, mb, bb.mixed, parts, bb.block_out, stream, err)) return false; st_end(layer, 5, stream); dump_half(bb, g, layer, bb.block_out, (uint64_t) g.n_embd, g.n_embd, stream); dump_half(bb, g, layer, bb.inject, (uint64_t) 2 * g.n_embd + g.hc, g.hc, stream); st_begin(layer, 6, stream); - const bool steer = strata::kernels::cvec().covers(layer); // --control-vector-scaled: after this write - try { - if (!(g_fused_gr && strata::kernels::fused_gr_supported(g.n_embd, g.hc, g.hc_lr))) { - gr_write(bb.R, bb.block_out, bb.inject, gs, bb.R, stream); - if (steer) strata::kernels::cvec_apply(bb.R, layer, 1, g.hc * g.n_embd, nullptr, 0, nullptr, 0, false, stream); - } else if (layer == g.n_layers - 1) { - gr_write(bb.R, bb.block_out, bb.inject2, gs, bb.R, stream); // materialise R for the head - if (steer) strata::kernels::cvec_apply(bb.R, layer, 1, g.hc * g.n_embd, nullptr, 0, nullptr, 0, false, stream); - } else if (steer) { - // the write the next layer's fused read would have folded, then the vector - strata::kernels::cvec_apply(bb.R, layer, 1, g.hc * g.n_embd, bb.block_out, g.n_embd, bb.inject2, g.hc, - true, stream); - } - } catch (const std::exception& e) { - err = std::string("block_layer_post: ") + e.what(); - return false; - } + if (!(g_fused_gr && strata::kernels::fused_gr_supported(g.n_embd, g.hc, g.hc_lr))) + gr_write(bb.R, bb.block_out, bb.inject, gs, bb.R, stream); + else if (layer == g.n_layers - 1) + gr_write(bb.R, bb.block_out, bb.inject2, gs, bb.R, stream); // materialise R for the head st_end(layer, 6, stream); return true;} bool block_layer(const WeightTable& tables, const ModelGeometry& g, int64_t layer, int64_t pos, int32_t pos_base, const GdnBuffers& gb, const QsaState& qst, const QsaBuffers& qb, const MoEBuffers& mb, int64_t k, const BlockBuffers& bb, const float* parts, void* stream, std::string& err, const Doorbell* db, const PleRun* ple) { // The two halves back to back. `parts` must be THIS layer's experts and must be ready before the call - diff --git a/src/core/mtp.cpp b/src/core/mtp.cpp index 11d16f79c79..5b443cdbf7a 100644 --- a/src/core/mtp.cpp +++ b/src/core/mtp.cpp @@ -1,6 +1,5 @@ // src/core/mtp.cpp - see include/strata/core/mtp.hpp. #include "strata/core/mtp.hpp" -#include "strata/core/on_device.hpp" #include "strata/core/native_head.hpp" #include "strata/kernels/bf16_gemv.hpp" @@ -8,7 +7,6 @@ #include "strata/kernels/elementwise.hpp" #include "strata/kernels/fused_gr.hpp" #include "strata/kernels/gr.hpp" -#include "strata/kernels/kv_q4.hpp" #include "strata/kernels/kv_q8.hpp" #include "strata/kernels/native_moe.hpp" #include "strata/kernels/native_mmvq.hpp" @@ -29,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -69,43 +68,25 @@ strata::kernels::QsaShapes shapes_of(const ModelGeometry& g) { return s; } -// 64-bit seek/tell on a `FILE*`: `fseek`/`ftell` take a 32-bit `long` on Windows and would wrap past 2 GiB. -#if defined(_WIN32) -#define STRATA_FILE_SEEK64(f, o, w) _fseeki64((f), (long long) (o), (w)) -#define STRATA_FILE_TELL64(f) _ftelli64(f) -#else -#define STRATA_FILE_SEEK64(f, o, w) fseeko((f), (off_t) (o), (w)) -#define STRATA_FILE_TELL64(f) ftello(f) -#endif - bool read_file(const std::string& path, std::vector& out) { - // Loader fix (0.1.15+loaderfix.2): `ifstream::read` reaches the disk as 4095-byte reads on MSVC - the same - // split `load_experts_ranges` had - so the drafter's 111 MiB dense blob paid 4 KiB per operation on a cold - // start. `fread` passes a request bigger than the stream buffer straight to `ReadFile()`. - FILE* f = std::fopen(path.c_str(), "rb"); - if (f == nullptr) return false; - // A `FILE*` has no destructor and the short-read path below returns early: the guard closes on every path. - struct Closer { - FILE* f; - ~Closer() { if (f != nullptr) std::fclose(f); } - } closer{f}; - if (STRATA_FILE_SEEK64(f, 0, SEEK_END) != 0) return false; - const long long n = (long long) STRATA_FILE_TELL64(f); - if (n < 0 || STRATA_FILE_SEEK64(f, 0, SEEK_SET) != 0) return false; + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) return false; + const std::streamsize n = f.tellg(); + f.seekg(0); out.resize((size_t) n); - return n == 0 || std::fread(out.data(), 1, (size_t) n, f) == (size_t) n; + return (bool) f.read((char*) out.data(), n); } } // namespace MtpDrafter::~MtpDrafter() { if (cs_) cudaStreamSynchronize(cs_); - for (auto& e : prefill_exec_) if (e) cudaGraphExecDestroy(e); - for (auto& e : prefill_dev_exec_) if (e) cudaGraphExecDestroy(e); - if (pf_dev_) cudaFree(pf_dev_); for (auto& e : round_exec_) if (e) cudaGraphExecDestroy(e); for (auto& e : step_exec_) if (e) cudaGraphExecDestroy(e); if (cs_) cudaStreamDestroy(cs_); + if (side_) cudaStreamDestroy(side_); + if (sh_fork_) cudaEventDestroy(sh_fork_); + if (sh_join_) cudaEventDestroy(sh_join_); if (dense_) cudaFree(dense_); if (experts_) cudaFree(experts_); if (state_arena_) cudaFree(state_arena_); @@ -132,15 +113,11 @@ const void* MtpDrafter::q8(const char* name) const { bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, SessionState& ss, int max_t, std::string& err, int64_t window) { - cudaGetDevice(&device_); // a layer split's last stage on another GPU: the drafter lives there g_ = &g; ss_ = &ss; max_t_ = max_t; rt_dir_ = rt_dir; if (max_t < 1 || max_t > strata::kernels::kVerifyMaxT) { err = "mtp: max_t out of range"; return false; } - // Loader fix (0.1.15+loaderfix.2): the two reads below are the whole “drafter files” cost; reporting - // them apart from the rest of the stage is what makes the next regression visible. - const auto t_files = std::chrono::steady_clock::now(); // ---- the index and the dense weights { std::ifstream idx(rt_dir + "/dense.txt"); @@ -156,34 +133,34 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session } std::vector blob; if (!read_file(rt_dir + "/dense.bin", blob)) { err = "mtp: cannot read dense.bin"; return false; } - const cudaError_t alloc = cudaMalloc((void**) &dense_, blob.size()); - if (alloc != cudaSuccess) { - size_t free_bytes = 0, total_bytes = 0; - const cudaError_t info = cudaMemGetInfo(&free_bytes, &total_bytes); - err = "mtp: dense weights allocation failed (" + std::string(cudaGetErrorString(alloc)) + - "), requested " + std::to_string(blob.size() >> 20) + " MiB, CUDA0 free " + - (info == cudaSuccess ? std::to_string(free_bytes >> 20) + " MiB" : "unknown"); - return false; - } + if (cudaMalloc((void**) &dense_, blob.size()) != cudaSuccess) { err = "mtp: dense weights do not fit"; return false; } cudaMemcpy(dense_, blob.data(), blob.size(), cudaMemcpyHostToDevice); vram_ += blob.size(); } - // ---- the 512 routed experts, one blob each + // ---- the 512 routed experts, one blob each: raw GGUF blocks [gate | up | down] in the formats experts.txt names + // (the ggml type ids of gate/up and down), or without it the Q2_0 blob layout of cpu/expert.hpp { - const uint64_t bytes = (uint64_t) g.n_expert * strata::kernels::cpu::BLOB; - // Loader fix (0.1.15+loaderfix.2): each 64 MiB read below reached the disk as ~16k 4095-byte reads under - // MSVC's `basic_filebuf::xsgetn`, which is what made 675 MiB of drafter experts take minutes. - FILE* f = std::fopen((rt_dir + "/experts.bin").c_str(), "rb"); - if (f == nullptr) { err = "mtp: cannot open experts.bin"; return false; } - struct Closer { - FILE* f; - ~Closer() { if (f != nullptr) std::fclose(f); } - } closer{f}; + blob_ = strata::kernels::cpu::BLOB; + if (std::ifstream fmt(rt_dir + "/experts.txt"); fmt) { + int gu = -1, dn = -1; + fmt >> gu >> dn; + if (!strata::kernels::native_expert_supported(gu, dn, g.n_embd, g.n_ff)) { + err = "mtp: experts.txt names formats the GPU expert kernels lack"; + return false; + } + nat_ = strata::kernels::native_expert_layout(gu, dn, g.n_embd, g.n_ff); + blob_ = nat_.bytes; + } + const uint64_t bytes = (uint64_t) g.n_expert * blob_; + std::ifstream f(rt_dir + "/experts.bin", std::ios::binary | std::ios::ate); + if (!f) { err = "mtp: cannot open experts.bin"; return false; } + if ((uint64_t) f.tellg() != bytes) { err = "mtp: experts.bin does not match its format"; return false; } + f.seekg(0); if (cudaMalloc((void**) &experts_, bytes) != cudaSuccess) { err = "mtp: the 512 experts do not fit in VRAM"; return false; } std::vector chunk(64u << 20); for (uint64_t off = 0; off < bytes;) { const uint64_t n = std::min(chunk.size(), bytes - off); - if (std::fread(chunk.data(), 1, (size_t) n, f) != (size_t) n) { err = "mtp: experts.bin is truncated"; return false; } + if (!f.read((char*) chunk.data(), (std::streamsize) n)) { err = "mtp: experts.bin is truncated"; return false; } cudaMemcpy(experts_ + off, chunk.data(), n, cudaMemcpyHostToDevice); off += n; } @@ -197,31 +174,9 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session // ---- the layer's own K/V (dense attention: no indexer state is read) const strata::kernels::QsaShapes s = shapes_of(g); const int64_t max_cells = ss.qsa_states[0].max_cells; - // KV streaming: the drafter only reads its last `window` cells, so with streaming on its K/V is a ring of the - // window (plus the cells a round writes ahead of its queries) over a host copy, refilled on a resume. The host copy - // is pinned after the expert arena has pinned what it could: if it does not fit, the K/V stays whole in VRAM. - int64_t ring = (window > 0 && window < max_cells) ? window + 4 * (int64_t) max_t + 64 : 0; - // K8V4 never applies to the drafter: its own attention paths (below, and verify.cpp) handle whole formats - // only, whatever ring shape it takes (0, a window, or the -1 fully-resident fallback). - const bool kv_hybrid_was = qsa_kv_hybrid(); - const bool kv_int8_was = qsa_kv_int8(); - qsa_set_kv_hybrid(false); - if (kv_hybrid_was) qsa_set_kv_int8(true); // the drafter under --kv k8v4: plain INT8 - uint64_t sb = qsa_state_bytes(g, max_cells, false, ring); + const uint64_t sb = qsa_state_bytes(g, max_cells, false); if (cudaMalloc(&state_arena_, sb) != cudaSuccess) { err = "mtp: the K/V state does not fit"; return false; } - if (qsa_state_init(g, max_cells, state_arena_, st_, &ss.qsa_states[0], ring) == 0) { - if (st_.kv_mode == 0) { err = "mtp: state init failed"; return false; } - std::fprintf(stderr, "strata mtp: no pinned RAM left for the draft layer's K/V copy; keeping it in VRAM\n"); - cudaGetLastError(); - cudaFree(state_arena_); - st_ = QsaState{}; - ring = -1; // fully resident - sb = qsa_state_bytes(g, max_cells, false, ring); - if (cudaMalloc(&state_arena_, sb) != cudaSuccess) { err = "mtp: the K/V state does not fit"; return false; } - if (qsa_state_init(g, max_cells, state_arena_, st_, &ss.qsa_states[0], ring) == 0) { err = "mtp: state init failed"; return false; } - } - qsa_set_kv_int8(kv_int8_was); - qsa_set_kv_hybrid(kv_hybrid_was); + if (qsa_state_init(g, max_cells, state_arena_, st_, &ss.qsa_states[0]) == 0) { err = "mtp: state init failed"; return false; } qsa_state_zero(st_, g, nullptr); cudaDeviceSynchronize(); vram_ += sb; @@ -248,11 +203,11 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session hn_ = b.take(T * HC * N); h2_ = b.take(T * HC * N); mixed_ = b.take(T * N); inj_ = b.take(T * HC); inj2_ = b.take(T * HC); lo_ = b.take(T * (uint64_t) g.hc_lr); rs_ = b.take(T * HC); bo_ = b.take(T * N); - xn_ = b.take(T * HC * N); + grs_ = (float*) b.take(strata::kernels::fused_gr_scratch_bytes()); xq_ = b.take(strata::kernels::native_q8_1_bytes((int) (NH * HD), 8)); qfull_ = b.take(T * NH * 2 * HD); qcur_ = b.take(T * NH * HD); kcur_ = b.take(T * NKV * HD); vcur_ = b.take(T * NKV * HD); - attn_ = b.take(T * NH * HD); attn32_ = b.take(T * NH * HD); + attn32_ = b.take(T * NH * HD); attn_scratch_ = b.take((uint64_t) attn_scratch_floats_); // the full layer runs one row at a time logits_ = b.take(T * (uint64_t) g.n_expert); w_ = b.take(T * K); ids_ = b.take(T * K); shared_ = b.take(T * N); parts_ = b.take(T * K * N); y_ = b.take(T * N); @@ -261,11 +216,18 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session grp_ptr_ = b.take(T * K); grp_start_ = b.take(T * K + 1); grp_counts_ = b.take(4); hit_xq_ = b.take(T * (N / 32) * 34); hit_xs_ = b.take(T * (N / 32)); - hit_scratch_ = b.take(strata::kernels::moe_hit_grouped_scratch_bytes((int64_t) (T * K), g.n_embd, g.n_ff)); + nat_xq_ = b.take(T * (N / 32) * 36); + rseq_ = b.take(4); + rcount_ = b.take(4); + hit_scratch_ = b.take(std::max( + strata::kernels::moe_hit_grouped_scratch_bytes((int64_t) (T * K), g.n_embd, g.n_ff), + strata::kernels::native_expert_scratch_bytes((int64_t) (T * K), g.n_ff))); sh_scratch_ = (float*) b.take(strata::kernels::shared_expert_scratch_bytes(g.n_ff)); x_bf16_ = b.take(N); out_ids_ = b.take(T + 4); probs_ = b.take(T + 4); + arg_scratch_ = b.take(strata::kernels::argmax_rows_scratch_bytes((int) T)); + top_scratch_ = b.take(strata::kernels::row_top_prob_scratch_bytes((int) T)); dummy_inj_ = b.take(HC); }; Bump count; @@ -282,18 +244,22 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session for (int64_t i = 0; i < cap_; ++i) id[(size_t) (t * (uint64_t) cap_ + (uint64_t) i)] = (int32_t) i; cudaMemcpy(ident_, id.data(), id.size() * 4, cudaMemcpyHostToDevice); } - if (cudaStreamCreateWithFlags(&cs_, cudaStreamNonBlocking) != cudaSuccess) { err = "mtp: stream"; return false; } - const double files_s = std::chrono::duration(std::chrono::steady_clock::now() - t_files).count(); - std::fprintf(stderr, "strata mtp: draft layer loaded, %.0f MiB of VRAM (experts %.0f, dense %.0f), files read in %.2f s (%.0f MiB/s)\n", - (double) vram_ / 1048576.0, (double) g.n_expert * strata::kernels::cpu::BLOB / 1048576.0, - (double) tensors_.back().off / 1048576.0, files_s, - files_s > 0 ? ((double) g.n_expert * strata::kernels::cpu::BLOB + (double) tensors_.back().off) / - 1048576.0 / files_s : 0.0); + if (cudaStreamCreateWithFlags(&cs_, cudaStreamNonBlocking) != cudaSuccess || + cudaStreamCreateWithFlags(&side_, cudaStreamNonBlocking) != cudaSuccess || + cudaEventCreateWithFlags(&sh_fork_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&sh_join_, cudaEventDisableTiming) != cudaSuccess) { + err = "mtp: streams"; + return false; + } + char fmts[48] = "Q2_0 blobs"; + if (nat_.gu_type >= 0) std::snprintf(fmts, sizeof fmts, "ggml types %d/%d", nat_.gu_type, nat_.d_type); + std::fprintf(stderr, "strata mtp: draft layer loaded, %.0f MiB of VRAM (experts %.0f as %s, dense %.0f)\n", + (double) vram_ / 1048576.0, (double) g.n_expert * blob_ / 1048576.0, fmts, + (double) tensors_.back().off / 1048576.0); return true; } bool MtpDrafter::bind(const WeightTable& wt, const NativeHead* head, const float* window_R, std::string& err) { - const OnDevice on_device(device_); wt_ = &wt; head_ = head; window_R_ = window_R; @@ -301,11 +267,24 @@ bool MtpDrafter::bind(const WeightTable& wt, const NativeHead* head, const float if (!wo) { err = "mtp: output.weight is missing"; return false; } n_vocab_ = wo->ne1; if (head == nullptr || !head->loaded()) { err = "mtp: the draft layer needs the native head (--native)"; return false; } + if (dhead_ == nullptr && !bind_head(head, err)) return false; + if (dhead_ == nullptr && head->rows() != n_vocab_) { + err = "mtp: the main head holds part of the vocabulary and the draft layer has no draft head of its own"; + return false; + } if (head_logits_ == nullptr && cudaMalloc((void**) &head_logits_, (size_t) max_t_ * (size_t) n_vocab_ * sizeof(float)) != cudaSuccess) { err = "mtp: the draft logits do not fit"; return false; } + return true; +} + +bool MtpDrafter::bind_head(const NativeHead* head, std::string& err) { + if (head == nullptr || !head->loaded()) { + err = "mtp: the draft layer needs the native head (--native)"; + return false; + } // the draft head's token subset, when tools/draft_vocab.py wrote one if (dhead_ == nullptr) { std::vector raw; @@ -330,31 +309,23 @@ bool MtpDrafter::bind(const WeightTable& wt, const NativeHead* head, const float // The layer for T rows. full = false stops after the K/V append (the prompt only needs the cache). bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::string& err) { - using namespace strata::kernels; - const ModelGeometry& g = *g_; - SessionState& ss = *ss_; // step_row0 >= 0: the full layer on step rows [step_row0, +T); step_row0 < 0: K/V only on rows [-1 - step_row0, +T) const bool full = step_row0 >= 0; - const int row0 = full ? step_row0 : -1 - step_row0; if (full && T != 1) { err = "mtp: the full layer runs one row at a time (its attention scratch is sized for one)"; return false; } - const int64_t N = g.n_embd, HC = g.hc, K = ss.k, NH = g.n_head, HD = g.head_dim, NKV = g.n_head_kv; + if (!record_front(T, full ? step_row0 : -1 - step_row0, cs, err)) return false; + return !full || record_rest(step_row0, cs, err); +} + +bool MtpDrafter::record_front(int T, int row0, cudaStream_t cs, std::string& err) { + using namespace strata::kernels; + const ModelGeometry& g = *g_; + const int64_t N = g.n_embd, HC = g.hc, NH = g.n_head, HD = g.head_dim, NKV = g.n_head_kv; const QsaShapes s = shapes_of(g); - const GrShapes gs{g.n_embd, g.hc, g.hc_lr}; const int32_t* step = step_ + row0 * 4; const int32_t* pos = pos_ + row0 * NH; try { // ---- the two input branches - const WeightRef* we = wt_->find("token_embd.weight"); - if (!we) { err = "mtp: token_embd.weight is missing"; return false; } - if (const NativeEmbed* ne = native_embed()) { // plan v0.3 P6: the GGUF-form table - ne->gather_dev(tok_, T, emb_, cs); - } else { - const auto* codes = (const uint8_t*) we->data; - const auto* scales = (const float*) (codes + we->codes_bytes); - const auto* offsets = we->has_offset ? (const float*) (codes + we->codes_bytes + we->scales_bytes) : nullptr; - embedding_gather_dev(codes, scales, offsets, tok_, T, we->ne0, we->code_bits, we->code_bias, we->group_elems, - (uint64_t) (we->ne0 / (8 / we->code_bits)), (uint64_t) (we->ne0 / we->group_elems), emb_, cs); - } + if (!embed_rows(*wt_, g, tok_, T, emb_, cs, err)) return false; native_qsa_rms_norm_weighted(emb_, f32("pre_fc_norm_embedding.weight"), en_, (int) N, T, EPS, cs); native_quantize_q8_1(en_, xq_, (int) N, T, cs); native_mmvq(GGML_Q8_0, q8("fc_embedding.weight"), xq_, e2_, (int) N, (int) N, T, cs); @@ -377,32 +348,48 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri fa[t].eps = EPS; fa[t].lo = lo_ + t * g.hc_lr; fa[t].rs = rs_ + t * HC; fa[t].inject_out = inj_ + t * HC; fa[t].mixed = mixed_ + t * N; } - fused_gr_read_multi(fa, T, xn_, cs); + fused_gr_read_multi(fa, T, grs_, cs); } - // ---- attention: K/V into the layer's own cache, then (full) dense attention over every cell - auto norm_rope = [&](float* data, const float* gamma, int rows, int cols, const int32_t* p) { - native_qsa_rms_norm_weighted(data, gamma, data, cols, rows, EPS, cs); - if (native_rope_enabled()) native_rope_apply(data, data, rows, cols, (int) s.n_rot, (float) qsa_freq_base(), p, cs); - else rope_neox_apply(data, data, rows, cols, (int) s.n_rot, st_.cos_tab, st_.sin_tab, p, cs); - }; + // ---- attention: K/V into the layer's own cache native_quantize_q8_1(mixed_, xq_, (int) N, T, cs); native_mmvq(GGML_Q8_0, q8("self_attn.k_proj.weight"), xq_, kcur_, (int) N, (int) (NKV * HD), T, cs); native_mmvq(GGML_Q8_0, q8("self_attn.v_proj.weight"), xq_, vcur_, (int) N, (int) (NKV * HD), T, cs); for (int t = 0; t < T; ++t) { - norm_rope(kcur_ + t * NKV * HD, f32("self_attn.k_norm.weight"), (int) NKV, (int) HD, pos + t * NH); - if (st_.kv_q4) { // Q4_0 KV (kv_q4.hpp): rotated K and V - fwht256_inplace_cuda(kcur_ + t * NKV * HD, NKV, cs); - fwht256_inplace_cuda(vcur_ + t * NKV * HD, NKV, cs); - kv_append_q4_step(st_.k_q4, st_.v_q4, st_.page_table, step + t * 4, kcur_ + t * NKV * HD, - vcur_ + t * NKV * HD, s, cs, &st_.host); - } else if (st_.kv_int8) + norm_rope(kcur_ + t * NKV * HD, f32("self_attn.k_norm.weight"), (int) NKV, (int) HD, pos + t * NH, cs); + if (st_.kv_int8) kv_append_q8_step(st_.k_q, st_.v_q, st_.k_scale, st_.v_scale, st_.page_table, step + t * 4, - kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs, &st_.host); + kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs); else kv_append_step(st_.k_pool, st_.v_pool, st_.page_table, step + t * 4, kcur_ + t * NKV * HD, - vcur_ + t * NKV * HD, s, cs, &st_.host); + vcur_ + t * NKV * HD, s, cs); } - if (!full) return true; + } catch (const std::exception& e) { + err = std::string("mtp: ") + e.what(); + return false; + } + return true; +} + +void MtpDrafter::norm_rope(float* data, const float* gamma, int rows, int cols, const int32_t* p, cudaStream_t cs) { + using namespace strata::kernels; + const QsaShapes s = shapes_of(*g_); + native_qsa_rms_norm_weighted(data, gamma, data, cols, rows, EPS, cs); + if (native_rope_enabled()) native_rope_apply(data, data, rows, cols, (int) s.n_rot, (float) qsa_freq_base(), p, cs); + else rope_neox_apply(data, data, rows, cols, (int) s.n_rot, st_.cos_tab, st_.sin_tab, p, cs); +} + +bool MtpDrafter::record_rest(int step_row, cudaStream_t cs, std::string& err) { + using namespace strata::kernels; + const ModelGeometry& g = *g_; + SessionState& ss = *ss_; + const int T = 1; + const int64_t N = g.n_embd, HC = g.hc, K = ss.k, NH = g.n_head, HD = g.head_dim; + const QsaShapes s = shapes_of(g); + const GrShapes gs{g.n_embd, g.hc, g.hc_lr}; + const int32_t* step = step_ + step_row * 4; + const int32_t* pos = pos_ + step_row * NH; + try { + // ---- dense attention over every cell native_mmvq(GGML_Q8_0, q8("self_attn.q_proj.weight"), xq_, qfull_, (int) N, (int) (NH * 2 * HD), T, cs); for (int t = 0; t < T; ++t) { float* qc = qcur_ + t * NH * HD; @@ -411,15 +398,14 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri err = "mtp: q split failed"; return false; } - norm_rope(qc, f32("self_attn.q_norm.weight"), (int) NH, (int) HD, pos + t * NH); - if (st_.kv_q4) fwht256_inplace_cuda(qc, NH, cs); + norm_rope(qc, f32("self_attn.q_norm.weight"), (int) NH, (int) HD, pos + t * NH, cs); } - const QsaAttnPools pools = qsa_attn_pools(st_); + QsaAttnPools pools; + pools.page_table = st_.page_table; + if (st_.kv_int8) { pools.k_q = st_.k_q; pools.v_q = st_.v_q; pools.k_scale = st_.k_scale; pools.v_scale = st_.v_scale; } + else { pools.k_pool = st_.k_pool; pools.v_pool = st_.v_pool; } if (window_ > 0) window_ids(const_cast(step), T, (int) window_, ident_, cap_, cs); - qsa_decode_attn_batch(qcur_, pools, ident_, step, cap_, s, attn_scratch_, attn_, T, cs); - if (st_.kv_q4) fwht256_inplace_cuda(attn_, (int64_t) T * NH, cs); - for (int t = 0; t < T; ++t) - native_qsa_gate_apply(attn_ + t * NH * HD, qfull_ + t * NH * 2 * HD, attn32_ + t * NH * HD, (int) NH, (int) HD, cs); + qsa_decode_attn_batch(qcur_, pools, ident_, step, cap_, s, attn_scratch_, attn32_, T, cs, qfull_); native_quantize_q8_1(attn32_, xq_, (int) (NH * HD), T, cs); native_mmvq(GGML_Q8_0, q8("self_attn.o_proj.weight"), xq_, bo_, (int) (NH * HD), (int) N, T, cs); // ---- the MLP hyper-connection (the attention write folded in) @@ -435,30 +421,56 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri fa[t].eps = EPS; fa[t].lo = lo_ + t * g.hc_lr; fa[t].rs = rs_ + t * HC; fa[t].inject_out = inj2_ + t * HC; fa[t].mixed = mixed_ + t * N; } - fused_gr_read_multi(fa, T, xn_, cs); + fused_gr_read_multi(fa, T, grs_, cs); } - // ---- MoE: router, the 512 resident experts, the shared expert, the combine, the write - for (int t = 0; t < T; ++t) { - bf16_gemv_fp32_mmvf(mixed_ + t * N, bf16("mlp.gate.weight"), logits_ + t * g.n_expert, (int) N, (int) g.n_expert, cs); - if (native_router_enabled()) native_router_top10(logits_ + t * g.n_expert, ids_ + t * K, w_ + t * K, cs); - else router_top10(logits_ + t * g.n_expert, 1, (int) g.n_expert, (int) K, ids_ + t * K, w_ + t * K, cs); - } - moe_group_resident(ids_, (int) (T * K), (int) K, experts_, (int64_t) strata::kernels::cpu::BLOB, grp_ptr_, - grp_start_, grp_counts_, hit_dst_, hit_slot_, cs); - quantize_q8_0_scaled(mixed_, hit_xq_, hit_xs_, (int64_t) T * N, cs); - moe_grouped_s2(grp_ptr_, grp_start_, grp_counts_, hit_dst_, hit_slot_, (int64_t) T * K, (int64_t) T * K, hit_xq_, - hit_xs_, hit_scratch_, parts_, cs); + // ---- MoE: the shared expert on a branch beside the router and the 512 resident experts, then the combine + // and the write NativeSharedWeights nsw; nsw.gate_type = GGML_Q8_0; nsw.gate_data = q8("mlp.shared_expert.gate_proj.weight"); nsw.up_type = GGML_Q8_0; nsw.up_data = q8("mlp.shared_expert.up_proj.weight"); nsw.down_type = GGML_Q8_0; nsw.down_data = q8("mlp.shared_expert.down_proj.weight"); - nsw.q8_1 = xq_; + nsw.q8_1 = xq_; // the routed experts read nat_xq_ / hit_xq_ const SForm none{}; + if (cudaEventRecord(sh_fork_, cs) != cudaSuccess || cudaStreamWaitEvent(side_, sh_fork_, 0) != cudaSuccess) { + err = "mtp: the shared expert's branch"; + return false; + } for (int t = 0; t < T; ++t) { - f32_to_bf16_bulk(mixed_ + t * N, x_bf16_, N, cs); + f32_to_bf16_bulk(mixed_ + t * N, x_bf16_, N, side_); shared_expert(nullptr, nullptr, x_bf16_, none, nullptr, nullptr, nullptr, none, nullptr, nullptr, nullptr, none, nullptr, nullptr, nullptr, bf16("mlp.shared_expert_gate.weight"), sh_scratch_, shared_ + t * N, - N, g.n_ff, 32, cs, mixed_ + t * N, &nsw); + N, g.n_ff, 32, side_, mixed_ + t * N, &nsw); + } + if (cudaEventRecord(sh_join_, side_) != cudaSuccess) { err = "mtp: the shared expert's branch"; return false; } + bool xq1 = false; // the router kernel wrote the native experts' q8_1 rows + if (native_router_enabled()) { + VerifyRouterArgs ra; + ra.x = mixed_; ra.w = bf16("mlp.gate.weight"); ra.logits = logits_; ra.ids = ids_; ra.weights = w_; + ra.xq1 = nat_.gu_type >= 0 ? nat_xq_ : nullptr; + ra.seq = rseq_; ra.counter = rcount_; + ra.n_tok = T; ra.n_embd = (int) N; ra.n_expert = (int) g.n_expert; + verify_router(ra, cs); + xq1 = ra.xq1 != nullptr; + } else { + for (int t = 0; t < T; ++t) { + bf16_gemv_fp32_mmvf(mixed_ + t * N, bf16("mlp.gate.weight"), logits_ + t * g.n_expert, (int) N, + (int) g.n_expert, cs); + router_top10(logits_ + t * g.n_expert, 1, (int) g.n_expert, (int) K, ids_ + t * K, w_ + t * K, cs); + } + } + moe_group_resident(ids_, (int) (T * K), (int) K, experts_, (int64_t) blob_, grp_ptr_, grp_start_, grp_counts_, + hit_dst_, hit_slot_, cs); + if (nat_.gu_type >= 0) { + if (!xq1) quantize_q8_1_rows(mixed_, T, N, nat_xq_, cs); + native_expert_grouped(nat_, grp_ptr_, grp_start_, grp_counts_, hit_dst_, hit_slot_, (int64_t) T * K, + (int64_t) T * K, nat_xq_, hit_scratch_, parts_, cs); + } else { + quantize_q8_0_scaled(mixed_, hit_xq_, hit_xs_, (int64_t) T * N, cs); + moe_grouped_s2(grp_ptr_, grp_start_, grp_counts_, hit_dst_, hit_slot_, (int64_t) T * K, (int64_t) T * K, + hit_xq_, hit_xs_, hit_scratch_, parts_, cs); + } + if (cudaStreamWaitEvent(cs, sh_join_, 0) != cudaSuccess) { err = "mtp: the shared expert's join"; return false; } + for (int t = 0; t < T; ++t) { if (native_moe_combine_enabled()) native_moe_combine(parts_ + (size_t) t * K * N, w_ + t * K, shared_ + t * N, y_ + t * N, N, K, cs); else @@ -466,20 +478,23 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri gr_write(R_ + (size_t) t * HC * N, y_ + t * N, inj2_ + t * HC, gs, R_ + (size_t) t * HC * N, cs); } // ---- the final mixer and the main model's head - for (int t = 0; t < T; ++t) - gr_read(R_ + (size_t) t * HC * N, f32("hyper_connection_mixer.hc_norm.weight"), - bf16("hyper_connection_mixer.input_mix_weight_down.weight"), - bf16("hyper_connection_mixer.input_mix_weight_up.weight"), nullptr, EPS, gs, ss.block.gr, - sample_ + t * N, dummy_inj_, cs); + { + FusedGrArgs fa[kFusedGrMaxT]; + for (int t = 0; t < T; ++t) { + fa[t].R = R_ + (size_t) t * HC * N; fa[t].R_out = R_ + (size_t) t * HC * N; fa[t].apply = false; + fa[t].w_norm = f32("hyper_connection_mixer.hc_norm.weight"); + fa[t].w_down = bf16("hyper_connection_mixer.input_mix_weight_down.weight"); + fa[t].w_up = bf16("hyper_connection_mixer.input_mix_weight_up.weight"); + fa[t].eps = EPS; fa[t].lo = lo_ + t * g.hc_lr; fa[t].rs = rs_ + t * HC; fa[t].mixed = sample_ + t * N; + } + fused_gr_read_multi(fa, T, grs_, cs); + } native_quantize_q8_1(sample_, xq_, (int) N, T, cs); const bool sub = dhead_ != nullptr; const int64_t nv = sub ? n_dvocab_ : n_vocab_; native_mmvq(head_->type(), sub ? dhead_ : head_->weights(), xq_, head_logits_, (int) N, (int) nv, T, cs); - SamplerParams sp; - sp.greedy = true; - sp.temperature = 0.0f; - sample_tokens(head_logits_, T, (int) nv, nullptr, 0, sp, out_ids_, cs); - row_top_prob(head_logits_, T, (int) nv, out_ids_, probs_, cs); + argmax_rows(head_logits_, T, (int) nv, arg_scratch_, out_ids_, cs); + row_top_prob_split(head_logits_, T, (int) nv, out_ids_, probs_, top_scratch_, cs); if (sub) map_ids(out_ids_, dvocab_, T, cs); } catch (const std::exception& e) { err = std::string("mtp: ") + e.what(); @@ -509,45 +524,28 @@ bool finish_capture(cudaStream_t cs, bool ok, cudaGraphExec_t& exec, const char* } } // namespace -bool MtpDrafter::capture_prefill(int T, std::string& err) { - if (prefill_exec_[T]) return true; - using namespace strata::kernels; - if (cudaStreamBeginCapture(cs_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) { err = "mtp: begin capture"; return false; } - copy_i32_from_mapped(tok_, m_tok_, T, cs_); - copy_i32_from_mapped(step_, m_step_, (int64_t) T * 4, cs_); - copy_i32_from_mapped(pos_, m_pos_, (int64_t) T * g_->n_head, cs_); - const bool ok = record_forward(T, -1, cs_, err); // K/V only, rows [0, T) - return finish_capture(cs_, ok, prefill_exec_[T], "prefill", err); -} - -bool MtpDrafter::capture_prefill_dev(int T, std::string& err) { - if (prefill_dev_exec_[T]) return true; - if (cudaStreamBeginCapture(cs_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) { err = "mtp: begin capture"; return false; } - const bool ok = record_forward(T, -1, cs_, err); // K/V only, rows [0, T); tok_/step_/pos_ filled before launch - return finish_capture(cs_, ok, prefill_dev_exec_[T], "prefill (device inputs)", err); -} - bool MtpDrafter::capture_round(int T, std::string& err) { if (round_exec_[T]) return true; using namespace strata::kernels; const int64_t HCN = g_->hc * g_->n_embd; if (cudaStreamBeginCapture(cs_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) { err = "mtp: begin capture"; return false; } bool ok = true; - copy_i32_from_mapped(tok_, m_tok_, T, cs_); - copy_i32_from_mapped(step_, m_step_, (int64_t) 2 * T * 4, cs_); - copy_i32_from_mapped(pos_, m_pos_, (int64_t) 2 * T * g_->n_head, cs_); - copy_i32_from_mapped(row_, m_row_, 2, cs_); - copy_from_mapped(Rin_, window_R_, (int64_t) T * HCN, cs_); - // the catch-up: K/V for the window's T cells, then the full layer for row a only (its cell's K/V is written - // again, identically), staged by the host in step row 2*max_t - 1; the draft chain is one graph per step - // (`capture_step`) so the host can stop it when a draft is unlikely + // the catch-up: the layer's front for the window's T cells (their K/V), then its rest for row a only, on row + // a's intermediates copied to row 0, at the cell the host staged in step row 2*max_t - 1; the draft chain is + // one graph per step (`capture_step`) so the host can stop it when a draft is unlikely const int ra = 2 * max_t_ - 1; - ok = record_forward(T, -1, cs_, err); - if (ok) mtp_select(Rin_, HCN, tok_, row_, Rin_, tok_, nullptr, 0, cs_); + const int64_t N = g_->n_embd, NH = g_->n_head; + const MappedCopy in[7] = {{tok_, m_tok_, T}, {step_, m_step_, (int64_t) 2 * T * 4}, {pos_, m_pos_, 2 * T * NH}, + {row_, m_row_, 2}, {Rin_, window_R_, T * HCN}, {step_ + ra * 4, m_step_ + ra * 4, 4}, + {pos_ + ra * NH, m_pos_ + ra * NH, NH}}; + copy_from_mapped_multi(in, 7, cs_); + ok = record_front(T, 0, cs_, err); if (ok) { - copy_i32_from_mapped(step_ + ra * 4, m_step_ + ra * 4, 4, cs_); - copy_i32_from_mapped(pos_ + ra * g_->n_head, m_pos_ + ra * g_->n_head, g_->n_head, cs_); - ok = record_forward(1, ra, cs_, err); + if (T > 1) { + copy_row_to_first(row_, R_, HCN, inj_, g_->hc, mixed_, N, cs_); + native_quantize_q8_1(mixed_, xq_, (int) N, 1, cs_); + } + ok = record_rest(ra, cs_, err); } if (ok) mtp_select(R_, HCN, out_ids_, row_ + 1, Rin_, tok_, m_out_, 0, cs_, probs_, m_prob_); return finish_capture(cs_, ok, round_exec_[T], "round", err); @@ -561,115 +559,33 @@ bool MtpDrafter::capture_step(int j, std::string& err) { const int64_t HCN = g_->hc * g_->n_embd; const int row = max_t_ + j - 1; if (cudaStreamBeginCapture(cs_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) { err = "mtp: begin capture"; return false; } - copy_i32_from_mapped(step_ + row * 4, m_step_ + row * 4, 4, cs_); - copy_i32_from_mapped(pos_ + row * g_->n_head, m_pos_ + row * g_->n_head, g_->n_head, cs_); + const MappedCopy in[2] = {{step_ + row * 4, m_step_ + row * 4, 4}, + {pos_ + row * g_->n_head, m_pos_ + row * g_->n_head, g_->n_head}}; + copy_from_mapped_multi(in, 2, cs_); bool ok = record_forward(1, row, cs_, err); if (ok) mtp_select(R_, HCN, out_ids_, row_ + 1, Rin_, tok_, m_out_, j, cs_, probs_, m_prob_); return finish_capture(cs_, ok, step_exec_[j], "step", err); } -void MtpDrafter::kv_restore(int64_t upto) { - const OnDevice on_device(device_); - if (st_.kv_mode != 2 || upto <= 0) return; - // the ring's blocks below `upto`, from the host copy: a checkpoint resume may have left later cells in them - const strata::kernels::QsaShapes s = shapes_of(*g_); - const int64_t b1 = (upto + s.page_size - 1) / s.page_size, b0 = std::max(0, b1 - st_.n_slots); - strata::kernels::kv_ring_restore(qsa_attn_pools(st_), st_.host, qsa_kv_format(st_), b0, b1, st_.n_slots, s, cs_); - cudaStreamSynchronize(cs_); -} - -bool MtpDrafter::prefill(const float* R_rows, const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err) { - const OnDevice on_device(device_); - const Clock::time_point t0 = Clock::now(); - const int64_t HCN = g_->hc * g_->n_embd; - // cells the window can never reach again need no K/V - const int64_t first_needed = (window_ > 0 && prompt_len_ > 0) ? prompt_len_ - window_ - 64 : 0; - // E-4: every group's token / step / position records uploaded at once; each group is then device copies and a - // graph on the one stream, with a single sync at the end (a group of <= max_t rows used to be staged in mapped - // memory and synced before the next: ~5,500 host round trips on a 32K prompt). The same work in the same - // order. STRATA_MTP_PREFILL_SYNC=1 keeps the old loop. - static const bool per_group_sync = std::getenv("STRATA_MTP_PREFILL_SYNC") != nullptr; - const int64_t NHp = g_->n_head, per_row = 1 + 4 + NHp; - if (!per_group_sync && n > 0) { - if (pf_cap_ < n * per_row) { - if (pf_dev_) cudaFree(pf_dev_); - pf_dev_ = nullptr; - pf_cap_ = 0; - if (cudaMalloc((void**) &pf_dev_, (size_t) (n * per_row) * sizeof(int32_t)) != cudaSuccess) { - err = "mtp prefill: the input records do not fit"; - return false; - } - pf_cap_ = n * per_row; - } - std::vector rec((size_t) (n * per_row)); - int32_t* tk = rec.data(); - int32_t* stp = tk + n; - int32_t* ps = stp + 4 * n; - for (int64_t i = 0; i < n; ++i) { - const int64_t cell = cell0 + i; - tk[i] = next_tokens[i]; - stp[i * 4 + 0] = (int32_t) cell; - stp[i * 4 + 1] = (int32_t) (cell + 1); - stp[i * 4 + 2] = (int32_t) ((cell + 1) / 4); - stp[i * 4 + 3] = (int32_t) (cell + 1); - for (int64_t h = 0; h < NHp; ++h) ps[i * NHp + h] = (int32_t) cell; - } - if (cudaMemcpyAsync(pf_dev_, rec.data(), rec.size() * sizeof(int32_t), cudaMemcpyHostToDevice, cs_) != cudaSuccess) { - err = "mtp prefill: the input records upload failed"; - return false; - } - const int32_t* d_tk = pf_dev_; - const int32_t* d_stp = d_tk + n; - const int32_t* d_ps = d_stp + 4 * n; - for (int64_t c = 0; c < n; c += max_t_) { - const int T = (int) std::min(max_t_, n - c); - if (cell0 + c + T <= first_needed) continue; - if (!capture_prefill_dev(T, err)) return false; - if (cudaMemcpyAsync(tok_, d_tk + c, (size_t) T * 4, cudaMemcpyDeviceToDevice, cs_) != cudaSuccess || - cudaMemcpyAsync(step_, d_stp + c * 4, (size_t) T * 16, cudaMemcpyDeviceToDevice, cs_) != cudaSuccess || - cudaMemcpyAsync(pos_, d_ps + c * NHp, (size_t) (T * NHp) * 4, cudaMemcpyDeviceToDevice, cs_) != cudaSuccess || - cudaMemcpyAsync(Rin_, R_rows + (size_t) c * HCN, (size_t) T * HCN * sizeof(float), cudaMemcpyDeviceToDevice, - cs_) != cudaSuccess || - cudaGraphLaunch(prefill_dev_exec_[T], cs_) != cudaSuccess) { - err = std::string("mtp prefill: ") + cudaGetErrorString(cudaGetLastError()); - return false; - } - } - if (cudaStreamSynchronize(cs_) != cudaSuccess) { - err = std::string("mtp prefill: ") + cudaGetErrorString(cudaGetLastError()); - return false; - } - ms_prefill += ms_since(t0); - return true; - } - for (int64_t c = 0; c < n; c += max_t_) { - const int T = (int) std::min(max_t_, n - c); - if (cell0 + c + T <= first_needed) continue; - if (!capture_prefill(T, err)) return false; - for (int t = 0; t < T; ++t) { - const int64_t cell = cell0 + c + t; - h_tok_[t] = next_tokens[c + t]; - h_step_[t * 4 + 0] = (int32_t) cell; - h_step_[t * 4 + 1] = (int32_t) (cell + 1); - h_step_[t * 4 + 2] = (int32_t) ((cell + 1) / 4); - h_step_[t * 4 + 3] = (int32_t) (cell + 1); - for (int64_t h = 0; h < g_->n_head; ++h) h_pos_[t * g_->n_head + h] = (int32_t) cell; - } - if (cudaMemcpyAsync(Rin_, R_rows + (size_t) c * HCN, (size_t) T * HCN * sizeof(float), cudaMemcpyDeviceToDevice, - cs_) != cudaSuccess || - cudaGraphLaunch(prefill_exec_[T], cs_) != cudaSuccess || - cudaStreamSynchronize(cs_) != cudaSuccess) { - err = std::string("mtp prefill: ") + cudaGetErrorString(cudaGetLastError()); - return false; - } - } - ms_prefill += ms_since(t0); - return true; +MtpPromptKv MtpDrafter::prompt_kv() { + MtpPromptKv k; + k.norm_emb = f32("pre_fc_norm_embedding.weight"); + k.norm_hidden = f32("pre_fc_norm_hidden.weight"); + k.fc_emb = q8("fc_embedding.weight"); + k.fc_hidden = q8("fc_hidden.weight"); + k.hc_norm = f32("attn_hyper_connection.hc_norm.weight"); + k.hc_down = bf16("attn_hyper_connection.input_mix_weight_down.weight"); + k.hc_up = bf16("attn_hyper_connection.input_mix_weight_up.weight"); + k.k_proj = q8("self_attn.k_proj.weight"); + k.v_proj = q8("self_attn.v_proj.weight"); + k.k_norm = f32("self_attn.k_norm.weight"); + k.kv = &st_; + k.first_cell = (window_ > 0 && prompt_len_ > 0) ? prompt_len_ - window_ - 64 : 0; + return k; } bool MtpDrafter::draft(int T, const int32_t* tokens, int64_t p, int a, int32_t* drafts, std::string& err, float* probs, float min_p, int* n_drafts) { - const OnDevice on_device(device_); if (T < 1 || T > max_t_ || a < 0 || a >= T) { err = "mtp: draft arguments out of range"; return false; } if (!capture_round(T, err)) return false; const Clock::time_point t0 = Clock::now(); @@ -696,9 +612,11 @@ bool MtpDrafter::draft(int T, const int32_t* tokens, int64_t p, int a, int32_t* drafts[0] = ((volatile int32_t*) h_out_)[0]; float pj = ((volatile float*) h_prob_)[0]; if (probs) probs[0] = pj; + if (on_draft) on_draft(0, drafts[0], pj); int n = 1; // the chain continues while the last draft is likely enough to be verified - for (int j = 1; j < std::min(max_t_ - 1, max_drafts_) && pj >= min_p; ++j) { + const int chain = max_drafts_ > 0 ? std::min(max_drafts_, max_t_ - 1) : max_t_ - 1; + for (int j = 1; j < chain && pj >= min_p; ++j) { if (!capture_step(j, err)) return false; put(max_t_ + j - 1, p + a + j); std::atomic_thread_fence(std::memory_order_seq_cst); @@ -709,6 +627,7 @@ bool MtpDrafter::draft(int T, const int32_t* tokens, int64_t p, int a, int32_t* drafts[j] = ((volatile int32_t*) h_out_)[j]; pj = ((volatile float*) h_prob_)[j]; if (probs) probs[j] = pj; + if (on_draft) on_draft(j, drafts[j], pj); ++n; } for (int j = n; j < max_t_ - 1; ++j) { drafts[j] = 0; if (probs) probs[j] = 0.0f; } @@ -720,7 +639,6 @@ bool MtpDrafter::draft(int T, const int32_t* tokens, int64_t p, int a, int32_t* bool MtpDrafter::draft_first(int T, const float* R_row, int32_t token, int64_t cell, int32_t* drafts, std::string& err, float* probs, float min_p, int* n_drafts) { - const OnDevice on_device(device_); // row 0 is the real pair; rows 1.. repeat it and only write cells the next round overwrites const int64_t HCN = g_->hc * g_->n_embd; for (int t = 0; t < T; ++t) diff --git a/src/core/native_dense.cpp b/src/core/native_dense.cpp index ba3099897d5..980ec434dd5 100644 --- a/src/core/native_dense.cpp +++ b/src/core/native_dense.cpp @@ -13,12 +13,9 @@ namespace strata::core { namespace { -bool eligible(const strata::TensorInfo& tensor, bool include_ple_key) { - const auto& name = tensor.name; +bool eligible(const std::string& name, bool include_ple_key) { if (name.rfind("blk.", 0) != 0) return false; - // Match the native PLE kernel: Q2_0, IQ3_XXS and IQ4_XS. Other keys retain the packed BF16 fallback. - if (name == "blk.1.ple_key.weight") - return include_ple_key && (tensor.type == 42 || tensor.type == 18 || tensor.type == 23); + if (include_ple_key && name == "blk.1.ple_key.weight") return true; static const char* suffixes[] = {".attn_qkv.weight", ".attn_gate.weight", ".ssm_out.weight", ".attn_q.weight", ".attn_k.weight", ".attn_v.weight", ".attn_output.weight", ".ffn_gate_shexp.weight", ".ffn_up_shexp.weight", ".ffn_down_shexp.weight"}; @@ -41,7 +38,7 @@ bool NativeDense::served_names(const std::vector& shards, bool incl for (const auto& path : shards) { strata::GgufFile gguf(path); for (const auto& tensor : gguf.tensors()) - if (eligible(tensor, include_ple_key) && strata::kernels::native_mmvq_supported(tensor.type) && + if (eligible(tensor.name, include_ple_key) && strata::kernels::native_mmvq_supported(tensor.type) && tensor.shape.size() == 2) out.insert(tensor.name); } @@ -127,7 +124,7 @@ bool NativeDense::load(const std::vector& shards, WeightTable& tabl } } for (const auto& tensor : gguf.tensors()) { - if (!eligible(tensor, include_ple_key)) continue; + if (!eligible(tensor.name, include_ple_key)) continue; if (!seen.insert(tensor.name).second) { err = "native dense: duplicate tensor " + tensor.name; return false; } diff --git a/src/core/native_head.cpp b/src/core/native_head.cpp index fc9f694ffcc..e77dff0d753 100644 --- a/src/core/native_head.cpp +++ b/src/core/native_head.cpp @@ -15,22 +15,37 @@ NativeHead::~NativeHead() { if (weights_) cudaFree(weights_); } -bool NativeHead::load(const std::string& path, int64_t n_in, int64_t n_out, std::string& err) { +bool NativeHead::keep_rows(int64_t rows, std::string& err) { + if (!loaded() || rows <= 0 || rows >= n_out_) { err = "native head: keep_rows out of range"; return false; } + const uint64_t bytes = (uint64_t) rows * row_bytes(); + void* w = nullptr; + if (cudaMalloc(&w, bytes) != cudaSuccess || + cudaMemcpy(w, weights_, bytes, cudaMemcpyDeviceToDevice) != cudaSuccess) { + if (w) cudaFree(w); + err = "native head: keep_rows: " + std::string(cudaGetErrorString(cudaGetLastError())); + return false; + } + cudaFree(weights_); + weights_ = w; + bytes_ = bytes; + n_out_ = (int) rows; + part_ = true; + return true; +} + +bool NativeHead::load(const std::vector& shards, int64_t n_in, int64_t n_out, std::string& err) { if (loaded()) { err = "native head is already loaded"; return false; } if (n_in <= 0 || n_out <= 0 || n_in > INT_MAX || n_out > INT_MAX || n_in % 256) { err = "native head requires positive int32 dimensions and whole 256-value rows"; return false; } try { - strata::GgufFile gguf(path); - err = strata::check_architecture(gguf); + const strata::GgufModel model(shards); + err = strata::check_architecture(model.shard(0)); if (!err.empty()) return false; - const strata::TensorInfo* tensor = nullptr; - for (const auto& candidate : gguf.tensors()) { - if (candidate.name != "output.weight") continue; - if (tensor) { err = "native head: duplicate output.weight"; return false; } - tensor = &candidate; - } + size_t at = 0; + const strata::TensorInfo* tensor = model.find("output.weight", &at); + const strata::GgufFile& gguf = model.shard(at); if (!tensor || !strata::kernels::native_mmvq_supported((int) tensor->type) || tensor->shape.size() != 2 || tensor->shape[0] != (uint64_t) n_in || tensor->shape[1] != (uint64_t) n_out) { err = "native head: expected a natively supported output.weight with the canonical head dimensions"; @@ -69,6 +84,7 @@ bool NativeHead::load(const std::string& path, int64_t n_in, int64_t n_out, std: } bool NativeHead::run(const float* mixed, float* logits, void* stream, std::string& err) const { + if (part_) { err = "native head: it holds part of the vocabulary (the rest is on the second GPU)"; return false; } if (!loaded() || !mixed || !logits || !stream) { err = "native head requires loaded weights, device buffers and an explicit stream"; return false; @@ -104,12 +120,12 @@ NativeEmbed::~NativeEmbed() { if (host_) cudaFreeHost(host_); } -bool NativeEmbed::load(const std::string& path, int64_t n_embd, int64_t n_vocab, std::string& err) { +bool NativeEmbed::load(const std::vector& shards, int64_t n_embd, int64_t n_vocab, std::string& err) { try { - strata::GgufFile gguf(path); - const strata::TensorInfo* t = nullptr; - for (const auto& c : gguf.tensors()) - if (c.name == "token_embd.weight") t = &c; + const strata::GgufModel model(shards); + size_t at = 0; + const strata::TensorInfo* t = model.find("token_embd.weight", &at); + const strata::GgufFile& gguf = model.shard(at); if (!t || t->shape.size() != 2 || t->shape[0] != (uint64_t) n_embd || t->shape[1] != (uint64_t) n_vocab || !strata::kernels::iq_supported((int) t->type) || n_embd % 256) { err = "native embedding: token_embd.weight is absent, of another shape, or of a type without a GPU " diff --git a/src/core/pinned.cu b/src/core/pinned.cu index 0e3e6b4a057..3d012df2d53 100644 --- a/src/core/pinned.cu +++ b/src/core/pinned.cu @@ -12,14 +12,6 @@ #include #include -// Loader fix: `fseek`/`ftell` are 32-bit on Windows by default (and the pack is 42.9 GB), and the 64-bit -// spelling is not the same on the two platforms the engine builds for. -#ifdef _WIN32 -#define STRATA_FSEEK64(f, o) _fseeki64((f), (long long) (o), SEEK_SET) -#else -#define STRATA_FSEEK64(f, o) fseeko((f), (off_t) (o), SEEK_SET) -#endif - #ifdef _WIN32 #include #else @@ -37,46 +29,20 @@ namespace { // A 2 MB-aligned reservation. Large pages first, then the largest alignment the OS will give us for free. void* reserve(uint64_t bytes, PageBacking& got, std::string& note) { #ifdef _WIN32 - // MEM_LARGE_PAGES needs SeLockMemoryPrivilege. Having it assigned to the account is not enough: the - // PROCESS must enable it in its own token (AdjustTokenPrivileges) before VirtualAlloc, or the call fails. - // An account without the assignment, or a failure to enable, leaves the process as it was: VirtualAlloc - // then refuses and the 4 KB fallback below runs - that is the EXPECTED outcome on a desktop. - { - HANDLE tok = nullptr; - if (OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &tok)) { - TOKEN_PRIVILEGES tp{}; - tp.PrivilegeCount = 1; - if (LookupPrivilegeValueW(nullptr, L"SeLockMemoryPrivilege", &tp.Privileges[0].Luid)) { - tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED; - if (!AdjustTokenPrivileges(tok, FALSE, &tp, 0, nullptr, nullptr) && GetLastError() != ERROR_NOT_ALL_ASSIGNED) - (void) 0; // nothing actionable: the large-page attempt below reports the outcome - } - CloseHandle(tok); - } - } // MEM_LARGE_PAGES needs SeLockMemoryPrivilege; a normal account does not have it and VirtualAlloc then // fails with ERROR_PRIVILEGE_NOT_HELD. That is the EXPECTED outcome on a desktop, not an error. SIZE_T large = GetLargePageMinimum(); - // A/B switch: STRATA_NO_LARGEPAGES=1 skips the large-page attempt, same run, same boot. - if (large > 0 && std::getenv("STRATA_NO_LARGEPAGES") == nullptr) { - // MEM_LARGE_PAGES requires the allocation size to be an exact multiple of the large page size - - // anything else is ERROR_INVALID_PARAMETER (87), which reads like a privilege problem but is not. - // Round up: the slack is under 2 MB and the tail stays unused. - const SIZE_T lbytes = (SIZE_T) (((SIZE_T) bytes + large - 1) / large * large); - void* p = VirtualAlloc(nullptr, lbytes, MEM_RESERVE | MEM_COMMIT | MEM_LARGE_PAGES, + if (large > 0) { + void* p = VirtualAlloc(nullptr, (SIZE_T) bytes, MEM_RESERVE | MEM_COMMIT | MEM_LARGE_PAGES, PAGE_READWRITE); if (p) { got = PageBacking::LargePages; note = "large pages (" + std::to_string((unsigned long long) large) + " B)"; return p; } - // 1450 (ERROR_NO_SYSTEM_RESOURCES) is the large-page pool saying no, 87 is a size that is not a - // multiple of the minimum, 1314 is the privilege: without the byte count the three read as one bug. - note = "large pages refused for " + std::to_string((unsigned long long) lbytes) + " B (GetLargePageMinimum=" + - std::to_string((unsigned long long) large) + ", VirtualAlloc error " + - std::to_string((unsigned long long) GetLastError()) + "); using 4 KB pages"; - } else if (std::getenv("STRATA_NO_LARGEPAGES") != nullptr) { - note = "large pages skipped (STRATA_NO_LARGEPAGES); using 4 KB pages"; + note = "large pages refused (GetLargePageMinimum=" + std::to_string((unsigned long long) large) + + ", VirtualAlloc error " + std::to_string((unsigned long long) GetLastError()) + + " - needs SeLockMemoryPrivilege); using 4 KB pages"; } else { note = "this system has no large-page minimum; using 4 KB pages"; } @@ -137,27 +103,23 @@ PinnedArena::PinnedArena(uint64_t bytes, uint64_t slice) : PinnedArena(bytes, un if (slice_bytes) slice_bytes = slice; // sliced registration: record the uniform size } -PinnedArena::PinnedArena(uint64_t bytes, const std::vector& bounds, - uint64_t max_pinned_bytes) : capacity(bytes) { +PinnedArena::PinnedArena(uint64_t bytes, const std::vector& bounds) : capacity(bytes) { if (bytes == 0) return; base = reserve(bytes, backing, note); // Register with CUDA BEFORE any page is touched: cudaHostRegister pins what is resident now, and a region // that has already been faulted in page by page is far more expensive to register and may fail outright. if (base) { - const bool capped = max_pinned_bytes > 0 && max_pinned_bytes < bytes && bounds.size() >= 2; - const cudaError_t e = capped ? cudaSuccess : - cudaHostRegister(base, (size_t) bytes, cudaHostRegisterPortable | cudaHostRegisterMapped); - if (!capped && e == cudaSuccess) { + const cudaError_t e = cudaHostRegister(base, (size_t) bytes, cudaHostRegisterPortable | cudaHostRegisterMapped); + if (e == cudaSuccess) { note = "cudaHostRegister PORTABLE ok; " + note; registered_bytes = bytes; - } else if (bounds.size() >= 2 && (capped || clear_error())) { + } else if (bounds.size() >= 2 && clear_error()) { // Plan v0.3 P5: the whole range is refused, so pin it slice by slice from the start. The rest stays // resident through the working-set lock below. (P6: slices may differ in size, one per layer.) slice_bytes = 1; // sliced; the uniform constructor records the size for (size_t i = 0; i + 1 < bounds.size(); ++i) { const uint64_t off = bounds[i], n = bounds[i + 1] - bounds[i]; - if (capped && (off > max_pinned_bytes || n > max_pinned_bytes - off)) break; if (cudaHostRegister((uint8_t*) base + off, (size_t) n, cudaHostRegisterPortable | cudaHostRegisterMapped) != cudaSuccess) { (void) cudaGetLastError(); break; @@ -166,9 +128,7 @@ PinnedArena::PinnedArena(uint64_t bytes, const std::vector& bounds, registered_bytes = off + n; ++registered_slices; } - note = (capped ? "cudaHostRegister limited to " + std::to_string(max_pinned_bytes >> 30) + - " GiB for CUDA1; " : - "cudaHostRegister of the whole arena FAILED (" + std::string(cudaGetErrorString(e)) + "); ") + + note = "cudaHostRegister of the whole arena FAILED (" + std::string(cudaGetErrorString(e)) + "); " + std::to_string(registered_slices) + " slices pinned (" + std::to_string(registered_bytes >> 30) + " GiB); " + note; if (registered_bytes < bytes) { @@ -232,6 +192,14 @@ LoadStats load_experts(const std::string& path, uint8_t* dst, uint64_t blob_byte LoadStats load_experts_ranges(const std::string& path, uint8_t* dst, const std::vector& layer_off, const std::vector& layer_bytes, int threads, uint64_t chunk) { + std::vector layer_dst; + for (uint64_t off : layer_off) layer_dst.push_back(dst + off); + return load_experts_ranges(path, layer_dst, layer_off, layer_bytes, threads, chunk); +} + +LoadStats load_experts_ranges(const std::string& path, const std::vector& layer_dst, + const std::vector& layer_off, const std::vector& layer_bytes, int threads, + uint64_t chunk) { LoadStats st; const uint64_t layers = (uint64_t) layer_off.size(); st.layers = layers; @@ -242,33 +210,11 @@ LoadStats load_experts_ranges(const std::string& path, uint8_t* dst, const std:: const auto t0 = std::chrono::steady_clock::now(); std::vector layer_hash((size_t) layers, 1469598103934665603ull); std::atomic next_layer{0}; - std::atomic read_ns_sum{0}; // summed over the threads: see LoadStats::read_seconds - std::atomic copy_ns_sum{0}; std::mutex err_mu; std::string err; auto worker = [&]() { std::vector buf((size_t) chunk); - // One handle per thread, seeked once per layer: a shared handle would need a lock around the seek and - // would serialise the very thing the threads are here to parallelise. `fread` on a `FILE*` rather than - // `std::ifstream`: see the header - MSVC's `basic_filebuf::xsgetn` splits any request larger than - // `_INTERNAL_BUFSIZ - 1` into 4095-byte freads, which turned one 8 MiB chunk into ~2048 4 KiB reads. - // `fread` sees a request bigger than the stream buffer and passes it to `_read()`/ - // `ReadFile()` unchanged, so the chunk size reaches the disk. Buffered, not `FILE_FLAG_NO_BUFFERING`: - // the cache should still hold what it can. - FILE* f = std::fopen(path.c_str(), "rb"); - if (f == nullptr) { - std::lock_guard g(err_mu); - err = "cannot open " + path; - return; - } - // A `FILE*` has no destructor that closes it, and this function has early returns below (open, seek and - // short-read failures), so the guard is what keeps the closing correct on every path. - struct Closer { - FILE* f; - ~Closer() { if (f != nullptr) std::fclose(f); } - } closer{f}; - uint64_t read_ns = 0, copy_ns = 0; for (;;) { const uint64_t L = next_layer.fetch_add(1); if (L >= layers) break; @@ -276,39 +222,30 @@ LoadStats load_experts_ranges(const std::string& path, uint8_t* dst, const std:: uint64_t remaining = layer_bytes[(size_t) L]; uint64_t pos = 0; uint64_t h = 1469598103934665603ull; - // 64-bit seek: the pack is 42.9 GB, so the 32-bit `fseek` would wrap past 4 GiB - if (STRATA_FSEEK64(f, off) != 0) { + // one handle per thread, seeked once per layer: a shared handle would need a lock around the seek + // and would serialise the very thing the threads are here to parallelise + std::ifstream f(path, std::ios::binary); + if (!f) { std::lock_guard g(err_mu); - err = "seek to " + std::to_string(off) + " B failed in layer " + std::to_string(L); + err = "cannot open " + path; return; } + f.seekg((std::streamoff) off); while (remaining > 0) { const uint64_t n = remaining < chunk ? remaining : chunk; - const auto t_read = std::chrono::steady_clock::now(); - const size_t got = std::fread(buf.data(), 1, (size_t) n, f); - read_ns += (uint64_t) std::chrono::duration_cast( - std::chrono::steady_clock::now() - t_read).count(); - // A short read is EOF or an I/O error, never a silent zero fill: say WHERE and HOW SHORT, and - // keep going no further - the caller turns this into a refused load, not a wrong answer. - if (got != (size_t) n) { + f.read((char*) buf.data(), (std::streamsize) n); + if ((uint64_t) f.gcount() != n) { std::lock_guard g(err_mu); - err = "short read in layer " + std::to_string(L) + ": got " + std::to_string(got) + " of " - + std::to_string(n) + " B at offset " + std::to_string(off + pos) - + (std::ferror(f) != 0 ? " (ferror set)" : ""); + err = "short read in layer " + std::to_string(L); return; } - const auto t_copy = std::chrono::steady_clock::now(); - std::memcpy(dst + off + pos, buf.data(), (size_t) n); + std::memcpy(layer_dst[(size_t) L] + pos, buf.data(), (size_t) n); h = fnv1a64(buf.data(), n, h); - copy_ns += (uint64_t) std::chrono::duration_cast( - std::chrono::steady_clock::now() - t_copy).count(); pos += n; remaining -= n; } layer_hash[(size_t) L] = h; } - read_ns_sum.fetch_add(read_ns); - copy_ns_sum.fetch_add(copy_ns); }; std::vector pool; @@ -319,12 +256,8 @@ LoadStats load_experts_ranges(const std::string& path, uint8_t* dst, const std:: if (!err.empty()) { std::fprintf(stderr, "load_experts: %s\n", err.c_str()); st.seconds = -1.0; - st.ok = false; - st.error = err; return st; } - st.read_seconds = (double) read_ns_sum.load() / 1e9; - st.copy_seconds = (double) copy_ns_sum.load() / 1e9; st.layer_checksums = std::move(layer_hash); st.seconds = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); return st; diff --git a/src/core/prompt_cache.cpp b/src/core/prompt_cache.cpp new file mode 100644 index 00000000000..01af3bda2d3 --- /dev/null +++ b/src/core/prompt_cache.cpp @@ -0,0 +1,285 @@ +// src/core/prompt_cache.cpp - the --serve prompt cache (see include/strata/core/prompt_cache.hpp). +#include "strata/core/prompt_cache.hpp" + +#include "strata/kernels/kv_q8.hpp" +#include "strata/kernels/qsa.hpp" + +#include + +#include +#include +#include +#include + +namespace strata::core { + +namespace { + +/// A replaced sequence is stashed when the switch would drop at least this many of its cells. +constexpr int64_t kStashMin = 1024; + +} // namespace + +PromptCache::PromptCache(const ModelGeometry& g, SessionState& ss, std::vector kv, void* stream, + int max_ckpts, uint64_t stash_bytes) + : g_(g), ss_(ss), kv_(std::move(kv)), stream_(stream), max_ckpts_(max_ckpts), stash_budget_(stash_bytes), + ckpt_bytes_(session_ckpt_bytes(g)) {} + +PromptCache::~PromptCache() { + for (const auto& [p, pinned] : owned_) { + if (pinned) cudaFreeHost(p); + else std::free(p); + } +} + +int64_t PromptCache::reusable(const Seq& seq, const std::vector& ids, bool live) const { + const int64_t m = std::min((int64_t) seq.tokens.size(), (int64_t) ids.size() - 1); + int64_t lcp = 0; + while (lcp < m && seq.tokens[(size_t) lcp] == ids[(size_t) lcp]) ++lcp; + if (live && lcp > 0 && lcp == (int64_t) seq.tokens.size()) return lcp; + int64_t best = -1; + for (const Ckpt& c : seq.ckpts) + if (c.pos <= lcp && c.pos > best) best = c.pos; + return best; +} + +PromptCache::Ckpt* PromptCache::find(Seq& seq, int64_t pos) { + for (Ckpt& c : seq.ckpts) + if (c.pos == pos) return &c; + return nullptr; +} + +void* PromptCache::take_buffer() { + if (!free_.empty()) { + void* h = free_.back(); + free_.pop_back(); + return h; + } + if (allocated_ < max_ckpts_) { + void* h = nullptr; + const bool pinned = cudaMallocHost(&h, (size_t) ckpt_bytes_) == cudaSuccess; + if (!pinned) { + (void) cudaGetLastError(); + h = std::malloc((size_t) ckpt_bytes_); // pageable: slower copies, same result + } + if (h != nullptr) { + owned_.emplace_back(h, pinned); + ++allocated_; + } + return h; + } + // the pool is full: the least recently used checkpoint of any sequence goes + Seq* owner = nullptr; + size_t at = 0; + uint64_t oldest = UINT64_MAX; + auto scan = [&](Seq& s) { + for (size_t i = 0; i < s.ckpts.size(); ++i) + if (s.ckpts[i].stamp < oldest) { + oldest = s.ckpts[i].stamp; + owner = &s; + at = i; + } + }; + scan(live_); + for (Seq& s : stash_) scan(s); + if (owner == nullptr) return nullptr; + void* h = owner->ckpts[at].host; + owner->ckpts.erase(owner->ckpts.begin() + (ptrdiff_t) at); + if (owner != &live_ && owner->ckpts.empty()) // a stashed sequence without a checkpoint cannot be resumed + stash_.erase(stash_.begin() + (owner - stash_.data())); + return h; +} + +void PromptCache::release(Seq& seq, int64_t keep_upto) { + for (size_t i = 0; i < seq.ckpts.size();) { + if (seq.ckpts[i].pos > keep_upto) { + free_.push_back(seq.ckpts[i].host); + seq.ckpts.erase(seq.ckpts.begin() + (ptrdiff_t) i); + } else { + ++i; + } + } +} + +std::vector PromptCache::kv_parts(int64_t cells) const { + // [page][kv_head][page_size][head_dim] with an identity page table: the first cells are one prefix of each + // pool. The indexer's pooled rows are one per idx_block cells. + const strata::kernels::QsaShapes sh = strata::kernels::qsa_real_shapes(); + const uint64_t hd = (uint64_t) g_.head_dim, row = (uint64_t) g_.idx_key_dim * sizeof(float); + std::vector parts; + for (QsaState* st : kv_) { + const uint64_t pages = (uint64_t) std::min(st->n_pages, (cells + sh.page_size - 1) / sh.page_size); + const uint64_t c = pages * (uint64_t) g_.n_head_kv * (uint64_t) sh.page_size; + if (st->kv_int8) { + const uint64_t scales = c * (hd / strata::kernels::KV_Q8_GROUP) * sizeof(uint16_t); + parts.push_back({st->k_q, c * hd}); + parts.push_back({st->v_q, c * hd}); + parts.push_back({st->k_scale, scales}); + parts.push_back({st->v_scale, scales}); + } else { + parts.push_back({st->k_pool, c * hd * sizeof(uint16_t)}); + parts.push_back({st->v_pool, c * hd * sizeof(uint16_t)}); + } + const int64_t rows = std::min(cells / sh.idx_block + 2, st->max_cells / sh.idx_block + 2); + parts.push_back({st->idx_pooled, (uint64_t) rows * row}); + } + return parts; +} + +bool PromptCache::kv_copy(const std::vector& parts, uint8_t* host, bool to_host) { + cudaStream_t cs = (cudaStream_t) stream_; + for (const Part& p : parts) { + const cudaError_t e = to_host ? cudaMemcpyAsync(host, p.dev, (size_t) p.bytes, cudaMemcpyDeviceToHost, cs) + : cudaMemcpyAsync(p.dev, host, (size_t) p.bytes, cudaMemcpyHostToDevice, cs); + if (e != cudaSuccess) return false; + host += p.bytes; + } + return cudaStreamSynchronize(cs) == cudaSuccess; +} + +void PromptCache::stash_live(int64_t keep_upto) { + const int64_t L = (int64_t) live_.tokens.size(); + if (stash_budget_ == 0 || L < kStashMin) return; + const std::vector parts = kv_parts(L); + uint64_t bytes = 0; + for (const Part& p : parts) bytes += p.bytes; + if (bytes > stash_budget_) return; + uint64_t used = 0; + for (const Seq& s : stash_) used += s.kv_bytes; + while (!stash_.empty() && used + bytes > stash_budget_) { // the least recently used sequences go + auto lru = std::min_element(stash_.begin(), stash_.end(), + [](const Seq& a, const Seq& b) { return a.stamp < b.stamp; }); + used -= lru->kv_bytes; + release(*lru, 0); + stash_.erase(lru); + } + Seq s; + s.kv.reset(new (std::nothrow) uint8_t[bytes]); + if (!s.kv || !kv_copy(parts, s.kv.get(), true)) return; + checkpoint(L); // the end state: the session is still at L + for (size_t i = 0; i < live_.ckpts.size();) { + if (live_.ckpts[i].pos > keep_upto) { + s.ckpts.push_back(live_.ckpts[i]); + live_.ckpts.erase(live_.ckpts.begin() + (ptrdiff_t) i); + } else { + ++i; + } + } + if (s.ckpts.empty()) return; + s.tokens = live_.tokens; + s.kv_cells = L; + s.kv_bytes = bytes; + s.stamp = ++clock_; + stash_.push_back(std::move(s)); +} + +bool PromptCache::restore(Seq& seq, int64_t pos) { + Ckpt* ck = find(seq, pos); + if (ck == nullptr) return false; + if (seq.kv && !kv_copy(kv_parts(seq.kv_cells), seq.kv.get(), false)) return false; + ck->stamp = ++clock_; + return session_ckpt_load(ss_, g_, pos, ck->host, stream_); +} + +int64_t PromptCache::begin(const std::vector& ids, bool cacheable) { + const int64_t L = (int64_t) live_.tokens.size(); + auto zero = [&]() -> int64_t { + release(live_, 0); + live_.tokens.clear(); + session_zero(ss_, g_, nullptr, stream_); + cudaStreamSynchronize((cudaStream_t) stream_); + return 0; + }; + if (!enabled()) return zero(); + if (!cacheable) { + stash_live(0); + return zero(); + } + int64_t best = reusable(live_, ids, true); + int from = -1; // -1: the live sequence, else a stashed one + for (size_t i = 0; i < stash_.size(); ++i) { + const int64_t r = reusable(stash_[i], ids, false); + if (r > best) { + best = r; + from = (int) i; + } + } + if (best <= 0) { + stash_live(0); + return zero(); + } + if (from < 0) { + // The live sequence continues. When the prompt drops most of it (another conversation that shares a + // prefix), the rest goes to the stash; an edit near the end just overwrites it. + const int64_t dropped = L - best; + if (dropped >= kStashMin && dropped > best) { + if (Ckpt* c = find(live_, best)) c->stamp = ++clock_; // keeps it through the stash's evictions + stash_live(best); + } + if (best < L && !restore(live_, best)) return zero(); + release(live_, best); + live_.tokens.resize((size_t) best); + return best; + } + // Another sequence continues: the live one is stashed and the chosen one becomes live. + Seq chosen = std::move(stash_[(size_t) from]); + stash_.erase(stash_.begin() + from); + stash_live(0); + release(live_, 0); + live_.tokens.clear(); + if (!restore(chosen, best)) { + release(chosen, 0); + return zero(); + } + live_.tokens.assign(chosen.tokens.begin(), chosen.tokens.begin() + best); + live_.ckpts = std::move(chosen.ckpts); + release(live_, best); + return best; +} + +void PromptCache::set(int64_t pos, const int64_t* tokens, int64_t n) { + if (!enabled()) return; + release(live_, pos); + live_.tokens.resize((size_t) pos); + live_.tokens.insert(live_.tokens.end(), tokens, tokens + n); +} + +void PromptCache::set(int64_t pos, const int32_t* tokens, int64_t n) { + if (!enabled()) return; + release(live_, pos); + live_.tokens.resize((size_t) pos); + for (int64_t i = 0; i < n; ++i) live_.tokens.push_back(tokens[i]); +} + +void PromptCache::checkpoint(int64_t pos) { + if (!enabled() || pos <= 0 || pos != (int64_t) live_.tokens.size() || find(live_, pos) != nullptr) return; + void* h = take_buffer(); + if (h == nullptr) return; + cudaDeviceSynchronize(); + if (!session_ckpt_save(ss_, g_, pos, h, stream_)) { + free_.push_back(h); + return; + } + live_.ckpts.push_back({pos, h, ++clock_}); +} + +void PromptCache::forget() { + release(live_, 0); + live_.tokens.clear(); +} + +std::string PromptCache::summary() const { + size_t n = live_.ckpts.size(); + uint64_t kv = 0; + for (const Seq& s : stash_) { + n += s.ckpts.size(); + kv += s.kv_bytes; + } + char buf[192]; + std::snprintf(buf, sizeof buf, "live %lld cells; %zu checkpoints (%.0f MiB); %zu stashed sequences (%.0f MiB)", + (long long) live_.tokens.size(), n, (double) (n * ckpt_bytes_) / 1048576.0, stash_.size(), + (double) kv / 1048576.0); + return buf; +} + +} // namespace strata::core diff --git a/src/core/second_gpu.cpp b/src/core/second_gpu.cpp new file mode 100644 index 00000000000..4cfb793aa4a --- /dev/null +++ b/src/core/second_gpu.cpp @@ -0,0 +1,253 @@ +// src/core/second_gpu.cpp - see include/strata/core/second_gpu.hpp. +#include "strata/core/second_gpu.hpp" + +#include "strata/kernels/cpu/expert_layout.hpp" +#include "strata/kernels/elementwise.hpp" +#include "strata/kernels/iq_kernels.hpp" + +#include + +#include +#include + +namespace strata::core { +namespace { + +// the plan block: [groups of part 0, of part 1, ring, pad] [starts, part 0: cap+1] [part 1: cap+1] [dst: cap] +// [tok: cap] (8-byte aligned) [slot addresses, part 0: cap] [part 1: cap] [out] [flag after part 0] [after part 1] +struct PlanLayout { + int64_t starts[2], dst, tok, ptr; // int32 offsets; `ptr` even + size_t bytes; + explicit PlanLayout(int64_t cap) + : starts{4, 4 + cap + 1}, dst(4 + 2 * (cap + 1)), tok(dst + cap), ptr((tok + cap + 1) & ~1ll), + bytes((size_t) ptr * 4 + (size_t) (2 * cap + 3) * 8) {} +}; + +// the grouped kernels' grids are sized for the bucket a layer's group count falls in +int bucket(int groups) { + for (int b : {1, 2, 3, 4, 6, 8, 12, 16, 24, 32}) if (groups <= b) return b; + return groups <= 48 ? 48 : groups; +} + +struct DeviceScope { // the second GPU current for one call, the main one again afterwards + int main; + DeviceScope(int dev, int main_dev) : main(main_dev) { cudaSetDevice(dev); } + ~DeviceScope() { cudaSetDevice(main); } +}; + +} // namespace + +SecondGpu::~SecondGpu() { + if (dev_ < 0) return; + DeviceScope scope(dev_, main_); + if (s_) cudaStreamSynchronize(s_); + if (pre_s_) cudaStreamSynchronize(pre_s_); + for (auto& gr : graphs_) cudaGraphExecDestroy(gr.exec); + cache_.close(); + void* dev[] = {d_xq_, d_plan_, d_scratch_, d_rows_, d_count_, d_pre_}; + for (void* p : dev) if (p) cudaFree(p); + void* host[] = {h_x_, h_plan_}; + for (void* p : host) if (p) cudaFreeHost(p); + if (ev_) cudaEventDestroy(ev_); + if (pre_ev_) cudaEventDestroy(pre_ev_); + if (s_) cudaStreamDestroy(s_); + if (pre_s_) cudaStreamDestroy(pre_s_); +} + +bool SecondGpu::init(int device, int main_device, int64_t n_embd, int64_t n_ff, int max_t, int64_t k, std::string& err) { + dev_ = device; + main_ = main_device; + max_t_ = max_t; + n_embd_ = n_embd; + n_ff_ = n_ff; + cap_ = (int64_t) max_t * k; + DeviceScope scope(dev_, main_); + const size_t xb = (size_t) max_t * (size_t) (n_embd / 32) * 36; // q8_1 rows + const size_t ob = (size_t) cap_ * (size_t) n_embd * sizeof(float); + const size_t pb = PlanLayout(cap_).bytes; + const unsigned pm = cudaHostAllocPortable | cudaHostAllocMapped; + const char* step = nullptr; + cudaError_t e = cudaSuccess; + int unified = 0; + auto run = [&](const char* what, cudaError_t r) { if (e == cudaSuccess && r != cudaSuccess) { e = r; step = what; } }; + run("context", cudaFree(nullptr)); + run("unified addressing", cudaDeviceGetAttribute(&unified, cudaDevAttrUnifiedAddressing, dev_)); + run("stream", cudaStreamCreateWithFlags(&s_, cudaStreamNonBlocking)); + run("event", cudaEventCreateWithFlags(&ev_, cudaEventDisableTiming)); + run("pinned staging", cudaHostAlloc((void**) &h_x_, xb, pm)); + run("pinned staging", cudaHostAlloc((void**) &h_plan_, pb, pm)); + run("mapped staging", cudaHostGetDevicePointer((void**) &m_x_, h_x_, 0)); + run("mapped staging", cudaHostGetDevicePointer((void**) &m_plan_, h_plan_, 0)); + run("device buffers", cudaMalloc((void**) &d_xq_, xb)); + run("device buffers", cudaMalloc((void**) &d_plan_, pb)); + run("device buffers", cudaMalloc((void**) &d_scratch_, strata::kernels::native_expert_scratch_bytes(cap_, n_ff))); + run("device buffers", cudaMalloc((void**) &d_rows_, ob)); + run("device buffers", cudaMalloc((void**) &d_count_, sizeof(unsigned))); + run("device buffers", cudaMemset(d_count_, 0, sizeof(unsigned))); + if (e != cudaSuccess) { + err = std::string("second GPU: ") + step + ": " + cudaGetErrorString(e); + return false; + } + if (!unified) { // it writes its rows through the host pointers of the first GPU's mapped memory + err = "second GPU: no unified addressing"; + return false; + } + std::memset(h_x_, 0, xb); + std::memset(h_plan_, 0, pb); + return true; +} + +bool SecondGpu::init_prefetch(int slots, uint64_t blob_bytes, std::string& err) { + slots = (std::min)(slots, kPrefetchMax); + if (slots <= 0) return true; + DeviceScope scope(dev_, main_); + pre_cap_ = (blob_bytes + 255) / 256 * 256; + if (cudaStreamCreateWithFlags(&pre_s_, cudaStreamNonBlocking) != cudaSuccess || + cudaEventCreateWithFlags(&pre_ev_, cudaEventDisableTiming) != cudaSuccess || + cudaMalloc((void**) &d_pre_, (size_t) slots * pre_cap_) != cudaSuccess) { + err = std::string("second GPU: prefetch slots: ") + cudaGetErrorString(cudaGetLastError()); + return false; + } + pre_max_ = slots; + return true; +} + +bool SecondGpu::prefetch(int64_t layer, const int32_t* ids, const uint8_t* const* src, int n, uint64_t bytes, + std::string& err) { + if (n <= 0) return true; + if (n > pre_max_ || bytes > pre_cap_) { err = "second GPU: a prefetch does not fit its slots"; return false; } + DeviceScope scope(dev_, main_); + void* dst[kPrefetchMax]; + size_t size[kPrefetchMax]; + for (int i = 0; i < n; ++i) { + dst[i] = d_pre_ + (size_t) i * pre_cap_; + size[i] = (size_t) bytes; + } + if (cudaStreamWaitEvent(pre_s_, ev_, 0) != cudaSuccess || // the last submitted layer may read the slots + !copy_blobs(dst, (const void* const*) src, size, (size_t) n, pre_s_) || + cudaEventRecord(pre_ev_, pre_s_) != cudaSuccess) { + err = std::string("second GPU: prefetch: ") + cudaGetErrorString(cudaGetLastError()); + return false; + } + (void) cudaStreamQuery(pre_s_); // submit now, not at the next driver call + pre_layer_ = layer; + pre_n_ = n; + for (int i = 0; i < n; ++i) pre_ids_[i] = ids[i]; + prefetch_copied += n; + return true; +} + +bool SecondGpu::graph_for(int gu_type, int d_type, int groups, int part, bool prep, cudaGraphExec_t& exec, + std::string& err) { + for (const auto& gr : graphs_) + if (gr.gu == gu_type && gr.d == d_type && gr.groups == groups && gr.part == part && gr.prep == prep) { + exec = gr.exec; + return true; + } + const strata::kernels::NativeExpertLayout L = strata::kernels::native_expert_layout(gu_type, d_type, n_embd_, n_ff_); + const PlanLayout P(cap_); + const auto* pi = (const int32_t*) d_plan_; + const auto* p64 = (const unsigned long long*) (pi + P.ptr); + cudaGraph_t graph = nullptr; + if (cudaStreamBeginCapture(s_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) { + err = "second GPU: cannot capture"; + return false; + } + try { + if (prep) { + // kernels copy the plan and the activations (a copy node costs ~15 us on this link); the grouped kernels + // read the group count from the plan and skip the groups past it + strata::kernels::copy_i32_from_mapped((int32_t*) d_plan_, (const int32_t*) m_plan_, (int64_t) P.bytes / 4, + s_); + strata::kernels::copy_from_mapped((float*) d_xq_, (const float*) m_x_, max_t_ * (n_embd_ / 32) * 9, s_); + } + strata::kernels::native_expert_grouped(L, p64 + part * cap_, pi + P.starts[part], pi + part, pi + P.dst, + pi + P.tok, groups, cap_, d_xq_, d_scratch_, d_rows_, s_); + strata::kernels::native_expert_rows_out(d_rows_, pi + part, pi + P.starts[part], pi + P.dst, n_embd_, cap_, + (float* const*) (p64 + 2 * cap_), + (uint32_t* const*) (p64 + 2 * cap_ + 1 + part), pi + 2, d_count_, s_); + } catch (const std::exception& ex) { + cudaStreamEndCapture(s_, &graph); + if (graph) cudaGraphDestroy(graph); + err = std::string("second GPU: ") + ex.what(); + return false; + } + const cudaError_t ce = cudaStreamEndCapture(s_, &graph); + if (ce != cudaSuccess || cudaGraphInstantiate(&exec, graph, 0) != cudaSuccess) { + if (graph) cudaGraphDestroy(graph); + err = std::string("second GPU: graph capture: ") + cudaGetErrorString(ce); + return false; + } + cudaGraphDestroy(graph); + cudaGraphUpload(exec, s_); + cudaStreamSynchronize(s_); + graphs_.push_back({gu_type, d_type, groups, part, prep, exec}); + return true; +} + +bool SecondGpu::submit(int64_t layer, const uint8_t* x, int n_tok, int64_t k, const int32_t* slots, + const int32_t* starts, const int32_t* entries, int n_groups, float* out, uint32_t* flag, + uint32_t ring, std::string& err) { + const int n = starts[n_groups]; + if (n_groups <= 0 || n > cap_ || n_tok > max_t_) { err = "second GPU: a layer's share is out of range"; return false; } + cudaError_t q; + while ((q = cudaEventQuery(ev_)) == cudaErrorNotReady) _mm_pause(); + if (q != cudaSuccess) { + err = std::string("second GPU: ") + cudaGetErrorString(q); + return false; + } + // the plan: the experts in its cache first (part 0), then the prefetched ones (part 1), their entries in that order + const PlanLayout P(cap_); + auto* pi = (int32_t*) h_plan_; + auto* p64 = (unsigned long long*) (pi + P.ptr); + int ng[2] = {0, 0}, e = 0; + for (int part = 0; part < 2; ++part) { + int32_t* st = pi + P.starts[part]; + for (int g = 0; g < n_groups; ++g) { + const bool pre = slots[g] <= -2; + if (pre != (part == 1)) continue; + p64[(size_t) part * (size_t) cap_ + (size_t) ng[part]] = + pre ? (unsigned long long) (d_pre_ + (size_t) (-2 - slots[g]) * pre_cap_) + : (unsigned long long) cache_.device_slot(slots[g]); + st[ng[part]++] = e; + for (int j = starts[g]; j < starts[g + 1]; ++j, ++e) { + pi[P.dst + e] = entries[j]; + pi[P.tok + e] = (int32_t) (entries[j] / k); + } + } + st[ng[part]] = e; + pi[part] = ng[part]; + } + pi[2] = (int32_t) ring; + p64[2 * cap_] = (unsigned long long) out; + p64[2 * cap_ + 1] = ng[1] == 0 ? (unsigned long long) flag : 0; // the last launch raises it + p64[2 * cap_ + 2] = ng[1] > 0 ? (unsigned long long) flag : 0; + std::memcpy(h_x_, x, (size_t) n_tok * (size_t) (n_embd_ / 32) * 36); + const auto& f = strata::kernels::cpu::expert_layout().fmt[(size_t) layer]; + DeviceScope scope(dev_, main_); + cudaGraphExec_t own = nullptr, pre = nullptr; + if ((ng[0] > 0 && !graph_for(f.gu_type, f.d_type, bucket(ng[0]), 0, true, own, err)) || + (ng[1] > 0 && !graph_for(f.gu_type, f.d_type, bucket(ng[1]), 1, ng[0] == 0, pre, err))) + return false; + if ((own != nullptr && cudaGraphLaunch(own, s_) != cudaSuccess) || + (pre != nullptr && + (cudaStreamWaitEvent(s_, pre_ev_, 0) != cudaSuccess || cudaGraphLaunch(pre, s_) != cudaSuccess)) || + cudaEventRecord(ev_, s_) != cudaSuccess) { + err = std::string("second GPU: ") + cudaGetErrorString(cudaGetLastError()); + return false; + } + ++layers; + experts += n_groups; + entries_done += n; + prefetch_used += ng[1]; + return true; +} + +bool SecondGpu::healthy(std::string& err) const { + const cudaError_t q = cudaEventQuery(ev_); + if (q == cudaSuccess || q == cudaErrorNotReady) return true; + err = std::string("second GPU: ") + cudaGetErrorString(q); + return false; +} + +} // namespace strata::core diff --git a/src/core/session.cpp b/src/core/session.cpp index d18a5ec38a5..5035c7d667f 100644 --- a/src/core/session.cpp +++ b/src/core/session.cpp @@ -4,6 +4,7 @@ #include "strata/kernels/qsa.hpp" #include "strata/kernels/elementwise.hpp" +#include "strata/kernels/fused_gr.hpp" #include "strata/kernels/quantize_act.hpp" #include "strata/kernels/s2_expert_grouped.hpp" #include "strata/kernels/cpu/pool.hpp" @@ -118,6 +119,8 @@ void session_zero(SessionState& s, const ModelGeometry& g, const float* R_init, } else { cudaMemsetAsync(s.block.R, 0, (size_t) g.hc * g.n_embd * 4, cs); } + // the fused hyper-connection read's scratch: its counters start at zero + cudaMemsetAsync(s.block.gr_scratch, 0, strata::kernels::fused_gr_scratch_bytes(), cs); // every GDN layer's recurrence and conv history cudaMemsetAsync(s.gdn_state, 0, (size_t) g.n_gdn_layers() * gdn_state_floats(g) * 4, cs); // and every QSA layer's cache and indexer @@ -133,6 +136,69 @@ void session_zero(SessionState& s, const ModelGeometry& g, const float* R_init, s.ple_token = -1; } +namespace { + +/// One region of a checkpoint: `bytes` of device memory at `dev`, stored back to back in the host buffer. +struct CkptPart { void* dev; uint64_t bytes; }; + +/// The parts of a checkpoint at `pos`, in storage order (the PLE token window follows them on the host). +std::vector ckpt_parts(const SessionState& s, const ModelGeometry& g, int64_t pos) { + strata::kernels::QsaShapes sh = strata::kernels::qsa_real_shapes(); + sh.idx_dim = g.idx_key_dim; + const uint64_t row = (uint64_t) sh.idx_dim * sizeof(float); + std::vector parts; + parts.push_back({s.gdn_state, (uint64_t) g.n_gdn_layers() * gdn_state_floats(g) * sizeof(float)}); + parts.push_back({s.ple_hist, ple_hist_bytes()}); + for (int64_t i = 0; i < g.n_qsa_layers(); ++i) { + const QsaState& st = s.qsa_states[i]; + parts.push_back({st.idx_tail, (uint64_t) (sh.idx_block - 1) * row}); + parts.push_back({st.idx_dead, row}); + parts.push_back({(uint8_t*) st.idx_pooled + (uint64_t) (pos / sh.idx_block) * row, row}); + parts.push_back({st.idx_block_pos, sizeof(int32_t)}); + } + return parts; +} + +constexpr uint64_t kCkptHostInts = 3; // ple_prev[0], ple_prev[1], ple_token + +} // namespace + +uint64_t session_ckpt_bytes(const ModelGeometry& g) { + SessionState probe; // only the sizes are read: null device pointers are never dereferenced here + std::vector qsa((size_t) g.n_qsa_layers()); + probe.qsa_states = qsa.data(); + uint64_t n = kCkptHostInts * sizeof(int32_t); + for (const CkptPart& p : ckpt_parts(probe, g, 0)) n += p.bytes; + return n; +} + +bool session_ckpt_save(const SessionState& s, const ModelGeometry& g, int64_t pos, void* host, void* stream) { + cudaStream_t cs = (cudaStream_t) stream; + uint8_t* h = (uint8_t*) host; + for (const CkptPart& p : ckpt_parts(s, g, pos)) { + if (cudaMemcpyAsync(h, p.dev, (size_t) p.bytes, cudaMemcpyDeviceToHost, cs) != cudaSuccess) return false; + h += p.bytes; + } + const int32_t ints[kCkptHostInts] = {s.ple_prev[0], s.ple_prev[1], s.ple_token}; + std::memcpy(h, ints, sizeof ints); + return cudaStreamSynchronize(cs) == cudaSuccess; +} + +bool session_ckpt_load(SessionState& s, const ModelGeometry& g, int64_t pos, const void* host, void* stream) { + cudaStream_t cs = (cudaStream_t) stream; + const uint8_t* h = (const uint8_t*) host; + for (const CkptPart& p : ckpt_parts(s, g, pos)) { + if (cudaMemcpyAsync(p.dev, h, (size_t) p.bytes, cudaMemcpyHostToDevice, cs) != cudaSuccess) return false; + h += p.bytes; + } + int32_t ints[kCkptHostInts]; + std::memcpy(ints, h, sizeof ints); + s.ple_prev[0] = ints[0]; + s.ple_prev[1] = ints[1]; + s.ple_token = ints[2]; + return cudaStreamSynchronize(cs) == cudaSuccess; +} + /// Sets `s.gdn.state`/`conv_state` for `layer`, which is what makes one layer's GDN state its own. Shared by /// the direct and captured paths so the two cannot disagree about which slice a layer owns. void gdn_point_at(const ModelGeometry& g, int64_t layer, SessionState& s) { @@ -490,7 +556,7 @@ void session_graphs_free(SessionGraphs& gr) { gr.captured = false; } -bool SessionLoopScratch::init(size_t parts_bytes_in, std::string& err) { +bool SessionLoopScratch::init(size_t parts_bytes_in, std::string& err, int host) { if (y_miss != nullptr || probe != nullptr) { err = "SessionLoopScratch::init: already initialised"; return false; @@ -507,15 +573,18 @@ bool SessionLoopScratch::init(size_t parts_bytes_in, std::string& err) { free(); return false; } - // **PIN THE HOST ONCE, NOT ONCE PER TOKEN.** `ExpertPool` builds its workers from `physical_cores(true)`, - // which drops the first physical core so the host loop can spin without taking a worker's cycles - and + // **PIN THE HOST ONCE, NOT ONCE PER TOKEN.** `ExpertPool` builds its workers from `CorePlacement`, which + // keeps them off the host's core so the host loop can spin without taking a worker's cycles - and // nothing in the pool can pin the host, so if this does not happen the spin is free to land on a worker's // core or its SMT sibling. The symptom is not an error: it is a CPU path at 26.9 GB/s where the same pool // runs at 36.32. It was being done and undone on EVERY token, which is a syscall pair on the critical path // for a property that wants to hold for the whole session. - const std::vector cores = strata::kernels::cpu::physical_cores(false); - if (!cores.empty()) { - pinned_core = strata::kernels::cpu::pin_current_thread(cores[0]); + if (host < 0) { // the last physical core, CorePlacement's start + const std::vector> cores = strata::kernels::cpu::physical_cores(); + host = cores.back().front(); + } + if (host >= 0) { + pinned_core = strata::kernels::cpu::pin_current_thread(host); pinned = true; } return true; diff --git a/src/core/split_head.cpp b/src/core/split_head.cpp new file mode 100644 index 00000000000..db45c9e7e72 --- /dev/null +++ b/src/core/split_head.cpp @@ -0,0 +1,164 @@ +// src/core/split_head.cpp - see include/strata/core/split_head.hpp. +#include "strata/core/split_head.hpp" + +#include "strata/artifact/gguf_reader.hpp" +#include "strata/kernels/elementwise.hpp" +#include "strata/kernels/native_mmvq.hpp" +#include "strata/kernels/verify_kernels.hpp" + +#include + +#include +#include + +namespace strata::core { +namespace { + +struct DeviceScope { // this GPU current for one call, the main one again afterwards + int main; + DeviceScope(int dev, int main_dev) : main(main_dev) { cudaSetDevice(dev); } + ~DeviceScope() { cudaSetDevice(main); } +}; + +} // namespace + +SplitHead::~SplitHead() { + if (dev_ < 0) return; + DeviceScope scope(dev_, main_); + if (s_) cudaStreamSynchronize(s_); + for (auto& e : exec_) + for (auto& x : e) if (x) cudaGraphExecDestroy(x); + void* dev[] = {w_, mixed_, out_, xq_, xil_, scratch_}; + for (void* p : dev) if (p) cudaFree(p); + void* host[] = {h_in_, h_hand_, h_done_, h_idx_, h_val_, h_logits_}; + for (void* p : host) if (p) cudaFreeHost(p); + if (s_) cudaStreamDestroy(s_); +} + +bool SplitHead::init(int device, int main_device, const std::vector& shards, int64_t n_in, + int64_t n_vocab, int64_t split, int max_t, std::string& err) { + if (split <= 0 || split >= n_vocab || max_t < 1 || max_t > strata::kernels::kVerifyMaxT) { + err = "split head: bad arguments"; + return false; + } + dev_ = device; + main_ = main_device; + max_t_ = max_t; + n_in_ = n_in; + split_ = split; + rows_ = n_vocab - split; + DeviceScope scope(dev_, main_); + try { + const strata::GgufModel model(shards); + size_t at = 0; + const strata::TensorInfo* t = model.find("output.weight", &at); + if (!t || !strata::kernels::native_mmvq_supported((int) t->type)) { + err = "split head: output.weight is missing or not a native format"; + return false; + } + type_ = (int) t->type; + const uint64_t row = strata::kernels::native_mmvq_weight_bytes(type_, (int) n_in, 1); + bytes_ = row * (uint64_t) rows_; + const unsigned pm = cudaHostAllocPortable | cudaHostAllocMapped; + const size_t T = (size_t) max_t; + cudaError_t e = cudaSuccess; + const char* step = nullptr; + auto run = [&](const char* what, cudaError_t r) { + if (e == cudaSuccess && r != cudaSuccess) { e = r; step = what; } + }; + run("context", cudaFree(nullptr)); + run("stream", cudaStreamCreateWithFlags(&s_, cudaStreamNonBlocking)); + run("its rows", cudaMalloc(&w_, bytes_)); + run("its rows", cudaMemcpy(w_, model.shard(at).tensor_data(*t) + (uint64_t) split * row, bytes_, + cudaMemcpyHostToDevice)); + run("buffers", cudaMalloc((void**) &mixed_, T * (size_t) n_in * 4)); + run("buffers", cudaMalloc((void**) &out_, T * (size_t) rows_ * 4)); + run("buffers", cudaMalloc((void**) &xq_, strata::kernels::native_q8_1_bytes((int) n_in, max_t))); + run("buffers", cudaMalloc((void**) &xil_, strata::kernels::native_q8_1_il_bytes((int) n_in, max_t))); + run("buffers", cudaMalloc((void**) &scratch_, strata::kernels::argmax_rows_scratch_bytes(max_t))); + run("buffers", cudaMemset(scratch_, 0, strata::kernels::argmax_rows_scratch_bytes(max_t))); + run("mapped", cudaHostAlloc((void**) &h_in_, T * (size_t) n_in * 4, pm)); + run("mapped", cudaHostAlloc((void**) &h_hand_, 64, pm)); + run("mapped", cudaHostAlloc((void**) &h_done_, 64, pm)); + run("mapped", cudaHostAlloc((void**) &h_idx_, T * 4 + 64, pm)); + run("mapped", cudaHostAlloc((void**) &h_val_, T * 4 + 64, pm)); + run("mapped", cudaHostAlloc((void**) &h_logits_, T * (size_t) rows_ * 4, pm)); + if (e != cudaSuccess) { + err = std::string("split head: ") + step + ": " + cudaGetErrorString(e); + return false; + } + *h_hand_ = 0; + *h_done_ = 0; + } catch (const std::exception& x) { + err = std::string("split head: ") + x.what(); + return false; + } + return true; +} + +// One window's part: the input from mapped memory, quantized as the main GPU's head does, the rows' projection, the +// per-token argmax (and logit) to mapped memory, the count raised. +bool SplitHead::capture(int T, bool logits, std::string& err) { + using namespace strata::kernels; + if (exec_[T][logits]) return true; + if (cudaStreamBeginCapture(s_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) { + err = "split head: begin capture"; + return false; + } + try { + copy_from_mapped(mixed_, h_in_, (int64_t) T * n_in_, s_); + if (T >= 2) { + native_quantize_q8_1_il(mixed_, xq_, xil_, (int) n_in_, T, s_); + native_mmvq_il(type_, w_, xq_, xil_, out_, (int) n_in_, (int) rows_, T, s_); + } else { + native_quantize_q8_1(mixed_, xq_, (int) n_in_, T, s_); + native_mmvq(type_, w_, xq_, out_, (int) n_in_, (int) rows_, T, s_); + } + argmax_rows(out_, T, (int) rows_, scratch_, h_idx_, s_, h_val_); + if (logits) copy_from_mapped(h_logits_, out_, (int64_t) T * rows_, s_); + mapped_bump(h_done_, s_); + } catch (const std::exception& x) { + cudaGraph_t g = nullptr; + cudaStreamEndCapture(s_, &g); + if (g) cudaGraphDestroy(g); + err = std::string("split head: ") + x.what(); + return false; + } + cudaGraph_t g = nullptr; + if (cudaStreamEndCapture(s_, &g) != cudaSuccess || cudaGraphInstantiate(&exec_[T][logits], g, 0) != cudaSuccess) { + if (g) cudaGraphDestroy(g); + err = "split head: capture failed"; + return false; + } + cudaGraphDestroy(g); + cudaGraphUpload(exec_[T][logits], s_); + cudaStreamSynchronize(s_); + return true; +} + +bool SplitHead::submit(int T, bool logits, std::string& err) { + if (T < 1 || T > max_t_) { err = "split head: window size out of range"; return false; } + DeviceScope scope(dev_, main_); + if (!capture(T, logits, err)) return false; + if (cudaGraphLaunch(exec_[T][logits], s_) != cudaSuccess) { + err = std::string("split head: launch: ") + cudaGetErrorString(cudaGetLastError()); + return false; + } + (void) cudaStreamQuery(s_); // submit now + return true; +} + +bool SplitHead::wait(uint32_t n, std::string& err) { + const auto t0 = std::chrono::steady_clock::now(); + while (*(volatile uint32_t*) h_done_ < n) { + _mm_pause(); + if (std::chrono::steady_clock::now() - t0 > std::chrono::seconds(5)) { + DeviceScope scope(dev_, main_); + err = std::string("split head: its part never finished (") + cudaGetErrorString(cudaStreamQuery(s_)) + ")"; + return false; + } + } + return true; +} + +} // namespace strata::core diff --git a/src/core/verify.cpp b/src/core/verify.cpp index 54df72b4857..ddbcd88906b 100644 --- a/src/core/verify.cpp +++ b/src/core/verify.cpp @@ -5,22 +5,20 @@ #endif #include "strata/core/native_head.hpp" -#include "strata/core/on_device.hpp" #include "strata/kernels/iq_kernels.hpp" #include "strata/kernels/cpu/expert_layout.hpp" #include "strata/kernels/bf16_gemv.hpp" -#include "strata/kernels/native_router.hpp" -#include "strata/kernels/native_moe.hpp" #include "strata/kernels/cpu/expert.hpp" #include "strata/kernels/elementwise.hpp" #include "strata/kernels/fused_gr.hpp" -#include "strata/kernels/cvec.hpp" #include "strata/kernels/gr.hpp" -#include "strata/kernels/kv_q4.hpp" #include "strata/kernels/kv_q8.hpp" #include "strata/kernels/native_mmvq.hpp" +#include "strata/kernels/native_moe.hpp" +#include "strata/kernels/native_ple_postops.hpp" #include "strata/kernels/native_qsa.hpp" #include "strata/kernels/native_qsa_indexer.hpp" +#include "strata/kernels/native_router.hpp" #include "strata/kernels/native_rope.hpp" #include "strata/kernels/ngram.hpp" #include "strata/kernels/ple.hpp" @@ -28,20 +26,13 @@ #include "strata/kernels/qsa_decode_attn.hpp" #include "strata/kernels/qsa_select.hpp" #include "strata/kernels/quantize_act.hpp" -#include "strata/kernels/rope.hpp" #include "strata/kernels/s2_expert_grouped.hpp" #include "strata/kernels/sampler.hpp" -#include "strata/core/progress.hpp" #include "strata/kernels/shared_expert.hpp" #include "strata/kernels/verify_kernels.hpp" -#include #include -#include -#include -#include #include -#include #include #include #include @@ -67,8 +58,8 @@ struct Bump { } }; -bool mapped(size_t bytes, void** h, void** d) { - if (cudaHostAlloc(h, bytes, cudaHostAllocMapped) != cudaSuccess) return false; +bool mapped(size_t bytes, void** h, void** d, unsigned flags = cudaHostAllocMapped) { + if (cudaHostAlloc(h, bytes, flags) != cudaSuccess) return false; std::memset(*h, 0, bytes); return cudaHostGetDevicePointer(d, *h, 0) == cudaSuccess; } @@ -100,54 +91,46 @@ bool native_of(const WeightRef* w, const std::string& name, std::string& err) { } // namespace -namespace { -std::atomic g_diag_verifier{nullptr}; -void diag_active_verifier(std::FILE* f) { - if (const Verifier* v = g_diag_verifier.load()) v->diag(f); -} -} // namespace - -void Verifier::diag(std::FILE* f) const { - auto rd = [](const uint32_t* p) { return p ? *(const volatile uint32_t*) p : 0u; }; - std::fprintf(f, " verify window: %d tokens at position %lld, host at layer step %u; the GPU rang %u; flags: " - "served %u, plan (A) %u, copies (B) %u\n", last_t_, (long long) last_pos0_, cur_layer_ + 1, - rd(h_seq_), rd(h_flag_), rd(h_flagA_), rd(h_flagB_)); -} - Verifier::~Verifier() { - const Verifier* self = this; - g_diag_verifier.compare_exchange_strong(self, nullptr); if (cs_) cudaStreamSynchronize(cs_); for (auto& e : exec_) if (e) cudaGraphExecDestroy(e); if (commit_exec_) cudaGraphExecDestroy(commit_exec_); if (cs_) cudaStreamDestroy(cs_); if (copy_) { cudaStreamSynchronize(copy_); cudaStreamDestroy(copy_); } + if (side_) cudaStreamDestroy(side_); + if (shs_) cudaStreamDestroy(shs_); + if (g2s_) cudaStreamDestroy(g2s_); + if (cps_) cudaStreamDestroy(cps_); + if (kst_) cudaStreamDestroy(kst_); + if (sgs_) cudaStreamDestroy(sgs_); + for (cudaEvent_t e : {join_, bfork_, bjoin_, pfork_, pjoin_, shfork_, shjoin_, res_ready_, done_, g2fork_, + g2join_[0], g2join_[1], cpfork_, cpjoin_[0], cpjoin_[1], efork_, rdone_, kjoin_, iproj_, + sgjoin_}) + if (e) cudaEventDestroy(e); if (arena_) cudaFree(arena_); - void* hosts[] = {h_tok_, h_step_, h_pos_, h_commit_, h_ple_, h_out_, h_x_, h_ids_, h_w_, h_seq_, h_flag_, h_ymiss_, - h_flagA_, h_plan_, h_flagB_}; + if (stamps_) cudaFree(stamps_); + void* hosts[] = {h_tok_, h_val_, h_step_, h_pos_, h_commit_, h_ple_, h_out_, h_x_, h_ids_, h_w_, h_seq_, h_flag_, + h_ymiss_, + h_flagA_, h_plan_, h_flagB_, h_pids_, h_pw_, h_pseq_, h_pleflag_, h_rows_, h_res_, h_flag2_, + h_list2_, h_xk_, h_xq1_}; for (void* h : hosts) if (h) cudaFreeHost(h); } bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& ss, const VerifyHits& hits, const NativeHead* head, int max_t, std::string& err) { - g_diag_verifier.store(this); - diag_verify_fn().store(&diag_active_verifier); - cudaGetDevice(&device_); // a layer split's stage on another GPU: its streams, graphs and buffers live there wt_ = &wt; g_ = &g; ss_ = &ss; hits_ = hits; head_ = head; max_t_ = max_t; - sampling_.greedy = true; // a fresh verifier samples greedily until set_sampling says otherwise - sampling_.temperature = 0.0f; if (max_t < 2 || max_t > strata::kernels::kVerifyMaxT || max_t > strata::kernels::cpu::MAXT) { err = "verify: the window must hold 2.." + std::to_string(strata::kernels::kVerifyMaxT) + " tokens"; return false; } - if (hits.d_res == nullptr || hits.cache_base == nullptr || hits.blob <= 0) { + if (hits.res == nullptr || hits.cache_base == nullptr || hits.blob <= 0 || hits.slots <= 0) { err = "verify: needs the profile-filled VRAM expert tier (--expert-profile and --expert-cache)"; return false; } @@ -161,10 +144,9 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& err = "verify: geometry differs from the artifact's"; return false; } - if (le_ < 0) le_ = g.n_layers; - if (lb_ < 0 || lb_ >= le_ || le_ > g.n_layers || (lb_ > 0 && hand_in_ == nullptr) || - (le_ < g.n_layers && hand_out_ == nullptr)) { - err = "verify: the stage's layer range or its hand-off buffers are wrong"; + if (ss.ple.ready() && (!strata::kernels::ple_native_postops_enabled() || !strata::kernels::ple_native_bf16() || + (ss.ple.w.key_bf16 == nullptr && ss.ple.w.key_native_data == nullptr))) { + err = "verify: the PLE block needs the native postops and value and a BF16 or native key"; return false; } const WeightRef* wo = wt.find("output.weight"); @@ -177,6 +159,11 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& attn_scratch_floats_ = (int64_t) strata::kernels::qsa_decode_attn_scratch_floats(cap_, s); const uint64_t T = (uint64_t) max_t, N = (uint64_t) g.n_embd, HC = (uint64_t) g.hc, K = (uint64_t) ss.k; + head_rows_ = (head != nullptr && head->loaded()) ? head->rows() : n_vocab_; + if (head_rows_ != n_vocab_ && (shead_ == nullptr || shead_->split() != head_rows_)) { + err = "verify: the head holds part of the vocabulary without its other part"; + return false; + } const uint64_t C = (uint64_t) g.ssm_conv_channels, ZV = (uint64_t) g.ssm_value_dim, HV = (uint64_t) g.ssm_v_heads; const uint64_t NH = (uint64_t) g.n_head, HD = (uint64_t) g.head_dim, NKV = (uint64_t) g.n_head_kv; const uint64_t IQ = (uint64_t) g.idx_q_heads, ID = (uint64_t) g.idx_key_dim; @@ -185,13 +172,22 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& const uint64_t TS = (uint64_t) (s.idx_block - 1) * ID; const int max_in = (int) std::max(std::max(N, ZV), NH * HD); - // ---- mapped staging + // ---- mapped staging (portable where the second GPU writes: its rows and flags) + const unsigned portable = cudaHostAllocMapped | cudaHostAllocPortable; + list2_stride_ = (4 + (int64_t) (T * K) + 3) & ~3ll; + const auto& lay = strata::kernels::cpu::expert_layout(); + quant_router_ = strata::kernels::native_router_enabled() && lay.native; + xk_layer_.assign((size_t) g.n_layers, 0); + for (int64_t l = 0; quant_router_ && l < g.n_layers; ++l) + xk_layer_[(size_t) l] = pool_takes_q8k(lay.fmt[(size_t) l]); + xk_stride_ = ((int64_t) (N / 256) * 292 + 15) & ~15ll; bool ok = mapped(T * 4, (void**) &h_tok_, (void**) &m_tok_) && mapped(T * strata::kernels::kStepCount * 4, (void**) &h_step_, (void**) &m_step_) && - mapped(T * (NH + NKV + IQ) * 4, (void**) &h_pos_, (void**) &m_pos_) && + mapped(T * NH * 4, (void**) &h_pos_, (void**) &m_pos_) && mapped((2 + T) * 4 + 16, (void**) &h_commit_, (void**) &m_commit_) && mapped(T * N * 4, (void**) &h_ple_, (void**) &m_ple_) && mapped(T * 4 + 16, (void**) &h_out_, (void**) &m_out_) && + mapped(T * 4 + 16, (void**) &h_val_, (void**) &m_val_) && mapped(T * N * 4, (void**) &h_x_, (void**) &m_x_) && mapped(T * K * 4, (void**) &h_ids_, (void**) &m_ids_) && mapped(T * K * 4, (void**) &h_w_, (void**) &m_w_) && @@ -199,40 +195,60 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& mapped(64, (void**) &h_flag_, (void**) &m_flag_) && mapped(64, (void**) &h_flagA_, (void**) &m_flagA_) && mapped(64, (void**) &h_flagB_, (void**) &m_flagB_) && - mapped(T * K * N * 4, (void**) &h_ymiss_, (void**) &m_ymiss_); + mapped(T * K * N * 4, (void**) &h_ymiss_, (void**) &m_ymiss_, portable) && + mapped(128, (void**) &h_flag2_, (void**) &m_flag2_, portable) && + mapped((size_t) (2 * list2_stride_) * 4, (void**) &h_list2_, (void**) &m_list2_) && + mapped(T * (uint64_t) xk_stride_, (void**) &h_xk_, (void**) &m_xk_) && + mapped(T * (N / 32) * 36, (void**) &h_xq1_, (void**) &m_xq1_) && + mapped(2 * T * K * 4, (void**) &h_pids_, (void**) &m_pids_) && + mapped(2 * T * K * 4, (void**) &h_pw_, (void**) &m_pw_) && + mapped(64, (void**) &h_pseq_, (void**) &m_pseq_) && + mapped(64, (void**) &h_pleflag_, (void**) &m_pleflag_) && + mapped(T * HC * N * 4, (void**) &h_rows_, (void**) &m_rows_); if (!ok) { err = "verify: mapped staging allocation failed"; return false; } - // the GPU plan: counts(4) | start(cap+1) | dst(cap) | tok(cap) | pad | ptr(cap u64) | ptr2(cap u64) | start2(cap+1) + // A plan block, one per token group: counts(4) | start(cap+1) | dst(cap) | tok(cap) | pad | ptr(cap u64) | + // ptr2(cap u64) | start2(cap+1). The main GPU's hits (`verify_hit_plan`, device) fill counts [groups, entries], + // start, dst, tok and ptr; the pool's PCIe share (mapped, `sink_`) counts [-, entries, groups], dst, tok, ptr2 and + // start2. { const int64_t cap = (int64_t) (T * K); - const int64_t i32 = 4 + (cap + 1) + cap + cap; - const int64_t ptr_off = (i32 + 1) & ~1ll; - plan_i32_ = ptr_off + 4 * cap + (cap + 1) + 1; + plan_ptr_off_ = (4 + (cap + 1) + cap + cap + 1) & ~1ll; + plan_i32_ = (plan_ptr_off_ + 4 * cap + (cap + 1) + 2) & ~1ll; + plan_stride_ = plan_i32_ + 16; if (!mapped((size_t) plan_i32_ * 4 * 2 + 64, (void**) &h_plan_, (void**) &m_plan_)) { err = "verify: mapped plan allocation failed"; return false; } sink_.counts = h_plan_; - sink_.start = h_plan_ + 4; - sink_.dst = sink_.start + cap + 1; + sink_.dst = h_plan_ + 4 + cap + 1; sink_.tok = sink_.dst + cap; - sink_.ptr = (unsigned long long*) (h_plan_ + ptr_off); - sink_.ptr2 = sink_.ptr + cap; - sink_.start2 = h_plan_ + ptr_off + 4 * cap; + sink_.ptr2 = (unsigned long long*) (h_plan_ + plan_ptr_off_) + cap; + sink_.start2 = h_plan_ + plan_ptr_off_ + 4 * cap; sink_.cap = cap; sink_.publish = &Verifier::publish_plan; + sink_.host_rows_only = true; // the gather-combine (post) reads the GPUs' from VRAM sink_.fetch = &Verifier::fetch_dma; sink_.ctx = this; } + // the residency snapshot, rounded up to whole 16-byte copies + res_words_ = (g.n_layers * g.n_expert + 3) & ~3ll; + if (!mapped((size_t) res_words_ * 4, (void**) &h_res_, (void**) &m_res_)) { + err = "verify: mapped residency allocation failed"; + return false; + } + std::memcpy(h_res_, hits.res, (size_t) (g.n_layers * g.n_expert) * sizeof(int32_t)); // ---- the device arena: the same sequence counted, then carved auto carve = [&](Bump& b) { tok_ = b.take(T); step_ = b.take(T * strata::kernels::kStepCount); - pos_ = b.take(T * (NH + NKV + IQ)); commit_ = b.take(2 + T); + pos_ = b.take(T * NH); commit_ = b.take(2 + T); ple_ = b.take(T * N); emb_ = b.take(T * N); R_ = b.take(T * HC * N); mixed_ = b.take(T * N); bo_ = b.take(T * N); inj_ = b.take(T * HC); inj2_ = b.take(T * HC); - lo_ = b.take(T * (uint64_t) g.hc_lr); rs_ = b.take(T * HC); xn_ = b.take(T * HC * N); + lo_ = b.take(T * (uint64_t) g.hc_lr); rs_ = b.take(T * HC); + grs_ = (float*) b.take(strata::kernels::fused_gr_scratch_bytes()); xq_ = b.take(strata::kernels::native_q8_1_bytes(max_in, (int) T)); + xil_ = b.take(strata::kernels::native_q8_1_il_bytes(max_in, (int) T)); qkv_L_ = b.take(nG * T * C); h_L_ = b.take(nG * T * C); gate_L_ = b.take(nG * T * HV); beta_L_ = b.take(nG * T * HV); z_ = b.take(T * ZV); y_ = b.take(T * ZV); y_dummy_ = b.take(T * ZV); @@ -240,13 +256,20 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& kcur_ = b.take(T * NKV * HD); vcur_ = b.take(T * NKV * HD); idx_raw_L_ = b.take(nQ * T * ID); qidx_ = b.take(T * IQ * ID); scores_ = b.take(T * (uint64_t) max_blocks_); sel_ = b.take(T * (uint64_t) cap_); - attn_ = b.take(T * NH * HD); attn32_ = b.take(T * NH * HD); + attn32_ = b.take(T * NH * HD); attn_scratch_ = b.take(T * (uint64_t) attn_scratch_floats_); tail_snap_ = b.take(nQ * TS); logits_ = b.take(T * (uint64_t) g.n_expert); w_ = b.take(T * K); ids_ = b.take(T * K); - shared_ = b.take(T * N); parts_ = b.take(T * K * N); hit_out_ = b.take(T * K * N); + plogits_ = b.take(2 * T * (uint64_t) g.n_expert); pw_ = b.take(T * K); + pids_ = b.take(T * K); + pmixed_ = b.take(T * N); gate_ema_ = b.take((uint64_t) g.n_layers * HC * N); + pbias_ = b.take((uint64_t) (g.n_layers * g.n_expert)); + shared_ = b.take(T * N); hit_out_ = b.take(T * K * N); hit_slot_ = b.take(T * K); hit_dst_ = b.take(T * K); hit_count_ = b.take(4); - plan_ = b.take(2 * ((uint64_t) plan_i32_ + 16)); + list2_ = b.take(2 * (uint64_t) list2_stride_); + ring_count_ = b.take(2); + plan_ = b.take(2 * (uint64_t) plan_stride_); dplan_ = b.take(2 * (uint64_t) plan_stride_); + res_ = b.take((uint64_t) res_words_); slot_ptr_ = b.take((uint64_t) hits.slots); staging_ = b.take((uint64_t) kStagingBlobs * strata::kernels::cpu::expert_layout().max_blob); hit_xq_ = b.take(T * (N / 32) * 34); hit_xs_ = b.take(T * (N / 32)); nat_xq_ = b.take(T * (N / 32) * 36); @@ -254,10 +277,16 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& strata::kernels::moe_hit_grouped_scratch_bytes((int64_t) (T * K), g.n_embd, g.n_ff), strata::kernels::native_expert_scratch_bytes((int64_t) (T * K), g.n_ff))); head_mixed_ = b.take(T * N); head_inj_ = b.take(HC); - sh_bf16_ = b.take(T * N); sh_gate_ = b.take(T * (uint64_t) g.n_ff); + sh_gate_ = b.take(T * (uint64_t) g.n_ff); sh_up_ = b.take(T * (uint64_t) g.n_ff); sh_g_ = b.take(T + 4); head_logits_ = b.take(T * (uint64_t) n_vocab_); + arg_scratch_ = b.take(strata::kernels::argmax_rows_scratch_bytes((int) T)); + if (shead_ != nullptr) head_full_ = b.take(T * (uint64_t) n_vocab_); + one_ = b.take(4); hist_snap_ = b.take(T * HS); + const uint64_t PD = (uint64_t) strata::kernels::NG_HC_DIM; + ple_key_ = b.take(T * PD); ple_val_ = b.take(T * N); ple_nkey_ = b.take(T * PD); + ple_norm_ = b.take(T * PD); ple_gated_ = b.take(T * PD); ple_gate_ = b.take(T * HC); }; Bump count; carve(count); @@ -266,41 +295,70 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& return false; } cudaMemset(arena_, 0, count.used); - prof_on_ = std::getenv("STRATA_VERIFY_PROFILE") != nullptr; - if (prof_on_) { - const size_t np = (size_t) g.n_layers * kProfPer + 4; - if (cudaMalloc((void**) &prof_, np * 8) != cudaSuccess) { prof_on_ = false; prof_ = nullptr; cudaGetLastError(); } - else { cudaMemset(prof_, 0, np * 8); prof_h_.assign(np, 0); } - } Bump real; real.base = (uint8_t*) arena_; carve(real); + { + const int32_t one = 1; + if (cudaMemcpy(one_, &one, sizeof one, cudaMemcpyHostToDevice) != cudaSuccess) { + err = "verify: the arena could not be set"; + return false; + } + } sink_.staging = (unsigned long long) staging_; sink_.staging_cap = kStagingBlobs; (void) TS; + { // the gates' running averages start at a gate's middle + const std::vector half((size_t) (g.n_layers * g.hc * g.n_embd), 0.5f); + if (cudaMemcpy(gate_ema_, half.data(), half.size() * sizeof(float), cudaMemcpyHostToDevice) != cudaSuccess) { + err = "verify: the gates' averages could not be set"; + return false; + } + } + { + std::vector sp((size_t) hits.slots); + for (int64_t i = 0; i < hits.slots; ++i) + sp[(size_t) i] = (unsigned long long) (hits.cache_base + (hits.slot_off != nullptr ? hits.slot_off[i] + : (uint64_t) i * (uint64_t) hits.blob)); + if (cudaMemcpy(slot_ptr_, sp.data(), sp.size() * sizeof(unsigned long long), cudaMemcpyHostToDevice) != cudaSuccess) { + err = "verify: the slot table could not be uploaded"; + return false; + } + } if (cudaStreamCreateWithFlags(©_, cudaStreamNonBlocking) != cudaSuccess) { err = "verify: copy stream create failed"; return false; } - if (cudaStreamCreateWithFlags(&cs_, cudaStreamNonBlocking) != cudaSuccess) { + if (cudaStreamCreateWithFlags(&cs_, cudaStreamNonBlocking) != cudaSuccess || + cudaStreamCreateWithFlags(&side_, cudaStreamNonBlocking) != cudaSuccess || + cudaStreamCreateWithFlags(&shs_, cudaStreamNonBlocking) != cudaSuccess || + cudaStreamCreateWithFlags(&g2s_, cudaStreamNonBlocking) != cudaSuccess || + cudaStreamCreateWithFlags(&cps_, cudaStreamNonBlocking) != cudaSuccess || + cudaStreamCreateWithFlags(&kst_, cudaStreamNonBlocking) != cudaSuccess || + cudaEventCreateWithFlags(&kjoin_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&iproj_, cudaEventDisableTiming) != cudaSuccess || + cudaStreamCreateWithFlags(&sgs_, cudaStreamNonBlocking) != cudaSuccess || + cudaEventCreateWithFlags(&sgjoin_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&cpfork_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&cpjoin_[0], cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&cpjoin_[1], cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&efork_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&rdone_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&g2fork_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&g2join_[0], cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&g2join_[1], cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&join_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&bfork_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&bjoin_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&pfork_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&pjoin_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&shfork_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&shjoin_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&res_ready_, cudaEventDisableTiming) != cudaSuccess || + cudaEventCreateWithFlags(&done_, cudaEventDisableTiming) != cudaSuccess) { err = "verify: stream create failed"; return false; } - // E-6: a layer whose routed experts are all resident is planned on the device (STRATA_VERIFY_DEVICE_PLAN=1: on; - // exact, but neutral on RIBPC 1-2 GPUs: off by default) - { - const char* v = std::getenv("STRATA_VERIFY_DEVICE_PLAN"); - device_plan_ = v != nullptr && std::atoi(v) != 0; - } - if (device_plan_) { - bool ok2 = cudaMalloc((void**) &skip_, 64) == cudaSuccess && cudaMemset(skip_, 0, 64) == cudaSuccess; - if (ok2 && hits.slot_off != nullptr && hits.n_slots > 0) { - ok2 = cudaMalloc((void**) &slot_off_d_, (size_t) hits.n_slots * sizeof(unsigned long long)) == cudaSuccess && - cudaMemcpy(slot_off_d_, hits.slot_off, (size_t) hits.n_slots * sizeof(unsigned long long), - cudaMemcpyHostToDevice) == cudaSuccess; - } - if (!ok2) { cudaGetLastError(); device_plan_ = false; } - } std::fprintf(stderr, "strata verify: window up to %d tokens, %.1f MiB of device buffers\n", max_t, (double) count.used / 1048576.0); return true; @@ -332,33 +390,56 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { (uint64_t) g.ssm_conv_channels * (g.ssm_d_conv - 1); const int64_t HS = (int64_t) NG_HIST * NG_HC_DIM; const int64_t TS = (s.idx_block - 1) * ID; - const bool ple_on = ss.ple.ready() && ple_stage(); + const bool ple_on = ss.ple.ready(); auto Rt = [&](int t) { return R_ + (size_t) t * HC * N; }; const int G = (split_ && T >= 2) ? 2 : 1; - static const bool dec_batch = [] { const char* v = std::getenv("STRATA_DEC_BATCH"); return v == nullptr || std::atoi(v) != 0; }(); - auto stamp = [&](int64_t l, int i, int grp) { if (prof_on_ && grp == 0) gpu_stamp(prof_, (int) (l * kProfPer + i), cs); }; const int tb_[2] = {0, (T + 1) / 2}, te_[2] = {G == 2 ? (T + 1) / 2 : T, T}; groups_[T] = G; + // --window-profile: row l < n_layers holds layer l's stages, row n_layers the window's own stamps + auto stamp = [&](int64_t row, int i, cudaStream_t st) { + if (stamps_ != nullptr && G == 1) + gpu_stamp(stamps_ + ((size_t) T * (size_t) (g.n_layers + 1) + (size_t) row) * kStamps + i, st); + }; + stamp(g.n_layers, 0, cs); + if (gpu2_ && !quant_router_) { + err = "verify: the second GPU takes the native router's q8_1 rows (a native pack)"; + return false; + } + // the layer's short branches: `fork` starts stream `to` from the main stream's work so far, `mark` records what + // `from` has queued, `wait` makes the main stream wait for a mark + auto fork = [&](cudaStream_t to, cudaEvent_t ev) { + if (cudaEventRecord(ev, cs) == cudaSuccess && cudaStreamWaitEvent(to, ev, 0) == cudaSuccess) return true; + err = "verify: a branch could not fork"; + return false; + }; + auto mark = [&](cudaStream_t from, cudaEvent_t ev) { + if (cudaEventRecord(ev, from) == cudaSuccess) return true; + err = "verify: a branch could not join"; + return false; + }; + auto wait = [&](cudaEvent_t ev) { + if (cudaStreamWaitEvent(cs, ev, 0) == cudaSuccess) return true; + err = "verify: a branch could not join"; + return false; + }; - // ---- the window's inputs, from mapped staging + // ---- the window's inputs, from mapped staging; the residency snapshot on the shared expert's branch (the first + // hit plan waits for it) copy_i32_from_mapped(tok_, m_tok_, T, cs); copy_i32_from_mapped(step_, m_step_, (int64_t) T * kStepCount, cs); - copy_i32_from_mapped(pos_, m_pos_, (int64_t) MT * (NH + NKV + IQ), cs); - // per-ROW positions of the K rows [t][NKV] and the indexer query rows [t][IQ] (for batched RoPE) - const int32_t* pos_k = pos_ + MT * NH; - const int32_t* pos_i = pos_ + MT * (NH + NKV); - if (ple_on) copy_from_mapped(ple_, m_ple_, (int64_t) T * N, cs); + copy_i32_from_mapped(pos_, m_pos_, (int64_t) T * NH, cs); + if (cudaEventRecord(shfork_, cs) != cudaSuccess || cudaStreamWaitEvent(shs_, shfork_, 0) != cudaSuccess) { + err = "verify: the residency copy could not fork"; + return false; + } + copy_from_mapped((float*) res_, (const float*) m_res_, res_words_, shs_); // int32 words, copied as they are + if (cudaEventRecord(res_ready_, shs_) != cudaSuccess) { + err = "verify: the residency copy could not join"; + return false; + } - // ---- the embeddings, broadcast to the hc streams - or, in a later stage of a layer split, the previous stage's - // residual, pending write and inject (see set_stage) - const int64_t HB = Verifier::handoff_floats(g); - if (lb_ > 0) { - for (int t = 0; t < T; ++t) { - copy_from_mapped(Rt(t), hand_in_ + (size_t) t * HB, HC * N, cs); - copy_from_mapped(bo_ + (size_t) t * N, hand_in_ + (size_t) t * HB + HC * N, N, cs); - copy_from_mapped(inj2_ + (size_t) t * HC, hand_in_ + (size_t) t * HB + HC * N + N, HC, cs); - } - } else if (const NativeEmbed* ne = native_embed()) { // plan v0.3 P6: the GGUF-form table + // ---- the embeddings, broadcast to the hc streams + if (const NativeEmbed* ne = native_embed()) { // plan v0.3 P6: the GGUF-form table ne->gather_dev(tok_, T, emb_, cs); broadcast_streams(emb_, R_, N, (int) HC, T, cs); } else { @@ -377,6 +458,16 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { broadcast_streams(emb_, R_, N, (int) HC, T, cs); } + // set_predict: this layer's FFN read also estimates the next layer's FFN input (that layer's read of this state, + // its gates their running average over the windows), and on a side branch forked after the read the next layer's + // router runs on it with that layer's logit bias. When learning, the FFN read of a layer the previous one + // predicts keeps its gates' average, and the layer corrects its bias by its own logits on that branch. + auto predicts = [&](int64_t l) { + return predict_ != nullptr && G == 1 && l + 1 < g.n_layers && native_router_enabled(); + }; + auto corrects = [&](int64_t l) { return learn_ && l > 0 && predicts(l - 1); }; + auto branches = [&](int64_t l) { return predicts(l) || corrects(l); }; + // per-layer state indices (GDN and QSA layers are numbered separately) std::vector gdn_idx((size_t) g.n_layers, -1), qsa_idx((size_t) g.n_layers, -1); { @@ -387,10 +478,56 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { } } + // the shared expert: in a window of one token group on a branch of its own forked after the FFN read (`efork_`), + // beside the router and the routed experts (joined before the combine); in a split window inline. It ends with + // its down projection: the combine applies its scalar gate, whose logits come from a branch of their own. + const bool beside = G == 1; + auto shared = [&](int64_t l, int grp) -> bool { + const int tb = tb_[grp], te = te_[grp], n = te - tb; + const LayerView v(wt, l); + const WeightRef *wgi = need(v, "ffn_gate_inp_shexp.weight", err), *wsg = need(v, "ffn_gate_shexp.weight", err), + *wsu = need(v, "ffn_up_shexp.weight", err), *wsd = need(v, "ffn_down_shexp.weight", err); + if (!wgi || !wsg || !wsu || !wsd) return false; + if (!native_of(wsg, v.name("ffn_gate_shexp.weight"), err) || !native_of(wsu, v.name("ffn_up_shexp.weight"), err) || + !native_of(wsd, v.name("ffn_down_shexp.weight"), err)) + return false; + NativeSharedWeights nsw; + nsw.gate_type = wsg->native_type; nsw.gate_data = wsg->native_data; + nsw.up_type = wsu->native_type; nsw.up_data = wsu->native_data; + nsw.down_type = wsd->native_type; nsw.down_data = wsd->native_data; + nsw.q8_1 = xq_; // the main stream's next use is the next layer's mixer + nsw.q8_1_il = xil_; + cudaStream_t ws = cs, gs = cs; + if (beside) { + if (cudaStreamWaitEvent(shs_, efork_, 0) != cudaSuccess || + cudaStreamWaitEvent(sgs_, efork_, 0) != cudaSuccess) { + err = "verify: the shared expert's branch could not fork"; + return false; + } + ws = shs_; + gs = sgs_; + } + try { + shared_expert_multi(n, mixed_ + tb * N, nsw, (const uint16_t*) wgi->data, sh_gate_ + (size_t) tb * g.n_ff, + sh_up_ + (size_t) tb * g.n_ff, sh_g_ + tb, shared_ + tb * N, N, g.n_ff, ws, gs); + } catch (const std::exception& e) { + err = std::string("verify shared expert: ") + e.what(); + return false; + } + if (beside) { + stamp(l, kLayerStamps, shs_); + if (cudaEventRecord(shjoin_, shs_) != cudaSuccess || cudaEventRecord(sgjoin_, sgs_) != cudaSuccess) { + err = "verify: the shared expert's branch could not join"; + return false; + } + } + return true; + }; + // ---------------------------------------------------------------- pre(l, group): up to the ring auto pre = [&](int64_t l, int grp) -> bool { const int tb = tb_[grp], te = te_[grp], n = te - tb; - stamp(l, 0, grp); + stamp(l, 0, cs); const LayerView v(wt, l); const char* pfx[2] = {"hc_attn_", "hc_ffn_"}; const WeightRef *wn[2], *wd[2], *wu[2], *wi[2]; @@ -401,28 +538,51 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { wi[h] = need(v, (std::string(pfx[h]) + "inject.weight").c_str(), err); if (!wn[h] || !wd[h] || !wu[h] || !wi[h]) return false; } - // the previous layer's FFN write, folded into this layer's first read (a control vector after it has - // already applied it) - bool pending = l > 0 && !cvec().covers(l - 1); + bool pending = l > 0; // the previous layer's FFN write, folded into this layer's first read + // the projections of the group's tokens: 2+ tokens are quantized with the interleaved copy the multi-token + // kernels read + auto quant = [&](const float* x, int64_t n_in) { + if (n >= 2) native_quantize_q8_1_il(x, xq_, xil_, (int) n_in, n, cs); + else native_quantize_q8_1(x, xq_, (int) n_in, n, cs); + }; + auto mm = [&](int type, const void* w, float* y, int64_t n_in, int64_t n_out, cudaStream_t ps) { + native_mmvq_il(type, w, xq_, xil_, y, (int) n_in, (int) n_out, n, ps); + }; + auto proj = [&](const WeightRef* w, float* y, int64_t n_in, int64_t n_out, cudaStream_t ps) { + mm(w->native_type, w->native_data, y, n_in, n_out, ps); + }; if (l == 1 && ple_on) { - float* normalized = (float*) ((uint8_t*) ss.ple.scratch + ple_block_scratch_bytes()); - for (int t = tb; t < te; ++t) { - gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs); - PleOut po; - po.normalized = normalized; - po.result = Rt(t); - try { - ple_block(ple_ + t * N, Rt(t), ss.ple.hist, ss.ple.w, po, ss.ple.scratch, cs); - ple_history_advance(ss.ple.hist, normalized, cs); - } catch (const std::exception& e) { - err = std::string("verify PLE: ") + e.what(); - return false; + if (grp == 0) { // the host reads the rows while layer 0 runs + wait_flag_ge(m_pleflag_, 1, cs); + copy_from_mapped(ple_, m_ple_, (int64_t) T * N, cs); + } + for (int t = tb; t < te; ++t) gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs); + // the PLE block: the group's projections with each weight read once, then the rest of the block over its + // tokens in order, which leaves the history after each of them for the commit + const PleWeights& pw = ss.ple.w; + const float* emb = ple_ + tb * N; + float* key = ple_key_ + (size_t) tb * NG_HC_DIM; + float* norm = ple_norm_ + (size_t) tb * NG_HC_DIM; + try { + if (pw.key_bf16 != nullptr) { + bf16_gemv_fp32_mmvf_multi(emb, pw.key_bf16, key, N, NG_HC_DIM, n, cs); + } else { + quant(emb, N); + mm(pw.key_native_type, pw.key_native_data, key, N, NG_HC_DIM, cs); } - copy_from_mapped(hist_snap_ + (size_t) t * HS, ss.ple.hist, HS, cs); + bf16_gemv_fp32_mmvf_multi(emb, pw.value_bf16, ple_val_ + tb * N, N, N, n, cs); + const NativePleTokensBuffers pb{ple_nkey_ + (size_t) tb * NG_HC_DIM, norm, ple_gate_ + tb * HC, + ple_gated_ + (size_t) tb * NG_HC_DIM, norm, Rt(tb)}; + native_ple_postops_tokens(key, Rt(tb), ple_val_ + tb * N, ss.ple.hist, pw, pb, n, cs, + hist_snap_ + (size_t) tb * HS); + } catch (const std::exception& e) { + err = std::string("verify PLE: ") + e.what(); + return false; } pending = false; } - auto gr_read_group = [&](int half, bool apply, float* inj_prev, float* inj_out) { + auto gr_read_group = [&](int half, bool apply, float* inj_prev, float* inj_out, float* gate_ema, + const float* est_norm, const float* est_gates) { FusedGrArgs fa[kFusedGrMaxT]; for (int t = tb; t < te; ++t) { FusedGrArgs& a = fa[t - tb]; @@ -431,13 +591,13 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { a.w_norm = (const float*) wn[half]->data; a.w_down = (const uint16_t*) wd[half]->data; a.w_up = (const uint16_t*) wu[half]->data; a.w_inject = (const uint16_t*) wi[half]->data; a.eps = EPS; a.lo = lo_ + t * g.hc_lr; a.rs = rs_ + t * HC; - a.inject_out = inj_out + t * HC; a.mixed = mixed_ + t * N; + a.inject_out = inj_out + t * HC; a.mixed = mixed_ + t * N; a.gate_ema = gate_ema; + a.est_norm = est_norm; a.est_gates = est_gates; a.est = pmixed_ + t * N; } - fused_gr_read_multi(fa, n, xn_ + (size_t) tb * HC * N, cs, (prof_on_ && grp == 0) ? prof_ : nullptr, - (int) (l * kProfPer + (half == 0 ? 27 : 30))); + fused_gr_read_multi(fa, n, grs_, cs); }; - gr_read_group(0, pending, inj2_, inj_); - stamp(l, 1, grp); + gr_read_group(0, pending, inj2_, inj_, nullptr, nullptr, nullptr); + stamp(l, 1, cs); float* xm = mixed_ + tb * N; try { if (!is_qsa_layer(g, l)) { @@ -458,23 +618,28 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { float* hb = h_L_ + (size_t) gi * MT * C; float* gate = gate_L_ + (size_t) gi * MT * HV; float* beta = beta_L_ + (size_t) gi * MT * HV; - native_quantize_q8_1(xm, xq_, (int) N, n, cs); - native_mmvq(wqkv->native_type, wqkv->native_data, xq_, qkv + (size_t) tb * C, (int) N, (int) C, n, cs); - stamp(l, 2, grp); - gdn_conv_l2_multi(conv, qkv, (const float*) wc->data, hb, (int) C, (int) (2 * HK), EPS, n, cs, tb); - stamp(l, 3, grp); + // beside qkv and the conv: alpha/beta (a few blocks, latency-bound) and the gate projection, joined + // before the recurrence. qkv's grid is recorded first, so it is dispatched first and the conv runs + // beside the gate projection's last blocks. + if (!fork(side_, bfork_)) return false; gdn_ab_multi(xm, (const uint16_t*) wa->data, (const uint16_t*) wb->data, (const float*) wdt->data, (const float*) wsa->data, gate + (size_t) tb * HV, beta + (size_t) tb * HV, (int) N, (int) HV, - n, cs); - stamp(l, 4, grp); - native_mmvq(wg->native_type, wg->native_data, xq_, z_ + (size_t) tb * ZV, (int) N, (int) ZV, n, cs); - stamp(l, 5, grp); + n, side_); + if (!mark(side_, bjoin_)) return false; + quant(xm, N); + if (!fork(shs_, pfork_)) return false; + proj(wqkv, qkv + (size_t) tb * C, N, C, cs); + proj(wg, z_ + (size_t) tb * ZV, N, ZV, shs_); + if (!mark(shs_, pjoin_)) return false; + // a one-token window keeps its token: it advances the conv history and the state itself (commit) + gdn_conv_l2_multi(conv, qkv, (const float*) wc->data, hb, (int) C, (int) (2 * HK), EPS, n, cs, tb, + T == 1); + if (!wait(bjoin_) || !wait(pjoin_)) return false; // the recurrence from the untouched state over tokens [0, te); outputs only for this group's gdn_step_norm_multi(state, hb, (int) C, gate, beta, z_, (const float*) wnm->data, EPS, y_, (int) HK, - (int) HV, te, nullptr, cs, tb); - stamp(l, 6, grp); - native_quantize_q8_1(y_ + (size_t) tb * ZV, xq_, (int) ZV, n, cs); - native_mmvq(wout->native_type, wout->native_data, xq_, bo_ + tb * N, (int) ZV, (int) N, n, cs); + (int) HV, te, T == 1 ? one_ : nullptr, cs, tb); + quant(y_ + (size_t) tb * ZV, ZV); + proj(wout, bo_ + tb * N, ZV, N, cs); } else { // ======================= QSA ======================= const int64_t qi = qsa_idx[(size_t) l]; @@ -488,337 +653,363 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) { if (!native_of(wq, v.name("attn_q.weight"), err) || !native_of(wk, v.name("attn_k.weight"), err) || !native_of(wv, v.name("attn_v.weight"), err) || !native_of(wo, v.name("attn_output.weight"), err)) return false; - auto norm_rope = [&](float* data, const WeightRef* norm, int rows, int cols, const int32_t* pos) { - if (native_qsa_enabled()) native_qsa_rms_norm_weighted(data, (const float*) norm->data, data, cols, rows, EPS, cs); - else rms_norm_weighted(data, (const float*) norm->data, rows, cols, EPS, cs); - if (native_rope_enabled()) native_rope_apply(data, data, rows, cols, (int) s.n_rot, (float) qsa_freq_base(), pos, cs); - else rope_neox_apply(data, data, rows, cols, (int) s.n_rot, st.cos_tab, st.sin_tab, pos, cs); + // the rms norm and rope of the group's heads in one launch: rows of `cols` read at in + r * stride, + // token t's positions at pos_ + t * NH + auto norm_rope = [&](const float* in, int stride, float* out, const WeightRef* norm, int heads, + int cols, cudaStream_t ns) { + native_qsa_norm_rope_tokens(in, stride, (const float*) norm->data, out, cols, n * heads, EPS, + (int) s.n_rot, (float) qsa_freq_base(), pos_ + tb * NH, heads, (int) NH, + ns); }; + // three branches beside the values, joined before the attention: the indexer (its key and query + // projections in one launch, the keys appended, the queries' block scores and top-k), the keys and the + // queries. A grid is dispatched whole before a later kernel's blocks, so the keys' and values' + // projections go ahead of the queries' grid, and that grid waits for the indexer's projections: the + // append is dispatched ahead of it, and the indexer's last kernels end about when the queries' do. + if (!fork(side_, bfork_)) return false; float* idx_raw = idx_raw_L_ + (size_t) qi * MT * ID; - // the per-token GEMVs / norms / RoPEs / copies of this layer as one launch over the - // window's rows each - row-wise identical arithmetic (STRATA_DEC_BATCH=0: token by token) - const bool qb = dec_batch && n > 1 && native_qsa_enabled() && native_rope_enabled() && !st.kv_q4; - native_quantize_q8_1(xm, xq_, (int) N, n, cs); - if (qb) bf16_gemv_fp32_mmvf_multi(mixed_ + tb * N, N, (const uint16_t*) wik->data, idx_raw + tb * ID, ID, N, ID, n, cs); - else for (int t = tb; t < te; ++t) - bf16_gemv_fp32_mmvf(mixed_ + t * N, (const uint16_t*) wik->data, idx_raw + t * ID, (int) N, (int) ID, cs); - stamp(l, 7, grp); - native_mmvq(wk->native_type, wk->native_data, xq_, kcur_ + tb * NKV * HD, (int) N, (int) (NKV * HD), n, cs); - native_mmvq(wv->native_type, wv->native_data, xq_, vcur_ + tb * NKV * HD, (int) N, (int) (NKV * HD), n, cs); - if (qb) norm_rope(kcur_ + tb * NKV * HD, wkn, (int) (n * NKV), (int) HD, pos_k + tb * NKV); - else for (int t = tb; t < te; ++t) norm_rope(kcur_ + t * NKV * HD, wkn, (int) NKV, (int) HD, pos_ + t * NH); - if (st.kv_q4) { // Q4_0 KV (kv_q4.hpp): K and V rotated before they are stored - fwht256_inplace_cuda(kcur_ + tb * NKV * HD, (int64_t) n * NKV, cs); - fwht256_inplace_cuda(vcur_ + tb * NKV * HD, (int64_t) n * NKV, cs); - } else if (st.kv_hybrid) { // K8V4: only V is rotated - fwht256_inplace_cuda(vcur_ + tb * NKV * HD, (int64_t) n * NKV, cs); - } - stamp(l, 8, grp); - if (grp == 0) copy_from_mapped(tail_snap_ + (size_t) qi * TS, st.idx_tail, TS, cs); - for (int t = tb; t < te; ++t) { - const int32_t* step_t = step_ + t * kStepCount; - if (st.kv_hybrid) { // K8V4: the unused half's lanes folded onto the used pool (layer.cpp) - kv_append_q8_step(st.k_q, st.k_q, st.k_scale, st.k_scale, st.page_table, step_t, - kcur_ + t * NKV * HD, kcur_ + t * NKV * HD, s, cs, nullptr); - kv_append_q4_step(st.v_q4, st.v_q4, st.page_table, step_t, vcur_ + t * NKV * HD, - vcur_ + t * NKV * HD, s, cs, nullptr); - } else if (st.kv_q4) - kv_append_q4_step(st.k_q4, st.v_q4, st.page_table, step_t, kcur_ + t * NKV * HD, - vcur_ + t * NKV * HD, s, cs, &st.host); - else if (st.kv_int8) - kv_append_q8_step(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, step_t, - kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs, &st.host); - else - kv_append_step(st.k_pool, st.v_pool, st.page_table, step_t, kcur_ + t * NKV * HD, - vcur_ + t * NKV * HD, s, cs, &st.host); - } + bf16_gemv_fp32_mmvf_multi(xm, (const uint16_t*) wik->data, idx_raw + tb * ID, N, ID, n, side_, + (const uint16_t*) wiq->data, qidx_ + tb * IQ * ID, IQ * ID); + if (!mark(side_, iproj_)) return false; const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; - for (int t = tb; t < te; ++t) - native_qsa_indexer_append(idx_raw + t * ID, step_ + t * kStepCount + kStepPos, 0, - (const float*) wikn->data, EPS, ib, s, st.max_cells, - (float) qsa_freq_base(), cs); - stamp(l, 9, grp); - native_mmvq(wq->native_type, wq->native_data, xq_, qfull_ + tb * NH * 2 * HD, (int) N, (int) (NH * 2 * HD), - n, cs); - if (qb) { - if (cudaMemcpy2DAsync(qcur_ + tb * NH * HD, (size_t) HD * 4, qfull_ + tb * NH * 2 * HD, (size_t) HD * 2 * 4, - (size_t) HD * 4, (size_t) (n * NH), cudaMemcpyDeviceToDevice, cs) != cudaSuccess) { - err = "verify: the q/gate split failed"; - return false; - } - norm_rope(qcur_ + tb * NH * HD, wqn, (int) (n * NH), (int) HD, pos_ + tb * NH); - bf16_gemv_fp32_mmvf_multi(mixed_ + tb * N, N, (const uint16_t*) wiq->data, qidx_ + tb * IQ * ID, IQ * ID, - N, IQ * ID, n, cs); - norm_rope(qidx_ + tb * IQ * ID, wiqn, (int) (n * IQ), (int) ID, pos_i + tb * IQ); - } else { - for (int t = tb; t < te; ++t) { - float* qc = qcur_ + t * NH * HD; - if (cudaMemcpy2DAsync(qc, (size_t) HD * 4, qfull_ + t * NH * 2 * HD, (size_t) HD * 2 * 4, - (size_t) HD * 4, (size_t) NH, cudaMemcpyDeviceToDevice, cs) != cudaSuccess) { - err = "verify: the q/gate split failed"; - return false; - } - norm_rope(qc, wqn, (int) NH, (int) HD, pos_ + t * NH); - if (st.kv_q4) fwht256_inplace_cuda(qc, NH, cs); // = - } - for (int t = tb; t < te; ++t) { - float* qx = qidx_ + t * IQ * ID; - bf16_gemv_fp32_mmvf(mixed_ + t * N, (const uint16_t*) wiq->data, qx, (int) N, (int) (IQ * ID), cs); - norm_rope(qx, wiqn, (int) IQ, (int) ID, pos_ + t * NH); - } - } - stamp(l, 10, grp); + native_qsa_indexer_append_multi(idx_raw + tb * ID, step_ + tb * kStepCount + kStepPos, (int) kStepCount, n, + 0, (const float*) wikn->data, EPS, ib, s, st.max_cells, + (float) qsa_freq_base(), side_, + grp == 0 ? tail_snap_ + (size_t) qi * TS : nullptr); + norm_rope(qidx_ + tb * IQ * ID, (int) ID, qidx_ + tb * IQ * ID, wiqn, (int) IQ, (int) ID, side_); qsa_block_scores(st.idx_pooled, st.idx_dead, qidx_ + tb * IQ * ID, step_ + tb * kStepCount, n, max_blocks_, - s, scores_ + (size_t) tb * max_blocks_, cs); + s, scores_ + (size_t) tb * max_blocks_, side_, qsa_score_grid_blocks); qsa_block_topk(scores_ + (size_t) tb * max_blocks_, step_ + tb * kStepCount, n, max_blocks_, cap_, s, - sel_ + (size_t) tb * cap_, cs); - stamp(l, 11, grp); - // KV streaming: the n selections' blocks resident (device-side, inside the graph) - qsa_kv_resolve(st, *g_, sel_ + (size_t) tb * cap_, step_ + tb * kStepCount, n, cap_, cs); - stamp(l, 12, grp); - const QsaAttnPools pools = qsa_attn_pools(st); - qsa_decode_attn_batch(qcur_ + tb * NH * HD, pools, sel_ + (size_t) tb * cap_, step_ + tb * kStepCount, cap_, - s, attn_scratch_ + (size_t) tb * attn_scratch_floats_, attn_ + tb * NH * HD, n, cs); - stamp(l, 13, grp); - if (st.kv_q4 || st.kv_hybrid) fwht256_inplace_cuda(attn_ + tb * NH * HD, (int64_t) n * NH, cs); // back: H^-1 = H - if (qb) native_qsa_gate_apply(attn_ + tb * NH * HD, qfull_ + tb * NH * 2 * HD, attn32_ + tb * NH * HD, - (int) (n * NH), (int) HD, cs); - else - for (int t = tb; t < te; ++t) { - if (native_qsa_enabled()) - native_qsa_gate_apply(attn_ + t * NH * HD, qfull_ + t * NH * 2 * HD, attn32_ + t * NH * HD, - (int) NH, (int) HD, cs); - else - qsa_gate_apply_f32(attn_ + t * NH * HD, qfull_ + t * NH * 2 * HD, s, attn32_ + t * NH * HD, cs); + sel_ + (size_t) tb * cap_, side_); + if (!mark(side_, bjoin_)) return false; + quant(xm, N); + if (cudaEventRecord(pfork_, cs) != cudaSuccess || cudaStreamWaitEvent(kst_, pfork_, 0) != cudaSuccess || + cudaStreamWaitEvent(shs_, pfork_, 0) != cudaSuccess || + cudaStreamWaitEvent(shs_, iproj_, 0) != cudaSuccess) { + err = "verify: a branch could not fork"; + return false; } - stamp(l, 14, grp); - native_quantize_q8_1(attn32_ + tb * NH * HD, xq_, (int) (NH * HD), n, cs); - native_mmvq(wo->native_type, wo->native_data, xq_, bo_ + tb * N, (int) (NH * HD), (int) N, n, cs); + proj(wk, kcur_ + tb * NKV * HD, N, NKV * HD, kst_); + norm_rope(kcur_ + tb * NKV * HD, (int) HD, kcur_ + tb * NKV * HD, wkn, (int) NKV, (int) HD, kst_); + if (!mark(kst_, kjoin_)) return false; + proj(wv, vcur_ + tb * NKV * HD, N, NKV * HD, cs); + proj(wq, qfull_ + tb * NH * 2 * HD, N, NH * 2 * HD, shs_); + // the query halves of every token's heads (rows of 2 * HD: query, then gate) + norm_rope(qfull_ + tb * NH * 2 * HD, (int) (2 * HD), qcur_ + tb * NH * HD, wqn, (int) NH, (int) HD, + shs_); + if (!mark(shs_, pjoin_) || !wait(kjoin_)) return false; + if (st.kv_int8) + kv_append_q8_steps(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, step_ + tb * kStepCount, + (int) kStepCount, kcur_ + tb * NKV * HD, vcur_ + tb * NKV * HD, n, s, cs); + else + for (int t = tb; t < te; ++t) + kv_append_step(st.k_pool, st.v_pool, st.page_table, step_ + t * kStepCount, kcur_ + t * NKV * HD, + vcur_ + t * NKV * HD, s, cs); + if (!wait(bjoin_) || !wait(pjoin_)) return false; + QsaAttnPools pools; + pools.page_table = st.page_table; + if (st.kv_int8) { pools.k_q = st.k_q; pools.v_q = st.v_q; pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; } + else { pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; } + qsa_decode_attn_batch(qcur_ + tb * NH * HD, pools, sel_ + (size_t) tb * cap_, step_ + tb * kStepCount, cap_, + s, attn_scratch_ + (size_t) tb * attn_scratch_floats_, attn32_ + tb * NH * HD, + n, cs, qfull_ + tb * NH * 2 * HD); + quant(attn32_ + tb * NH * HD, NH * HD); + proj(wo, bo_ + tb * N, NH * HD, N, cs); } } catch (const std::exception& e) { err = "verify layer " + std::to_string(l) + ": " + e.what(); return false; } - stamp(l, 16, grp); - gr_read_group(1, true, inj_, inj2_); - // the window's rows routed in 2 launches (one router GEMV reading the weight once, one - // top-10) instead of 2 per token; every row's arithmetic is the single-token call's (STRATA_DEC_BATCH=0: old) - const WeightRef* w_router = v.get("ffn_gate_inp.weight"); - if (dec_batch && n > 1 && w_router != nullptr && native_router_enabled() && NE == 512 && K == 10) { + stamp(l, 2, cs); + // the FFN read, and the next layer's estimated FFN input for the prediction (set_predict) + const WeightRef* wn1 = predicts(l) ? need(LayerView(wt, l + 1), "hc_ffn_norm.weight", err) : nullptr; + if (predicts(l) && !wn1) return false; + gr_read_group(1, true, inj_, inj2_, corrects(l) ? gate_ema_ + (size_t) l * HC * N : nullptr, + wn1 != nullptr ? (const float*) wn1->data : nullptr, + wn1 != nullptr ? gate_ema_ + (size_t) (l + 1) * HC * N : nullptr); + stamp(l, 3, cs); + // the routed experts' activations: from the router kernel, else beside it (joined before the experts) + if (!quant_router_) { + if (!fork(side_, bfork_)) return false; + if (strata::kernels::cpu::expert_layout().native) + quantize_q8_1_rows(xm, n, N, nat_xq_ + (size_t) tb * (N / 32) * 36, side_); + else + quantize_q8_0_scaled(xm, hit_xq_ + (size_t) tb * (N / 32) * 34, hit_xs_ + (size_t) tb * (N / 32), + (int64_t) n * N, side_); + if (!mark(side_, bjoin_)) return false; + } + // the branches that need only the FFN read fork here; they are recorded after the router, whose blocks are + // then dispatched first, and run beside its routing, the ring and the hit plan + if (beside && cudaEventRecord(efork_, cs) != cudaSuccess) { + err = "verify: the FFN read's branches could not fork"; + return false; + } + // the router, the doorbell and the main GPU's routed experts - the ones its VRAM tier holds, decided here from + // the residency snapshot the pool reads too (it leaves their rows to the GPU) + if (l == 0 && grp == 0 && cudaStreamWaitEvent(cs, res_ready_, 0) != cudaSuccess) { + err = "verify: the residency copy could not join"; + return false; + } + int32_t* hit_plan = dplan_ + (size_t) grp * (size_t) plan_stride_; + if (native_router_enabled()) { // one kernel: logits, top 10, ring, hit plan + const WeightRef* wr = need(v, "ffn_gate_inp.weight", err); + if (!wr) return false; + VerifyRouterArgs ra; + ra.x = xm; ra.w = (const uint16_t*) wr->data; ra.logits = logits_ + tb * NE; + ra.ids = ids_ + tb * K; ra.weights = w_ + tb * K; + ra.x_out = xk_layer_[(size_t) l] ? nullptr : m_x_ + tb * N; + ra.ids_out = m_ids_ + tb * K; ra.w_out = m_w_ + tb * K; ra.seq = m_seq_; + if (quant_router_) { // the q8_1 rows for both GPUs' experts, the Q8_K rows for the CPU's + ra.xq1 = nat_xq_ + (size_t) tb * (N / 32) * 36; + ra.xq1_out = gpu2_ ? m_xq1_ + (size_t) tb * (N / 32) * 36 : nullptr; + ra.xk_out = xk_layer_[(size_t) l] ? m_xk_ + (size_t) tb * (size_t) xk_stride_ : nullptr; + ra.xk_stride = (int) xk_stride_; + } + ra.ring = (uint32_t) (l * G + grp + 1); // the rings so far (the host waits for this one) + ra.res = res_ + (size_t) l * NE; ra.slot_ptr = slot_ptr_; ra.plan = hit_plan; + ra.cap = (int) (MT * K); ra.ptr_off = (int) plan_ptr_off_; + ra.counter = ring_count_; ra.n_tok = n; ra.n_embd = (int) N; ra.n_expert = (int) NE; try { - bf16_gemv_fp32_mmvf_multi(mixed_ + tb * N, N, (const uint16_t*) w_router->data, logits_ + tb * NE, NE, N, - NE, n, cs); - native_router_top10_multi(logits_ + tb * NE, ids_ + tb * K, w_ + tb * K, n, cs); - } catch (const std::exception& e) { err = "verify router: " + std::string(e.what()); return false; } - } else - for (int t = tb; t < te; ++t) { - MoEBuffers mb = ss.moe; - mb.logits = logits_ + t * NE; mb.ids = ids_ + t * K; mb.weights = w_ + t * K; - if (!moe_route(wt, g, l, K, mb, mixed_ + t * N, cs, err, nullptr)) return false; + verify_router(ra, cs); + } catch (const std::exception& e) { + err = "verify layer " + std::to_string(l) + " router: " + e.what(); + return false; + } + } else { + for (int t = tb; t < te; ++t) { + MoEBuffers mb = ss.moe; + mb.logits = logits_ + t * NE; mb.ids = ids_ + t * K; mb.weights = w_ + t * K; + if (!moe_route(wt, g, l, K, mb, mixed_ + t * N, cs, err, nullptr)) return false; + } + doorbell_publish(xm, ids_ + tb * K, w_ + tb * K, (int64_t) n * N, (int64_t) n * K, m_x_ + tb * N, + m_ids_ + tb * K, m_w_ + tb * K, m_seq_, cs); + verify_hit_plan(ids_ + tb * K, (int) (n * K), (int) K, res_ + (size_t) l * NE, (int) NE, slot_ptr_, hit_plan, + MT * K, plan_ptr_off_, cs); + } + stamp(l, 4, cs); + // after the doorbell: the second GPU's rows of this token group, taken into VRAM once their flag rises, and + // the wait for the CPU's rows, each on a branch beside the routed experts; then the branches forked after the + // FFN read: the shared expert, and the next layer's prediction + if (gpu2_) { + if (!fork(g2s_, g2fork_)) return false; + wait_flag_ge(m_flag2_ + (size_t) grp * 16, (uint32_t) (l * G + grp + 1), g2s_); + fetch_listed_rows(m_list2_ + grp * list2_stride_, m_ymiss_ + (size_t) tb * K * N, + hit_out_ + (size_t) tb * K * N, list2_ + grp * list2_stride_, (int) (n * K), N, g2s_); + stamp(l, kLayerStamps + 1, g2s_); + if (!mark(g2s_, g2join_[grp])) return false; + } + if (!fork(cps_, cpfork_)) return false; + wait_flag_ge(m_flag_, (uint32_t) (l * G + grp + 1), cps_); // the CPU's share is in the mapped rows + if (!mark(cps_, cpjoin_[grp])) return false; + if (beside && cudaEventRecord(rdone_, cs) != cudaSuccess) { + err = "verify: the router's end could not be marked"; + return false; } - if (device_plan_) // E-6: every routed expert resident: this group's plan without the host - resident_plan(ids_ + tb * K, n * (int) K, (int) K, hits_.d_res + l * g.n_expert, (int) g.n_expert, - hits_.cache_base, slot_off_d_, (long long) hits_.blob, - plan_ + (size_t) grp * (size_t) (plan_i32_ + 16), (long long) max_t_ * K, skip_ + grp, - (uint32_t) ((l - lb_) * G + grp + 1), cs); - doorbell_publish(xm, ids_ + tb * K, w_ + tb * K, (int64_t) n * N, (int64_t) n * K, m_x_ + tb * N, - m_ids_ + tb * K, m_w_ + tb * K, m_seq_, cs); - stamp(l, 17, grp); - { - const WeightRef *wgi = need(v, "ffn_gate_inp_shexp.weight", err), *wsg = need(v, "ffn_gate_shexp.weight", err), - *wsu = need(v, "ffn_up_shexp.weight", err), *wsd = need(v, "ffn_down_shexp.weight", err); - if (!wgi || !wsg || !wsu || !wsd) return false; - if (!native_of(wsg, v.name("ffn_gate_shexp.weight"), err) || !native_of(wsu, v.name("ffn_up_shexp.weight"), err) || - !native_of(wsd, v.name("ffn_down_shexp.weight"), err)) + if (beside && !shared(l, grp)) return false; + if (branches(l) && cudaStreamWaitEvent(side_, efork_, 0) != cudaSuccess) { + err = "verify: the prediction branch could not fork"; + return false; + } + if (predicts(l)) { + const WeightRef* wr1 = need(LayerView(wt, l + 1), "ffn_gate_inp.weight", err); + if (!wr1) return false; + const size_t half = (size_t) ((l + 1) & 1) * (size_t) (MT * K); + VerifyRouterArgs ra; + ra.x = pmixed_; ra.w = (const uint16_t*) wr1->data; + ra.logits = plogits_ + (size_t) ((l + 1) & 1) * MT * NE; ra.ids = pids_; ra.weights = pw_; + ra.ids_out = m_pids_ + half; ra.w_out = m_pw_ + half; ra.seq = m_pseq_; + ra.ring = (uint32_t) (l + 1); // the predictions published so far + ra.counter = ring_count_ + 1; ra.n_tok = n; ra.n_embd = (int) N; ra.n_expert = (int) NE; + if (learn_) ra.bias = pbias_ + (size_t) (l + 1) * NE; + try { + verify_router(ra, side_); + } catch (const std::exception& e) { + err = "verify layer " + std::to_string(l) + " prediction: " + e.what(); return false; - NativeSharedWeights nsw; - nsw.gate_type = wsg->native_type; nsw.gate_data = wsg->native_data; - nsw.up_type = wsu->native_type; nsw.up_data = wsu->native_data; - nsw.down_type = wsd->native_type; nsw.down_data = wsd->native_data; - nsw.q8_1 = xq_; - if (dec_batch) f32_to_bf16_bulk(mixed_ + tb * N, sh_bf16_ + tb * N, (int64_t) n * N, cs); // contiguous rows - else for (int t = tb; t < te; ++t) f32_to_bf16_bulk(mixed_ + t * N, sh_bf16_ + t * N, N, cs); + } + } + if (corrects(l)) { // after the prediction and the router, from the prediction the previous layer made + if (cudaStreamWaitEvent(side_, rdone_, 0) != cudaSuccess) { + err = "verify: the prediction's bias could not wait for the router"; + return false; + } try { - shared_expert_multi(n, xm, sh_bf16_ + tb * N, nsw, (const uint16_t*) wgi->data, sh_gate_ + (size_t) tb * g.n_ff, - sh_up_ + (size_t) tb * g.n_ff, sh_g_ + tb, shared_ + tb * N, N, g.n_ff, cs); + verify_router_bias_update(logits_ + tb * NE, plogits_ + (size_t) (l & 1) * MT * NE, n, + pbias_ + (size_t) l * NE, side_); } catch (const std::exception& e) { - err = std::string("verify shared expert: ") + e.what(); + err = "verify layer " + std::to_string(l) + " prediction bias: " + e.what(); return false; } } - if (strata::kernels::cpu::expert_layout().native) - quantize_q8_1_rows(xm, n, N, nat_xq_ + (size_t) tb * (N / 32) * 36, cs); - else - quantize_q8_0_scaled(xm, hit_xq_ + (size_t) tb * (N / 32) * 34, hit_xs_ + (size_t) tb * (N / 32), (int64_t) n * N, cs); - stamp(l, 18, grp); - return true; + if (branches(l) && cudaEventRecord(join_, side_) != cudaSuccess) { + err = "verify: the prediction branch could not join"; + return false; + } + if (!quant_router_ && !wait(bjoin_)) return false; + stamp(l, 5, cs); + return beside || shared(l, grp); }; // ---------------------------------------------------------------- post(l, group): experts, combine auto post = [&](int64_t l, int grp) -> bool { const int tb = tb_[grp], te = te_[grp], n = te - tb; - const uint32_t ring = (uint32_t) ((l - lb_) * G + grp + 1); + const uint32_t ring = (uint32_t) (l * G + grp + 1); const int64_t cap = (int64_t) n * K, capx = (int64_t) max_t_ * K; - int32_t* pl = plan_ + (size_t) grp * (size_t) (plan_i32_ + 16); - if (device_plan_) { // E-6: skipped when the device planned this group (all its experts resident) - wait_flag_ge_or(m_flagA_, ring, skip_ + grp, cs); - copy_i32_from_mapped_unless(pl, m_plan_ + (size_t) grp * (size_t) plan_i32_, plan_i32_, skip_ + grp, ring, cs); - } else { - wait_flag_ge(m_flagA_, ring, cs); // the pool published this group's GPU plan - copy_i32_from_mapped(pl, m_plan_ + (size_t) grp * (size_t) plan_i32_, plan_i32_, cs); - } - stamp(l, 19, grp); - const int32_t* p_counts = pl; - const int32_t* p_start = pl + 4; - const int32_t* p_dst = p_start + capx + 1; - const int32_t* p_tok = p_dst + capx; - const int64_t ptr_off = ((4 + (capx + 1) + 2 * capx) + 1) & ~1ll; - const unsigned long long* p_ptr = (const unsigned long long*) (pl + ptr_off); - const unsigned long long* p_ptr2 = p_ptr + capx; - const int32_t* p_start2 = pl + ptr_off + 4 * capx; + struct Plan { // a plan block's fields (`init`) + const int32_t *counts = nullptr, *start = nullptr, *dst = nullptr, *tok = nullptr, *start2 = nullptr; + const unsigned long long *ptr = nullptr, *ptr2 = nullptr; + }; + auto plan_at = [&](const int32_t* pl) { + Plan p; + p.counts = pl; + p.start = pl + 4; + p.dst = p.start + capx + 1; + p.tok = p.dst + capx; + p.ptr = (const unsigned long long*) (pl + plan_ptr_off_); + p.ptr2 = p.ptr + capx; + p.start2 = pl + plan_ptr_off_ + 4 * capx; + return p; + }; float* hit_out = hit_out_ + (size_t) tb * K * N; const auto& lay = strata::kernels::cpu::expert_layout(); - // plan v0.3 P6: the VRAM groups now; the PCIe groups once the copy engine has landed them in staging - auto grouped = [&](const unsigned long long* gp, const int32_t* gs, const int32_t* gn) { + auto grouped = [&](const Plan& p, const unsigned long long* gp, const int32_t* gs, const int32_t* gn) { if (lay.native) { // the layer's GGUF formats (i-quant gate/up, Q2_0 / IQ4_NL down) const auto& f = lay.fmt[(size_t) l]; const NativeExpertLayout L = native_expert_layout(f.gu_type, f.d_type, f.n_embd, f.n_ff); - native_expert_grouped(L, gp, gs, gn, p_dst, p_tok, cap, cap, + native_expert_grouped(L, gp, gs, gn, p.dst, p.tok, cap, cap, nat_xq_ + (size_t) tb * (N / 32) * 36, hit_scratch_, hit_out, cs); } else { - moe_grouped_s2(gp, gs, gn, p_dst, p_tok, cap, cap, hit_xq_ + (size_t) tb * (N / 32) * 34, + moe_grouped_s2(gp, gs, gn, p.dst, p.tok, cap, cap, hit_xq_ + (size_t) tb * (N / 32) * 34, hit_xs_ + (size_t) tb * (N / 32), hit_scratch_, hit_out, cs); } }; - grouped(p_ptr, p_start, p_counts); - stamp(l, 20, grp); - if (device_plan_) wait_flag_ge_or(m_flagB_, ring, skip_ + grp, cs); - else wait_flag_ge(m_flagB_, ring, cs); // the PCIe share is in staging (DMA) or mapped - if (sink_.pcie_mode == 2) { // stage it with a copy kernel, then point at staging - const int64_t per = G == 2 ? kStagingBlobs / 2 : kStagingBlobs; - uint8_t* stage = staging_ + (size_t) (grp * per) * lay.max_blob; - fetch_blobs(p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), (int) per, cs); - rebase_ptrs((unsigned long long*) p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), cs); - } - stamp(l, 21, grp); - grouped(p_ptr2, p_start2, p_counts + 2); - stamp(l, 22, grp); - if (device_plan_) { // no CPU share when the device planned the group: its rows are zeros - wait_flag_ge_or(m_flag_, ring, skip_ + grp, cs); - copy_or_zero_from_mapped(parts_ + (size_t) tb * K * N, m_ymiss_ + (size_t) tb * K * N, (long long) n * K * N, - skip_ + grp, ring, cs); + // plan v0.3 P6: the VRAM groups at once; the PCIe groups once the pool has published them and the copy engine + // has landed them in staging + const Plan vram = plan_at(dplan_ + (size_t) grp * (size_t) plan_stride_); + grouped(vram, vram.ptr, vram.start, vram.counts); + stamp(l, 6, cs); + Plan pcie; + if (pcie_share_) { + wait_flag_ge(m_flagA_, ring, cs); // the pool published this group's PCIe share + stamp(l, 7, cs); + int32_t* pl = plan_ + (size_t) grp * (size_t) plan_stride_; + copy_i32_from_mapped(pl, m_plan_ + (size_t) grp * (size_t) plan_i32_, plan_i32_, cs); + pcie = plan_at(pl); + wait_flag_ge(m_flagB_, ring, cs); // it is in staging (DMA) or mapped + if (sink_.pcie_mode == 2) { // stage it with a copy kernel, then point at staging + const int64_t per = G == 2 ? kStagingBlobs / 2 : kStagingBlobs; + uint8_t* stage = staging_ + (size_t) (grp * per) * lay.max_blob; + fetch_blobs(pcie.ptr2, pcie.counts + 2, stage, (int64_t) lay.blob_bytes(l), (int) per, cs); + rebase_ptrs((unsigned long long*) pcie.ptr2, pcie.counts + 2, stage, (int64_t) lay.blob_bytes(l), cs); + } + grouped(pcie, pcie.ptr2, pcie.start2, pcie.counts + 2); } else { - wait_flag_ge(m_flag_, ring, cs); // the CPU's share is in the mapped rows - stamp(l, 23, grp); - if (dec_batch) // only the CPU rows cross PCIe (p_dst[0, counts[1]) = the GPU's own rows) - copy_rows_from_mapped(parts_ + (size_t) tb * K * N, m_ymiss_ + (size_t) tb * K * N, (int64_t) n * K, N, - p_dst, p_counts + 1, cs); - else - copy_from_mapped(parts_ + (size_t) tb * K * N, m_ymiss_ + (size_t) tb * K * N, (int64_t) n * K * N, cs); + stamp(l, 7, cs); } - moe_hit_add(parts_ + (size_t) tb * K * N, hit_out, p_dst, p_counts + 1, cap, N, cs); - if (dec_batch && n > 1 && native_moe_combine_enabled()) { // one launch for the window's rows - try { - native_moe_combine_multi(parts_ + (size_t) tb * K * N, w_ + tb * K, shared_ + tb * N, bo_ + tb * N, N, K, n, cs); - } catch (const std::exception& e) { err = "verify combine: " + std::string(e.what()); return false; } - } else - for (int t = tb; t < te; ++t) { - MoEBuffers mb = ss.moe; - mb.weights = w_ + t * K; mb.shared = shared_ + t * N; - if (!moe_combine_parts(g, l, K, mb, parts_ + (size_t) t * K * N, bo_ + t * N, cs, err)) return false; + stamp(l, 8, cs); + if (!wait(cpjoin_[grp])) return false; // the CPU's share is in the mapped rows + stamp(l, 9, cs); + if (gpu2_ && !wait(g2join_[grp])) return false; // the second GPU's, in `hit_out` + stamp(l, 10, cs); + if (beside && (cudaStreamWaitEvent(cs, shjoin_, 0) != cudaSuccess || + cudaStreamWaitEvent(cs, sgjoin_, 0) != cudaSuccess)) { + err = "verify: the shared expert's branch could not join"; + return false; } - stamp(l, 24, grp); - if (l == g.n_layers - 1) { - for (int t = tb; t < te; ++t) gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs); - if (cvec().covers(l)) cvec_apply(Rt(tb), l, n, HC * N, nullptr, 0, nullptr, 0, false, cs); - } else if (cvec().covers(l)) { - cvec_apply(Rt(tb), l, n, HC * N, bo_ + tb * N, N, inj2_ + tb * HC, HC, true, cs); + // only the CPU's rows cross PCIe here; the GPUs' come from `hit_out` (the pool leaves their rows unwritten) + const int32_t* list2 = gpu2_ ? list2_ + grp * list2_stride_ : nullptr; + try { + native_moe_gather_combine(hit_out, m_ymiss_ + (size_t) tb * K * N, vram.dst, vram.counts + 1, pcie.dst, + pcie.counts != nullptr ? pcie.counts + 1 : nullptr, + list2 != nullptr ? list2 + 4 : nullptr, list2, w_ + tb * K, shared_ + tb * N, + sh_g_ + tb, bo_ + tb * N, N, K, n, cs); + } catch (const std::exception& e) { + err = "verify layer " + std::to_string(l) + " combine: " + e.what(); + return false; } + stamp(l, 11, cs); + if (branches(l) && cudaStreamWaitEvent(cs, join_, 0) != cudaSuccess) { + err = "verify: the prediction branch could not join"; + return false; + } + if (l == g.n_layers - 1) + for (int t = tb; t < te; ++t) gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs); + stamp(l, 12, cs); return true; }; + stamp(g.n_layers, 1, cs); for (int grp = 0; grp < G; ++grp) - if (!pre(lb_, grp)) return false; - for (int64_t l = lb_; l < le_; ++l) + if (!pre(0, grp)) return false; + for (int64_t l = 0; l < g.n_layers; ++l) for (int grp = 0; grp < G; ++grp) { if (!post(l, grp)) return false; - if (l + 1 < le_ && !pre(l + 1, grp)) return false; - } - if (le_ < g.n_layers) { // a layer split's earlier stage: hand the residual on, no head - for (int t = 0; t < T; ++t) { - copy_from_mapped(hand_out_ + (size_t) t * HB, Rt(t), HC * N, cs); - copy_from_mapped(hand_out_ + (size_t) t * HB + HC * N, bo_ + (size_t) t * N, N, cs); - copy_from_mapped(hand_out_ + (size_t) t * HB + HC * N + N, inj2_ + (size_t) t * HC, HC, cs); + if (l + 1 < g.n_layers && !pre(l + 1, grp)) return false; } - return true; - } + stamp(g.n_layers, 2, cs); // ---- the head, T columns, and the argmax of each - stamp(g.n_layers, 0, 0); { const WeightRef *hn = wt.find("output_hc_norm.weight"), *hd = wt.find("output_hc_down.weight"), *hu = wt.find("output_hc_up.weight"); if (!hn || !hd || !hu) { err = "verify: an output_hc_* weight is missing"; return false; } - for (int t = 0; t < T; ++t) { - BlockBuffers bb = ss.block; - bb.R = Rt(t); - bb.mixed = head_mixed_ + t * N; - if (head_ != nullptr && head_->loaded()) { - if (!lm_head_mix(wt, g, bb, cs, err)) return false; - } else if (!lm_head(wt, g, bb, head_logits_ + (size_t) t * n_vocab_, cs, err)) { + if (head_ != nullptr && head_->loaded()) { // the final mixer, the window's tokens in one read + if (hn->kind != WeightKind::F32 || hd->kind != WeightKind::Bf16InF32 || hu->kind != WeightKind::Bf16InF32) { + err = "verify: the output_hc_* weights have the wrong engine forms"; return false; } + FusedGrArgs fa[kFusedGrMaxT]; + for (int t = 0; t < T; ++t) { + fa[t].R = Rt(t); fa[t].R_out = Rt(t); fa[t].apply = false; + fa[t].w_norm = (const float*) hn->data; fa[t].w_down = (const uint16_t*) hd->data; + fa[t].w_up = (const uint16_t*) hu->data; fa[t].eps = EPS; + fa[t].lo = lo_ + t * g.hc_lr; fa[t].rs = rs_ + t * HC; fa[t].mixed = head_mixed_ + t * N; + } + fused_gr_read_multi(fa, T, grs_, cs); + } else { + for (int t = 0; t < T; ++t) { + BlockBuffers bb = ss.block; + bb.R = Rt(t); + bb.mixed = head_mixed_ + t * N; + if (!lm_head(wt, g, bb, head_logits_ + (size_t) t * n_vocab_, cs, err)) return false; + } } if (head_ != nullptr && head_->loaded()) { try { - native_quantize_q8_1(head_mixed_, xq_, (int) N, T, cs); - native_mmvq(head_->type(), head_->weights(), xq_, head_logits_, (int) N, (int) n_vocab_, T, cs); + if (shead_ != nullptr) { // the second GPU's part starts once the host sees the count + copy_from_mapped(shead_->input(), head_mixed_, (int64_t) T * N, cs); + mapped_bump(shead_->handoff(), cs); + } + if (T >= 2) { + native_quantize_q8_1_il(head_mixed_, xq_, xil_, (int) N, T, cs); + native_mmvq_il(head_->type(), head_->weights(), xq_, xil_, head_logits_, (int) N, (int) head_rows_, + T, cs); + } else { + native_quantize_q8_1(head_mixed_, xq_, (int) N, T, cs); + native_mmvq(head_->type(), head_->weights(), xq_, head_logits_, (int) N, (int) head_rows_, T, cs); + } } catch (const std::exception& e) { err = std::string("verify head: ") + e.what(); return false; } } - // Greedy, the default, is recorded here as before (no extra launch or sync per window). A request that - // samples or penalizes is sampled again host-side after the replay (run()) with its own parameters and a - // fresh draw counter: a captured sampler would bake them in and replay the same draws forever. - SamplerParams sp; - sp.greedy = true; - sp.temperature = 0.0f; - sample_tokens(head_logits_, T, (int) n_vocab_, nullptr, 0, sp, m_out_, cs); + argmax_rows(head_logits_, T, (int) head_rows_, arg_scratch_, m_out_, cs, m_val_); } - stamp(g.n_layers, 1, 0); + if (host_rows_) copy_from_mapped(m_rows_, R_, (int64_t) T * HC * N, cs); // a plain copy kernel, here to host + stamp(g.n_layers, 3, cs); + stamp(g.n_layers, 4, cs); // right after the last: one stamp's own cost return true; } -std::string Verifier::profile_report() { - if (!prof_on_ || prof_windows_ == 0) return std::string(); - static const char* names[kProfPer] = {"-", "hc-read0", "q8+qkv/q-idx gemv", "conv", "ab", "z", "rec", "q8+kv-idx", - "k/v+norm-rope", "kv+idx append", "q+q-idx", "scores+topk", "kv-resolve", - "attention", "gate", "", "out-proj", "hc-read1+router", "shared+quant", - "waitA", "VRAM hits", "waitB", "PCIe grp", "waitCPU", "copy+combine", - "(gap)", "head", " hc0 norm", " hc0 down", " hc0 up", "", ""}; - std::string out; - char b[80]; - double total = 0; - for (int k = 0; k < 2; ++k) { - out += k == 0 ? " GDN layers:" : " | QSA layers:"; - for (int i = 0; i < kProfPer; ++i) { - if (prof_sum_[k][i] <= 0) continue; - total += prof_sum_[k][i]; - std::snprintf(b, sizeof b, " %s %.2f", names[i], prof_sum_[k][i] / 1e6 / (double) prof_windows_); - out += b; - } - } - std::snprintf(b, sizeof b, " | total %.2f ms/window over %lld windows", total / 1e6 / (double) prof_windows_, (long long) prof_windows_); - out += b; - for (auto& r : prof_sum_) for (double& d : r) d = 0; - prof_windows_ = 0; - return out; -} - bool Verifier::capture(int T, std::string& err) { if (exec_[T] != nullptr) return true; + if (profile_ && stamps_ == nullptr && + cudaMalloc(&stamps_, (size_t) (strata::kernels::kVerifyMaxT + 1) * (size_t) (g_->n_layers + 1) * kStamps * + sizeof(unsigned long long)) != cudaSuccess) { + err = "verify: the profile's stamp buffer does not fit"; + return false; + } if (cudaStreamBeginCapture(cs_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) { err = "verify: begin capture failed"; return false; @@ -836,35 +1027,6 @@ bool Verifier::capture(int T, std::string& err) { err = std::string("verify: end capture: ") + cudaGetErrorString(ce); return false; } -#if !defined(STRATA_USE_HIP) // a CUDA debug listing (node types, kernel names) - if (std::getenv("STRATA_VERIFY_NODES") != nullptr) { // what the window graph holds - size_t nn = 0; - cudaGraphGetNodes(graph, nullptr, &nn); - std::vector nodes(nn); - cudaGraphGetNodes(graph, nodes.data(), &nn); - std::map kinds; - for (cudaGraphNode_t nd : nodes) { - cudaGraphNodeType ty; - cudaGraphNodeGetType(nd, &ty); - std::string name = "type" + std::to_string((int) ty); - if (ty == cudaGraphNodeTypeKernel) { - cudaKernelNodeParams kp{}; - if (cudaGraphKernelNodeGetParams(nd, &kp) == cudaSuccess) { - const char* fn = nullptr; - if (cudaFuncGetName(&fn, kp.func) == cudaSuccess && fn) name = fn; - } - } else if (ty == cudaGraphNodeTypeMemcpy) name = "memcpy"; - else if (ty == cudaGraphNodeTypeMemset) name = "memset"; - ++kinds[name]; - } - std::vector> v; - for (auto& [k2, c] : kinds) v.push_back({c, k2}); - std::sort(v.rbegin(), v.rend()); - std::fprintf(stderr, "strata verify: the %d-token window graph has %zu nodes:", T, nn); - for (size_t i = 0; i < v.size() && i < 40; ++i) std::fprintf(stderr, " %d x %.60s;", v[i].first, v[i].second.c_str()); - std::fprintf(stderr, "\n"); - } -#endif const cudaError_t ie = cudaGraphInstantiate(&exec_[T], graph, 0); cudaGraphDestroy(graph); if (ie != cudaSuccess) { @@ -897,8 +1059,7 @@ bool Verifier::capture_commit(std::string& err) { try { copy_i32_from_mapped(commit_, m_commit_, 2 + MT, cs_); int64_t qsa_index = 0, gdn_index = 0; - for (int64_t l = 0; l < lb_; ++l) (is_qsa_layer(g, l) ? qsa_index : gdn_index) += 1; - for (int64_t l = lb_; l < le_ && ok; ++l) { + for (int64_t l = 0; l < g.n_layers && ok; ++l) { const LayerView v(*wt_, l); if (!is_qsa_layer(g, l)) { const WeightRef* wnm = need(v, "ssm_norm.weight", err); @@ -917,14 +1078,13 @@ bool Verifier::capture_commit(std::string& err) { if (!wikn) { ok = false; break; } copy_from_mapped(st.idx_tail, tail_snap_ + (size_t) qsa_index * TS, TS, cs_); const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; - for (int64_t t = 0; t < MT; ++t) - native_qsa_indexer_append(idx_raw_L_ + (size_t) (qsa_index * MT + t) * ID, commit_ + 2 + t, 0, - (const float*) wikn->data, EPS, ib, s, st.max_cells, - (float) qsa_freq_base(), cs_); + native_qsa_indexer_append_multi(idx_raw_L_ + (size_t) (qsa_index * MT) * ID, commit_ + 2, 1, (int) MT, 0, + (const float*) wikn->data, EPS, ib, s, st.max_cells, + (float) qsa_freq_base(), cs_); ++qsa_index; } } - if (ok && ss.ple.ready() && ple_stage()) copy_indexed(ss.ple.hist, hist_snap_, HS, commit_ + 1, HS, cs_); + if (ok && ss.ple.ready()) copy_indexed(ss.ple.hist, hist_snap_, HS, commit_ + 1, HS, cs_); } catch (const std::exception& e) { err = std::string("verify commit: ") + e.what(); ok = false; @@ -947,7 +1107,6 @@ bool Verifier::capture_commit(std::string& err) { bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, void* user, int32_t* out, std::string& err) { using namespace strata::kernels; - const OnDevice on_device(device_); if (T < 1 || T > max_t_) { err = "verify: window size out of range"; return false; } const ModelGeometry& g = *g_; SessionState& ss = *ss_; @@ -960,25 +1119,28 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, h_tok_[t] = tokens[t]; qsa_step_fill(h_step_ + t * kStepCount, pos0 + t, s); for (int64_t h = 0; h < g.n_head; ++h) h_pos_[t * g.n_head + h] = (int32_t) (pos0 + t); - int32_t* pk = h_pos_ + (size_t) max_t_ * g.n_head; - int32_t* pi = pk + (size_t) max_t_ * g.n_head_kv; - for (int64_t h = 0; h < g.n_head_kv; ++h) pk[t * g.n_head_kv + h] = (int32_t) (pos0 + t); - for (int64_t h = 0; h < g.idx_q_heads; ++h) pi[t * g.idx_q_heads + h] = (int32_t) (pos0 + t); } - if (ss.ple.ready() && ple_stage()) { - uint32_t rows[kVerifyMaxT * PLE_N_HEADS]; + const bool ple_on = ss.ple.ready(); + if (ple_on) { // the reads start here unless ple_ahead started them int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]}; for (int t = 0; t < T; ++t) { - ngram_rows(&tokens[t], prev, 1, ss.ple.consts, rows + t * PLE_N_HEADS); + uint32_t rows[PLE_N_HEADS]; + ngram_rows(&tokens[t], prev, 1, ss.ple.consts, rows); + if (!ss.ple.table->ahead_holds(t, rows)) ss.ple.table->ahead(t, rows); prev[0] = prev[1]; prev[1] = tokens[t]; } - if (!ss.ple.table->gather_batch(rows, (size_t) T, h_ple_, err)) return false; } + // the residency this window decides the main GPU's hits by, and the pool reads (`residency`) + std::memcpy(h_res_, hits_.res, (size_t) (g.n_layers * g.n_expert) * sizeof(int32_t)); *(volatile uint32_t*) h_seq_ = 0; *(volatile uint32_t*) h_flag_ = 0; *(volatile uint32_t*) h_flagA_ = 0; *(volatile uint32_t*) h_flagB_ = 0; + *(volatile uint32_t*) h_pseq_ = 0; + *(volatile uint32_t*) h_pleflag_ = 0; + h_flag2_[0] = h_flag2_[16] = 0; + sink_.gpu2_ring = 0; std::atomic_thread_fence(std::memory_order_seq_cst); last_t_ = T; last_pos0_ = pos0; @@ -986,22 +1148,54 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, ms_host += ms_since(t0); VDBG("staged; launching\n"); const cudaError_t le = cudaGraphLaunch(exec_[T], cs_); - if (le != cudaSuccess) { err = std::string("verify: launch: ") + cudaGetErrorString(le); return false; } + if (le != cudaSuccess || cudaEventRecord(done_, cs_) != cudaSuccess) { + err = std::string("verify: launch: ") + cudaGetErrorString(cudaGetLastError()); + return false; + } (void) cudaStreamQuery(cs_); VDBG("launched\n"); + // on a failure every wait of the graph passes, so it runs to its end before the next window resets the flags + auto release = [&] { + h_list2_[0] = h_list2_[list2_stride_] = 0; + std::atomic_thread_fence(std::memory_order_seq_cst); + for (uint32_t* f : {h_flag_, h_flagA_, h_flagB_, h_pleflag_, h_flag2_, h_flag2_ + 16}) + *(volatile uint32_t*) f = 0xffffffffu; + cudaStreamSynchronize(cs_); + return false; + }; + if (ple_on) { // before layer 1: the graph waits for the flag, raised even when a read failed + const Clock::time_point tp = Clock::now(); + bool ok = true; + for (int t = 0; t < T && ok; ++t) ok = ss.ple.table->ahead_collect(t, h_ple_ + (size_t) t * g.n_embd, err); + std::atomic_thread_fence(std::memory_order_seq_cst); + _mm_sfence(); + *(volatile uint32_t*) h_pleflag_ = 1; + ms_host += ms_since(tp); + if (!ok) return release(); + } volatile uint32_t* const seq = h_seq_; volatile uint32_t* const flag = h_flag_; const int G = groups_[T] > 0 ? groups_[T] : 1; const int gtb[2] = {0, (T + 1) / 2}, gte[2] = {G == 2 ? (T + 1) / 2 : T, T}; - for (int64_t k = 0; k < (le_ - lb_) * G; ++k) { - const int64_t l = lb_ + k / G; + // a pool whose second-GPU rows came after the CPU's: the host time until they were in + const volatile uint32_t* late_flag = nullptr; + uint32_t late_want = 0; + Clock::time_point late_at{}; + auto settle = [&](bool now) { + if (late_flag != nullptr && (now || *late_flag >= late_want)) { + ms_late2 += ms_since(late_at); + late_flag = nullptr; + } + }; + for (int64_t k = 0; k < g.n_layers * G; ++k) { + const int64_t l = k / G; const int grp = (int) (k % G); const uint32_t want = (uint32_t) (k + 1); const Clock::time_point a = Clock::now(); auto last_flush = a; uint32_t spins = 0; - progress_at("verify window: waiting for the GPU to reach layer", l); while (*seq < want) { + settle(false); _mm_pause(); if ((++spins & 1023u) != 0) continue; const auto now = Clock::now(); @@ -1011,123 +1205,220 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, if (q != cudaErrorNotReady && *seq < want) { err = "verify: layer " + std::to_string(l) + " never rang (" + (q == cudaSuccess ? std::string("graph finished") : std::string(cudaGetErrorString(q))) + ")"; - return false; + return release(); } + if (watch_ != nullptr && !watch_(watch_user_, err)) return release(); + } + if (now - a > std::chrono::seconds(20)) { + err = "verify: timed out at layer " + std::to_string(l); + return release(); } - if (now - a > std::chrono::seconds(20)) { err = "verify: timed out at layer " + std::to_string(l); return false; } } const Clock::time_point b = Clock::now(); VDBG("layer %lld rang\n", (long long) l); cur_layer_ = want - 1; set_plan_slot(grp); const int tb = gtb[grp], n = gte[grp] - gtb[grp]; - progress_at("verify window: the CPU experts of layer", l); + sink_.xk = xk_layer_[(size_t) l] ? h_xk_ + (size_t) tb * (size_t) xk_stride_ : nullptr; + sink_.xk_stride = xk_stride_; + sink_.x1 = h_xq1_ + (size_t) tb * (size_t) (g.n_embd / 32) * 36; if (pool != nullptr) pool(user, h_x_ + (size_t) tb * g.n_embd, h_ids_ + (size_t) tb * ss.k, n, ss.k, h_ymiss_ + (size_t) tb * ss.k * g.n_embd, l); VDBG("layer %lld served\n", (long long) l); - progress_tick(); std::atomic_thread_fence(std::memory_order_seq_cst); _mm_sfence(); - if (*(volatile uint32_t*) h_flagA_ != want) { // the pool did not publish a plan: an empty one + if (pcie_share_ && *(volatile uint32_t*) h_flagA_ != want) { // no PCIe share published: an empty one sink_.counts[0] = 0; sink_.counts[1] = 0; sink_.counts[2] = 0; - sink_.start[0] = 0; sink_.start2[0] = 0; std::atomic_thread_fence(std::memory_order_seq_cst); *(volatile uint32_t*) h_flagA_ = want; raise_flag(h_flagB_, want); } + if (gpu2_ && sink_.gpu2_ring != want) { // no second-GPU share went out (the pool failed): an empty one + sink_.gpu2_list[0] = 0; + std::atomic_thread_fence(std::memory_order_seq_cst); + *(volatile uint32_t*) sink_.gpu2_flag = want; + } *flag = want; ms_wait += std::chrono::duration(b - a).count(); ms_pool += ms_since(b); + if (gpu2_ && *(volatile uint32_t*) sink_.gpu2_flag < want) { + settle(true); // a split window's other token group + late_flag = sink_.gpu2_flag; + late_want = want; + late_at = Clock::now(); + ++late2; + } + if (tail_ != nullptr && k + 1 == g.n_layers * G) tail_(tail_user_); + // the next layer's prediction, published during this layer's pool, else microseconds later (its branch is + // joined before the next layer): waited for, so the experts the second GPU takes never depend on timing + if (predict_ != nullptr && G == 1 && l + 1 < g.n_layers && native_router_enabled()) { + const Clock::time_point c = Clock::now(); + if (*(volatile uint32_t*) h_pseq_ < (uint32_t) (l + 1)) { + ++predict_late; + while (*(volatile uint32_t*) h_pseq_ < (uint32_t) (l + 1) && + Clock::now() - c < std::chrono::milliseconds(50)) { + settle(false); + _mm_pause(); + } + } + if (*(volatile uint32_t*) h_pseq_ >= (uint32_t) (l + 1)) { + std::atomic_thread_fence(std::memory_order_acquire); + const size_t half = (size_t) ((l + 1) & 1) * (size_t) max_t_ * (size_t) ss.k; + predict_(predict_user_, l + 1, h_pids_ + half, h_pw_ + half, T, ss.k); + } + ms_predict += ms_since(c); + } + if (rows_in_ != nullptr && k + 1 < g.n_layers * G) rows_in_(rows_user_, l); + } + while (late_flag != nullptr && ms_since(late_at) < 50.0) { // the last layer's (the graph ends after them) + settle(false); + _mm_pause(); + } + settle(true); + const bool sampled = head_sampling_ && ((!sampling_.greedy && sampling_.temperature > 0.0f) || hist_d_ != nullptr); + const bool whole = shead_ != nullptr && (logits_wanted_ || sampled); + if (shead_ != nullptr) { // the head's input is handed over after the last layer: the second GPU's part then + ++hand_n_; + const Clock::time_point h0 = Clock::now(); + while (*(volatile uint32_t*) shead_->handoff() < hand_n_) { + _mm_pause(); + if (ms_since(h0) > 5000.0) { + err = std::string("verify: the head's input never came (") + cudaGetErrorString(cudaStreamQuery(cs_)) + + ")"; + return false; + } + } + if (!shead_->submit(T, whole, err)) return false; + ++done_n_; } - progress_at("verify window: waiting for the GPU to finish the window (flags A/B/M raised)", (int64_t) T); const cudaError_t se = cudaStreamSynchronize(cs_); if (se != cudaSuccess) { err = std::string("verify: ") + cudaGetErrorString(se); return false; } - progress_at("verify window: waiting for the expert copies", (int64_t) T); cudaStreamSynchronize(copy_); // no host function of this window may raise flag B in the next one - if (prof_on_ && G == 1) { // the window's GPU stage stamps - cudaMemcpy(prof_h_.data(), prof_, prof_h_.size() * 8, cudaMemcpyDeviceToHost); - const int64_t L = g.n_layers; - auto at = [&](int64_t l, int i) { return prof_h_[(size_t) (l * kProfPer + i)]; }; - for (int64_t l = 0; l < L; ++l) { - const int kind = is_qsa_layer(g, l) ? 1 : 0; - unsigned long long prev = at(l, 0); - for (int i = 1; i <= 24; ++i) { - const unsigned long long x = at(l, i); - if (x == 0 || x < prev) continue; - prof_sum_[kind][i] += (double) (x - prev); - prev = x; - } - if (l + 1 < L) prof_sum_[kind][25] += (double) (at(l + 1, 0) - at(l, 24)); - prof_sum_[kind][27] += (double) (at(l, 27) - at(l, 0)); // hc-read0: norm - prof_sum_[kind][28] += (double) (at(l, 28) - at(l, 27)); // down - prof_sum_[kind][29] += (double) (at(l, 1) - at(l, 28)); // up - prof_sum_[kind][1] -= (double) (at(l, 1) - at(l, 0)); // (hc-read0 shown split) + if (shead_ != nullptr) { // each token's larger logit; on equality this GPU's, the lower index + if (!shead_->wait(done_n_, err)) return false; + for (int t = 0; t < T; ++t) { + const float v1 = ((volatile float*) h_val_)[t], v2 = ((volatile const float*) shead_->val())[t]; + if (v2 > v1) ((volatile int32_t*) h_out_)[t] = (int32_t) (shead_->split() + shead_->idx()[t]); } - prof_sum_[0][26] += (double) (at(L, 1) - at(L, 0)); - ++prof_windows_; + if (whole) join_rows(head_full_, n_vocab_, head_logits_, head_rows_, shead_->logits(), T, cs_); } - // ---- a sampled or penalized request: the head's sampling again, host-side so its parameters are this call's - // own (a captured kernel would replay the same draws forever). Row t's draw is Philox(seed, pos0 + t): tied to - // the POSITION it samples, not to how the text was cut into windows, so a seed replays the same text whatever - // the drafts were. Exact: a rejected row's draw is discarded, and no kept decision depends on a reused draw. - if (le_ < g.n_layers) { // a layer split's earlier stage: the hand-off is written (synced above) - ++windows; - return next_ == nullptr || next_->run(T, tokens, pos0, pool, next_user_, out, err); - } - const bool sampled = !sampling_.greedy && sampling_.temperature > 0.0f; - if (head_sampling_ && (sampled || hist_d_ != nullptr)) { + // a sampled or penalized request: the head's sampling again, after the graph (see set_sampling) + if (sampled) { SamplerParams sp = sampling_; sp.counter = (uint64_t) pos0; - sample_tokens(head_logits_, T, (int) n_vocab_, hist_d_, hist_len_, sp, m_out_, cs_); - if (cudaStreamSynchronize(cs_) != cudaSuccess) { // m_out_ is the mapped h_out_: synced, it is readable - err = "verify: the head sampling failed"; - return false; - } + sample_tokens(shead_ != nullptr ? head_full_ : head_logits_, T, (int) n_vocab_, hist_d_, hist_len_, sp, m_out_, + cs_); + if (cudaStreamSynchronize(cs_) != cudaSuccess) { err = "verify: the head sampling failed"; return false; } } for (int t = 0; t < T; ++t) out[t] = ((volatile int32_t*) h_out_)[t]; - if (static const bool dbg = std::getenv("STRATA_DBG_NAN") != nullptr; dbg) { // debug: the first non-finite head - static bool reported = false; - if (!reported) { - std::vector h((size_t) T * (size_t) n_vocab_); - cudaMemcpy(h.data(), head_logits_, h.size() * 4, cudaMemcpyDeviceToHost); - for (int t = 0; t < T && !reported; ++t) { - int64_t bad = 0; - for (int64_t v = 0; v < n_vocab_; ++v) bad += !std::isfinite(h[(size_t) t * n_vocab_ + v]); - if (bad) { - reported = true; - std::fprintf(stderr, "strata dbg: verify window at position %lld, row %d: %lld of %lld logits non-finite " - "(token out %d)\n", (long long) pos0, t, (long long) bad, (long long) n_vocab_, out[t]); - } + if (stamps_ != nullptr && G == 1) { + const int64_t NL = g.n_layers; + h_stamps_.resize((size_t) (NL + 1) * kStamps); + cudaMemcpy(h_stamps_.data(), stamps_ + (size_t) T * (size_t) (NL + 1) * kStamps, + h_stamps_.size() * sizeof(unsigned long long), cudaMemcpyDeviceToHost); + if (prof_ns_.empty()) prof_ns_.assign(2 * kStamps, 0.0); + const unsigned long long* w = h_stamps_.data() + (size_t) NL * kStamps; + for (int64_t l = 0; l < NL; ++l) { + const unsigned long long* st = h_stamps_.data() + (size_t) l * kStamps; + const int kind = is_qsa_layer(g, l) ? 1 : 0; + const unsigned long long before = l == 0 ? w[1] : (st - kStamps)[kLayerStamps - 1]; + prof_ns_[(size_t) kind * kStamps] += (double) (st[0] - before); + for (int i = 1; i < kLayerStamps; ++i) prof_ns_[(size_t) kind * kStamps + i] += (double) (st[i] - st[i - 1]); + // the shared expert's branch and the second GPU's rows', from the doorbell (they fork there) to their ends + prof_ns_[(size_t) kind * kStamps + kLayerStamps] += (double) (st[kLayerStamps] - st[4]); + if (gpu2_) prof_ns_[(size_t) kind * kStamps + kLayerStamps + 1] += (double) (st[kLayerStamps + 1] - st[4]); + { // the wait for the CPU's rows, by bucket + const double us = (double) (st[9] - st[8]) / 1e3; + int b = 0; + while (b < kWaitBuckets - 1 && us >= kWaitEdges[b]) ++b; + ++prof_wait_hist_[kind][b]; } + ++prof_layers_[kind]; } + prof_window_ns_[0] += (double) (w[1] - w[0]); + prof_window_ns_[1] += (double) (w[3] - w[2]); + prof_window_ns_[2] += (double) (w[3] - w[0]); + prof_window_ns_[3] += (double) (w[4] - w[3]); + ++prof_windows_; } VDBG("window done\n"); ++windows; - progress_at("decode"); - progress_beat(); return true; } +bool Verifier::copy_logits(int T, float* out, std::string& err) const { + if (shead_ != nullptr && !logits_wanted_) { err = "verify: the logits need set_logits_wanted"; return false; } + const float* src = shead_ != nullptr ? head_full_ : head_logits_; + if (cudaMemcpyAsync(out, src, (size_t) T * (size_t) n_vocab_ * sizeof(float), cudaMemcpyDeviceToHost, + cs_) != cudaSuccess || + cudaStreamSynchronize(cs_) != cudaSuccess) { + err = std::string("verify: reading the logits back: ") + cudaGetErrorString(cudaGetLastError()); + return false; + } + return true; +} + +void Verifier::print_profile() const { + if (prof_windows_ == 0) return; + static const char* const names[kStamps] = {"stamp gap", "HC read, mixer", "mixer", "HC read, FFN", + "router, ring, hits", "expert inputs", "VRAM experts", + "wait for PCIe plan", "PCIe experts", "wait for CPU", + "wait for 2nd GPU", "rows + combine", "last HC write", + "shared expert, beside", "2nd GPU rows, beside"}; + const double nw = (double) prof_windows_; + std::printf("%-24s GPU time over %lld windows, ms per window: inputs %.3f, head %.3f, whole window %.3f; " + "one stamp %.2f us (each stage below includes one; the last two run from the ring beside the " + "six before \"rows + combine\")\n", "window profile", (long long) prof_windows_, + prof_window_ns_[0] / nw / 1e6, prof_window_ns_[1] / nw / 1e6, prof_window_ns_[2] / nw / 1e6, + prof_window_ns_[3] / nw / 1e3); + std::printf(" %-22s %12s %12s %14s\n", "stage", "GDN us/layer", "QSA us/layer", "ms per window"); + for (int i = 0; i < kStamps; ++i) { + const double a = prof_ns_[(size_t) i], b = prof_ns_[(size_t) kStamps + i]; + std::printf(" %-22s %12.1f %12.1f %14.3f\n", names[i], + prof_layers_[0] ? a / (double) prof_layers_[0] / 1e3 : 0.0, + prof_layers_[1] ? b / (double) prof_layers_[1] / 1e3 : 0.0, (a + b) / nw / 1e6); + } + std::printf(" %-22s", "wait for CPU, layers"); + for (int b = 0; b < kWaitBuckets; ++b) { + const double f = 100.0 * (double) (prof_wait_hist_[0][b] + prof_wait_hist_[1][b]) / + (double) std::max(1, prof_layers_[0] + prof_layers_[1]); + if (b < kWaitBuckets - 1) std::printf(" <%.0f us %.0f%%,", kWaitEdges[b], f); + else std::printf(" more %.0f%%\n", f); + } +} + +void Verifier::ple_ahead(int t, int32_t token) { + using namespace strata::kernels; + SessionState& ss = *ss_; + if (!ss.ple.ready() || t < 0 || t >= max_t_) return; + ahead_tok_[t] = token; + int32_t prev[2] = {t >= 2 ? ahead_tok_[t - 2] : ss.ple_prev[t], t >= 1 ? ahead_tok_[t - 1] : ss.ple_prev[1]}; + uint32_t rows[PLE_N_HEADS]; + ngram_rows(&token, prev, 1, ss.ple.consts, rows); + ss.ple.table->ahead(t, rows); +} + void Verifier::set_plan_slot(int grp) { const int64_t cap = sink_.cap; int32_t* base = h_plan_ + (size_t) grp * (size_t) plan_i32_; - const int64_t i32 = 4 + (cap + 1) + cap + cap; - const int64_t ptr_off = (i32 + 1) & ~1ll; sink_.counts = base; - sink_.start = base + 4; - sink_.dst = sink_.start + cap + 1; + sink_.dst = base + 4 + cap + 1; sink_.tok = sink_.dst + cap; - sink_.ptr = (unsigned long long*) (base + ptr_off); - sink_.ptr2 = sink_.ptr + cap; - sink_.start2 = base + ptr_off + 4 * cap; + sink_.ptr2 = (unsigned long long*) (base + plan_ptr_off_) + cap; + sink_.start2 = base + plan_ptr_off_ + 4 * cap; const int G = groups_[last_t_] > 0 ? groups_[last_t_] : 1; const int64_t per = G == 2 ? kStagingBlobs / 2 : kStagingBlobs; sink_.staging = (unsigned long long) (staging_ + (size_t) (grp * per) * strata::kernels::cpu::expert_layout().max_blob); sink_.staging_cap = per; + if (gpu2_) { + sink_.gpu2_list = h_list2_ + (size_t) grp * (size_t) list2_stride_; + sink_.gpu2_flag = h_flag2_ + (size_t) grp * 16; + sink_.ring = cur_layer_ + 1; + } } // Flag B only rises: a host function of an earlier layer may run after a later layer already raised it directly. @@ -1167,25 +1458,35 @@ void Verifier::publish_plan(void* ctx) { *(volatile uint32_t*) v->h_flagA_ = v->cur_layer_ + 1; } -bool Verifier::commit(int n_keep, std::string& err) { - const OnDevice on_device(device_); +bool Verifier::commit(int n_keep, std::string& err, bool wait) { if (n_keep < 1 || n_keep > last_t_) { err = "verify: commit count out of range"; return false; } const Clock::time_point t0 = Clock::now(); h_commit_[0] = n_keep; h_commit_[1] = n_keep - 1; for (int t = 0; t < max_t_; ++t) h_commit_[2 + t] = t < n_keep ? (int32_t) (last_pos0_ + t) : -1; - std::atomic_thread_fence(std::memory_order_seq_cst); - const cudaError_t le = cudaGraphLaunch(commit_exec_, cs_); - if (le != cudaSuccess) { err = std::string("verify: commit launch: ") + cudaGetErrorString(le); return false; } + if (last_t_ > 1) { // a one-token window has advanced the state itself (record_window) + std::atomic_thread_fence(std::memory_order_seq_cst); + const cudaError_t le = cudaGraphLaunch(commit_exec_, cs_); + if (le != cudaSuccess) { err = std::string("verify: commit launch: ") + cudaGetErrorString(le); return false; } + (void) cudaStreamQuery(cs_); // submit now, not at the next driver call + commit_pending_ = true; + } + for (int t = 0; t < n_keep; ++t) { + ss_->ple_prev[0] = ss_->ple_prev[1]; + ss_->ple_prev[1] = last_tokens_[t]; + } + ms_commit += ms_since(t0); + return !wait || wait_commit(err); +} + +bool Verifier::wait_commit(std::string& err) { + if (!commit_pending_) return true; + const Clock::time_point t0 = Clock::now(); + commit_pending_ = false; const cudaError_t se = cudaStreamSynchronize(cs_); - if (se != cudaSuccess) { err = std::string("verify: commit: ") + cudaGetErrorString(se); return false; } - if (ple_stage()) // stages that share one session must advance it once - for (int t = 0; t < n_keep; ++t) { - ss_->ple_prev[0] = ss_->ple_prev[1]; - ss_->ple_prev[1] = last_tokens_[t]; - } ms_commit += ms_since(t0); - return next_ == nullptr || next_->commit(n_keep, err); + if (se != cudaSuccess) { err = std::string("verify: commit: ") + cudaGetErrorString(se); return false; } + return true; } } // namespace strata::core diff --git a/src/kernels/cpu/expert_layout.cpp b/src/kernels/cpu/expert_layout.cpp index ec78774cf7e..4c9b92305a8 100644 --- a/src/kernels/cpu/expert_layout.cpp +++ b/src/kernels/cpu/expert_layout.cpp @@ -92,6 +92,7 @@ bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n L.bytes.assign((size_t) n_layers, 0); L.max_blob = 0; std::string line; + // `layer gu_type d_type offset blob_bytes`; older packs' extra columns (shard-1 offsets) are not read while (std::getline(in, line)) { if (line.empty() || line[0] == '#') { if (!line.empty() && line[0] == '#') { @@ -105,7 +106,7 @@ bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n } std::istringstream ss(line); long long l = -1, gt = -1, dt = -1; - unsigned long long off = 0, blob = 0, go = 0, uo = 0, dox = 0; + unsigned long long off = 0, blob = 0; if (!(ss >> l >> gt >> dt >> off >> blob) || l < 0 || l >= n_layers) { err = "native_experts.txt: a malformed line: " + line; return false; @@ -117,17 +118,6 @@ bool expert_layout_load(const std::string& pack_dir, int64_t n_layers, int64_t n " B but its formats make " + std::to_string(f.bytes); return false; } - if (ss >> go >> uo >> dox) { // v2 lines: the GGUF offsets - if (L.gguf_off.empty()) L.gguf_off.assign((size_t) (3 * n_layers), 0); - L.gguf_off[(size_t) (3 * l)] = go; - L.gguf_off[(size_t) (3 * l + 1)] = uo; - L.gguf_off[(size_t) (3 * l + 2)] = dox; - std::string file; // v3: the shard that holds this layer (a file name beside --native) - if (ss >> file) { - if (L.gguf_file.empty()) L.gguf_file.assign((size_t) n_layers, std::string()); - L.gguf_file[(size_t) l] = file; - } - } L.fmt[(size_t) l] = f; L.offset[(size_t) l] = off; L.bytes[(size_t) l] = blob; diff --git a/src/kernels/cpu/pool.cpp b/src/kernels/cpu/pool.cpp index c5f4075d643..85810a2203d 100644 --- a/src/kernels/cpu/pool.cpp +++ b/src/kernels/cpu/pool.cpp @@ -1,6 +1,5 @@ // src/kernels/cpu/pool.cpp - P2.S3: the CPU expert pool. Read pool.hpp first; it explains the protocol. #include "strata/kernels/cpu/pool.hpp" -#include "strata/core/progress.hpp" #include "strata/kernels/cpu/expert_layout.hpp" #include @@ -9,7 +8,6 @@ #include #include -#include #if defined(_WIN32) #define WIN32_LEAN_AND_MEAN @@ -21,80 +19,155 @@ namespace strata::kernels::cpu { -namespace { -constexpr uint64_t pack_head(uint32_t epoch, uint32_t n, uint32_t i) { - return ((uint64_t) epoch << 32) | ((uint64_t) n << 16) | (uint64_t) i; -} -} // namespace - -std::vector physical_cores(bool skip_first) { - std::vector cores; +std::vector> physical_cores() { + std::vector> cores; #if defined(_WIN32) // Ask the OS rather than assuming a layout. `hardware_concurrency()` returns LOGICAL processors, and on // every SMT machine half of them are siblings - pinning one worker to each of the first N would put two // workers on each physical core and halve the bandwidth the expert kernel is bound by. DWORD len = 0; GetLogicalProcessorInformationEx(RelationProcessorCore, nullptr, &len); - if (len == 0) { - for (unsigned i = 0; i < std::thread::hardware_concurrency(); ++i) cores.push_back((int) i); - } else { - std::vector buf(len); - if (GetLogicalProcessorInformationEx(RelationProcessorCore, - (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX) buf.data(), &len)) { - const char* p = buf.data(); - const char* end = p + len; - while (p < end) { - const auto* e = (const SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX*) p; - if (e->Relationship == RelationProcessorCore) { - const GROUP_AFFINITY& g = e->Processor.GroupMask[0]; - for (int bit = 0; bit < 64; ++bit) - if (g.Mask & (1ull << bit)) { cores.push_back((int) (g.Group * 64 + bit)); break; } - } - p += e->Size; + std::vector buf(len); + if (len > 0 && GetLogicalProcessorInformationEx(RelationProcessorCore, + (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX) buf.data(), &len)) { + for (const char* p = buf.data(); p < buf.data() + len;) { + const auto* e = (const SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX*) p; + if (e->Relationship == RelationProcessorCore) { + const GROUP_AFFINITY& g = e->Processor.GroupMask[0]; + std::vector core; + for (int bit = 0; bit < 64; ++bit) + if (g.Mask & (1ull << bit)) core.push_back((int) (g.Group * 64 + bit)); + if (!core.empty()) cores.push_back(core); } + p += e->Size; } } #else - // The logical CPUs this process may run on, ONE PER PHYSICAL CORE (issue #40): SMT siblings share a core's - // load/store bandwidth, so a worker on each would put two workers on one core, as the Windows branch above - // explains. sysfs names each CPU's (package, core); the first allowed CPU of each pair is kept, so a taskset - // that leaves out the first sibling still gets its core. Without sysfs every allowed CPU counts, as before. - auto topo = [](int cpu, const char* what) -> long { - char path[96]; - std::snprintf(path, sizeof path, "/sys/devices/system/cpu/cpu%d/topology/%s", cpu, what); - long v = -1; - if (std::FILE* f = std::fopen(path, "r")) { - if (std::fscanf(f, "%ld", &v) != 1) v = -1; - std::fclose(f); - } - return v; - }; - std::vector allowed; - cpu_set_t set; + cpu_set_t set; // each logical processor as a core of its own CPU_ZERO(&set); if (sched_getaffinity(0, sizeof set, &set) == 0) { for (int i = 0; i < CPU_SETSIZE; ++i) - if (CPU_ISSET(i, &set)) allowed.push_back(i); - } else { - for (unsigned i = 0; i < std::thread::hardware_concurrency(); ++i) allowed.push_back((int) i); + if (CPU_ISSET(i, &set)) cores.push_back({i}); } - std::vector> seen; - for (int cpu : allowed) { - const long pkg = topo(cpu, "physical_package_id"), core = topo(cpu, "core_id"); - if (pkg >= 0 && core >= 0) { - const std::pair key{pkg, core}; - if (std::find(seen.begin(), seen.end(), key) != seen.end()) continue; // an SMT sibling - seen.push_back(key); +#endif + if (cores.empty()) + for (unsigned i = 0; i < std::thread::hardware_concurrency(); ++i) cores.push_back({(int) i}); + return cores; +} + +namespace { + +#if defined(_WIN32) +// Each logical processor's time in interrupts and DPCs so far (100 ns units), from NtQuerySystemInformation's +// processor performance information (winternl.h names its DpcTime and InterruptTime Reserved1); the kernel adds to +// them at each clock tick, so over a second they resolve ~1.6%. +std::vector interrupt_times() { + struct Info { LARGE_INTEGER idle, kernel, user, dpc, interrupt; ULONG interrupts; }; + using Query = LONG(WINAPI*)(int, void*, ULONG, ULONG*); + static const Query query = reinterpret_cast( + reinterpret_cast(GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtQuerySystemInformation"))); + SYSTEM_INFO si; + GetSystemInfo(&si); + std::vector v(si.dwNumberOfProcessors); + ULONG len = 0; + std::vector t; + if (query == nullptr || query(8, v.data(), (ULONG) (v.size() * sizeof(Info)), &len) != 0) return t; + for (const Info& i : v) t.push_back((uint64_t) (i.dpc.QuadPart + i.interrupt.QuadPart)); + return t; +} +#endif + +constexpr double kBusy = 0.05; // a processor over this share of a second in interrupts and DPCs takes them +constexpr double kQuiet = 0.02; // one below it may take a thread + +} // namespace + +CorePlacement::CorePlacement(bool spare_first) : cores_(physical_cores()) { + for (size_t c = 0; c < cores_.size(); ++c) + for (int p : cores_[c]) { + if (p >= (int) core_of_.size()) core_of_.resize((size_t) p + 1, -1); + core_of_[(size_t) p] = (int) c; + } + host_ = cores_.back().front(); + for (size_t c = spare_first && cores_.size() > 4 ? 1 : 0; c + 1 < cores_.size(); ++c) + workers_.push_back(cores_[c].front()); +} + +// A quiet processor for a thread on `from`: the first of the quietest free core whose processors are all quiet, else +// the other processor of `from`'s own core, or -1. +int CorePlacement::quiet_processor(int from, const std::vector& load) const { + auto quiet = [&](int p) { return p < (int) load.size() && load[(size_t) p] < kQuiet && hot_[(size_t) p] == 0; }; + std::vector taken(cores_.size(), 0); + for (int p : workers_) taken[(size_t) core_of_[(size_t) p]] = 1; + taken[(size_t) core_of_[(size_t) host_]] = 1; + int best = -1; + double best_load = 0; + for (size_t c = 0; c < cores_.size(); ++c) { + if (taken[c]) continue; + bool ok = true; + double sum = 0; + for (int p : cores_[c]) { + ok = ok && quiet(p); + sum += p < (int) load.size() ? load[(size_t) p] : 0.0; + } + if (ok && (best < 0 || sum < best_load)) { + best = cores_[c].front(); + best_load = sum; } - cores.push_back(cpu); } + if (best >= 0) return best; + for (int p : cores_[(size_t) core_of_[(size_t) from]]) + if (p != from && quiet(p)) return p; + return -1; +} + +std::string CorePlacement::tick(ExpertPool& pool) { +#if defined(_WIN32) + const auto now = std::chrono::steady_clock::now(); + const double s = std::chrono::duration(now - at_).count(); + if (!busy_.empty() && s < 1.0) return {}; + std::vector cur = interrupt_times(); + // a first sample, or the engine idle since the last one (between requests): a new start, nothing to judge + const bool fresh = busy_.empty() || s > 3.0 || cur.size() != busy_.size(); + std::vector load(cur.size(), 0.0); + if (!fresh) + for (size_t p = 0; p < cur.size(); ++p) load[p] = (double) (cur[p] - busy_[p]) / 1e7 / s; + busy_ = std::move(cur); + at_ = now; + hot_.resize(busy_.size(), 0); + for (size_t p = 0; p < hot_.size(); ++p) hot_[p] = !fresh && load[p] > kBusy ? hot_[p] + 1 : 0; + std::string moved; + for (int t = -1; t < (int) workers_.size(); ++t) { // the host, then the workers + int& at = t < 0 ? host_ : workers_[(size_t) t]; + if (at >= (int) hot_.size() || hot_[(size_t) at] < 2) continue; + const int to = quiet_processor(at, load); + if (to < 0) continue; + if (t < 0) pin_current_thread(to); + else pool.pin_worker(t, to); + char line[160]; + std::snprintf(line, sizeof line, + "%slogical processor %d spent %.0f%% of the last second in interrupts and DPCs: " + "the %s moved to %d", moved.empty() ? "" : "; ", at, 100.0 * load[(size_t) at], + t < 0 ? "host thread" : "pool worker", to); + moved += line; + hot_[(size_t) at] = 0; + at = to; + ++moves; + } + return moved; +#else + (void) pool; + return {}; #endif - if (skip_first && !cores.empty()) cores.erase(cores.begin()); - return cores; } namespace { +// `state_`: the epoch in the high half, then the closed bit, then the parked workers +constexpr uint64_t kClosed = 1ull << 31; +constexpr uint64_t kParked = kClosed - 1; +constexpr uint64_t kEpoch = 1ull << 32; + void pin_this_thread(int core) { if (core < 0) return; #if defined(_WIN32) @@ -141,50 +214,12 @@ void restore_thread_affinity(long long previous) { #endif } -namespace { -std::atomic g_diag_pool{nullptr}; -void diag_active_pool(std::FILE* f) { - if (const ExpertPool* p = g_diag_pool.load()) p->diag(f); -} -int64_t now_ms() { - return std::chrono::duration_cast(std::chrono::steady_clock::now().time_since_epoch()) - .count(); -} -} // namespace - -void ExpertPool::diag(std::FILE* f) const { - const uint64_t h = head_.load(); - std::fprintf(f, " expert pool: epoch %u, batch epoch %u: %u of %u jobs claimed, %u done; %u of %d workers parked, " - "%u sleeping; mode %d\n", epoch_.load(), (uint32_t) (h >> 32), (uint32_t) h & 0xffffu, - (uint32_t) (h >> 16) & 0xffffu, done_.load(), parked_.load(), n_, sleepers_.load(), mode_); - std::fprintf(f, " expert pool threads:"); - for (int i = 0; i < n_; ++i) { - const int32_t s = wstate_[(size_t) i].load(); - if (s == kParked) std::fprintf(f, " w%d=parked", i); - else if (s == kSleeping) std::fprintf(f, " w%d=sleeping", i); - else if (s == kBetween) std::fprintf(f, " w%d=draining", i); - else std::fprintf(f, " w%d=job%d", i, s); - } - const int32_t hs = hstate_.load(); - const char* hn = hs == kIdle ? "idle" : hs == kWaitParked ? "waiting for the workers to park" - : hs == kWaitDone ? "waiting for the jobs to finish" : "running a job"; - std::fprintf(f, "; host %s", hn); - if (hs >= 0) std::fprintf(f, " %d", hs); - std::fprintf(f, " for %lld ms\n", (long long) (now_ms() - hstate_ms_.load())); -} - -ExpertPool::ExpertPool(int n_workers, bool pin, bool host_works) : host_works_(host_works) { - if (const char* e = std::getenv("STRATA_POOL_SPIN_US")) // a test knob; see kSpinBeforeSleep - spin_before_sleep_ = std::chrono::microseconds((std::max)(0, std::atoi(e))); - const std::vector cores = physical_cores(true); +ExpertPool::ExpertPool(int n_workers, bool pin, bool host_works, std::vector cores) : host_works_(host_works) { + if (cores.empty()) cores = CorePlacement(false).workers(); n_ = n_workers > 0 ? n_workers : (int) cores.size(); if (n_ < 1) n_ = 1; + state_.store(kClosed | (uint64_t) n_, std::memory_order_relaxed); // epoch 0, closed, everyone parked scratch_.resize((size_t) n_); - wstate_.reset(new std::atomic[(size_t) n_]); - for (int i = 0; i < n_; ++i) wstate_[(size_t) i].store(kParked); - hstate_ms_.store(now_ms()); - g_diag_pool.store(this); - strata::core::diag_pool_fn().store(&diag_active_pool); split_.resize((size_t) kMaxSplit); split_multi_.resize((size_t) kMaxSplitMulti); threads_.reserve((size_t) n_); @@ -197,153 +232,84 @@ ExpertPool::ExpertPool(int n_workers, bool pin, bool host_works) : host_works_(h } } -ExpertPool::~ExpertPool() { - const ExpertPool* self = this; - g_diag_pool.compare_exchange_strong(self, nullptr); - stop_.store(true, std::memory_order_release); - // Bump the epoch so a PARKED worker notices the stop flag rather than sleeping through it. - publish(); - for (auto& t : threads_) t.join(); +void ExpertPool::pin_worker(int i, int core) { + if (i < 0 || i >= (int) threads_.size() || core < 0) return; +#if defined(_WIN32) + SetThreadAffinityMask((HANDLE) threads_[(size_t) i].native_handle(), (DWORD_PTR) 1 << (core & 63)); +#else + cpu_set_t set; + CPU_ZERO(&set); + CPU_SET(core, &set); + pthread_setaffinity_np(threads_[(size_t) i].native_handle(), sizeof set, &set); +#endif } -void ExpertPool::publish() { - // Both sides are seq_cst, and that is the whole lost-wakeup argument: a worker going to sleep does - // `sleepers_++` and then reads `epoch_`, the host does `epoch_++` and then reads `sleepers_`. In one total - // order at least one of them sees the other's write - the worker sees the new epoch and does not sleep, or - // the host sees the sleeper and notifies under the mutex the worker holds until it is inside `wait`. - // On x86 the fetch_add is a locked xadd either way, so this costs the token path nothing. - epoch_.fetch_add(1, std::memory_order_seq_cst); - if (sleepers_.load(std::memory_order_seq_cst) != 0) { - std::lock_guard lk(sleep_mu_); - sleep_cv_.notify_all(); - } +ExpertPool::~ExpertPool() { + stop_.store(true, std::memory_order_seq_cst); + // a new value, so a sleeping worker's wait returns and it sees the stop flag + state_.fetch_add(kEpoch, std::memory_order_seq_cst); + state_.notify_all(); + for (auto& t : threads_) t.join(); } void ExpertPool::worker(int id) { - uint32_t seen = 0; - // ARRIVE at the park before the first wait, so `parked_ == n_` is true from construction. Counting only - // on the RETURN from a drain leaves `parked_` at 0 until each worker has finished one batch, and the first - // `run()` - which waits for `parked_ == n_` before publishing - then deadlocks. It deadlocks on the very - // first call, which is the good case; a version that deadlocked on the second would be far worse. - parked_.fetch_add(1, std::memory_order_acq_rel); + uint32_t seen = 0; // the last epoch this worker joined for (;;) { - // Park: wait for work. `_mm_pause` rather than a bare spin because it yields the pipeline to the - // sibling hyperthread; `epoch_` is bumped once per LAYER, not once per expert, so most of these - // iterations are spent here with nothing to do. - // - // **AND NOTHING ELSE HAPPENS IN HERE.** This loop used to do `pauses_.fetch_add(1)` on every iteration - // - a locked read-modify-write, five workers against one cache line - so the workers spent their wait - // invalidating each other's caches and the very line the host writes to publish work. The counter was - // diagnostic and nothing branched on it. See the note on the atomics in pool.hpp. - // - // After `kSpinBeforeSleep` with no work the worker sleeps instead (issue #4). The clock is read once - // every 1024 pauses, so the spin itself is unchanged. - const auto parked_at = std::chrono::steady_clock::now(); - uint32_t spins = 0; - while (epoch_.load(std::memory_order_acquire) == seen) { + // Park: wait for an open phase this worker has not joined. `_mm_pause` rather than a bare spin because + // it yields the pipeline to the sibling hyperthread. The loop writes nothing shared (see the note on + // the atomics in pool.hpp) until it joins, or goes to sleep after `kIdleSpin` without work. + uint64_t s = state_.load(std::memory_order_acquire); + auto idle_from = std::chrono::steady_clock::now(); + for (uint32_t spins = 1;; ++spins) { if (stop_.load(std::memory_order_relaxed)) return; + if (!(s & kClosed) && (uint32_t) (s >> 32) != seen) { + // join; fails when the host has closed this phase or other workers joined since `s` was read + if (state_.compare_exchange_weak(s, s - 1, std::memory_order_acq_rel, std::memory_order_acquire)) break; + continue; + } _mm_pause(); - if ((++spins & 1023u) != 0) continue; - if (std::chrono::steady_clock::now() - parked_at < spin_before_sleep_) continue; - std::unique_lock lk(sleep_mu_); - wstate_[(size_t) id].store(kSleeping, std::memory_order_relaxed); - sleepers_.fetch_add(1, std::memory_order_seq_cst); - sleep_cv_.wait(lk, [&] { - return epoch_.load(std::memory_order_seq_cst) != seen || stop_.load(std::memory_order_relaxed); - }); - sleepers_.fetch_sub(1, std::memory_order_relaxed); - wstate_[(size_t) id].store(kParked, std::memory_order_relaxed); + if (spins % 1024 == 0 && std::chrono::steady_clock::now() - idle_from > kIdleSpin) { + // seq_cst, against open_phase's store then load: the host sees this sleeper or this worker's + // wait sees the new epoch + sleepers_.fetch_add(1, std::memory_order_seq_cst); + sleeps_.fetch_add(1, std::memory_order_relaxed); + state_.wait(s, std::memory_order_seq_cst); + sleepers_.fetch_sub(1, std::memory_order_seq_cst); + idle_from = std::chrono::steady_clock::now(); + } + s = state_.load(std::memory_order_acquire); } - if (stop_.load(std::memory_order_acquire)) return; - // acquire: the batch this epoch published (`head`, and the description before it) is visible from here - seen = epoch_.load(std::memory_order_acquire); - parked_.fetch_sub(1, std::memory_order_acq_rel); // leaving the park + seen = (uint32_t) (s >> 32); // Drain: one claim per iteration, so a slow worker takes fewer experts and a fast one takes more. - // Every job is the same size (all experts are 1,382,400 bytes), so there is nothing to schedule. Only - // this epoch's jobs: if the host has already moved on, the claims fail and the worker parks again. - wstate_[(size_t) id].store(kBetween, std::memory_order_relaxed); - drain(id, scratch_[(size_t) id], seen); - wstate_[(size_t) id].store(kParked, std::memory_order_relaxed); - parked_.fetch_add(1, std::memory_order_acq_rel); // back at the park - } -} - -int ExpertPool::claim(uint32_t epoch) { - uint64_t h = head_.load(std::memory_order_acquire); - for (;;) { - if ((uint32_t) (h >> 32) != epoch) return -1; // not the batch this thread woke for - const uint32_t n = (uint32_t) (h >> 16) & 0xffffu, i = (uint32_t) h & 0xffffu; - if (i >= n) return -1; // exhausted - if (head_.compare_exchange_weak(h, h + 1, std::memory_order_acq_rel, std::memory_order_acquire)) - return (int) i; + // Every job is the same size (all experts are 1,382,400 bytes), so there is nothing to schedule. + drain(id, scratch_[(size_t) id]); + state_.fetch_add(1, std::memory_order_release); // parked again } } -uint32_t ExpertPool::begin_batch(int n) { - if (n < 0 || n > 0xffff) { - std::fprintf(stderr, "strata: expert pool batch of %d jobs is out of range\n", n); - std::abort(); - } - // Every job of the previous batch has completed (`wait_done`), and a claim of it can no longer succeed, so - // nothing adds to `done` until this batch's first claim - which the release below orders after the reset. - done_.store(0, std::memory_order_relaxed); - const uint32_t e = epoch_.load(std::memory_order_relaxed) + 1; // only the host bumps the epoch - head_.store(pack_head(e, (uint32_t) n, 0), std::memory_order_release); - publish(); - return e; -} - -void ExpertPool::wait_parked(const char* what) { - hstate_.store(kWaitParked, std::memory_order_relaxed); - hstate_ms_.store(now_ms(), std::memory_order_relaxed); - uint32_t spins = 0; - std::chrono::steady_clock::time_point t0{}; - while (parked_.load(std::memory_order_acquire) != (uint32_t) n_) { - _mm_pause(); - if ((++spins & 1023u) != 0) continue; - const auto now = std::chrono::steady_clock::now(); - if (spins == 1024u) t0 = now; - else if (now - t0 > kStall) { - std::fprintf(stderr, "strata: the CPU expert pool stalled %s (%u of %d workers parked) - stopping the engine " - "so the server can start it again (issue #29)\n", - what, parked_.load(), n_); - std::fflush(stderr); - std::abort(); - } - } +void ExpertPool::open_phase() { + // The previous phase is closed, so no worker is inside a drain and none can join until this store. + const uint64_t s = state_.load(std::memory_order_relaxed); + state_.store(((s >> 32) + 1) << 32 | (uint64_t) n_, std::memory_order_seq_cst); + if (sleepers_.load(std::memory_order_seq_cst) != 0) state_.notify_all(); } -void ExpertPool::wait_done(int n) { - hstate_.store(kWaitDone, std::memory_order_relaxed); - hstate_ms_.store(now_ms(), std::memory_order_relaxed); - uint32_t spins = 0, seen = 0; - std::chrono::steady_clock::time_point t0{}; - for (;;) { - const uint32_t d = done_.load(std::memory_order_acquire); - if (d >= (uint32_t) n) return; // `>=`: never a wait that an overshoot outlives +void ExpertPool::close_phase() { + // Every job is done; wait for the workers that joined to park again, and shut the phase in the same step, + // so a worker that notices it late cannot join it. + const uint64_t open = (state_.load(std::memory_order_relaxed) & ~kParked) | (uint64_t) n_; + for (uint64_t s = open; + !state_.compare_exchange_weak(s, open | kClosed, std::memory_order_acq_rel, std::memory_order_relaxed); + s = open) _mm_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 - else if (now - t0 > kStall) { - std::fprintf(stderr, "strata: the CPU expert pool stalled: %u of %d jobs done, %u of %d workers parked - " - "stopping the engine so the server can start it again (issue #29)\n", - d, n, parked_.load(), n_); - std::fflush(stderr); - std::abort(); - } - } } -void ExpertPool::drain(int id, ExpertScratch& scratch, uint32_t epoch) { +void ExpertPool::drain(int id, ExpertScratch& scratch) { (void) id; for (;;) { - const int ci = claim(epoch); - if (ci < 0) break; - const uint32_t i = (uint32_t) ci; - if (id >= 0) wstate_[(size_t) id].store(ci, std::memory_order_relaxed); - else { hstate_.store(ci, std::memory_order_relaxed); hstate_ms_.store(now_ms(), std::memory_order_relaxed); } + const uint32_t i = head_.fetch_add(1, std::memory_order_relaxed); + if (i >= (uint32_t) njobs_) break; if (mode_ == 0) { const ExpertJob& j = jobs_[i]; s2_expert_vnni_q(j.blob, *j.act, j.out, scratch); @@ -416,16 +382,16 @@ void ExpertPool::drain(int id, ExpertScratch& scratch, uint32_t epoch) { } } -void ExpertPool::run_phase(int mode, int n_tasks) { - wait_parked("before a phase"); +void ExpertPool::run_phase(int mode, int n_tasks, void (*first)(void*), void* ctx) { mode_ = mode; njobs_ = n_tasks; - const uint32_t e = begin_batch(n_tasks); - if (host_works_) drain(-1, host_scratch_, e); - wait_done(n_tasks); - wait_parked("after a phase"); - hstate_.store(kIdle, std::memory_order_relaxed); - hstate_ms_.store(now_ms(), std::memory_order_relaxed); + head_.store(0, std::memory_order_relaxed); + done_.store(0, std::memory_order_relaxed); + open_phase(); + if (first != nullptr) first(ctx); + if (host_works_) drain(-1, host_scratch_); + while (done_.load(std::memory_order_acquire) != (uint32_t) n_tasks) _mm_pause(); + close_phase(); } void ExpertPool::run_split(ExpertJob* jobs, int n) { @@ -483,8 +449,12 @@ void ExpertPool::run_split_multi(ExpertJobMulti* jobs, int n) { ms_drain_ += std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); } -void ExpertPool::run_split_multi_native(const NativeFmt& f, ExpertJobMulti* jobs, int n) { - if (n <= 0) return; +void ExpertPool::run_split_multi_native(const NativeFmt& f, ExpertJobMulti* jobs, int n, void (*first)(void*), + void* ctx) { + if (n <= 0) { + if (first != nullptr) first(ctx); + return; + } const auto t0 = std::chrono::steady_clock::now(); // more distinct experts than buffers: run them in batches for (int b0 = 0; b0 < n; b0 += kMaxSplitMulti) { @@ -495,7 +465,7 @@ void ExpertPool::run_split_multi_native(const NativeFmt& f, ExpertJobMulti* jobs mtasks_ = 3 * threads; mrows_ = (int64_t) nb * FF; const auto a = std::chrono::steady_clock::now(); - run_phase(5, mtasks_); + run_phase(5, mtasks_, b0 == 0 ? first : nullptr, ctx); const auto b = std::chrono::steady_clock::now(); for (int e = 0; e < nb; ++e) for (int t = 0; t < mjobs_[e].nt; ++t) @@ -520,23 +490,22 @@ void ExpertPool::run(ExpertJob* jobs, int n) { for (int i = 0; i < n; ++i) s2_expert_vnni_q(jobs[i].blob, *jobs[i].act, jobs[i].out, scratch_[0]); return; } - // Wait for every worker to be parked BEFORE touching the batch, so the publish below is the only thing - // that can move a worker into the drain loop. + // The previous phase is closed, so no worker can be in the drain loop while the batch is written. // - // THE THREE PHASES ARE TIMED SEPARATELY. They were one number, which cannot distinguish a pool that is - // slow at the WORK from one that is slow at the SYNCHRONISATION - and those need opposite fixes. - const auto t_a = std::chrono::steady_clock::now(); - wait_parked("before a batch"); + // THE PHASES ARE TIMED SEPARATELY. They were one number, which cannot distinguish a pool that is slow at + // the WORK from one that is slow at the SYNCHRONISATION - and those need opposite fixes. const auto t_b = std::chrono::steady_clock::now(); jobs_ = jobs; njobs_ = n; mode_ = 0; - const uint32_t e = begin_batch(n); // done, then head (release), then the epoch: the batch is described first + head_.store(0, std::memory_order_relaxed); + done_.store(0, std::memory_order_relaxed); + open_phase(); // seq_cst: jobs_/njobs_ are visible to a worker that joins // ---- **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 - // `physical_cores(true)` deliberately keeps the five workers off - so for the whole drain that core was + // The loop below used to be `while (done_ != n) _mm_pause();`. The host is pinned to its own core - the + // one `CorePlacement` deliberately keeps the 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. // @@ -549,28 +518,23 @@ void ExpertPool::run(ExpertJob* jobs, int n) { // device after `run()` returns, so the write must be published, not merely performed. if (host_works_) { for (;;) { - const int ci = claim(e); - if (ci < 0) break; - hstate_.store(ci, std::memory_order_relaxed); - hstate_ms_.store(now_ms(), std::memory_order_relaxed); - const ExpertJob& j = jobs_[ci]; + const uint32_t i = head_.fetch_add(1, std::memory_order_relaxed); + if (i >= (uint32_t) n) break; + const ExpertJob& j = jobs_[i]; s2_expert_vnni_q(j.blob, *j.act, j.out, host_scratch_); done_.fetch_add(1, std::memory_order_release); } } - wait_done(n); - // And park again, so the next `run` starts from a known state. See the header for why `done` alone is - // not enough. + while (done_.load(std::memory_order_acquire) != (uint32_t) n) _mm_pause(); + // And close the phase, so the next `run` starts from a known state. See the header for why `done` alone + // is not enough. const auto t_c = std::chrono::steady_clock::now(); - wait_parked("after a batch"); - hstate_.store(kIdle, std::memory_order_relaxed); - hstate_ms_.store(now_ms(), std::memory_order_relaxed); + close_phase(); const auto t_d = std::chrono::steady_clock::now(); - ms_wait_park_ += std::chrono::duration(t_b - t_a).count(); ms_drain_ += std::chrono::duration(t_c - t_b).count(); - ms_repark_ += std::chrono::duration(t_d - t_c).count(); + ms_close_ += std::chrono::duration(t_d - t_c).count(); } } // namespace strata::kernels::cpu diff --git a/src/kernels/cpu/pool_test.cpp b/src/kernels/cpu/pool_test.cpp index dbc7288a1d5..50d59246ff2 100644 --- a/src/kernels/cpu/pool_test.cpp +++ b/src/kernels/cpu/pool_test.cpp @@ -13,8 +13,16 @@ // mode the park protocol exists to prevent, so `run()` is called many times in a row. // 4. A BATCH BIGGER AND SMALLER THAN THE WORKER COUNT, because `n < workers` leaves most workers claiming // nothing and `n > workers` is the real case (10 experts, 5 workers). +// 5. (--stress, synthetic experts, no file) SLEEPING AND LATE WORKERS. Idle workers sleep after `kIdleSpin`; +// a phase opened while they sleep is often finished by the host alone before they wake, and a worker that +// wakes then must not join it: the host may already be writing the next phase. Thousands of batches, with +// idle gaps either side of `kIdleSpin` and one-expert batches after the long ones, every output bitwise +// against a serial run of the same kernels. #include "strata/kernels/cpu/pool.hpp" #include "strata/kernels/cpu/expert.hpp" +#if defined(STRATA_NATIVE_EXPERTS) +#include "ggml.h" +#endif #include #include @@ -22,8 +30,10 @@ #include #include #include +#include #include #include +#include #include namespace cpu = strata::kernels::cpu; @@ -49,19 +59,224 @@ double now_ms() { return duration(steady_clock::now().time_since_epoch()).count(); } +// ---- 5: --stress +const float kSentinel = -1.2345e33f; + +std::vector gauss_vec(std::mt19937& rng, size_t n, float sd) { + std::normal_distribution g(0.0f, sd); + std::vector v(n); + for (auto& x : v) x = g(rng); + return v; +} + +double median(std::vector v) { + if (v.empty()) return 0.0; + std::sort(v.begin(), v.end()); + return v[v.size() / 2]; +} + +/// The pause before a batch: mostly none (generation's layers come back to back), sometimes up to 2 ms, and one +/// in 40 longer than `kIdleSpin`, so every worker is asleep when the batch is published. True for a long one. +bool gap(std::mt19937& rng) { + const unsigned r = rng() % 40; + if (r == 0) { + std::this_thread::sleep_for(cpu::ExpertPool::kIdleSpin + std::chrono::milliseconds(15)); + return true; + } + if (r < 5) std::this_thread::sleep_for(std::chrono::microseconds(100 + rng() % 1900)); + return false; +} + +struct StressStats { + long long batches = 0, outputs = 0, differ = 0; + int long_gaps = 0; + std::vector ms_spinning, ms_after_sleep; // one-expert batches +}; + +void stress_report(const char* what, const StressStats& s, uint32_t sleeps) { + std::printf(" %-34s %s (%lld of %lld outputs differ, %lld batches; %d long gaps, %u worker sleeps)\n", what, + s.differ ? "*** NO ***" : "identical to serial", s.differ, s.outputs, s.batches, s.long_gaps, sleeps); + std::printf(" %-34s %.3f ms while the workers spin, %.3f ms right after they slept (median)\n", + " a one-expert batch", median(s.ms_spinning), median(s.ms_after_sleep)); +} + +/// `run()`: whole Q2_0 experts. Random codes with finite scales: the protocol is under test, not the kernel. +StressStats stress_run(cpu::ExpertPool& pool, std::mt19937& rng, int iters) { + const int NB = 4, NA = 4, MAXJ = 16; + std::vector> blobs((size_t) NB, std::vector(cpu::BLOB)); + for (auto& b : blobs) { + for (auto& v : b) v = (uint8_t) rng(); + for (size_t o = cpu::O_GU_SCALES; o + 2 <= cpu::BLOB; o += 2) { + const uint16_t h = (uint16_t) (0x2000 | (rng() & 0x3FF)); // fp16 in [2^-7, 2^-6) + std::memcpy(b.data() + o, &h, 2); + } + } + std::vector acts((size_t) NA); + for (auto& a : acts) cpu::act_quant_q8_1(gauss_vec(rng, cpu::H, 1.0f).data(), cpu::H, a); + std::vector jobs((size_t) MAXJ); + std::vector out((size_t) MAXJ * cpu::H), ref(cpu::H); + cpu::ExpertScratch ws; + StressStats s; + for (int it = 0; it < iters; ++it) { + const bool slept = gap(rng); + s.long_gaps += slept; + const bool tiny = slept ? rng() % 2 == 0 : rng() % 20 == 0; + const int n = tiny ? 1 : 1 + (int) (rng() % MAXJ); + for (int j = 0; j < n; ++j) { + jobs[(size_t) j].blob = blobs[rng() % NB].data(); + jobs[(size_t) j].act = &acts[rng() % NA]; + jobs[(size_t) j].out = out.data() + (size_t) j * cpu::H; + std::fill(jobs[(size_t) j].out, jobs[(size_t) j].out + cpu::H, kSentinel); + } + const double t0 = now_ms(); + pool.run(jobs.data(), n); + const double ms = now_ms() - t0; + if (tiny) (slept ? s.ms_after_sleep : s.ms_spinning).push_back(ms); + for (int j = 0; j < n; ++j) { + cpu::s2_expert_vnni_q(jobs[(size_t) j].blob, *jobs[(size_t) j].act, ref.data(), ws); + s.differ += std::memcmp(ref.data(), jobs[(size_t) j].out, sizeof(float) * cpu::H) != 0; + } + s.outputs += n; + ++s.batches; + } + return s; +} + +#if defined(STRATA_NATIVE_EXPERTS) +/// `run_split_multi_native`: one layer's experts in GGUF formats, split by rows, several tokens per expert. +StressStats stress_native(cpu::ExpertPool& pool, std::mt19937& rng, int iters, ggml_type gu, ggml_type dn, + std::string& err) { + StressStats s; + cpu::NativeFmt f; + if (!cpu::native_fmt((int) gu, (int) dn, cpu::H, cpu::FF, f, err)) return s; + const int NB = 4, NA = 8, MAXE = 12, MAXNT = 4; + // ggml's quantizers on random weights (valid blocks), one thread per expert: the i-quant ones are slow + std::vector> blobs((size_t) NB, std::vector(f.bytes)); + { + std::vector th; + for (int b = 0; b < NB; ++b) + th.emplace_back([&f, &blobs, gu, dn, b, seed = rng()] { + std::mt19937 r(seed); + uint8_t* p = blobs[(size_t) b].data(); + ggml_quantize_chunk(gu, gauss_vec(r, (size_t) cpu::FF * cpu::H, 0.02f).data(), p, 0, cpu::FF, cpu::H, nullptr); + ggml_quantize_chunk(gu, gauss_vec(r, (size_t) cpu::FF * cpu::H, 0.02f).data(), p + f.up_off, 0, cpu::FF, + cpu::H, nullptr); + ggml_quantize_chunk(dn, gauss_vec(r, (size_t) cpu::H * cpu::FF, 0.02f).data(), p + f.down_off, 0, cpu::H, + cpu::FF, nullptr); + }); + for (auto& t : th) t.join(); + } + std::vector> acts((size_t) NA, std::vector(cpu::kNativeActBytes)); + for (auto& a : acts) cpu::native_quant_act(f, gauss_vec(rng, cpu::H, 1.0f).data(), a.data()); + std::vector jobs((size_t) MAXE); + std::vector out((size_t) MAXE * MAXNT * cpu::H), ref((size_t) MAXNT * cpu::H); + std::vector ff((size_t) MAXNT * cpu::FF); + std::vector hq((size_t) MAXNT * cpu::kNativeHBytes); + for (int it = 0; it < iters; ++it) { + const bool slept = gap(rng); + s.long_gaps += slept; + const bool tiny = slept ? rng() % 2 == 0 : rng() % 20 == 0; + const int n = tiny ? 1 : 1 + (int) (rng() % MAXE); + for (int e = 0; e < n; ++e) { + cpu::ExpertJobMulti& j = jobs[(size_t) e]; + j = cpu::ExpertJobMulti{}; + j.blob = blobs[rng() % NB].data(); + j.nt = tiny ? 1 : 1 + (int) (rng() % MAXNT); + int tok[NA]; + std::iota(tok, tok + NA, 0); + std::shuffle(tok, tok + NA, rng); + for (int t = 0; t < j.nt; ++t) { + j.nact[t] = acts[(size_t) tok[t]].data(); + j.out[t] = out.data() + ((size_t) e * MAXNT + (size_t) t) * cpu::H; + std::fill(j.out[t], j.out[t] + cpu::H, kSentinel); + } + } + const double t0 = now_ms(); + pool.run_split_multi_native(f, jobs.data(), n); + const double ms = now_ms() - t0; + if (tiny) (slept ? s.ms_after_sleep : s.ms_spinning).push_back(ms); + // the same kernels and token groups on this thread, every row at once + for (int e = 0; e < n; ++e) { + const cpu::ExpertJobMulti& j = jobs[(size_t) e]; + float* ffp[cpu::MAXT]; + float* rp[cpu::MAXT]; + const void* hqp[cpu::MAXT]; + for (int t = 0; t < j.nt; ++t) { + ffp[t] = ff.data() + (size_t) t * cpu::FF; + rp[t] = ref.data() + (size_t) t * cpu::H; + hqp[t] = hq.data() + (size_t) t * cpu::kNativeHBytes; + } + cpu::native_gu_rows(f, j.blob, j.nact, j.nt, ffp, 0, cpu::FF); + for (int t = 0; t < j.nt; ++t) cpu::native_quant_h(f, ffp[t], hq.data() + (size_t) t * cpu::kNativeHBytes); + cpu::native_down_rows(f, j.blob, hqp, j.nt, rp, 0, cpu::H); + for (int t = 0; t < j.nt; ++t) s.differ += std::memcmp(rp[t], j.out[t], sizeof(float) * cpu::H) != 0; + s.outputs += j.nt; + } + ++s.batches; + } + return s; +} +#endif + +int stress(int iters) { + std::mt19937 rng(20260927); + cpu::ExpertPool pool; + std::printf(" %-34s %d + the host thread; they sleep after %lld ms without work\n", "workers", pool.workers(), + (long long) cpu::ExpertPool::kIdleSpin.count()); + int bad = 0; + if (cpu::cpu_features().usable()) { + const uint32_t z0 = pool.sleeps(); + const StressStats s = stress_run(pool, rng, iters); + stress_report("run(), whole Q2_0 experts", s, pool.sleeps() - z0); + bad += s.differ != 0 || s.long_gaps == 0 || pool.sleeps() == z0; + } else { + std::printf(" run(): SKIPPED, the CPU lacks %s\n", cpu::cpu_features().reason()); + } +#if defined(STRATA_NATIVE_EXPERTS) + const struct { ggml_type gu, dn; const char* name; } fmts[] = { + {GGML_TYPE_Q4_K, GGML_TYPE_Q5_1, "native rows, Q4_K / Q5_1"}, // ggml-cpu's dot products + {GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ4_NL, "native rows, IQ3_XXS / IQ4_NL"}, // + the AVX-512 i-quant kernel + }; + for (const auto& ft : fmts) { + std::string err; + const uint32_t z0 = pool.sleeps(); + const StressStats s = stress_native(pool, rng, iters, ft.gu, ft.dn, err); + if (!err.empty()) { + std::printf(" %s: %s\n", ft.name, err.c_str()); + ++bad; + continue; + } + stress_report(ft.name, s, pool.sleeps() - z0); + bad += s.differ != 0 || s.long_gaps == 0 || pool.sleeps() == z0; + } +#else + std::printf(" native rows: SKIPPED, built without STRATA_NATIVE_EXPERTS\n"); +#endif + std::printf("\npool stress: %d failures\n", bad); + return bad ? 1 : 0; +} + } // namespace int main(int argc, char** argv) { bool selftest = false; + int stress_iters = 0; const char* path = "pack/full/experts.bin"; long long layer = 0; for (int i = 1; i < argc; ++i) { const std::string a = argv[i]; if (a == "--selftest") selftest = true; + else if (a == "--stress") stress_iters = 1500; + else if (a == "--iters" && i + 1 < argc) stress_iters = std::atoi(argv[++i]); else if (a == "--file" && i + 1 < argc) path = argv[++i]; else if (a == "--layer" && i + 1 < argc) layer = std::atoll(argv[++i]); - else { std::fprintf(stderr, "usage: pool_test [--selftest] [--file P] [--layer N]\n"); return 2; } + else { + std::fprintf(stderr, "usage: pool_test [--selftest] [--file P] [--layer N]\n" + " pool_test --stress [--iters N] (synthetic experts, no file)\n"); + return 2; + } } + if (stress_iters > 0) return stress(stress_iters); const cpu::CpuFeatures feat = cpu::cpu_features(); if (!feat.usable()) { @@ -99,9 +314,9 @@ int main(int argc, char** argv) { } // ---- the pool - const std::vector cores = cpu::physical_cores(true); + const std::vector cores = cpu::CorePlacement(false).workers(); const int hw = (int) std::thread::hardware_concurrency(); - std::printf(" %-44s %d logical, %d physical (skipping the first)\n", "cores the pool will use", + std::printf(" %-44s %d logical, %d physical (not the last, the host's)\n", "cores the pool will use", hw, (int) cores.size()); cpu::ExpertPool pool; diff --git a/src/kernels/cuda/elementwise.cu b/src/kernels/cuda/elementwise.cu index a2e9e53b5c7..60c26e8b186 100644 --- a/src/kernels/cuda/elementwise.cu +++ b/src/kernels/cuda/elementwise.cu @@ -293,6 +293,37 @@ __global__ void copy_i32_from_mapped_kernel(int32_t* __restrict__ dst, const vol for (int i = threadIdx.x; i < n; i += blockDim.x) dst[i] = src[i]; } +struct MappedCopies { + MappedCopy c[8]; +}; + +__global__ void copy_from_mapped_multi_kernel(MappedCopies a) { + const MappedCopy c = a.c[blockIdx.y]; + const int64_t i0 = (int64_t) blockIdx.x * blockDim.x + threadIdx.x, st = (int64_t) gridDim.x * blockDim.x; + if ((c.words & 3) == 0 && ((uintptr_t) c.dst & 15) == 0 && ((uintptr_t) c.src & 15) == 0) { + const volatile float4* s = (const volatile float4*) c.src; + for (int64_t i = i0; i < c.words / 4; i += st) ((float4*) c.dst)[i] = const_cast(s)[i]; + } else { + const volatile int32_t* s = (const volatile int32_t*) c.src; + for (int64_t i = i0; i < c.words; i += st) ((int32_t*) c.dst)[i] = s[i]; + } +} + +void copy_from_mapped_multi(const MappedCopy* copies, int n, void* stream) { + if (n <= 0) return; + if (n > 8) { std::fprintf(stderr, "copy_from_mapped_multi: at most 8 copies\n"); std::exit(1); } + MappedCopies a; + int64_t most = 0; + for (int i = 0; i < n; ++i) { + a.c[i] = copies[i]; + most = copies[i].words > most ? copies[i].words : most; + } + const int64_t units = (most + 3) / 4; + const unsigned bx = (unsigned) ((units + 255) / 256 < 32 ? (units + 255) / 256 : 32); + copy_from_mapped_multi_kernel<<>>(a); + check_launch("copy_from_mapped_multi"); +} + void copy_i32_from_mapped(int32_t* dst, const int32_t* src, int64_t n, void* stream) { if (n <= 0) return; copy_i32_from_mapped_kernel<<<1, 128, 0, (cudaStream_t) stream>>>(dst, (const volatile int32_t*) src, (int) n); diff --git a/src/kernels/cuda/fused_gr.cu b/src/kernels/cuda/fused_gr.cu index f56e152bc14..9c2fc805118 100644 --- a/src/kernels/cuda/fused_gr.cu +++ b/src/kernels/cuda/fused_gr.cu @@ -1,10 +1,10 @@ // src/kernels/cuda/fused_gr.cu - see include/strata/kernels/fused_gr.hpp. #include "strata/kernels/fused_gr.hpp" -#include "strata/kernels/bf16_bits.hpp" -#include "strata/kernels/verify_kernels.hpp" +#include "xor_scatter.cuh" #include +#include #include #include @@ -15,11 +15,23 @@ constexpr int N = 2560; // n_embd constexpr int HC = 4; // streams constexpr int D = N * HC; // 10240 constexpr int LR = 320; // hc_lr +constexpr int NR = LR + HC; // the down projection's rows: w_down's 320, then w_inject's 4 + +// down: the NR rows in DG groups of DRB = DW warps x DR rows, the input in DS slices of DCS columns (a slice lies in +// one stream); a block per (slice, group), up to DT tokens a launch +constexpr int DG = 4, DS = 20, DW = 9, DR = 9, DT = 4; +constexpr int DTH = DW * 32, DRB = DW * DR, DCS = D / DS, DSPC = DS / HC, DSTEPS = DCS / 256, DNF = DCS / 4; +static_assert(DRB * DG == NR && DCS % 256 == 0 && DS % HC == 0 && DNF <= DTH, "down geometry"); + +// up: UPM_COLS columns of every stream a block, UPM_ROWS rows a warp constexpr int THREADS = 256; constexpr int WARPS = THREADS / 32; -constexpr int DOWN_BLOCKS = LR / WARPS; // 40 blocks of 8 rows; one more for the inject rows -constexpr int UP_COLS = 32; // columns d per `up` block (x 4 streams = 128 rows) -constexpr int UP_BLOCKS = N / UP_COLS; // 80 +constexpr int UPM_COLS = 16; +constexpr int UPM_BLOCKS = N / UPM_COLS; // 160 +constexpr int UPM_ROWS = HC * UPM_COLS / WARPS; // 8 + +// the scratch (floats): DG counters, the slices' sums of squares [DG][DS][DT], the partial dots [DS][DT][NR] +constexpr int SC_CTR = 0, SC_SSP = 64, SC_PART = SC_SSP + DG * DS * DT, SC_FLOATS = SC_PART + DS * DT * NR; __device__ __forceinline__ float warp_sum(float v) { #pragma unroll @@ -28,281 +40,268 @@ __device__ __forceinline__ float warp_sum(float v) { } __device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f + __expf(-x)); } -// 8 bf16 packed in a uint4 against 8 floats. -__device__ __forceinline__ float dot8(const uint4 w, const float* x) { +// 8 bf16 packed in a uint4 against 8 floats held as two float4 +__device__ __forceinline__ float dot8h(const uint4 w, const float4 x0, const float4 x1) { float acc = 0.0f; - const uint32_t v[4] = {w.x, w.y, w.z, w.w}; + acc = fmaf(__uint_as_float(w.x << 16), x0.x, acc); + acc = fmaf(__uint_as_float(w.x & 0xffff0000u), x0.y, acc); + acc = fmaf(__uint_as_float(w.y << 16), x0.z, acc); + acc = fmaf(__uint_as_float(w.y & 0xffff0000u), x0.w, acc); + acc = fmaf(__uint_as_float(w.z << 16), x1.x, acc); + acc = fmaf(__uint_as_float(w.z & 0xffff0000u), x1.y, acc); + acc = fmaf(__uint_as_float(w.w << 16), x1.z, acc); + acc = fmaf(__uint_as_float(w.w & 0xffff0000u), x1.w, acc); + return acc; +} + +struct GrMulti { + FusedGrArgs a[kFusedGrMaxT]; + int T; +}; + +__host__ __device__ constexpr int pow2_floor(int n) { return n < 2 ? 1 : 2 * pow2_floor(n / 2); } + +// v[B, V) summed across the warp in power-of-two chunks of xor_scatter; emit(value, sum) on one lane per value +template +__device__ __forceinline__ void reduce_chunks(const float (&v)[VT], int lane, F& emit) { + if constexpr (B < V) { + constexpr int P = pow2_floor(V - B > 32 ? 32 : V - B); + float cv[P]; #pragma unroll - for (int j = 0; j < 4; ++j) { - acc = fmaf(__uint_as_float(v[j] << 16), x[2 * j], acc); - acc = fmaf(__uint_as_float(v[j] & 0xffff0000u), x[2 * j + 1], acc); + for (int e = 0; e < P; ++e) cv[e] = v[B + e]; + xor_scatter

(cv, lane); + if ((lane & (32 / P - 1)) == 0) emit(B + xor_scatter_index

(lane), cv[0]); + reduce_chunks(v, lane, emit); } - return acc; } -__global__ void __launch_bounds__(THREADS) gr_down_kernel(FusedGrArgs a) { - __shared__ __align__(16) float xn[D]; - __shared__ float part[WARPS][HC]; - __shared__ float s_rs[HC]; +// The down projection and the injection of TT tokens. Block (s, g): warp w takes rows g DRB + w DR + [0, DR) of +// [w_down; w_inject] over the slice's DCS columns, a lane 8 columns a step, each activation chunk loaded once for its +// DR rows. The slice's R' = R + bo_prev * 2 sigmoid(inj_prev / hc) times w_norm is staged in shared memory split into +// the halves of the 8-column chunks (conflict-free 16-byte reads), and the sums of squares of R' go to the scratch: +// a slice lies in one stream, so the stream's 1/rms multiplies the slice sums later. The group's last block (a +// counter) adds each row's partials in slice order, per stream times the stream's 1/rms, and writes lo = silu(y / hc) +// or the injection; group 0's also writes rs. Every token's arithmetic is the same whatever its place in the window +// and the window's size. +template +__global__ void __launch_bounds__(DTH) gr_down_split_kernel(GrMulti m, float* __restrict__ scratch) { + static_assert(TT >= 1 && TT <= DT, "tokens a launch"); + __shared__ __align__(16) float4 xs[TT][DNF]; + __shared__ float ssw[TT][DW]; + __shared__ float s_rs[TT][HC]; + __shared__ bool s_last; + unsigned* ctr = reinterpret_cast(scratch + SC_CTR); + float* ssp = scratch + SC_SSP; + float* part = scratch + SC_PART; const int t = threadIdx.x, lane = t & 31, warp = t >> 5; - float gw[HC]; + const int s = blockIdx.x, g = blockIdx.y; + const int c = s / DSPC, col0 = s * DCS; + const bool has_inj = m.a[0].w_inject != nullptr; + const int rb0 = g * DRB + warp * DR; + // the staging loads first, then the weights + float4 pr[TT], pb[TT], pw = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + if (t < DNF) { + pw = *reinterpret_cast(m.a[0].w_norm + col0 + 4 * t); #pragma unroll - for (int c = 0; c < HC; ++c) gw[c] = a.apply ? 2.0f * sigmoidf_(a.inj_prev[c] / (float) HC) : 0.0f; - // 1. R' * w_norm into shared memory, and the per-stream sums of squares of R'. - float ss[HC] = {0.0f, 0.0f, 0.0f, 0.0f}; - for (int i = t * 4; i < D; i += THREADS * 4) { - const int c = i / N, d = i - c * N; - float4 r = *reinterpret_cast(a.R + i); - if (a.apply) { - const float4 b = *reinterpret_cast(a.bo_prev + d); - r.x = fmaf(b.x, gw[c], r.x); r.y = fmaf(b.y, gw[c], r.y); - r.z = fmaf(b.z, gw[c], r.z); r.w = fmaf(b.w, gw[c], r.w); + for (int k = 0; k < TT; ++k) { + const FusedGrArgs& a = m.a[k]; + pr[k] = *reinterpret_cast(a.R + col0 + 4 * t); + if (a.apply) pb[k] = *reinterpret_cast(a.bo_prev + col0 - c * N + 4 * t); } - const float4 g = *reinterpret_cast(a.w_norm + i); - float sq = r.x * r.x + r.y * r.y + r.z * r.z + r.w * r.w; -#pragma unroll - for (int cc = 0; cc < HC; ++cc) if (cc == c) ss[cc] += sq; - *reinterpret_cast(xn + i) = make_float4(r.x * g.x, r.y * g.y, r.z * g.z, r.w * g.w); } + uint4 w[DR][DSTEPS]; #pragma unroll - for (int c = 0; c < HC; ++c) { - const float v = warp_sum(ss[c]); - if (lane == 0) part[warp][c] = v; - } - __syncthreads(); - if (t < HC) { - float s = 0.0f; - for (int w = 0; w < WARPS; ++w) s += part[w][t]; - s_rs[t] = rsqrtf(s / (float) N + a.eps); - if (blockIdx.x == 0) a.rs[t] = s_rs[t]; - } - __syncthreads(); - for (int i = t; i < D; i += THREADS) xn[i] *= s_rs[i / N]; - __syncthreads(); - // 2. one warp per output row: 10240 bf16 = 1280 chunks of 8, 40 per lane. - const bool inject_block = blockIdx.x == DOWN_BLOCKS; - const int row = inject_block ? warp : blockIdx.x * WARPS + warp; - if (inject_block && (a.w_inject == nullptr || warp >= HC)) return; - const uint16_t* wrow = (inject_block ? a.w_inject : a.w_down) + (size_t) row * D; - const uint4* w4 = reinterpret_cast(wrow); - float acc = 0.0f; -#pragma unroll 4 - for (int j = lane; j < D / 8; j += 32) acc += dot8(__ldg(w4 + j), xn + j * 8); - acc = warp_sum(acc); - if (lane != 0) return; - if (inject_block) { - a.inject_out[row] = acc; - } else { - const float x = acc / (float) HC; - a.lo[row] = x / (1.0f + __expf(-x)); + for (int i = 0; i < DR; ++i) { + const int rb = rb0 + i; + const bool ok = rb < LR || has_inj; + const uint16_t* base = + rb < LR ? m.a[0].w_down + (size_t) rb * D : m.a[0].w_inject + (size_t) (ok ? rb - LR : 0) * D; + const uint4* w4 = reinterpret_cast(base + col0) + lane; +#pragma unroll + for (int j = 0; j < DSTEPS; ++j) w[i][j] = ok ? __ldg(w4 + j * 32) : make_uint4(0, 0, 0, 0); } -} - -__global__ void __launch_bounds__(THREADS) gr_up_kernel(FusedGrArgs a) { - __shared__ __align__(16) float lo[LR]; - __shared__ float g[HC][UP_COLS]; - const int t = threadIdx.x, lane = t & 31, warp = t >> 5; - const int d0 = blockIdx.x * UP_COLS; - for (int k = t; k < LR; k += THREADS) lo[k] = a.lo[k]; - __syncthreads(); - // 128 rows (4 streams x 32 columns), 16 per warp: 320 bf16 = 40 chunks of 8. - for (int r = warp; r < HC * UP_COLS; r += WARPS) { - const int c = r / UP_COLS, dd = r - c * UP_COLS, i = c * N + d0 + dd; - const uint4* w4 = reinterpret_cast(a.w_up + (size_t) i * LR); - float acc = dot8(__ldg(w4 + lane), lo + lane * 8); - if (lane < LR / 8 - 32) acc += dot8(__ldg(w4 + 32 + lane), lo + (32 + lane) * 8); - acc = warp_sum(acc); - if (lane == 0) { - float rv = a.R[i]; +#pragma unroll + for (int k = 0; k < TT; ++k) { + const FusedGrArgs& a = m.a[k]; + const float gw = a.apply ? 2.0f * sigmoidf_(a.inj_prev[c] / (float) HC) : 0.0f; + float sq = 0.0f; + if (t < DNF) { + float4 r = pr[k]; if (a.apply) { - rv = fmaf(a.bo_prev[d0 + dd], 2.0f * sigmoidf_(a.inj_prev[c] / (float) HC), rv); - a.R_out[i] = rv; // this block owns column d0+dd of every stream + r.x = fmaf(pb[k].x, gw, r.x); r.y = fmaf(pb[k].y, gw, r.y); + r.z = fmaf(pb[k].z, gw, r.z); r.w = fmaf(pb[k].w, gw, r.w); } - const float x = rv * a.w_norm[i] * a.rs[c]; - g[c][dd] = x * sigmoidf_(acc); + sq = fmaf(r.x, r.x, sq); sq = fmaf(r.y, r.y, sq); sq = fmaf(r.z, r.z, sq); sq = fmaf(r.w, r.w, sq); + xs[k][(t & 1) * (DNF / 2) + (t >> 1)] = make_float4(r.x * pw.x, r.y * pw.y, r.z * pw.z, r.w * pw.w); } + sq = warp_sum(sq); + if (lane == 0) ssw[k][warp] = sq; } __syncthreads(); - if (t < UP_COLS) { - float s = 0.0f; + if (t < TT) { + float sq = 0.0f; #pragma unroll - for (int c = 0; c < HC; ++c) s += g[c][t]; - a.mixed[d0 + t] = s / (float) HC; + for (int q = 0; q < DW; ++q) sq += ssw[t][q]; + ssp[(g * DS + s) * DT + t] = sq; } -} - -// ================================ plan v0.3 P6: T tokens, one weight read ================================ -struct GrMulti { - FusedGrArgs a[kFusedGrMaxT]; - float* xn; - int T; -}; - -// Step 1 of `gr_down_kernel`, one block per token, same threads and reduction order: rs[t] and xn[t] to global. -__global__ void __launch_bounds__(THREADS) gr_norm_multi_kernel(GrMulti m) { - __shared__ float part[WARPS][HC]; - __shared__ float s_rs[HC]; - const FusedGrArgs& a = m.a[blockIdx.x]; - float* xn = m.xn + (size_t) blockIdx.x * D; - const int t = threadIdx.x, lane = t & 31, warp = t >> 5; - float gw[HC]; + float acc[DR * TT]; #pragma unroll - for (int c = 0; c < HC; ++c) gw[c] = a.apply ? 2.0f * sigmoidf_(a.inj_prev[c] / (float) HC) : 0.0f; - float ss[HC] = {0.0f, 0.0f, 0.0f, 0.0f}; - for (int i = t * 4; i < D; i += THREADS * 4) { - const int c = i / N, d = i - c * N; - float4 r = *reinterpret_cast(a.R + i); - if (a.apply) { - const float4 b = *reinterpret_cast(a.bo_prev + d); - r.x = fmaf(b.x, gw[c], r.x); r.y = fmaf(b.y, gw[c], r.y); - r.z = fmaf(b.z, gw[c], r.z); r.w = fmaf(b.w, gw[c], r.w); - } - const float4 g = *reinterpret_cast(a.w_norm + i); - float sq = r.x * r.x + r.y * r.y + r.z * r.z + r.w * r.w; + for (int e = 0; e < DR * TT; ++e) acc[e] = 0.0f; #pragma unroll - for (int cc = 0; cc < HC; ++cc) if (cc == c) ss[cc] += sq; - *reinterpret_cast(xn + i) = make_float4(r.x * g.x, r.y * g.y, r.z * g.z, r.w * g.w); - } + for (int j = 0; j < DSTEPS; ++j) { + float4 x0[TT], x1[TT]; +#pragma unroll + for (int k = 0; k < TT; ++k) { x0[k] = xs[k][j * 32 + lane]; x1[k] = xs[k][DNF / 2 + j * 32 + lane]; } +#pragma unroll + for (int i = 0; i < DR; ++i) #pragma unroll - for (int c = 0; c < HC; ++c) { - const float v = warp_sum(ss[c]); - if (lane == 0) part[warp][c] = v; + for (int k = 0; k < TT; ++k) acc[i * TT + k] += dot8h(w[i][j], x0[k], x1[k]); } + auto emit = [=](int vi, float sum) { + const int i = vi / TT, k = vi - i * TT, rb = rb0 + i; + if (rb < LR || has_inj) part[(s * DT + k) * NR + rb] = sum; + }; + reduce_chunks(acc, lane, emit); + // the group's last block + __threadfence(); __syncthreads(); - if (t < HC) { - float s = 0.0f; - for (int w = 0; w < WARPS; ++w) s += part[w][t]; - s_rs[t] = rsqrtf(s / (float) N + a.eps); - a.rs[t] = s_rs[t]; - } + if (t == 0) s_last = atomicAdd(ctr + g, 1u) == (unsigned) (DS - 1); __syncthreads(); - for (int i = t; i < D; i += THREADS) xn[i] *= s_rs[i / N]; -} - -#if defined(__HIPCC__) -constexpr int TILE = 1280; // eight-token tile fits gfx1100's 64 KiB LDS limit -#else -constexpr int TILE = 2560; // xn floats per token staged at a time: 320 chunks of 8, 10 per lane -#endif -constexpr int TQ = TILE / 8 / 32; // uint4 weight chunks per lane per tile - -// Step 2 of `gr_down_kernel` for T tokens. One warp per row (so each lane accumulates the same chunks in the -// same order as the single-token kernel); per tile the lane's 10 weight chunks are loaded BEFORE the activation -// tile is staged, so the DRAM and L2 traffic are in flight together. -__global__ void __launch_bounds__(THREADS) gr_down_multi_kernel(GrMulti m) { - extern __shared__ __align__(16) float tile[]; // [T][TILE] - const int t = threadIdx.x, lane = t & 31, warp = t >> 5; - const int T = m.T; - const bool inject_block = blockIdx.x == DOWN_BLOCKS; - const int row = inject_block ? warp : blockIdx.x * WARPS + warp; - const bool active = !(inject_block && (m.a[0].w_inject == nullptr || warp >= HC)); - const uint16_t* wrow = (inject_block ? m.a[0].w_inject : m.a[0].w_down) + (size_t) (active ? row : 0) * D; - const uint4* w4 = reinterpret_cast(wrow); - float acc[kFusedGrMaxT]; -#pragma unroll - for (int k = 0; k < kFusedGrMaxT; ++k) acc[k] = 0.0f; - for (int base = 0; base < D; base += TILE) { - uint4 wv[TQ]; - if (active) { -#pragma unroll - for (int q = 0; q < TQ; ++q) wv[q] = __ldg(w4 + base / 8 + lane + 32 * q); - } - __syncthreads(); // the previous tile is consumed - const float4* src4 = reinterpret_cast(m.xn); - float4* tile4 = reinterpret_cast(tile); - for (int i = t; i < T * (TILE / 4); i += THREADS) { - const int k = i / (TILE / 4), off = i - k * (TILE / 4); - tile4[i] = src4[((size_t) k * D + base) / 4 + off]; - } - __syncthreads(); - if (!active) continue; + if (!s_last) return; + __threadfence(); + if (t < TT * HC) { + const int k = t / HC, cc = t - k * HC; + float sq = 0.0f; #pragma unroll - for (int q = 0; q < TQ; ++q) { - const int j = lane + 32 * q; -#pragma unroll - for (int k = 0; k < kFusedGrMaxT; ++k) - if (k < T) acc[k] += dot8(wv[q], tile + k * TILE + j * 8); - } + for (int q = 0; q < DSPC; ++q) sq += __ldcg(ssp + (g * DS + cc * DSPC + q) * DT + k); + s_rs[k][cc] = rsqrtf(sq / (float) N + m.a[0].eps); } - if (!active) return; - float s[kFusedGrMaxT]; + __syncthreads(); + for (int e = t; e < TT * DRB; e += DTH) { + const int k = e / DRB, r = e - k * DRB, rb = g * DRB + r; + if (rb >= LR && !has_inj) continue; + float y = 0.0f; #pragma unroll - for (int k = 0; k < kFusedGrMaxT; ++k) s[k] = k < T ? warp_sum(acc[k]) : 0.0f; - // lane k writes token k (every lane holds every sum after the xor reduction) + for (int cc = 0; cc < HC; ++cc) { + float p = 0.0f; #pragma unroll - for (int k = 0; k < kFusedGrMaxT; ++k) { - if (k >= T || lane != k) continue; - if (inject_block) { - m.a[k].inject_out[row] = s[k]; + for (int q = 0; q < DSPC; ++q) p += __ldcg(part + ((cc * DSPC + q) * DT + k) * NR + rb); + y = fmaf(s_rs[k][cc], p, y); + } + if (rb >= LR) { + m.a[k].inject_out[rb - LR] = y; } else { - const float x = s[k] / (float) HC; - m.a[k].lo[row] = x / (1.0f + __expf(-x)); + const float x = y / (float) HC; + m.a[k].lo[rb] = x / (1.0f + __expf(-x)); } } + if (g == 0 && t < TT * HC) m.a[t / HC].rs[t % HC] = s_rs[t / HC][t % HC]; + if (t == 0) ctr[g] = 0u; } -constexpr int UPM_COLS = 16; // columns per block (x 4 streams = 64 rows, 8 per warp) -constexpr int UPM_BLOCKS = N / UPM_COLS; // 160 - -// `gr_up_kernel` for T tokens: each row of w_up read once; the T dots reduced by xor so every lane holds every -// sum, and lane k runs token k's epilogue - the T epilogues in parallel instead of one after another. +// The up projection of TT tokens: each row of w_up read once, RW rows a warp at a time; the TT x RW dots of those rows +// summed by one xor tree (xor_scatter) and every (row, token) epilogue on the lane holding its sum, its inputs fetched +// while the dots run: R <- R' for this block's columns (when apply), mixed[d] = mean_c xn[c,d] * sigmoid(u[c,d]), the +// gates' running average (when gate_ema) and the other read's estimate (when est_norm). The lo vectors are staged +// split into chunk halves. +template __global__ void __launch_bounds__(THREADS) gr_up_multi_kernel(GrMulti m) { - __shared__ __align__(16) float lo[kFusedGrMaxT][LR]; - __shared__ float g[kFusedGrMaxT][HC][UPM_COLS]; + constexpr int RW = TT <= 4 ? 4 : 2; + constexpr int NV = pow2_at_least(TT * RW); // sums a lane reduces, zero-padded + constexpr int LPV = 32 / NV; // lanes that end up holding each sum + __shared__ __align__(16) float4 lo4[TT][2][LR / 8]; + __shared__ float g[TT][HC][UPM_COLS]; + __shared__ float gs[TT][HC][UPM_COLS]; // the gates + __shared__ float ge[TT][HC][UPM_COLS]; // the estimate's terms + const bool est = m.a[0].est_norm != nullptr; const int t = threadIdx.x, lane = t & 31, warp = t >> 5; - const int T = m.T; const int d0 = blockIdx.x * UPM_COLS; - for (int i = t; i < T * LR; i += THREADS) lo[i / LR][i % LR] = m.a[i / LR].lo[i % LR]; + for (int i = t; i < TT * LR / 4; i += THREADS) { + const int k = i / (LR / 4), f = i - k * (LR / 4); + lo4[k][f & 1][f >> 1] = reinterpret_cast(m.a[k].lo)[f]; + } __syncthreads(); - for (int r = warp; r < HC * UPM_COLS; r += WARPS) { - const int c = r / UPM_COLS, dd = r - c * UPM_COLS, i = c * N + d0 + dd; - const uint4* w4 = reinterpret_cast(m.a[0].w_up + (size_t) i * LR); - const uint4 wa = __ldg(w4 + lane); - const uint4 wb = lane < LR / 8 - 32 ? __ldg(w4 + 32 + lane) : make_uint4(0, 0, 0, 0); - // the epilogue inputs of this lane's token, fetched while the dots run - float rv = 0.0f, wn = 0.0f, rsc = 0.0f, bo = 0.0f, ip = 0.0f; + const int vi = xor_scatter_index(lane); // the sum this lane ends with: row vi / TT, token vi % TT + const int vq = vi / TT, vk = vi - vq * TT; + const bool epi = (lane & (LPV - 1)) == 0 && vi < TT * RW; + for (int r0 = warp * UPM_ROWS; r0 < (warp + 1) * UPM_ROWS; r0 += RW) { + uint4 wa[RW], wb[RW]; +#pragma unroll + for (int q = 0; q < RW; ++q) { + const int r = r0 + q, c = r / UPM_COLS, dd = r - c * UPM_COLS, i = c * N + d0 + dd; + const uint4* w4 = reinterpret_cast(m.a[0].w_up + (size_t) i * LR); + wa[q] = __ldg(w4 + lane); + wb[q] = lane < LR / 8 - 32 ? __ldg(w4 + 32 + lane) : make_uint4(0, 0, 0, 0); + } + const int er = r0 + vq, ec = er / UPM_COLS, edd = er - ec * UPM_COLS, ei = ec * N + d0 + edd; + float rv = 0.0f, wn = 0.0f, rsc = 0.0f, bo = 0.0f, ip = 0.0f, en = 0.0f, eg = 0.0f; bool apply = false; - if (lane < T) { - const FusedGrArgs& a = m.a[lane]; - rv = a.R[i]; - wn = a.w_norm[i]; - rsc = a.rs[c]; + if (epi) { + const FusedGrArgs& a = m.a[vk]; + rv = a.R[ei]; + wn = a.w_norm[ei]; + rsc = a.rs[ec]; apply = a.apply; - if (apply) { bo = a.bo_prev[d0 + dd]; ip = a.inj_prev[c]; } + if (apply) { bo = a.bo_prev[d0 + edd]; ip = a.inj_prev[ec]; } + if (est) { en = a.est_norm[ei]; eg = a.est_gates[ei]; } } - float mine = 0.0f; + float v[NV]; #pragma unroll - for (int k = 0; k < kFusedGrMaxT; ++k) { - if (k >= T) break; - float acc = dot8(wa, lo[k] + lane * 8); - if (lane < LR / 8 - 32) acc += dot8(wb, lo[k] + (32 + lane) * 8); - acc = warp_sum(acc); - if (lane == k) mine = acc; - } - if (lane < T) { + for (int q = 0; q < RW; ++q) +#pragma unroll + for (int k = 0; k < TT; ++k) { + float acc = dot8h(wa[q], lo4[k][0][lane], lo4[k][1][lane]); + if (lane < LR / 8 - 32) acc += dot8h(wb[q], lo4[k][0][32 + lane], lo4[k][1][32 + lane]); + v[q * TT + k] = acc; + } +#pragma unroll + for (int e = TT * RW; e < NV; ++e) v[e] = 0.0f; + xor_scatter(v, lane); + if (epi) { + float x = rv; if (apply) { - rv = fmaf(bo, 2.0f * sigmoidf_(ip / (float) HC), rv); - m.a[lane].R_out[i] = rv; + x = fmaf(bo, 2.0f * sigmoidf_(ip / (float) HC), x); + m.a[vk].R_out[ei] = x; } - const float x = rv * wn * rsc; - g[lane][c][dd] = x * sigmoidf_(mine); + const float y = x * wn * rsc; + const float gate = sigmoidf_(v[0]); + g[vk][ec][edd] = y * gate; + gs[vk][ec][edd] = gate; + ge[vk][ec][edd] = x * rsc * en * eg; } } __syncthreads(); - for (int i = t; i < T * UPM_COLS; i += THREADS) { + for (int i = t; i < TT * UPM_COLS; i += THREADS) { const int k = i / UPM_COLS, col = i - k * UPM_COLS; float s = 0.0f; #pragma unroll for (int c = 0; c < HC; ++c) s += g[k][c][col]; m.a[k].mixed[d0 + col] = s / (float) HC; } + if (est) + for (int i = t; i < TT * UPM_COLS; i += THREADS) { + const int k = i / UPM_COLS, col = i - k * UPM_COLS; + float s = 0.0f; +#pragma unroll + for (int c = 0; c < HC; ++c) s += ge[k][c][col]; + m.a[k].est[d0 + col] = s / (float) HC; + } + if (m.a[0].gate_ema != nullptr && t < HC * UPM_COLS) { + const int c = t / UPM_COLS, col = t - c * UPM_COLS; + float s = 0.0f; +#pragma unroll + for (int k = 0; k < TT; ++k) s += gs[k][c][col]; + float* e = m.a[0].gate_ema + (size_t) c * N + d0 + col; + *e += 0.25f * (s / (float) TT - *e); + } } } // namespace -void fused_gr_read_multi(const FusedGrArgs* a, int n_tok, float* xn_scratch, void* stream, unsigned long long* stamp_buf, - int stamp_i0) { - if (n_tok < 1 || n_tok > kFusedGrMaxT || xn_scratch == nullptr) { +size_t fused_gr_scratch_bytes() { return (size_t) SC_FLOATS * sizeof(float); } + +void fused_gr_read_multi(const FusedGrArgs* a, int n_tok, float* scratch, void* stream) { + if (n_tok < 1 || n_tok > kFusedGrMaxT || scratch == nullptr) { std::fprintf(stderr, "fused_gr_read_multi: invalid arguments\n"); std::exit(1); } @@ -311,56 +310,32 @@ void fused_gr_read_multi(const FusedGrArgs* a, int n_tok, float* xn_scratch, voi m.a[t] = a[t]; const FusedGrArgs& x = a[t]; if (!x.R || !x.w_norm || !x.w_down || !x.w_up || !x.lo || !x.rs || !x.mixed || (x.w_inject && !x.inject_out) || - (x.apply && (!x.bo_prev || !x.inj_prev || !x.R_out)) || x.w_down != a[0].w_down || x.w_up != a[0].w_up || - x.w_inject != a[0].w_inject || x.w_norm != a[0].w_norm) { + (x.apply && (!x.bo_prev || !x.inj_prev || !x.R_out || x.inj_prev == x.inject_out)) || + x.w_down != a[0].w_down || x.w_up != a[0].w_up || x.w_inject != a[0].w_inject || + x.w_norm != a[0].w_norm || x.gate_ema != a[0].gate_ema || x.eps != a[0].eps || + x.est_norm != a[0].est_norm || x.est_gates != a[0].est_gates || (x.est_norm && (!x.est_gates || !x.est))) { std::fprintf(stderr, "fused_gr_read_multi: invalid arguments for token %d\n", t); std::exit(1); } } - m.xn = xn_scratch; m.T = n_tok; cudaStream_t st = (cudaStream_t) stream; - gr_norm_multi_kernel<<>>(m); - if (stamp_buf) gpu_stamp(stamp_buf, stamp_i0, stream); - // the shared-memory opt-in is a per-DEVICE setting: once per device, not once per process (a layer split - // runs this kernel on two cards) - static bool attr[64] = {}; - static int chunk[64] = {}; // Turing port: tokens the down kernel may carry in one launch on this card - int dev = 0; - cudaGetDevice(&dev); - if (dev >= 0 && dev < 64 && !attr[dev]) { - // at most what the card allows (Turing: 64 KB - enough for windows of up to 6 tokens) - int optin = 0; - cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - int want = (int) (kFusedGrMaxT * TILE * sizeof(float)); - if (optin > 0 && want > optin) want = optin; - cudaFuncSetAttribute(gr_down_multi_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, want); - cudaGetLastError(); // drop any error the attempt left behind - // Turing port: the down kernel stages n_tok*TILE floats of dynamic shared memory - 80 KB at the full - // 8 tokens. A card whose opt-in is below that (Turing: 64 KB, so 7+ tokens fail to launch as - // "invalid argument") processes the tokens in slices that fit; a card that reports no opt-in gets what - // fits the 48 KB default (4 tokens of the CUDA tile; all 8 of HIP's smaller tile). The down kernel's - // outputs (lo, inject_out) are strictly per-token, so the chunk boundaries are safe, and the up kernel - // below still sees every token of the batch in one launch. - const int capacity = (optin > 0 ? optin : 48 * 1024) / (int) (TILE * sizeof(float)); - chunk[dev] = capacity < 1 ? 1 : (capacity > kFusedGrMaxT ? kFusedGrMaxT : capacity); - attr[dev] = true; - } - const int chunk_tok = (dev >= 0 && dev < 64 && chunk[dev]) ? chunk[dev] : kFusedGrMaxT; - if (chunk_tok >= n_tok) { - gr_down_multi_kernel<<>>(m); - } else { - for (int c0 = 0; c0 < n_tok; c0 += chunk_tok) { - const int ct = n_tok - c0 < chunk_tok ? n_tok - c0 : chunk_tok; - GrMulti c{}; - c.xn = xn_scratch + (size_t) c0 * D; - c.T = ct; - for (int k = 0; k < ct; ++k) c.a[k] = a[c0 + k]; - gr_down_multi_kernel<<>>(c); + const dim3 down_grid(DS, DG); + for (int t0 = 0; t0 < n_tok; t0 += DT) { // DT tokens a down launch; each token's arithmetic is the same + GrMulti md; + md.T = std::min(DT, n_tok - t0); + for (int k = 0; k < md.T; ++k) md.a[k] = a[t0 + k]; + switch (md.T) { +#define STRATA_DOWN(T) case T: gr_down_split_kernel<<>>(md, scratch); break; + STRATA_DOWN(1) STRATA_DOWN(2) STRATA_DOWN(3) STRATA_DOWN(4) +#undef STRATA_DOWN } } - if (stamp_buf) gpu_stamp(stamp_buf, stamp_i0 + 1, stream); - gr_up_multi_kernel<<>>(m); + switch (n_tok) { +#define STRATA_UP(T) case T: gr_up_multi_kernel<<>>(m); break; + STRATA_UP(1) STRATA_UP(2) STRATA_UP(3) STRATA_UP(4) STRATA_UP(5) STRATA_UP(6) STRATA_UP(7) STRATA_UP(8) +#undef STRATA_UP + } const cudaError_t e = cudaGetLastError(); if (e != cudaSuccess) { std::fprintf(stderr, "fused_gr_read_multi: %s\n", cudaGetErrorString(e)); @@ -372,21 +347,4 @@ bool fused_gr_supported(int64_t n_embd, int64_t hc, int64_t hc_lr) { return n_embd == N && hc == HC && hc_lr == LR; } -void fused_gr_read(const FusedGrArgs& a, void* stream) { - if (!a.R || !a.w_norm || !a.w_down || !a.w_up || !a.lo || !a.rs || !a.mixed || - (a.w_inject && !a.inject_out) || (a.apply && (!a.bo_prev || !a.inj_prev || !a.R_out)) || - (a.apply && a.inj_prev == a.inject_out)) { - std::fprintf(stderr, "fused_gr_read: invalid arguments\n"); - std::exit(1); - } - cudaStream_t st = (cudaStream_t) stream; - gr_down_kernel<<>>(a); - gr_up_kernel<<>>(a); - const cudaError_t e = cudaGetLastError(); - if (e != cudaSuccess) { - std::fprintf(stderr, "fused_gr_read: %s\n", cudaGetErrorString(e)); - std::exit(1); - } -} - } // namespace strata::kernels diff --git a/src/kernels/cuda/gdn.cu b/src/kernels/cuda/gdn.cu index aad247a797e..5111afda106 100644 --- a/src/kernels/cuda/gdn.cu +++ b/src/kernels/cuda/gdn.cu @@ -31,8 +31,6 @@ constexpr int JTHREADS = 32; ///< threads along j, the state's fast axis /// The arithmetic does not change at all: the same columns are computed by the same code, in different blocks. constexpr int MAX_H = 1; ///< heads staged in shared memory per block -/// ggml_compute_softplus_f32: log1p(exp(x)), with the large-x branch that avoids overflow. -__device__ __forceinline__ float softplus_f(float x) { return x > 20.0f ? x : log1pf(expf(x)); } __device__ __forceinline__ float sigmoid_f(float x) { return 1.0f / (1.0f + expf(-x)); } /// One step of the recurrence. One thread per (h, j) column of the state. diff --git a/src/kernels/cuda/gr.cu b/src/kernels/cuda/gr.cu index 7eea7f49f43..24d615b2fb2 100644 --- a/src/kernels/cuda/gr.cu +++ b/src/kernels/cuda/gr.cu @@ -66,11 +66,6 @@ __device__ __forceinline__ void store_activation(uint16_t* dst, int i, float x) __device__ __forceinline__ float silu_f(float x) { return x / (1.0f + expf(-x)); } __device__ __forceinline__ float sigmoid_f(float x) { return 1.0f / (1.0f + expf(-x)); } -__device__ __forceinline__ double warp_sum(double v) { - for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xFFFFFFFFu, v, off); - return __shfl_sync(0xFFFFFFFFu, v, 0); -} - __device__ __forceinline__ float warp_sumf(float v) { for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xFFFFFFFFu, v, off); return __shfl_sync(0xFFFFFFFFu, v, 0); @@ -84,9 +79,9 @@ __device__ __forceinline__ float warp_sumf(float v) { /// the cast cannot overrun, and it keeps the change to one function instead of one per caller. __device__ float block_sumf(float v, double* scratch_raw) { float* scratch = (float*) scratch_raw; - // The leading barrier is not decoration - see the note on `block_sum` above: the result is read straight out - // of `scratch[0]` by every thread and a later call reuses the array, so without it a fast thread can - // overwrite `scratch[0]` before a slow one has read the previous result. + // The leading barrier is not decoration: the result is read straight out of `scratch[0]` by every thread and + // a later call reuses the array, so without it a fast thread can overwrite `scratch[0]` before a slow one has + // read the previous result - a race that is invisible in most runs. __syncthreads(); const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; v = warp_sumf(v); @@ -100,29 +95,6 @@ __device__ float block_sumf(float v, double* scratch_raw) { return scratch[0]; } -/// Block-wide sum in DOUBLE, broadcast to every thread. -/// -/// Double because `ref/gr.py` reduces in float64 for the norm and the bottleneck, and because a block-wide -/// f32 sum of 2560 squares is a different number from the reference's; the shuffle tree here also reorders -/// the additions, so the only defence is enough precision that the order stops mattering. -/// -/// The leading `__syncthreads()` is not decoration: the result is read straight out of `scratch[0]` by every -/// thread, and a later call reuses the same array. Without a barrier on entry a fast thread can overwrite -/// `scratch[0]` before a slow one has read the previous result - a race that is invisible in most runs. -__device__ double block_sum(double v, double* scratch) { - __syncthreads(); - const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; - v = warp_sum(v); - if (lane == 0) scratch[warp] = v; - __syncthreads(); - const int nw = ((int) blockDim.x + 31) >> 5; - v = (threadIdx.x < nw) ? scratch[threadIdx.x] : 0.0; - if (warp == 0) v = warp_sum(v); - if (threadIdx.x == 0) scratch[0] = v; - __syncthreads(); - return scratch[0]; -} - /// `hc` blocks, one per residual stream: per-stream RMSNorm, then the activated stream in both forms. /// /// ggml_rms_norm reduces over ne[0], and ne[0] is n_embd, so the reduction is PER STREAM. The `[n_embd, hc]` diff --git a/src/kernels/cuda/hit_plan.cuh b/src/kernels/cuda/hit_plan.cuh new file mode 100644 index 00000000000..56ef546e315 --- /dev/null +++ b/src/kernels/cuda/hit_plan.cuh @@ -0,0 +1,71 @@ +// src/kernels/cuda/hit_plan.cuh - the verify window's main-GPU hit plan (verify_kernels.hpp, verify_hit_plan) as the +// work of one block, for verify_kernels.cu's kernel and the router kernel that rings the doorbell (native_router.cu). +#pragma once + +#include + +#include + +namespace strata::kernels { +namespace { + +constexpr int kHitPlanMax = 128; + +// A block of at least kHitPlanMax threads, thread i = routed entry i (n <= kHitPlanMax). A group per distinct resident +// expert, in the order of their first entries; a group's entries in routing order. `ids` must be visible to the whole +// block. +__device__ __forceinline__ void hit_plan_block(const int32_t* __restrict__ ids, int n, int k, + const int32_t* __restrict__ res, int n_expert, + const unsigned long long* __restrict__ slot_ptr, + int32_t* __restrict__ plan, int cap, int ptr_off) { + __shared__ int32_t s_id[kHitPlanMax], s_size[kHitPlanMax], s_start[kHitPlanMax]; + const int i = threadIdx.x; + int32_t e = -1, slot = -1; + if (i < n) { + e = ids[i]; + if (e >= 0 && e < n_expert) slot = res[e]; + } + if (i < kHitPlanMax) s_id[i] = e; + __syncthreads(); + const bool hit = slot >= 0; + int first = i; + if (hit) + for (int j = 0; j < i; ++j) + if (s_id[j] == e) { first = j; break; } + const bool lead = hit && first == i; + int size = 0; + if (lead) + for (int j = i; j < n; ++j) size += s_id[j] == e; + if (i < kHitPlanMax) s_size[i] = size; // > 0 at a group's first entry only + const int groups = __syncthreads_count(lead); + const int entries = __syncthreads_count(hit); + int32_t* start = plan + 4; + int32_t* dst = start + cap + 1; + int32_t* tok = dst + cap; + auto* ptr = (unsigned long long*) (plan + ptr_off); + if (lead) { + int grp = 0, at = 0; + for (int j = 0; j < i; ++j) + if (s_size[j] > 0) { ++grp; at += s_size[j]; } + s_start[i] = at; + start[grp] = at; + ptr[grp] = slot_ptr[slot]; + } + __syncthreads(); + if (hit) { + int rank = 0; + for (int j = first; j < i; ++j) rank += s_id[j] == e; + dst[s_start[first] + rank] = i; + tok[s_start[first] + rank] = i / k; + } + if (i == 0) { + plan[0] = groups; + plan[1] = entries; + plan[2] = 0; + plan[3] = 0; + start[groups] = entries; + } +} + +} // namespace +} // namespace strata::kernels diff --git a/src/kernels/cuda/iq_dot.cuh b/src/kernels/cuda/iq_dot.cuh new file mode 100644 index 00000000000..c8fdb6879c4 --- /dev/null +++ b/src/kernels/cuda/iq_dot.cuh @@ -0,0 +1,610 @@ +// src/kernels/cuda/iq_dot.cuh - the GGUF formats' dot products against q8_1 activations (one call covers a slice of a +// block), for iq_kernels.cu and the parity test that keeps its earlier expert kernels. +// +// Transcribed from llama.cpp (ggml/src/ggml-cuda/vecdotq.cuh at the commit in third_party/ggml/VERSION.txt; MIT +// license, third_party/ggml/LICENSE); the i-quants apply their signs another way (`sign_masks`), with the same +// integer sums. The block structs and codebook grids come from its ggml-common.h, included unchanged. +#pragma once + +#include +#include + +#define GGML_COMMON_DECL_CUDA +#define GGML_COMMON_IMPL_CUDA +#include "ggml-common.h" + +namespace strata::kernels { +namespace { + +// ---------------------------------------------------------------- llama.cpp helpers (vecdotq.cuh) +__device__ __forceinline__ int get_int_b2(const void* x, const int& i32) { + const uint16_t* x16 = (const uint16_t*) x; + int x32 = x16[2 * i32 + 0] << 0; + x32 |= x16[2 * i32 + 1] << 16; + return x32; +} +__device__ __forceinline__ int get_int_b4(const void* x, const int& i32) { return ((const int*) x)[i32]; } +__device__ __forceinline__ uint32_t unpack_ksigns(const uint8_t v) { + const uint32_t p = __popc(v) & 1; + const uint32_t s = v ^ p << 7; + return s * 0x01010101; +} +__device__ __forceinline__ int2 get_int_from_table_16(const int& q4, const int8_t* table) { + const uint32_t* table32 = (const uint32_t*) table; + uint32_t tmp[2]; + const uint32_t low_high_selection_indices = (0x32103210 | ((q4 & 0x88888888) >> 1)); +#pragma unroll + for (uint32_t i = 0; i < 2; ++i) { + const uint32_t shift = 16 * i; + const uint32_t low = __byte_perm(table32[0], table32[1], q4 >> shift); + const uint32_t high = __byte_perm(table32[2], table32[3], q4 >> shift); + tmp[i] = __byte_perm(low, high, low_high_selection_indices >> shift); + } + return make_int2(__byte_perm(tmp[0], tmp[1], 0x6420), __byte_perm(tmp[0], tmp[1], 0x7531)); +} +#define ggml_cuda_dp4a(a, b, c) __dp4a((a), (b), (c)) + +// The i-quants' signs as byte masks (0xff where a value is negative), four values a word: the multiply puts sign bit +// i on bit 7 of byte i, and prmt's sign-replicate mode (bit 3 of a selector nibble, which __byte_perm drops) fills +// the byte with it. Their dot products take sum(s g u) = sum(g u) - 2 sum((g & m) u) over the grid's unsigned +// values: the integers llama.cpp's sign-applied grid gives, in fewer instructions than __vcmpne4 and __vsub4. +__device__ __forceinline__ uint32_t msb_bytes(uint32_t x) { + uint32_t d; + asm("prmt.b32 %0, %1, 0, 0xBA98;" : "=r"(d) : "r"(x)); + return d; +} +__device__ __forceinline__ uint2 sign_masks(uint32_t s8) { // bit i of s8 -> byte i of (x, y) + return make_uint2(msb_bytes((s8 & 0x0F) * 0x10204080u), msb_bytes((s8 >> 4 & 0x0F) * 0x10204080u)); +} +__device__ __forceinline__ uint2 ksign_masks(uint32_t v7) { // 7 sign bits, the eighth their parity (ksigns_iq2xs) + return sign_masks(v7 | (__popc(v7) & 1) << 7); +} + +// ---------------------------------------------------------------- the dot products (vecdotq.cuh) +__device__ __forceinline__ float vec_dot_q2_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_q2_0* bq2_0 = (const block_q2_0*) vbq + kbx; + const float d2 = bq2_0->d; + const int16_t* qs = (const int16_t*) bq2_0->qs + iqs * 4; + const block_q8_1* bq8_1_chunk = bq8_1 + iqs; + int sumi = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int q = qs[j]; + const int u = get_int_b4(bq8_1_chunk->qs, j * 2 + 0); + const int v = get_int_b4(bq8_1_chunk->qs, j * 2 + 1); + const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0); + const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2); + const int qx = __byte_perm(qe, qo, 0x5140); + const int qy = __byte_perm(qe, qo, 0x7362); + sumi = ggml_cuda_dp4a(u, qx, sumi); + sumi = ggml_cuda_dp4a(v, qy, sumi); + } + const float d8 = __low2float(bq8_1_chunk->ds); + return d2 * d8 * sumi; +} + +__device__ __forceinline__ float vec_dot_iq2_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq2_xxs* bq2 = (const block_iq2_xxs*) vbq + kbx; + const int q2 = get_int_b2(bq2->qs, iqs); + const uint8_t* aux8 = (const uint8_t*) &q2; + const uint32_t aux32 = get_int_b2(bq2->qs, iqs + 1); + int sp = 0, sn = 0; // sum(g u), and over the negative values +#pragma unroll + for (int k0 = 0; k0 < 8; k0 += 2) { + const uint2 grid_pos = ((const uint2*) iq2xxs_grid)[aux8[k0 / 2]]; + const uint2 m = ksign_masks((aux32 >> (7 * k0 / 2)) & 0x7F); + const int u0 = get_int_b4(bq8_1[iqs / 2].qs, k0 + 0); + const int u1 = get_int_b4(bq8_1[iqs / 2].qs, k0 + 1); + sp = ggml_cuda_dp4a((int) grid_pos.x, u0, sp); + sn = ggml_cuda_dp4a((int) (grid_pos.x & m.x), u0, sn); + sp = ggml_cuda_dp4a((int) grid_pos.y, u1, sp); + sn = ggml_cuda_dp4a((int) (grid_pos.y & m.y), u1, sn); + } + int sumi = sp - 2 * sn; + const int ls = aux32 >> 27 | 1; + sumi = sumi * ls / 8; + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs / 2].ds); + return d * sumi; +} + +__device__ __forceinline__ float vec_dot_iq2_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq2_xs* bq2 = (const block_iq2_xs*) vbq + kbx; + const int2 q2_packed = make_int2(get_int_b2(bq2->qs, iqs + 0), get_int_b2(bq2->qs, iqs + 1)); + const uint16_t* q2 = (const uint16_t*) &q2_packed; + const int ls0 = bq2->scales[iqs / 2] & 0x0F; + const int ls1 = bq2->scales[iqs / 2] >> 4; + int sp[2] = {0, 0}, sn[2] = {0, 0}; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const uint2 grid_pos = ((const uint2*) iq2xs_grid)[q2[l0 / 2] & 0x1FF]; + const uint2 m = ksign_masks(q2[l0 / 2] >> 9); + const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); + const int h = l0 / 4; + sp[h] = ggml_cuda_dp4a((int) grid_pos.x, u0, sp[h]); + sn[h] = ggml_cuda_dp4a((int) (grid_pos.x & m.x), u0, sn[h]); + sp[h] = ggml_cuda_dp4a((int) grid_pos.y, u1, sp[h]); + sn[h] = ggml_cuda_dp4a((int) (grid_pos.y & m.y), u1, sn[h]); + } + const int sumi0 = sp[0] - 2 * sn[0], sumi1 = sp[1] - 2 * sn[1]; + const int sumi = (sumi0 * ls0 + sumi1 * ls1 + (sumi0 + sumi1) / 2) / 4; + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs / 2].ds); + return d * sumi; +} + +__device__ __forceinline__ float vec_dot_iq2_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq2_s* bq2 = (const block_iq2_s*) vbq + kbx; + const int qs_packed = get_int_b2(bq2->qs, iqs / 2); + const uint8_t* qs = (const uint8_t*) &qs_packed; + const int qh = bq2->qh[iqs / 2]; + const int signs_packed_32 = get_int_b2(bq2->qs, QK_K / 32 + iqs / 2); + const uint8_t* signs_packed_8 = (const uint8_t*) &signs_packed_32; + const int ls0 = bq2->scales[iqs / 2] & 0x0F; + const int ls1 = bq2->scales[iqs / 2] >> 4; + int sp[2] = {0, 0}, sn[2] = {0, 0}; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const uint2 grid_pos = ((const uint2*) iq2s_grid)[qs[l0 / 2] | ((qh << (8 - l0)) & 0x300)]; + const uint2 m = sign_masks(signs_packed_8[l0 / 2]); + const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); + const int h = l0 / 4; + sp[h] = ggml_cuda_dp4a((int) grid_pos.x, u0, sp[h]); + sn[h] = ggml_cuda_dp4a((int) (grid_pos.x & m.x), u0, sn[h]); + sp[h] = ggml_cuda_dp4a((int) grid_pos.y, u1, sp[h]); + sn[h] = ggml_cuda_dp4a((int) (grid_pos.y & m.y), u1, sn[h]); + } + const int sumi0 = sp[0] - 2 * sn[0], sumi1 = sp[1] - 2 * sn[1]; + const int sumi = (sumi0 * ls0 + sumi1 * ls1 + (sumi0 + sumi1) / 2) / 4; + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs / 2].ds); + return d * sumi; +} + +__device__ __forceinline__ float vec_dot_iq3_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq3_xxs* bq3 = (const block_iq3_xxs*) vbq + kbx; + const int2 q3_packed = make_int2(get_int_b2(bq3->qs, iqs), get_int_b2(bq3->qs, iqs + 1)); + const uint8_t* q3 = (const uint8_t*) &q3_packed; + const uint32_t aux32 = get_int_b2(bq3->qs, QK_K / 16 + iqs / 2); + int sp = 0, sn = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const uint32_t gx = iq3xxs_grid[q3[l0 + 0]], gy = iq3xxs_grid[q3[l0 + 1]]; + const uint2 m = ksign_masks((aux32 >> (7 * l0 / 2)) & 0x7F); + const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); + sp = ggml_cuda_dp4a((int) gx, u0, sp); + sn = ggml_cuda_dp4a((int) (gx & m.x), u0, sn); + sp = ggml_cuda_dp4a((int) gy, u1, sp); + sn = ggml_cuda_dp4a((int) (gy & m.y), u1, sn); + } + int sumi = sp - 2 * sn; + const int ls = aux32 >> 28; + sumi = (ls * sumi + sumi / 2) / 2; + const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs / 2].ds); + return d * sumi; +} + +__device__ __forceinline__ float vec_dot_iq3_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq3_s* bq3 = (const block_iq3_s*) vbq + kbx; + const int2 qs_packed = make_int2(get_int_b2(bq3->qs, iqs + 0), get_int_b2(bq3->qs, iqs + 1)); + const uint8_t* qs = (const uint8_t*) &qs_packed; + const int qh = bq3->qh[iqs / 2]; + const int signs_packed_32 = get_int_b2(bq3->signs, iqs / 2); + const uint8_t* signs_packed_8 = (const uint8_t*) &signs_packed_32; + int sp = 0, sn = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const uint32_t gx = iq3s_grid[qs[l0 + 0] | ((qh << (8 - l0)) & 0x100)]; + const uint32_t gy = iq3s_grid[qs[l0 + 1] | ((qh << (7 - l0)) & 0x100)]; + const uint2 m = sign_masks(signs_packed_8[l0 / 2]); + const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); + sp = ggml_cuda_dp4a((int) gx, u0, sp); + sn = ggml_cuda_dp4a((int) (gx & m.x), u0, sn); + sp = ggml_cuda_dp4a((int) gy, u1, sp); + sn = ggml_cuda_dp4a((int) (gy & m.y), u1, sn); + } + int sumi = sp - 2 * sn; + sumi *= 1 + 2 * ((bq3->scales[iqs / 4] >> ((iqs << 1) & 0x04)) & 0x0F); + const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs / 2].ds); + return d * sumi; +} + +// vec_dot_iq3_s_q8_1 for several activation columns: the grid values, masks and scale decoded once (`load`), each +// column's integer sums and float expression as above (`apply`; X::u(b, p, u) gives int p of q8_1 block b of every +// column, X::scales(b, d) their __low2float scales). +struct IQ3SCols { + uint32_t g[8], gm[8]; // grid words, and their negative values' bytes + int sc; + float dw; + __device__ __forceinline__ void load(const void* __restrict__ vbq, const int& kbx, const int& iqs) { + const block_iq3_s* bq3 = (const block_iq3_s*) vbq + kbx; + const int2 qs_packed = make_int2(get_int_b2(bq3->qs, iqs + 0), get_int_b2(bq3->qs, iqs + 1)); + const uint8_t* qs = (const uint8_t*) &qs_packed; + const int qh = bq3->qh[iqs / 2]; + const int signs_packed_32 = get_int_b2(bq3->signs, iqs / 2); + const uint8_t* signs_packed_8 = (const uint8_t*) &signs_packed_32; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const uint32_t gx = iq3s_grid[qs[l0 + 0] | ((qh << (8 - l0)) & 0x100)]; + const uint32_t gy = iq3s_grid[qs[l0 + 1] | ((qh << (7 - l0)) & 0x100)]; + const uint2 m = sign_masks(signs_packed_8[l0 / 2]); + g[l0] = gx; + gm[l0] = gx & m.x; + g[l0 + 1] = gy; + gm[l0 + 1] = gy & m.y; + } + sc = 1 + 2 * ((bq3->scales[iqs / 4] >> ((iqs << 1) & 0x04)) & 0x0F); + dw = __half2float(bq3->d); + } + template + __device__ __forceinline__ void apply(const X& x, const int& b, float (&out)[NC]) const { + int u[8][NC]; // every position's activations first, so that their loads overlap +#pragma unroll + for (int l = 0; l < 8; ++l) x.u(b, l, u[l]); + int sp[NC], sn[NC]; +#pragma unroll + for (int c = 0; c < NC; ++c) sp[c] = sn[c] = 0; +#pragma unroll + for (int l = 0; l < 8; ++l) { +#pragma unroll + for (int c = 0; c < NC; ++c) { + sp[c] = ggml_cuda_dp4a((int) g[l], u[l][c], sp[c]); + sn[c] = ggml_cuda_dp4a((int) gm[l], u[l][c], sn[c]); + } + } + float d8[NC]; + x.scales(b, d8); +#pragma unroll + for (int c = 0; c < NC; ++c) { + int sumi = sp[c] - 2 * sn[c]; + sumi *= sc; + const float d = dw * d8[c]; + out[c] = d * sumi; + } + } +}; + +__device__ __forceinline__ float vec_dot_iq1_m_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq1_m* bq1 = (const block_iq1_m*) vbq + kbx; + const int qs_packed = get_int_b4(bq1->qs, iqs); + const uint8_t* qs = (const uint8_t*) &qs_packed; + int sumi[2] = {0, 0}; + float sumf[2] = {0.0f, 0.0f}; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int qhl = bq1->qh[2 * iqs + l0 / 4] >> (4 * ((l0 / 2) % 2)); + const int grid = iq1s_grid_gpu[qs[l0 / 2] | ((qhl & 0x07) << 8)]; + const int grid0 = (grid >> 0) & 0x0F0F0F0F; + const int grid1 = (grid >> 4) & 0x0F0F0F0F; + const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); + sumi[l0 / 4] = ggml_cuda_dp4a(grid0, u0, sumi[l0 / 4]); + sumi[l0 / 4] = ggml_cuda_dp4a(grid1, u1, sumi[l0 / 4]); + const float delta = -1.0f + IQ1M_DELTA - (qhl & 0x08) * (2.0f * IQ1M_DELTA / 0x08); + int sumy = 0; + sumy = ggml_cuda_dp4a(u0, 0x01010101, sumy); + sumy = ggml_cuda_dp4a(u1, 0x01010101, sumy); + sumf[l0 / 4] += delta * sumy; + } + const uint16_t* sc = (const uint16_t*) bq1->scales; + iq1m_scale_t scale; + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00F0) | ((sc[2] >> 4) & 0x0F00) | (sc[3] & 0xF000); + const float d = __half2float(scale.f16) * __low2float(bq8_1[iqs].ds); + const int tmp = sc[iqs / 2] >> (6 * (iqs % 2)); + const int sc0 = 2 * ((tmp >> 0) & 0x07) + 1; + const int sc1 = 2 * ((tmp >> 3) & 0x07) + 1; + return d * ((sumi[0] + sumf[0]) * sc0 + (sumi[1] + sumf[1]) * sc1); +} + +__device__ __forceinline__ float vec_dot_iq4_nl_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq4_nl* bq4 = (const block_iq4_nl*) vbq + kbx; + const int* q8 = (const int*) bq8_1->qs + iqs; + int sumi = 0; +#pragma unroll + for (int l = 0; l < 2; ++l) { + const int aux_q4 = get_int_b2(bq4->qs, iqs + l); + const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); + sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi); + sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi); + } + const float d = __half2float(bq4->d) * __low2float(bq8_1->ds); + return d * sumi; +} + +// Q4_K / Q5_K gate/up and Q5_1 / Q8_0 down (the unsloth UD-Q4_K_XL file's experts) +__device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_vmmq(const int* __restrict__ v, const int* __restrict__ u, + const uint8_t* __restrict__ sc, const uint8_t* __restrict__ m, + const half2& dm4, const float* __restrict__ d8) { + float sumf_d = 0.0f; + float sumf_m = 0.0f; +#pragma unroll + for (int i = 0; i < QR4_K; ++i) { + const int v0i = (v[0] >> (4 * i)) & 0x0F0F0F0F; + const int v1i = (v[1] >> (4 * i)) & 0x0F0F0F0F; + const int dot1 = ggml_cuda_dp4a(v1i, u[2 * i + 1], ggml_cuda_dp4a(v0i, u[2 * i + 0], 0)); + const int dot2 = ggml_cuda_dp4a(0x01010101, u[2 * i + 1], ggml_cuda_dp4a(0x01010101, u[2 * i + 0], 0)); + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); + } + const float2 dm4f = __half22float2(dm4); + return dm4f.x * sumf_d - dm4f.y * sumf_m; +} + +// the 6-bit scale and min of the 32-value group pair `bq8_offset / 2` (vecdotq.cuh, shared by Q4_K and Q5_K) +__device__ __forceinline__ void k_scales(const uint8_t* scales8, int bq8_offset, uint16_t aux[2]) { + const uint16_t* scales = (const uint16_t*) scales8; + const int j = bq8_offset / 2; + const int jm = j & 1; + const uint32_t s0 = scales[jm + 0]; + const uint32_t s2 = scales[jm + 2]; + const uint32_t s4 = scales[jm + 4]; + const uint32_t hi = (uint32_t) -(int32_t) (j >= 2); + aux[0] = (uint16_t) (((s0 & 0x3f3f) & ~hi) | ((((s4 >> 0) & 0x0f0f) | ((s0 & 0xc0c0) >> 2)) & hi)); + aux[1] = (uint16_t) (((s2 & 0x3f3f) & ~hi) | ((((s4 >> 4) & 0x0f0f) | ((s2 & 0xc0c0) >> 2)) & hi)); +} + +__device__ __forceinline__ float vec_dot_q4_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_q4_K* bq4_K = (const block_q4_K*) vbq + kbx; + int v[2]; + int u[2 * QR4_K]; + float d8[QR4_K]; + const int bq8_offset = QR4_K * ((iqs / 2) / (QI8_1 / 2)); + const int* q4 = (const int*) (bq4_K->qs + 16 * bq8_offset + 4 * ((iqs / 2) % 4)); + v[0] = q4[0]; + v[1] = q4[4]; + uint16_t aux[2]; + k_scales(bq4_K->scales, bq8_offset, aux); + const uint8_t* sc = (const uint8_t*) aux; + const uint8_t* m = sc + 2; + for (int i = 0; i < QR4_K; ++i) { + const block_q8_1* bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + const int* q8 = (const int*) bq8i->qs + ((iqs / 2) % 4); + u[2 * i + 0] = q8[0]; + u[2 * i + 1] = q8[4]; + } + return vec_dot_q4_K_q8_1_impl_vmmq(v, u, sc, m, bq4_K->dm, d8); +} + +__device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_vmmq(const int* __restrict__ vl, const int* __restrict__ vh, + const int* __restrict__ u, const uint8_t* __restrict__ sc, + const uint8_t* __restrict__ m, const half2& dm5, + const float* __restrict__ d8) { + float sumf_d = 0.0f; + float sumf_m = 0.0f; +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const int vl0i = (vl[0] >> (4 * i)) & 0x0F0F0F0F; + const int vl1i = (vl[1] >> (4 * i)) & 0x0F0F0F0F; + const int vh0i = ((vh[0] >> i) << 4) & 0x10101010; + const int vh1i = ((vh[1] >> i) << 4) & 0x10101010; + const int v0i = vl0i | vh0i; + const int v1i = vl1i | vh1i; + const int dot1 = ggml_cuda_dp4a(v0i, u[2 * i + 0], ggml_cuda_dp4a(v1i, u[2 * i + 1], 0)); + const int dot2 = ggml_cuda_dp4a(0x01010101, u[2 * i + 0], ggml_cuda_dp4a(0x01010101, u[2 * i + 1], 0)); + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); + } + const float2 dm5f = __half22float2(dm5); + return dm5f.x * sumf_d - dm5f.y * sumf_m; +} + +__device__ __forceinline__ float vec_dot_q5_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_q5_K* bq5_K = (const block_q5_K*) vbq + kbx; + int vl[2]; + int vh[2]; + int u[2 * QR5_K]; + float d8[QR5_K]; + const int bq8_offset = QR5_K * ((iqs / 2) / (QI8_1 / 2)); + const int* ql = (const int*) (bq5_K->qs + 16 * bq8_offset + 4 * ((iqs / 2) % 4)); + const int* qh = (const int*) (bq5_K->qh + 4 * ((iqs / 2) % 4)); + vl[0] = ql[0]; + vl[1] = ql[4]; + vh[0] = qh[0] >> bq8_offset; + vh[1] = qh[4] >> bq8_offset; + uint16_t aux[2]; + k_scales(bq5_K->scales, bq8_offset, aux); + const uint8_t* sc = (const uint8_t*) aux; + const uint8_t* m = sc + 2; +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const block_q8_1* bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + const int* q8 = (const int*) bq8i->qs + ((iqs / 2) % 4); + u[2 * i + 0] = q8[0]; + u[2 * i + 1] = q8[4]; + } + return vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, bq5_K->dm, d8); +} + +constexpr int VDR_Q5_1 = 2, VDR_Q8_0 = 2; // VDR_Q5_1_Q8_1_MMVQ, VDR_Q8_0_Q8_1_MMVQ + +// The weights are llama.cpp's; the min term is not. llama.cpp multiplies the Q5_1 min by the q8_1 block's sum of +// the ORIGINAL activations while the scaled term uses the quantized ones, which adds min * (sum x - sum x_q) to +// every block: 1.9% instead of 1.2% relative error per expert (native_expert_parity). Like ggml-cpu, and like +// llama.cpp's own Q4_K / Q5_K dots, the min multiplies the quantized sum here. +__device__ __forceinline__ float vec_dot_q5_1_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_q5_1* bq5_1 = (const block_q5_1*) vbq + kbx; + int sumi = 0, sumu = 0; +#pragma unroll + for (int i = 0; i < VDR_Q5_1; ++i) { + const int vl = get_int_b4(bq5_1->qs, iqs + i); + const int vh = get_int_b4(bq5_1->qh, 0) >> (4 * (iqs + i)); + const int u0 = get_int_b4(bq8_1->qs, iqs + i), u1 = get_int_b4(bq8_1->qs, iqs + i + QI5_1); + int vi0 = (vl >> 0) & 0x0F0F0F0F; + vi0 |= (vh << 4) & 0x00000010; + vi0 |= (vh << 11) & 0x00001000; + vi0 |= (vh << 18) & 0x00100000; + vi0 |= (vh << 25) & 0x10000000; + sumi = ggml_cuda_dp4a(vi0, u0, sumi); + int vi1 = (vl >> 4) & 0x0F0F0F0F; + vi1 |= (vh >> 12) & 0x00000010; + vi1 |= (vh >> 5) & 0x00001000; + vi1 |= (vh << 2) & 0x00100000; + vi1 |= (vh << 9) & 0x10000000; + sumi = ggml_cuda_dp4a(vi1, u1, sumi); + sumu = ggml_cuda_dp4a(0x01010101, u1, ggml_cuda_dp4a(0x01010101, u0, sumu)); + } + const float2 dm5 = __half22float2(bq5_1->dm); + const float d8 = __low2float(bq8_1->ds); + return sumi * (dm5.x * d8) + sumu * (dm5.y * d8); +} + +// Q6_K gate/up (UD-Q5_K_XL, UD-Q6_K_XL), VDR_Q6_K_Q8_1_MMVQ = 1 +__device__ __forceinline__ float vec_dot_q6_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_q6_K* bq6_K = (const block_q6_K*) vbq + kbx; + const int bq8_offset = 2 * QR6_K * (iqs / (QI6_K / 2)) + (iqs % (QI6_K / 2)) / (QI6_K / 4); + const int scale_offset = (QI6_K / 4) * (iqs / (QI6_K / 2)) + (iqs % (QI6_K / 2)) / (QI6_K / 8); + const int vh_shift = 2 * ((iqs % (QI6_K / 2)) / (QI6_K / 4)); + const int vl = get_int_b2(bq6_K->ql, iqs); + const int vh = get_int_b2(bq6_K->qh, (QI6_K / 4) * (iqs / (QI6_K / 2)) + iqs % (QI6_K / 4)) >> vh_shift; + const int8_t* scales = bq6_K->scales + scale_offset; + float sumf = 0.0f; +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + const int u = get_int_b4(bq8_1[bq8_offset + 2 * i].qs, iqs % QI8_1); + const float d8 = __low2float(bq8_1[bq8_offset + 2 * i].ds); + const int sc = scales[4 * i]; + const int vil = (vl >> (4 * i)) & 0x0F0F0F0F; + const int vih = ((vh >> (4 * i)) << 4) & 0x30303030; + const int vi = __vsubss4((vil | vih), 0x20202020); // (vil | vih) - 32 + sumf += d8 * (ggml_cuda_dp4a(vi, u, 0) * sc); + } + return __half2float(bq6_K->d) * sumf; +} + +// IQ1_S gate/up (UD-IQ1_S), VDR 1. As for Q5_1, the delta term takes the quantized activation sum (ggml-cpu's +// bsums) rather than llama.cpp's sum of the original values. +__device__ __forceinline__ float vec_dot_iq1_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq1_s* bq1 = (const block_iq1_s*) vbq + kbx; + const int qs_packed = get_int_b2(bq1->qs, iqs); + const uint8_t* qs = (const uint8_t*) &qs_packed; + const int qh = bq1->qh[iqs]; + int sumi = 0, sumy = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int grid = iq1s_grid_gpu[qs[l0 / 2] | (((qh >> 3 * (l0 / 2)) & 0x07) << 8)]; + const int grid0 = (grid >> 0) & 0x0F0F0F0F; + const int grid1 = (grid >> 4) & 0x0F0F0F0F; + const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); + sumi = ggml_cuda_dp4a(grid0, u0, sumi); + sumi = ggml_cuda_dp4a(grid1, u1, sumi); + sumy = ggml_cuda_dp4a(0x01010101, u1, ggml_cuda_dp4a(0x01010101, u0, sumy)); + } + const float d1q = __half2float(bq1->d) * (((qh >> 11) & 0x0E) + 1); + const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f * IQ1S_DELTA / 0x8000); + return d1q * __low2float(bq8_1[iqs].ds) * (sumi + sumy * delta); +} + +// IQ4_XS gate/up (UD-IQ4_XS, UD-Q3_K_XL), VDR 4 +__device__ __forceinline__ float vec_dot_iq4_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq4_xs* bq4 = (const block_iq4_xs*) vbq + kbx; + int sumi = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int aux_q4 = get_int_b4(bq4->qs, iqs + j); + const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); + const int u0 = get_int_b4(bq8_1[iqs / 4].qs, j + 0); + const int u1 = get_int_b4(bq8_1[iqs / 4].qs, j + 4); + sumi = ggml_cuda_dp4a(v.x, u0, sumi); + sumi = ggml_cuda_dp4a(v.y, u1, sumi); + } + const int ls = ((bq4->scales_l[iqs / 8] >> (iqs & 0x04)) & 0x0F) | (((bq4->scales_h >> (iqs / 2)) & 0x03) << 4); + sumi *= ls - 32; + const float d = __half2float(bq4->d) * __low2float(bq8_1[iqs / 4].ds); + return d * sumi; +} + +__device__ __forceinline__ float vec_dot_q8_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_q8_0* bq8_0 = (const block_q8_0*) vbq + kbx; + int sumi = 0; +#pragma unroll + for (int i = 0; i < VDR_Q8_0; ++i) + sumi = ggml_cuda_dp4a(get_int_b2(bq8_0->qs, iqs + i), get_int_b4(bq8_1->qs, iqs + i), sumi); + const float d8_0 = __half2float(bq8_0->d), d8_1 = __low2float(bq8_1->ds); + return d8_0 * d8_1 * ((float) sumi); +} + +// ---------------------------------------------------------------- the formats +// qk = values per block, ipb = dot calls per block (qi / vdr), step = the iqs stride between calls. +template struct Fmt; +template<> struct Fmt<16> { static constexpr int qk = 256, ipb = 8, step = 2; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq2_xxs_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<17> { static constexpr int qk = 256, ipb = 8, step = 2; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq2_xs_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<18> { static constexpr int qk = 256, ipb = 8, step = 2; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq3_xxs_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<20> { static constexpr int qk = 32, ipb = 2, step = 2; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq4_nl_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<21> { static constexpr int qk = 256, ipb = 8, step = 2; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq3_s_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<22> { static constexpr int qk = 256, ipb = 8, step = 2; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq2_s_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<29> { static constexpr int qk = 256, ipb = 8, step = 1; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq1_m_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<42> { static constexpr int qk = 64, ipb = 2, step = 1; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q2_0_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<12> { static constexpr int qk = 256, ipb = QI4_K / 2, step = 2; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q4_K_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<13> { static constexpr int qk = 256, ipb = QI5_K / 2, step = 2; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q5_K_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<7> { static constexpr int qk = 32, ipb = QI5_1 / VDR_Q5_1, step = VDR_Q5_1; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q5_1_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<8> { static constexpr int qk = 32, ipb = QI8_0 / VDR_Q8_0, step = VDR_Q8_0; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q8_0_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<14> { static constexpr int qk = 256, ipb = QI6_K, step = 1; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q6_K_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<19> { static constexpr int qk = 256, ipb = QI1_S, step = 1; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq1_s_q8_1(v, y, kbx, iqs); } }; +template<> struct Fmt<23> { static constexpr int qk = 256, ipb = QI4_XS / 4, step = 4; + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq4_xs_q8_1(v, y, kbx, iqs); } }; + +// Every format above, for the dispatch switches: X(type) once per Fmt. +#define STRATA_FMTS(X) X(7) X(8) X(12) X(13) X(14) X(16) X(17) X(18) X(19) X(20) X(21) X(22) X(23) X(29) X(42) + +inline int fmt_qk(int t) { + switch (t) { +#define STRATA_QK(T) case T: return Fmt::qk; + STRATA_FMTS(STRATA_QK) +#undef STRATA_QK + default: return 0; + } +} + +__device__ __forceinline__ float warp_sum(float v) { +#pragma unroll + for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +} + +// One row against one q8_1 activation, the whole warp: call k = (block, part) is lane-strided. +template +__device__ __forceinline__ float row_dot(const uint8_t* row, const block_q8_1* x, int nb, int lane) { + using F = Fmt; + float s = 0.0f; + for (int k = lane; k < nb * F::ipb; k += 32) { + const int kbx = k / F::ipb, iqs = F::step * (k % F::ipb); + s += F::dot(row, x + kbx * (F::qk / 32), kbx, iqs); + } + return warp_sum(s); +} + +} // namespace +} // namespace strata::kernels diff --git a/src/kernels/cuda/iq_kernels.cu b/src/kernels/cuda/iq_kernels.cu index ea05d9d230c..8b4621580c8 100644 --- a/src/kernels/cuda/iq_kernels.cu +++ b/src/kernels/cuda/iq_kernels.cu @@ -1,17 +1,11 @@ // src/kernels/cuda/iq_kernels.cu - see include/strata/kernels/iq_kernels.hpp. // -// The dot products (vec_dot_*_q8_1), the dequantizers and the q8_1 quantizer are transcribed from llama.cpp -// (ggml/src/ggml-cuda/vecdotq.cuh, dequantize.cuh, quantize.cu at the commit in third_party/ggml/VERSION.txt; -// MIT license, third_party/ggml/LICENSE). The block structs and codebook grids come from its ggml-common.h, -// included unchanged. +// The dequantizers and the q8_1 quantizer are transcribed from llama.cpp (ggml/src/ggml-cuda/dequantize.cuh, +// quantize.cu at the commit in third_party/ggml/VERSION.txt; MIT license, third_party/ggml/LICENSE), the dot products +// in iq_dot.cuh. #include "strata/kernels/iq_kernels.hpp" - -#include -#include - -#define GGML_COMMON_DECL_CUDA -#define GGML_COMMON_IMPL_CUDA -#include "ggml-common.h" +#include "iq_dot.cuh" +#include "q8_1_il.cuh" #include #include @@ -24,312 +18,6 @@ void check(const char* what) { if (e != cudaSuccess) { std::fprintf(stderr, "%s: %s\n", what, cudaGetErrorString(e)); std::exit(1); } } -// ---------------------------------------------------------------- llama.cpp helpers (vecdotq.cuh) -__device__ __forceinline__ int get_int_b2(const void* x, const int& i32) { - const uint16_t* x16 = (const uint16_t*) x; - int x32 = x16[2 * i32 + 0] << 0; - x32 |= x16[2 * i32 + 1] << 16; - return x32; -} -__device__ __forceinline__ int get_int_b4(const void* x, const int& i32) { return ((const int*) x)[i32]; } -__device__ __forceinline__ uint32_t unpack_ksigns(const uint8_t v) { - const uint32_t p = __popc(v) & 1; - const uint32_t s = v ^ p << 7; - return s * 0x01010101; -} -__device__ __forceinline__ int2 get_int_from_table_16(const int& q4, const int8_t* table) { - const uint32_t* table32 = (const uint32_t*) table; - uint32_t tmp[2]; - const uint32_t low_high_selection_indices = (0x32103210 | ((q4 & 0x88888888) >> 1)); -#pragma unroll - for (uint32_t i = 0; i < 2; ++i) { - const uint32_t shift = 16 * i; - const uint32_t low = __byte_perm(table32[0], table32[1], q4 >> shift); - const uint32_t high = __byte_perm(table32[2], table32[3], q4 >> shift); - tmp[i] = __byte_perm(low, high, low_high_selection_indices >> shift); - } - return make_int2(__byte_perm(tmp[0], tmp[1], 0x6420), __byte_perm(tmp[0], tmp[1], 0x7531)); -} -#define ggml_cuda_dp4a(a, b, c) __dp4a((a), (b), (c)) - -// ---------------------------------------------------------------- the dot products (vecdotq.cuh) -__device__ __forceinline__ float vec_dot_q2_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, - const int& kbx, const int& iqs) { - const block_q2_0* bq2_0 = (const block_q2_0*) vbq + kbx; - const float d2 = bq2_0->d; - const int16_t* qs = (const int16_t*) bq2_0->qs + iqs * 4; - const block_q8_1* bq8_1_chunk = bq8_1 + iqs; - int sumi = 0; -#pragma unroll - for (int j = 0; j < 4; ++j) { - const int q = qs[j]; - const int u = get_int_b4(bq8_1_chunk->qs, j * 2 + 0); - const int v = get_int_b4(bq8_1_chunk->qs, j * 2 + 1); - const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0); - const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2); - const int qx = __byte_perm(qe, qo, 0x5140); - const int qy = __byte_perm(qe, qo, 0x7362); - sumi = ggml_cuda_dp4a(u, qx, sumi); - sumi = ggml_cuda_dp4a(v, qy, sumi); - } - const float d8 = __low2float(bq8_1_chunk->ds); - return d2 * d8 * sumi; -} - -__device__ __forceinline__ float vec_dot_iq2_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, - const int& kbx, const int& iqs) { - const block_iq2_xxs* bq2 = (const block_iq2_xxs*) vbq + kbx; - const int q2 = get_int_b2(bq2->qs, iqs); - const uint8_t* aux8 = (const uint8_t*) &q2; - const uint32_t aux32 = get_int_b2(bq2->qs, iqs + 1); - int sumi = 0; -#pragma unroll - for (int k0 = 0; k0 < 8; k0 += 2) { - const uint2 grid_pos = ((const uint2*) iq2xxs_grid)[aux8[k0 / 2]]; - const uint32_t signs = unpack_ksigns(aux32 >> (7 * k0 / 2)); - const int signs0 = __vcmpne4(signs & 0x08040201, 0); - const int grid0 = __vsub4(grid_pos.x ^ signs0, signs0); - const int u0 = get_int_b4(bq8_1[iqs / 2].qs, k0 + 0); - sumi = ggml_cuda_dp4a(grid0, u0, sumi); - const int signs1 = __vcmpne4(signs & 0x80402010, 0); - const int grid1 = __vsub4(grid_pos.y ^ signs1, signs1); - const int u1 = get_int_b4(bq8_1[iqs / 2].qs, k0 + 1); - sumi = ggml_cuda_dp4a(grid1, u1, sumi); - } - const int ls = aux32 >> 27 | 1; - sumi = sumi * ls / 8; - const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs / 2].ds); - return d * sumi; -} - -__device__ __forceinline__ float vec_dot_iq2_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, - const int& kbx, const int& iqs) { - const block_iq2_xs* bq2 = (const block_iq2_xs*) vbq + kbx; - const int2 q2_packed = make_int2(get_int_b2(bq2->qs, iqs + 0), get_int_b2(bq2->qs, iqs + 1)); - const uint16_t* q2 = (const uint16_t*) &q2_packed; - const int ls0 = bq2->scales[iqs / 2] & 0x0F; - const int ls1 = bq2->scales[iqs / 2] >> 4; - int sumi0 = 0, sumi1 = 0; -#pragma unroll - for (int l0 = 0; l0 < 8; l0 += 2) { - const uint2 grid_pos = ((const uint2*) iq2xs_grid)[q2[l0 / 2] & 0x1FF]; - const uint32_t signs = unpack_ksigns(q2[l0 / 2] >> 9); - const int signs0 = __vcmpne4(signs & 0x08040201, 0); - const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); - const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); - const int signs1 = __vcmpne4(signs & 0x80402010, 0); - const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); - const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); - if (l0 < 4) { - sumi0 = ggml_cuda_dp4a(grid_l, u0, sumi0); - sumi0 = ggml_cuda_dp4a(grid_h, u1, sumi0); - } else { - sumi1 = ggml_cuda_dp4a(grid_l, u0, sumi1); - sumi1 = ggml_cuda_dp4a(grid_h, u1, sumi1); - } - } - const int sumi = (sumi0 * ls0 + sumi1 * ls1 + (sumi0 + sumi1) / 2) / 4; - const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs / 2].ds); - return d * sumi; -} - -__device__ __forceinline__ float vec_dot_iq2_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, - const int& kbx, const int& iqs) { - const block_iq2_s* bq2 = (const block_iq2_s*) vbq + kbx; - const int qs_packed = get_int_b2(bq2->qs, iqs / 2); - const uint8_t* qs = (const uint8_t*) &qs_packed; - const int qh = bq2->qh[iqs / 2]; - const int signs_packed_32 = get_int_b2(bq2->qs, QK_K / 32 + iqs / 2); - const uint8_t* signs_packed_8 = (const uint8_t*) &signs_packed_32; - const int ls0 = bq2->scales[iqs / 2] & 0x0F; - const int ls1 = bq2->scales[iqs / 2] >> 4; - int sumi0 = 0, sumi1 = 0; -#pragma unroll - for (int l0 = 0; l0 < 8; l0 += 2) { - const int* grid_pos = (const int*) (iq2s_grid + (qs[l0 / 2] | ((qh << (8 - l0)) & 0x300))); - const int signs0 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x03) << 7) | ((signs_packed_8[l0 / 2] & 0x0C) << 21), 0x00000000); - const int signs1 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x30) << 3) | ((signs_packed_8[l0 / 2] & 0xC0) << 17), 0x00000000); - const int grid_l = __vsub4(grid_pos[0] ^ signs0, signs0); - const int grid_h = __vsub4(grid_pos[1] ^ signs1, signs1); - const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); - const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); - if (l0 < 4) { - sumi0 = ggml_cuda_dp4a(grid_l, u0, sumi0); - sumi0 = ggml_cuda_dp4a(grid_h, u1, sumi0); - } else { - sumi1 = ggml_cuda_dp4a(grid_l, u0, sumi1); - sumi1 = ggml_cuda_dp4a(grid_h, u1, sumi1); - } - } - const int sumi = (sumi0 * ls0 + sumi1 * ls1 + (sumi0 + sumi1) / 2) / 4; - const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs / 2].ds); - return d * sumi; -} - -__device__ __forceinline__ float vec_dot_iq3_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, - const int& kbx, const int& iqs) { - const block_iq3_xxs* bq3 = (const block_iq3_xxs*) vbq + kbx; - const int2 q3_packed = make_int2(get_int_b2(bq3->qs, iqs), get_int_b2(bq3->qs, iqs + 1)); - const uint8_t* q3 = (const uint8_t*) &q3_packed; - const uint32_t aux32 = get_int_b2(bq3->qs, QK_K / 16 + iqs / 2); - int sumi = 0; -#pragma unroll - for (int l0 = 0; l0 < 8; l0 += 2) { - const int2 grid_pos = make_int2(iq3xxs_grid[q3[l0 + 0]], iq3xxs_grid[q3[l0 + 1]]); - const uint32_t signs = unpack_ksigns(aux32 >> (7 * l0 / 2)); - const int signs0 = __vcmpne4(signs & 0x08040201, 0); - const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); - const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); - const int signs1 = __vcmpne4(signs & 0x80402010, 0); - const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); - const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); - sumi = ggml_cuda_dp4a(grid_l, u0, sumi); - sumi = ggml_cuda_dp4a(grid_h, u1, sumi); - } - const int ls = aux32 >> 28; - sumi = (ls * sumi + sumi / 2) / 2; - const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs / 2].ds); - return d * sumi; -} - -__device__ __forceinline__ float vec_dot_iq3_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, - const int& kbx, const int& iqs) { - const block_iq3_s* bq3 = (const block_iq3_s*) vbq + kbx; - const int2 qs_packed = make_int2(get_int_b2(bq3->qs, iqs + 0), get_int_b2(bq3->qs, iqs + 1)); - const uint8_t* qs = (const uint8_t*) &qs_packed; - const int qh = bq3->qh[iqs / 2]; - const int signs_packed_32 = get_int_b2(bq3->signs, iqs / 2); - const uint8_t* signs_packed_8 = (const uint8_t*) &signs_packed_32; - int sumi = 0; -#pragma unroll - for (int l0 = 0; l0 < 8; l0 += 2) { - const int2 grid_pos = make_int2(iq3s_grid[qs[l0 + 0] | ((qh << (8 - l0)) & 0x100)], - iq3s_grid[qs[l0 + 1] | ((qh << (7 - l0)) & 0x100)]); - const int signs0 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x03) << 7) | ((signs_packed_8[l0 / 2] & 0x0C) << 21), 0x00000000); - const int signs1 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x30) << 3) | ((signs_packed_8[l0 / 2] & 0xC0) << 17), 0x00000000); - const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); - const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); - const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); - const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); - sumi = ggml_cuda_dp4a(grid_l, u0, sumi); - sumi = ggml_cuda_dp4a(grid_h, u1, sumi); - } - sumi *= 1 + 2 * ((bq3->scales[iqs / 4] >> ((iqs << 1) & 0x04)) & 0x0F); - const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs / 2].ds); - return d * sumi; -} - -__device__ __forceinline__ float vec_dot_iq1_m_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, - const int& kbx, const int& iqs) { - const block_iq1_m* bq1 = (const block_iq1_m*) vbq + kbx; - const int qs_packed = get_int_b4(bq1->qs, iqs); - const uint8_t* qs = (const uint8_t*) &qs_packed; - int sumi[2] = {0, 0}; - float sumf[2] = {0.0f, 0.0f}; -#pragma unroll - for (int l0 = 0; l0 < 8; l0 += 2) { - const int qhl = bq1->qh[2 * iqs + l0 / 4] >> (4 * ((l0 / 2) % 2)); - const int grid = iq1s_grid_gpu[qs[l0 / 2] | ((qhl & 0x07) << 8)]; - const int grid0 = (grid >> 0) & 0x0F0F0F0F; - const int grid1 = (grid >> 4) & 0x0F0F0F0F; - const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); - const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); - sumi[l0 / 4] = ggml_cuda_dp4a(grid0, u0, sumi[l0 / 4]); - sumi[l0 / 4] = ggml_cuda_dp4a(grid1, u1, sumi[l0 / 4]); - const float delta = -1.0f + IQ1M_DELTA - (qhl & 0x08) * (2.0f * IQ1M_DELTA / 0x08); - int sumy = 0; - sumy = ggml_cuda_dp4a(u0, 0x01010101, sumy); - sumy = ggml_cuda_dp4a(u1, 0x01010101, sumy); - sumf[l0 / 4] += delta * sumy; - } - const uint16_t* sc = (const uint16_t*) bq1->scales; - iq1m_scale_t scale; - scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00F0) | ((sc[2] >> 4) & 0x0F00) | (sc[3] & 0xF000); - const float d = __half2float(scale.f16) * __low2float(bq8_1[iqs].ds); - const int tmp = sc[iqs / 2] >> (6 * (iqs % 2)); - const int sc0 = 2 * ((tmp >> 0) & 0x07) + 1; - const int sc1 = 2 * ((tmp >> 3) & 0x07) + 1; - return d * ((sumi[0] + sumf[0]) * sc0 + (sumi[1] + sumf[1]) * sc1); -} - -__device__ __forceinline__ float vec_dot_iq4_nl_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, - const int& kbx, const int& iqs) { - const block_iq4_nl* bq4 = (const block_iq4_nl*) vbq + kbx; - const int* q8 = (const int*) bq8_1->qs + iqs; - int sumi = 0; -#pragma unroll - for (int l = 0; l < 2; ++l) { - const int aux_q4 = get_int_b2(bq4->qs, iqs + l); - const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); - sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi); - sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi); - } - const float d = __half2float(bq4->d) * __low2float(bq8_1->ds); - return d * sumi; -} - -// IQ4_XS: 256 values as 8 sub-blocks of 32 (6-bit scale each); one call covers one sub-block (iqs = 4 * sub-block), -// and `bq8_1` is the super-block's first q8_1 block, so the call's activation is bq8_1[iqs / 4]. The GSQ-RCO IQ3_S -// file keeps one layer's routed gate/up experts in this format. -__device__ __forceinline__ float vec_dot_iq4_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, - const int& kbx, const int& iqs) { - const block_iq4_xs* bq4 = (const block_iq4_xs*) vbq + kbx; - int sumi = 0; -#pragma unroll - for (int j = 0; j < 4; ++j) { - const int aux_q4 = get_int_b4(bq4->qs, iqs + j); - const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); - const int u0 = get_int_b4(bq8_1[iqs / 4].qs, j + 0); - const int u1 = get_int_b4(bq8_1[iqs / 4].qs, j + 4); - sumi = ggml_cuda_dp4a(v.x, u0, sumi); - sumi = ggml_cuda_dp4a(v.y, u1, sumi); - } - const int ls = ((bq4->scales_l[iqs / 8] >> (iqs & 0x04)) & 0x0F) | (((bq4->scales_h >> (iqs / 2)) & 0x03) << 4); - sumi *= ls - 32; - const float d = __half2float(bq4->d) * __low2float(bq8_1[iqs / 4].ds); - return d * sumi; -} - -// ---------------------------------------------------------------- the formats -// qk = values per block, ipb = dot calls per block (qi / vdr), step = the iqs stride between calls. -template struct Fmt; -template<> struct Fmt<16> { static constexpr int qk = 256, ipb = 8, step = 2; - __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq2_xxs_q8_1(v, y, kbx, iqs); } }; -template<> struct Fmt<17> { static constexpr int qk = 256, ipb = 8, step = 2; - __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq2_xs_q8_1(v, y, kbx, iqs); } }; -template<> struct Fmt<18> { static constexpr int qk = 256, ipb = 8, step = 2; - __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq3_xxs_q8_1(v, y, kbx, iqs); } }; -template<> struct Fmt<20> { static constexpr int qk = 32, ipb = 2, step = 2; - __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq4_nl_q8_1(v, y, kbx, iqs); } }; -template<> struct Fmt<21> { static constexpr int qk = 256, ipb = 8, step = 2; - __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq3_s_q8_1(v, y, kbx, iqs); } }; -template<> struct Fmt<23> { static constexpr int qk = 256, ipb = 8, step = 4; - __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq4_xs_q8_1(v, y, kbx, iqs); } }; -template<> struct Fmt<22> { static constexpr int qk = 256, ipb = 8, step = 2; - __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq2_s_q8_1(v, y, kbx, iqs); } }; -template<> struct Fmt<29> { static constexpr int qk = 256, ipb = 8, step = 1; - __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_iq1_m_q8_1(v, y, kbx, iqs); } }; -template<> struct Fmt<42> { static constexpr int qk = 64, ipb = 2, step = 1; - __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return vec_dot_q2_0_q8_1(v, y, kbx, iqs); } }; - -__device__ __forceinline__ float warp_sum(float v) { -#pragma unroll - for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); - return v; -} - -// One row against one q8_1 activation, the whole warp: call k = (block, part) is lane-strided. -template -__device__ __forceinline__ float row_dot(const uint8_t* row, const block_q8_1* x, int nb, int lane) { - using F = Fmt; - float s = 0.0f; - for (int k = lane; k < nb * F::ipb; k += 32) { - const int kbx = k / F::ipb, iqs = F::step * (k % F::ipb); - s += F::dot(row, x + kbx * (F::qk / 32), kbx, iqs); - } - return warp_sum(s); -} - template __global__ void __launch_bounds__(128) mmvq_kernel(const uint8_t* __restrict__ w, size_t row_bytes, const block_q8_1* __restrict__ x, float* __restrict__ y, int n_in, @@ -345,60 +33,179 @@ __global__ void __launch_bounds__(128) mmvq_kernel(const uint8_t* __restrict__ w } } +// mmvq_kernel for NC columns from their interleaved (position-major) copy, or one from its q8_1 blocks: a warp takes R +// rows, decodes each call's weights once and keeps row_dot's lane-strided sum per (row, column), so every value is +// bitwise mmvq_kernel's. +template +__global__ void __launch_bounds__(128) mmvq_il_iq3_s_kernel(const uint8_t* __restrict__ w, size_t row_bytes, + const void* __restrict__ xq, const float* __restrict__ xd, + float* __restrict__ y, int n_in, int n_out) { + using F = Fmt<21>; + const int lane = threadIdx.x & 31; + const auto x = q8_1_cols(xq, xd, n_in / 32); + const int nb = n_in / F::qk; + // a grid-stride loop, though the grid gives each warp one group: native_mmvq_il_kernel's form, which its + // rows-a-warp table is measured on + for (int grp = blockIdx.x * 4 + (threadIdx.x >> 5); grp * R < n_out; grp += gridDim.x * 4) { + const int row0 = grp * R; + float acc[R][NC]; +#pragma unroll + for (int i = 0; i < R; ++i) +#pragma unroll + for (int c = 0; c < NC; ++c) acc[i][c] = 0.0f; + for (int k = lane; k < nb * F::ipb; k += 32) { + const int kbx = k / F::ipb, iqs = F::step * (k % F::ipb); +#pragma unroll + for (int i = 0; i < R; ++i) { + const int row = min(row0 + i, n_out - 1); // a partial last group recomputes its last row + IQ3SCols wv; + wv.load(w + (size_t) row * row_bytes, kbx, iqs); + float o[NC]; + wv.apply(x, kbx * (F::qk / 32) + iqs / 2, o); +#pragma unroll + for (int c = 0; c < NC; ++c) acc[i][c] += o[c]; + } + } +#pragma unroll + for (int i = 0; i < R; ++i) +#pragma unroll + for (int c = 0; c < NC; ++c) { + const float s = warp_sum(acc[i][c]); + if (lane == 0 && row0 + i < n_out) y[(size_t) c * n_out + row0 + i] = s; + } + } +} + // ---------------------------------------------------------------- grouped native experts -constexpr int GU_ROWS = 8; // rows per block (one warp each) +// Two kernels a call. The first computes a group's gate and up rows for one chunk of 32 rows of n_ff - one q8_1 block +// of h - SwiGLU and that block for each entry; the second the down rows. A warp takes several rows at once, one +// lane-strided accumulator each, so every value's sums and their order are `row_dot`'s. +constexpr int GUS_ROWS = 2, GUS_WARPS = 16; // gate and up rows per warp; 8 warps x 4 rows = one q8_1 block +constexpr int GUS_BATCH = 8; // entries whose block of h is staged at once +constexpr int DOWN_ROWS = 4, DOWN_WARPS = 8; + +// The K-quants' gate/up rows need more registers than the i-quants': with one block per SM they do not spill (Q4_K, Q5_K +// 15-17% faster at 2-4 tokens), while the i-quants run best at the occupancy the compiler picks (0: no bound). +template constexpr int gus_min_blocks() { return T == 12 || T == 13 || T == 14 ? 1 : 0; } template -__global__ void __launch_bounds__(256) native_gu_kernel(const unsigned long long* __restrict__ grp_ptr, - const int32_t* __restrict__ grp_start, - const int32_t* __restrict__ n_groups, - const int32_t* __restrict__ ent_tok, - const block_q8_1* __restrict__ xq, NativeExpertLayout L, - float* __restrict__ gate, float* __restrict__ up) { +__global__ void __launch_bounds__(GUS_WARPS * 32, gus_min_blocks()) native_gus_kernel(const unsigned long long* __restrict__ grp_ptr, + const int32_t* __restrict__ grp_start, + const int32_t* __restrict__ n_groups, + const int32_t* __restrict__ ent_tok, + const block_q8_1* __restrict__ xq, + NativeExpertLayout L, block_q8_1* __restrict__ hq) { + using F = Fmt; const int g = blockIdx.y; if (g >= *n_groups) return; - const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; - const int row = blockIdx.x * GU_ROWS + warp; // 0 .. 2*n_ff - if (row >= 2 * L.n_ff) return; - const bool is_up = row >= L.n_ff; - const int r = is_up ? row - (int) L.n_ff : row; - const uint8_t* blob = (const uint8_t*) grp_ptr[g]; - const uint8_t* wr = blob + (is_up ? L.up_off : 0) + (size_t) r * L.gu_row; - const int nb = (int) (L.n_embd / Fmt::qk), xb = (int) (L.n_embd / 32); + __shared__ float sh[GUS_BATCH][32]; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, c = blockIdx.x; + const uint8_t* wg = (const uint8_t*) grp_ptr[g] + (size_t) (32 * c + GUS_ROWS * warp) * L.gu_row; + const uint8_t* wu = wg + L.up_off; + const int nb = (int) (L.n_embd / F::qk), xb = (int) (L.n_embd / 32), hb = (int) (L.n_ff / 32); const int e0 = grp_start[g], e1 = grp_start[g + 1]; - for (int e = e0; e < e1; ++e) { - const float s = row_dot(wr, xq + (size_t) ent_tok[e] * xb, nb, lane); - if (lane == 0) (is_up ? up : gate)[(size_t) e * L.n_ff + r] = s; + for (int b0 = e0; b0 < e1; b0 += GUS_BATCH) { + const int b1 = min(e1, b0 + GUS_BATCH); + for (int e = b0; e < b1; ++e) { + const block_q8_1* x = xq + (size_t) ent_tok[e] * xb; + float sg[GUS_ROWS], su[GUS_ROWS]; +#pragma unroll + for (int i = 0; i < GUS_ROWS; ++i) sg[i] = su[i] = 0.0f; + for (int k = lane; k < nb * F::ipb; k += 32) { + const int kbx = k / F::ipb, iqs = F::step * (k % F::ipb); + const block_q8_1* xk = x + kbx * (F::qk / 32); +#pragma unroll + for (int i = 0; i < GUS_ROWS; ++i) { + sg[i] += F::dot(wg + (size_t) i * L.gu_row, xk, kbx, iqs); + su[i] += F::dot(wu + (size_t) i * L.gu_row, xk, kbx, iqs); + } + } +#pragma unroll + for (int i = 0; i < GUS_ROWS; ++i) { + const float gv = warp_sum(sg[i]), uv = warp_sum(su[i]); + if (lane == 0) sh[e - b0][GUS_ROWS * warp + i] = (gv / (1.0f + __expf(-gv))) * uv; + } + } + __syncthreads(); + for (int e = b0 + warp; e < b1; e += GUS_WARPS) { // an entry's block of h, as quantize_q8_1_kernel + const float xi = sh[e - b0][lane]; + float amax = fabsf(xi), sum = xi; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, o)); + sum += __shfl_xor_sync(0xffffffffu, sum, o); + } + const float d = amax / 127.0f; + const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); + block_q8_1* y = hq + (size_t) e * hb + c; + y->qs[lane] = q; + if (lane == 0) y->ds = make_half2(d, sum); + } + __syncthreads(); } } -__global__ void swiglu_entries_kernel(const float* __restrict__ gate, const float* __restrict__ up, float* __restrict__ h, - long long n) { - const long long i = (long long) blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n) return; - const float g = gate[i]; - h[i] = (g / (1.0f + __expf(-g))) * up[i]; -} - template -__global__ void __launch_bounds__(256) native_down_kernel(const unsigned long long* __restrict__ grp_ptr, - const int32_t* __restrict__ grp_start, - const int32_t* __restrict__ n_groups, - const int32_t* __restrict__ ent_dst, - const block_q8_1* __restrict__ hq, NativeExpertLayout L, - float* __restrict__ out) { +__global__ void __launch_bounds__(DOWN_WARPS * 32) native_down_kernel(const unsigned long long* __restrict__ grp_ptr, + const int32_t* __restrict__ grp_start, + const int32_t* __restrict__ n_groups, + const int32_t* __restrict__ ent_dst, + const block_q8_1* __restrict__ hq, + NativeExpertLayout L, float* __restrict__ out) { + using F = Fmt; const int g = blockIdx.y; if (g >= *n_groups) return; const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; - const int r = blockIdx.x * 8 + warp; - if (r >= L.n_embd) return; - const uint8_t* blob = (const uint8_t*) grp_ptr[g]; - const uint8_t* wr = blob + L.down_off + (size_t) r * L.d_row; - const int nb = (int) (L.n_ff / Fmt::qk), hb = (int) (L.n_ff / 32); + const int r0 = (blockIdx.x * DOWN_WARPS + warp) * DOWN_ROWS; + const uint8_t* wr = (const uint8_t*) grp_ptr[g] + L.down_off + (size_t) r0 * L.d_row; + const int nb = (int) (L.n_ff / F::qk), hb = (int) (L.n_ff / 32); const int e0 = grp_start[g], e1 = grp_start[g + 1]; for (int e = e0; e < e1; ++e) { - const float s = row_dot(wr, hq + (size_t) e * hb, nb, lane); - if (lane == 0) out[(size_t) ent_dst[e] * L.n_embd + r] = s; + const block_q8_1* x = hq + (size_t) e * hb; + float sd[DOWN_ROWS]; +#pragma unroll + for (int i = 0; i < DOWN_ROWS; ++i) sd[i] = 0.0f; + for (int k = lane; k < nb * F::ipb; k += 32) { + const int kbx = k / F::ipb, iqs = F::step * (k % F::ipb); + const block_q8_1* xk = x + kbx * (F::qk / 32); +#pragma unroll + for (int i = 0; i < DOWN_ROWS; ++i) sd[i] += F::dot(wr + (size_t) i * L.d_row, xk, kbx, iqs); + } + float v[DOWN_ROWS]; +#pragma unroll + for (int i = 0; i < DOWN_ROWS; ++i) v[i] = warp_sum(sd[i]); + if (lane == 0) { + float* o = out + (size_t) ent_dst[e] * L.n_embd + r0; + if constexpr (DOWN_ROWS == 4) *(float4*) o = make_float4(v[0], v[1], v[2], v[3]); + else if constexpr (DOWN_ROWS == 2) *(float2*) o = make_float2(v[0], v[1]); + else +#pragma unroll + for (int i = 0; i < DOWN_ROWS; ++i) o[i] = v[i]; + } + } +} + +__global__ void rows_out_kernel(const float4* __restrict__ src, const int32_t* __restrict__ n_groups, + const int32_t* __restrict__ grp_start, const int32_t* __restrict__ ent_dst, + long long n4, float4* const* out, uint32_t* const* flag, const int32_t* ring, + unsigned* counter) { + const int e0 = grp_start[0]; + const long long total = (long long) (grp_start[*n_groups] - e0) * n4, step = (long long) gridDim.x * blockDim.x; + float4* o = *out; + for (long long i = (long long) blockIdx.x * blockDim.x + threadIdx.x; i < total; i += step) { + const long long at = (long long) ent_dst[e0 + i / n4] * n4 + i % n4; + o[at] = src[at]; + } + uint32_t* f = *flag; + if (f == nullptr) return; + __shared__ bool last; + __threadfence_system(); // this thread's rows before the block counts as done + __syncthreads(); + if (threadIdx.x == 0) last = atomicAdd(counter, 1u) == gridDim.x - 1; + __syncthreads(); + if (last && threadIdx.x == 0) { + *counter = 0u; // for the next launch + *(volatile uint32_t*) f = (uint32_t) *ring; } } @@ -567,6 +374,113 @@ __device__ void dq_q2_0(const void* vx, int64_t ibs, dst_t* yy, int tid) { yy[b * 64 + i] = cvt(d * (float) (code - 1)); } } +__device__ __forceinline__ void scale_min_k4(int j, const uint8_t* q, uint8_t& d, uint8_t& m) { + if (j < 4) { + d = q[j] & 63; m = q[j + 4] & 63; + } else { + d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4); + m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); + } +} +template +__device__ void dq_q4_k(const void* vx, int64_t ibs, dst_t* yy, int tid) { + const block_q4_K* x = (const block_q4_K*) vx; + const int64_t il = tid / 8, ir = tid % 8, is = 2 * il; + const int n = 4; + dst_t* y = yy + 64 * il + n * ir; + const float dall = __low2half(x[ibs].dm); + const float dmin = __high2half(x[ibs].dm); + const uint8_t* q = x[ibs].qs + 32 * il + n * ir; + uint8_t sc, m; + scale_min_k4((int) is + 0, x[ibs].scales, sc, m); + const float d1 = dall * sc, m1 = dmin * m; + scale_min_k4((int) is + 1, x[ibs].scales, sc, m); + const float d2 = dall * sc, m2 = dmin * m; + for (int l = 0; l < n; ++l) { + y[l + 0] = cvt(d1 * (q[l] & 0xF) - m1); + y[l + 32] = cvt(d2 * (q[l] >> 4) - m2); + } +} +// llama.cpp's dequantize_q5_K, its 64 threads folded onto 32 +template +__device__ void dq_q5_k(const void* vx, int64_t ibs, dst_t* yy, int tid) { + const block_q5_K* x = (const block_q5_K*) vx; + for (int tt = tid; tt < 64; tt += 32) { + const int il = tt / 16, ir = tt % 16, is = 2 * il; + dst_t* y = yy + 64 * il + 2 * ir; + const float dall = __low2half(x[ibs].dm); + const float dmin = __high2half(x[ibs].dm); + const uint8_t* ql = x[ibs].qs + 32 * il + 2 * ir; + const uint8_t* qh = x[ibs].qh + 2 * ir; + uint8_t sc, m; + scale_min_k4(is + 0, x[ibs].scales, sc, m); + const float d1 = dall * sc, m1 = dmin * m; + scale_min_k4(is + 1, x[ibs].scales, sc, m); + const float d2 = dall * sc, m2 = dmin * m; + uint8_t hm = (uint8_t) (1 << (2 * il)); + y[0] = cvt(d1 * ((ql[0] & 0xF) + (qh[0] & hm ? 16 : 0)) - m1); + y[1] = cvt(d1 * ((ql[1] & 0xF) + (qh[1] & hm ? 16 : 0)) - m1); + hm <<= 1; + y[32] = cvt(d2 * ((ql[0] >> 4) + (qh[0] & hm ? 16 : 0)) - m2); + y[33] = cvt(d2 * ((ql[1] >> 4) + (qh[1] & hm ? 16 : 0)) - m2); + } +} +// 32-value blocks: a "superblock" is 8 of them; thread tid writes 8 values of block tid % 8 +template +__device__ void dq_q5_1(const void* vx, int64_t ibs, dst_t* yy, int tid) { + const block_q5_1* x = (const block_q5_1*) vx + ibs * (QK_K / QK5_1); + const int ib = tid % 8, il = tid / 8; + const float2 dm = __half22float2(x[ib].dm); + uint32_t qh; + memcpy(&qh, x[ib].qh, sizeof(qh)); + dst_t* y = yy + 32 * ib; + for (int j = 0; j < 4; ++j) { + const int iqs = 4 * il + j; + const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10; + const int xh_1 = ((qh >> (iqs + 12))) & 0x10; + y[iqs] = cvt((float) ((x[ib].qs[iqs] & 0xf) | xh_0) * dm.x + dm.y); + y[iqs + 16] = cvt((float) ((x[ib].qs[iqs] >> 4) | xh_1) * dm.x + dm.y); + } +} +template +__device__ void dq_q8_0(const void* vx, int64_t ibs, dst_t* yy, int tid) { + const block_q8_0* x = (const block_q8_0*) vx + ibs * (QK_K / QK8_0); + const int ib = tid % 8, il = tid / 8; + const float d = __half2float(x[ib].d); + dst_t* y = yy + 32 * ib + 8 * il; + for (int j = 0; j < 8; ++j) y[j] = cvt((float) x[ib].qs[8 * il + j] * d); +} +// llama.cpp's dequantize_q6_K, its 64 threads folded onto 32 +template +__device__ void dq_q6_k(const void* vx, int64_t ibs, dst_t* yy, int tid) { + const block_q6_K* x = (const block_q6_K*) vx; + for (int tt = tid; tt < 64; tt += 32) { + const int ip = tt / 32, il = tt - 32 * ip, is = 8 * ip + il / 16; + dst_t* y = yy + 128 * ip + il; + const float d = __half2float(x[ibs].d); + const uint8_t* ql = x[ibs].ql + 64 * ip + il; + const uint8_t qh = x[ibs].qh[32 * ip + il]; + const int8_t* sc = x[ibs].scales + is; + y[0] = cvt(d * sc[0] * ((int8_t) ((ql[0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32)); + y[32] = cvt(d * sc[2] * ((int8_t) ((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32)); + y[64] = cvt(d * sc[4] * ((int8_t) ((ql[0] >> 4) | (((qh >> 4) & 3) << 4)) - 32)); + y[96] = cvt(d * sc[6] * ((int8_t) ((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32)); + } +} +template +__device__ void dq_iq1_s(const void* vx, int64_t ibs, dst_t* yy, int tid) { + const block_iq1_s* x = (const block_iq1_s*) vx; + const int64_t il = tid / 8, ib = tid % 8; + dst_t* y = yy + 32 * ib + 8 * il; + const float delta = x[ibs].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA; + const float d = (float) x[ibs].d * (2 * ((x[ibs].qh[ib] >> 12) & 7) + 1); + uint32_t grid32[2]; + const int8_t* q = (const int8_t*) grid32; + grid32[0] = iq1s_grid_gpu[x[ibs].qs[4 * ib + il] | (((x[ibs].qh[ib] >> 3 * il) & 7) << 8)]; + grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f; + grid32[0] &= 0x0f0f0f0f; + for (int j = 0; j < 8; ++j) y[j] = cvt(d * (q[j] + delta)); +} template __device__ __forceinline__ void dq_dispatch(int ty, const void* vx, int64_t ibs, dst_t* y, int tid) { @@ -581,6 +495,12 @@ __device__ __forceinline__ void dq_dispatch(int ty, const void* vx, int64_t ibs, case 23: dq_iq4_xs(vx, ibs, y, tid); break; case 11: dq_q3_k(vx, ibs, y, tid); break; case 42: dq_q2_0(vx, ibs, y, tid); break; + case 12: dq_q4_k(vx, ibs, y, tid); break; + case 13: dq_q5_k(vx, ibs, y, tid); break; + case 7: dq_q5_1(vx, ibs, y, tid); break; + case 8: dq_q8_0(vx, ibs, y, tid); break; + case 14: dq_q6_k(vx, ibs, y, tid); break; + case 19: dq_iq1_s(vx, ibs, y, tid); break; default: break; } } @@ -600,7 +520,29 @@ __global__ void dequant_gu_kernel(int ty, const void* __restrict__ gate, const v dq_dispatch<__half>(ty, parity ? up : gate, i, y + ((2 * r + parity) * per_row + c) * QK_K, threadIdx.x); } -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; } +// several experts at once, expert z's blob at blobs[z]: first its gate/up superblocks (as dequant_gu_kernel), then +// its down superblocks (as dequant_flat_kernel) +__global__ void dequant_experts_kernel(int gu_ty, int d_ty, const uint8_t* const* __restrict__ blobs, size_t up_off, + size_t down_off, int64_t per_row, int64_t gu_blocks, __half* __restrict__ gu, + __half* __restrict__ dn, int64_t gu_elems, int64_t dn_elems) { + const uint8_t* b = blobs[blockIdx.z]; + const int64_t i = blockIdx.x; + if (i < 2 * gu_blocks) { + const int parity = (int) (i / gu_blocks); + const int64_t j = i % gu_blocks, r = j / per_row, c = j % per_row; + dq_dispatch<__half>(gu_ty, parity ? b + up_off : b, j, + gu + blockIdx.z * gu_elems + ((2 * r + parity) * per_row + c) * QK_K, threadIdx.x); + } else { + const int64_t j = i - 2 * gu_blocks; + dq_dispatch<__half>(d_ty, b + down_off, j, dn + blockIdx.z * dn_elems + j * QK_K, threadIdx.x); + } +} + +// the formats dq_dispatch dequantizes +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 || + t == 12 || t == 13 || t == 14 || t == 19 || t == 7 || t == 8; +} } // namespace @@ -618,6 +560,12 @@ size_t iq_row_bytes(int t, int64_t n) noexcept { case 23: return (size_t) (n / 256) * sizeof(block_iq4_xs); case 11: return (size_t) (n / 256) * sizeof(block_q3_K); case 42: return (size_t) (n / 64) * sizeof(block_q2_0); + case 12: return (size_t) (n / 256) * sizeof(block_q4_K); + case 13: return (size_t) (n / 256) * sizeof(block_q5_K); + case 14: return (size_t) (n / 256) * sizeof(block_q6_K); + case 19: return (size_t) (n / 256) * sizeof(block_iq1_s); + case 7: return (size_t) (n / 32) * sizeof(block_q5_1); + case 8: return (size_t) (n / 32) * sizeof(block_q8_0); default: return 0; } } @@ -636,20 +584,52 @@ void iq_mmvq(int t, const void* w, const void* x_q8_1, float* y, int n_in, int n const auto* W = (const uint8_t*) w; const auto* X = (const block_q8_1*) x_q8_1; switch (t) { - case 16: mmvq_kernel<16><<>>(W, rb, X, y, n_in, n_out, ncols); break; - case 17: mmvq_kernel<17><<>>(W, rb, X, y, n_in, n_out, ncols); break; - case 18: mmvq_kernel<18><<>>(W, rb, X, y, n_in, n_out, ncols); break; - case 20: mmvq_kernel<20><<>>(W, rb, X, y, n_in, n_out, ncols); break; - case 21: mmvq_kernel<21><<>>(W, rb, X, y, n_in, n_out, ncols); break; - case 22: mmvq_kernel<22><<>>(W, rb, X, y, n_in, n_out, ncols); break; - case 23: mmvq_kernel<23><<>>(W, rb, X, y, n_in, n_out, ncols); break; - case 29: mmvq_kernel<29><<>>(W, rb, X, y, n_in, n_out, ncols); break; - case 42: mmvq_kernel<42><<>>(W, rb, X, y, n_in, n_out, ncols); break; +#define STRATA_MMVQ(T) case T: mmvq_kernel<<>>(W, rb, X, y, n_in, n_out, ncols); break; + STRATA_FMTS(STRATA_MMVQ) +#undef STRATA_MMVQ default: std::fprintf(stderr, "iq_mmvq: type %d is not supported\n", t); std::exit(1); } check("iq_mmvq"); } +bool iq_mmvq_il_supported(int t) noexcept { return t == 21; } + +namespace { +template +void launch_iq3_s_il(const uint8_t* w, size_t rb, const void* x_q8_1, const void* x_il, float* y, int n_in, + int n_out, int ncols, cudaStream_t s) { + const unsigned blocks = (unsigned) (((n_out + R - 1) / R + 3) / 4); + if (ncols == 1) { + mmvq_il_iq3_s_kernel<1, R><<>>(w, rb, x_q8_1, nullptr, y, n_in, n_out); + return; + } + const Q81IlParts parts = q8_1_il_parts(x_il, n_in, ncols); + switch (ncols) { +#define STRATA_IL(N) \ + case N: mmvq_il_iq3_s_kernel<<>>(w, rb, parts.pm, parts.d, y, n_in, n_out); break; + STRATA_IL(2) STRATA_IL(3) STRATA_IL(4) STRATA_IL(5) STRATA_IL(6) STRATA_IL(7) STRATA_IL(8) +#undef STRATA_IL + } +} +} // namespace + +void iq_mmvq_il(int t, const void* w, const void* x_q8_1, const void* x_il, float* y, int n_in, int n_out, int ncols, + void* stream, int rows) { + if (t != 21 || ncols < 1 || ncols > 8 || n_in % 256 != 0) { + std::fprintf(stderr, "iq_mmvq_il: type %d, %d columns are not supported\n", t, ncols); + std::exit(1); + } + const size_t rb = iq_row_bytes(t, n_in); + const auto* W = (const uint8_t*) w; + cudaStream_t s = (cudaStream_t) stream; + switch (rows) { + case 1: launch_iq3_s_il<1>(W, rb, x_q8_1, x_il, y, n_in, n_out, ncols, s); break; + case 4: launch_iq3_s_il<4>(W, rb, x_q8_1, x_il, y, n_in, n_out, ncols, s); break; + default: launch_iq3_s_il<2>(W, rb, x_q8_1, x_il, y, n_in, n_out, ncols, s); break; + } + check("iq_mmvq_il"); +} + void iq_dequant_f16(int t, const void* src, int64_t n, uint16_t* dst, void* stream) { if (n % 256 != 0 || !is_iq(t)) { std::fprintf(stderr, "iq_dequant_f16: bad arguments\n"); std::exit(1); } dequant_flat_kernel<__half><<<(unsigned) (n / 256), 32, 0, (cudaStream_t) stream>>>(t, src, (__half*) dst); @@ -688,6 +668,25 @@ void iq_dequant_gu_f16(int t, const void* gate, const void* up, int64_t n_ff, in check("iq_dequant_gu_f16"); } +void iq_dequant_experts_f16(int gu_type, int d_type, const uint8_t* const* blobs, int n, size_t up_off, size_t down_off, + int64_t n_ff, int64_t n_embd, uint16_t* gu, uint16_t* dn, void* stream) { + if (n <= 0) return; + if (n_embd % 256 != 0 || (n_embd * n_ff) % 256 != 0 || !is_iq(gu_type) || !is_iq(d_type) || n > 65535) { + std::fprintf(stderr, "iq_dequant_experts_f16: bad arguments\n"); + std::exit(1); + } + const int64_t per_row = n_embd / 256, gu_blocks = n_ff * per_row, d_blocks = n_embd * n_ff / 256; + dequant_experts_kernel<<>>( + gu_type, d_type, blobs, up_off, down_off, per_row, gu_blocks, (__half*) gu, (__half*) dn, 2 * n_ff * n_embd, + n_embd * n_ff); + check("iq_dequant_experts_f16"); +} + +bool native_expert_supported(int gu_type, int d_type, int64_t n_embd, int64_t n_ff) noexcept { + const int qg = fmt_qk(gu_type), qd = fmt_qk(d_type); + return qg > 0 && qd > 0 && is_iq(gu_type) && is_iq(d_type) && n_embd % qg == 0 && n_ff % qd == 0; +} + NativeExpertLayout native_expert_layout(int gu_type, int d_type, int64_t n_embd, int64_t n_ff) { NativeExpertLayout L; L.gu_type = gu_type; @@ -703,45 +702,47 @@ NativeExpertLayout native_expert_layout(int gu_type, int d_type, int64_t n_embd, } size_t native_expert_scratch_bytes(int64_t cap, int64_t n_ff) { - const size_t f = (size_t) cap * (size_t) n_ff * sizeof(float); - return 3 * ((f + 255) & ~(size_t) 255) + (((size_t) cap * (size_t) (n_ff / 32) * sizeof(block_q8_1) + 255) & ~(size_t) 255); + return ((size_t) cap * (size_t) (n_ff / 32) * sizeof(block_q8_1) + 255) & ~(size_t) 255; } void native_expert_grouped(const NativeExpertLayout& L, const unsigned long long* grp_ptr, const int32_t* grp_start, const int32_t* n_groups, const int32_t* ent_dst, const int32_t* ent_tok, int64_t cap_groups, int64_t cap_entries, const void* x_q8_1, void* scratch, float* out, void* stream) { if (cap_groups <= 0 || cap_entries <= 0) return; + if (L.n_ff % 32 != 0 || L.n_embd % (DOWN_WARPS * DOWN_ROWS) != 0) { + std::fprintf(stderr, "native_expert_grouped: n_ff %lld / n_embd %lld do not divide into the kernels' rows\n", + (long long) L.n_ff, (long long) L.n_embd); + std::exit(1); + } cudaStream_t s = (cudaStream_t) stream; - const size_t f = (size_t) cap_entries * (size_t) L.n_ff * sizeof(float), fa = (f + 255) & ~(size_t) 255; - float* gate = (float*) scratch; - float* up = (float*) ((uint8_t*) scratch + fa); - float* h = (float*) ((uint8_t*) scratch + 2 * fa); - block_q8_1* hq = (block_q8_1*) ((uint8_t*) scratch + 3 * fa); + auto* hq = (block_q8_1*) scratch; const auto* X = (const block_q8_1*) x_q8_1; - const dim3 ggu((unsigned) ((2 * L.n_ff + GU_ROWS - 1) / GU_ROWS), (unsigned) cap_groups); + const dim3 ggu((unsigned) (L.n_ff / 32), (unsigned) cap_groups); switch (L.gu_type) { - case 16: native_gu_kernel<16><<>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; - case 17: native_gu_kernel<17><<>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; - case 18: native_gu_kernel<18><<>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; - case 21: native_gu_kernel<21><<>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; - case 22: native_gu_kernel<22><<>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; - case 23: native_gu_kernel<23><<>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; - case 29: native_gu_kernel<29><<>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; - case 42: native_gu_kernel<42><<>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; +#define STRATA_GU(T) case T: native_gus_kernel<<>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, hq); break; + STRATA_FMTS(STRATA_GU) +#undef STRATA_GU default: std::fprintf(stderr, "native_expert_grouped: gate/up type %d\n", L.gu_type); std::exit(1); } check("native_expert_grouped/gu"); - const long long nh = (long long) cap_entries * L.n_ff; - swiglu_entries_kernel<<<(unsigned) ((nh + 255) / 256), 256, 0, s>>>(gate, up, h, nh); - quantize_q8_1_kernel<<<(unsigned) ((nh + 255) / 256), 256, 0, s>>>(h, hq, nh); - const dim3 gd((unsigned) ((L.n_embd + 7) / 8), (unsigned) cap_groups); + const dim3 gd((unsigned) (L.n_embd / (DOWN_WARPS * DOWN_ROWS)), (unsigned) cap_groups); switch (L.d_type) { - case 20: native_down_kernel<20><<>>(grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break; - case 23: native_down_kernel<23><<>>(grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break; - case 42: native_down_kernel<42><<>>(grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break; +#define STRATA_DOWN(T) case T: native_down_kernel<<>>(grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break; + STRATA_FMTS(STRATA_DOWN) +#undef STRATA_DOWN default: std::fprintf(stderr, "native_expert_grouped: down type %d\n", L.d_type); std::exit(1); } check("native_expert_grouped/down"); } +void native_expert_rows_out(const float* rows, const int32_t* n_groups, const int32_t* grp_start, + const int32_t* ent_dst, int64_t n_embd, int64_t max_rows, float* const* out, + uint32_t* const* flag, const int32_t* ring, unsigned* counter, void* stream) { + const long long n4 = n_embd / 4, most = max_rows * n4; + const unsigned blocks = (unsigned) std::min((most + 255) / 256, 256); + rows_out_kernel<<>>((const float4*) rows, n_groups, grp_start, ent_dst, n4, + (float4* const*) out, flag, ring, counter); + check("native_expert_rows_out"); +} + } // namespace strata::kernels diff --git a/src/kernels/cuda/kv_q8.cu b/src/kernels/cuda/kv_q8.cu index fb2ccec444e..78e9eed86b7 100644 --- a/src/kernels/cuda/kv_q8.cu +++ b/src/kernels/cuda/kv_q8.cu @@ -31,10 +31,15 @@ __global__ void kv_append_q8_kernel(int8_t* __restrict__ k_q, int8_t* __restrict uint16_t* __restrict__ k_scale, uint16_t* __restrict__ v_scale, const int32_t* __restrict__ table, const int32_t* __restrict__ step, const float* __restrict__ kcur, const float* __restrict__ vcur, int kv_heads, - int head_dim, int page_size, KvHostPools host) { + int head_dim, int page_size, int step_stride) { + // blockIdx.z = 2 * token + (0 K, 1 V) + const int tok = blockIdx.z >> 1; + step += (size_t) tok * step_stride; + kcur += (size_t) tok * kv_heads * head_dim; + vcur += (size_t) tok * kv_heads * head_dim; const long long pos = (long long) __ldg(step + kStepPos); const int h = blockIdx.x, g = blockIdx.y, t = threadIdx.x; - const bool is_v = blockIdx.z == 1; + const bool is_v = (blockIdx.z & 1) == 1; const int groups = head_dim / KV_Q8_GROUP; const float x = (is_v ? vcur : kcur)[h * head_dim + g * KV_Q8_GROUP + t]; // max |x| over the 64 values: two warps, then combine through shared memory in a fixed order @@ -51,18 +56,10 @@ __global__ void kv_append_q8_kernel(int8_t* __restrict__ k_q, int8_t* __restrict q = __float2int_rn(x / sf); q = q < -127 ? -127 : (q > 127 ? 127 : q); } - // KV streaming: the host copy (identity layout) always, the VRAM page only if the block is resident const long long page = (long long) table[pos / page_size]; - if (page >= 0) { - const long long row = (page * kv_heads + h) * page_size + (pos % page_size); - (is_v ? v_q : k_q)[row * head_dim + g * KV_Q8_GROUP + t] = (int8_t) q; - if (t == 0) (is_v ? v_scale : k_scale)[row * groups + g] = sbits; - } - if (host.k_q != nullptr) { - const long long row = ((pos / page_size) * kv_heads + h) * page_size + (pos % page_size); - (is_v ? host.v_q : host.k_q)[row * head_dim + g * KV_Q8_GROUP + t] = (int8_t) q; - if (t == 0) (is_v ? host.v_scale : host.k_scale)[row * groups + g] = sbits; - } + const long long row = (page * kv_heads + h) * page_size + (pos % page_size); + (is_v ? v_q : k_q)[row * head_dim + g * KV_Q8_GROUP + t] = (int8_t) q; + if (t == 0) (is_v ? v_scale : k_scale)[row * groups + g] = sbits; } // One thread = 4 consecutive values of one cell and head (as the FP16 gather does with uint2). @@ -101,13 +98,20 @@ __global__ void kv_gather_q8_kernel(const int8_t* __restrict__ k_q, const int8_t } // namespace void kv_append_q8_step(int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, const int32_t* page_table, - const int32_t* step, const float* kcur, const float* vcur, const QsaShapes& s, void* stream, - const KvHostPools* host) { + const int32_t* step, const float* kcur, const float* vcur, const QsaShapes& s, void* stream) { + validate(s, "kv_append_q8"); + kv_append_q8_steps(k_q, v_q, k_scale, v_scale, page_table, step, 0, kcur, vcur, 1, s, stream); +} + +void kv_append_q8_steps(int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, const int32_t* page_table, + const int32_t* step, int step_stride, const float* kcur, const float* vcur, int n_tok, + const QsaShapes& s, void* stream) { validate(s, "kv_append_q8"); - const dim3 grid((unsigned) s.n_head_kv, (unsigned) (s.head_dim / KV_Q8_GROUP), 2); + if (n_tok < 1) return; + const dim3 grid((unsigned) s.n_head_kv, (unsigned) (s.head_dim / KV_Q8_GROUP), (unsigned) (2 * n_tok)); kv_append_q8_kernel<<>>( k_q, v_q, k_scale, v_scale, page_table, step, kcur, vcur, (int) s.n_head_kv, (int) s.head_dim, - (int) s.page_size, host ? *host : KvHostPools{}); + (int) s.page_size, step_stride); check("kv_append_q8 launch"); } diff --git a/src/kernels/cuda/mmvf_multi.cuh b/src/kernels/cuda/mmvf_multi.cuh new file mode 100644 index 00000000000..b2ddc663d66 --- /dev/null +++ b/src/kernels/cuda/mmvf_multi.cuh @@ -0,0 +1,141 @@ +// src/kernels/cuda/mmvf_multi.cuh - one block's row, or rows, of the FP32-activation BF16 MMVF for up to 8 activation +// rows, for bf16_gemv_fp32_mmvf_multi (native_bf16.cu) and the verify window's router (native_router.cu); both +// translation units compile with --use_fast_math, the pinned contract's arithmetic. +// +// Adapted from llama.cpp 3cf03257f219afbe7334045ff7c6a06ac68c627d, ggml/src/ggml-cuda/{mmvf.cu,common.cuh}. +// MIT License, Copyright (c) 2023-2026 The ggml authors (see native_bf16.cu for the full notice). +#pragma once + +#include "strata/kernels/bf16_bits.hpp" + +#include + +#include + +namespace strata::kernels { +namespace { + +constexpr int kMmvfMaxRows = 8; + +__device__ __forceinline__ float mmvf_warp_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + value += __shfl_xor_sync(0xffffffffu, value, offset, 32); + return value; +} + +// `bf16_f32_mmvf_kernel` for up to kMmvfMaxRows activation rows (x: n_tok rows of n_in): each thread walks the same +// pairs in the same order for every row, and each row is reduced exactly as the single-row kernel reduces it; thread 0 +// ends with the sums in acc[0 .. n_tok). +template +__device__ __forceinline__ void mmvf_multi_row(const float* __restrict__ x, const uint16_t* __restrict__ row, int n_in, + int n_tok, float (&acc)[kMmvfMaxRows]) { + const int t = threadIdx.x; + const uint32_t* weights2 = reinterpret_cast(row); + __shared__ float partials[kMmvfMaxRows][32]; + if constexpr (BLOCK_SIZE > 32) { + if (t < 32) + for (int k = 0; k < kMmvfMaxRows; ++k) partials[k][t] = 0.0f; + __syncthreads(); + } +#pragma unroll + for (int k = 0; k < kMmvfMaxRows; ++k) acc[k] = 0.0f; + for (int pair = t; pair < n_in / 2; pair += BLOCK_SIZE) { + const uint32_t weight = weights2[pair]; + const float w0 = f32_from_bf16((uint16_t) weight), w1 = f32_from_bf16((uint16_t) (weight >> 16)); +#pragma unroll + for (int k = 0; k < kMmvfMaxRows; ++k) { + if (k >= n_tok) break; + const float2 input = reinterpret_cast(x + (size_t) k * n_in)[pair]; + acc[k] = __fmaf_rn(w0, input.x, acc[k]); + acc[k] = __fmaf_rn(w1, input.y, acc[k]); + } + } +#pragma unroll + for (int k = 0; k < kMmvfMaxRows; ++k) { + if (k >= n_tok) break; + acc[k] = mmvf_warp_sum(acc[k]); + } + if constexpr (BLOCK_SIZE > 32) { + if ((t & 31) == 0) + for (int k = 0; k < n_tok; ++k) partials[k][t / 32] = acc[k]; + __syncthreads(); + if (t < 32) { +#pragma unroll + for (int k = 0; k < kMmvfMaxRows; ++k) { + if (k >= n_tok) break; + acc[k] = mmvf_warp_sum(partials[k][t]); + } + } + } +} + +// mmvf_multi_row for R consecutive rows of one matrix (the first `rows` of them real): each pair's inputs are loaded +// once for the R rows, and every row's sums and their order are mmvf_multi_row's; thread 0 ends with acc[r][0 .. n_tok). +template +__device__ __forceinline__ void mmvf_multi_rows(const float* __restrict__ x, const uint16_t* __restrict__ row0, int n_in, + int n_tok, int rows, float (&acc)[R][kMmvfMaxRows]) { + const int t = threadIdx.x; + __shared__ float partials[R][kMmvfMaxRows][32]; + if constexpr (BLOCK_SIZE > 32) { + if (t < 32) +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int k = 0; k < kMmvfMaxRows; ++k) partials[r][k][t] = 0.0f; + __syncthreads(); + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int k = 0; k < kMmvfMaxRows; ++k) acc[r][k] = 0.0f; + for (int pair = t; pair < n_in / 2; pair += BLOCK_SIZE) { + float w0[R], w1[R]; +#pragma unroll + for (int r = 0; r < R; ++r) { + const uint32_t weight = r < rows ? reinterpret_cast(row0 + (size_t) r * n_in)[pair] : 0u; + w0[r] = f32_from_bf16((uint16_t) weight); + w1[r] = f32_from_bf16((uint16_t) (weight >> 16)); + } +#pragma unroll + for (int k = 0; k < kMmvfMaxRows; ++k) { + if (k >= n_tok) break; + const float2 input = reinterpret_cast(x + (size_t) k * n_in)[pair]; +#pragma unroll + for (int r = 0; r < R; ++r) { + acc[r][k] = __fmaf_rn(w0[r], input.x, acc[r][k]); + acc[r][k] = __fmaf_rn(w1[r], input.y, acc[r][k]); + } + } + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int k = 0; k < kMmvfMaxRows; ++k) { + if (k >= n_tok) break; + acc[r][k] = mmvf_warp_sum(acc[r][k]); + } + if constexpr (BLOCK_SIZE > 32) { + if ((t & 31) == 0) +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int k = 0; k < kMmvfMaxRows; ++k) { + if (k >= n_tok) break; + partials[r][k][t / 32] = acc[r][k]; + } + __syncthreads(); + if (t < 32) { +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int k = 0; k < kMmvfMaxRows; ++k) { + if (k >= n_tok) break; + acc[r][k] = mmvf_warp_sum(partials[r][k][t]); + } + } + } +} + +} // namespace +} // namespace strata::kernels diff --git a/src/kernels/cuda/native_bf16.cu b/src/kernels/cuda/native_bf16.cu index 1390a6abecf..069cc212099 100644 --- a/src/kernels/cuda/native_bf16.cu +++ b/src/kernels/cuda/native_bf16.cu @@ -1,5 +1,6 @@ #include "strata/kernels/bf16_gemv.hpp" #include "strata/kernels/bf16_bits.hpp" +#include "mmvf_multi.cuh" #include #include @@ -33,13 +34,6 @@ namespace { // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. -__device__ __forceinline__ float mmvf_warp_sum(float value) { -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - value += __shfl_xor_sync(0xffffffffu, value, offset, 32); - return value; -} - template __global__ void bf16_f32_mmvf_kernel(const float* __restrict__ x, const uint16_t* __restrict__ w, float* __restrict__ y, int n_in) { @@ -70,52 +64,25 @@ __global__ void bf16_f32_mmvf_kernel(const float* __restrict__ x, const uint16_t if (t == 0) y[blockIdx.x] = acc; } -// the same kernel for up to 8 activation rows - the weight row is read ONCE and every -// token keeps its own accumulator with exactly the single-row kernel's order (pairs, two ordered FMAs, the same warp -// and block reductions), so each output is bit-identical to a bf16_f32_mmvf_kernel launch of its own. -template -__global__ void bf16_f32_mmvf_multi_kernel(const float* __restrict__ x, int64_t ldx, const uint16_t* __restrict__ w, - float* __restrict__ y, int64_t ldy, int n_in, int n_tok) { - const int t = threadIdx.x; - const uint16_t* row = w + (size_t) blockIdx.x * n_in; - const uint32_t* weights2 = reinterpret_cast(row); - __shared__ float partials[NT][32]; - if constexpr (BLOCK_SIZE > 32) { - if (t < 32) -#pragma unroll - for (int k = 0; k < NT; ++k) partials[k][t] = 0.0f; - __syncthreads(); - } - float acc[NT]; -#pragma unroll - for (int k = 0; k < NT; ++k) acc[k] = 0.0f; - for (int pair = t; pair < n_in / 2; pair += BLOCK_SIZE) { - const uint32_t weight = weights2[pair]; - const float w0 = f32_from_bf16((uint16_t) weight), w1 = f32_from_bf16((uint16_t) (weight >> 16)); -#pragma unroll - for (int k = 0; k < NT; ++k) { - if (k < n_tok) { - const float2 input = reinterpret_cast(x + (size_t) k * ldx)[pair]; - acc[k] = __fmaf_rn(w0, input.x, acc[k]); - acc[k] = __fmaf_rn(w1, input.y, acc[k]); - } - } - } -#pragma unroll - for (int k = 0; k < NT; ++k) acc[k] = mmvf_warp_sum(acc[k]); - if constexpr (BLOCK_SIZE > 32) { - if ((t & 31) == 0) -#pragma unroll - for (int k = 0; k < NT; ++k) partials[k][t / 32] = acc[k]; - __syncthreads(); - if (t < 32) +// `bf16_f32_mmvf_kernel` for up to kMultiT activation rows, kRows weight rows a block (mmvf_multi_rows); the blocks +// past n_out's take the rows of a second matrix of the same input +constexpr int kMultiT = kMmvfMaxRows, kRows = 4; +template +__global__ void bf16_f32_mmvf_multi_kernel(const float* __restrict__ x, const uint16_t* __restrict__ w, + float* __restrict__ y, int n_in, int n_out, int n_tok, + const uint16_t* __restrict__ w2, float* __restrict__ y2, int n_out2) { + const int blocks1 = (n_out + kRows - 1) / kRows; + const bool first = (int) blockIdx.x < blocks1; + const int row0 = (first ? (int) blockIdx.x : (int) blockIdx.x - blocks1) * kRows, n = first ? n_out : n_out2; + const int rows = min(kRows, n - row0); + float acc[kRows][kMultiT]; + mmvf_multi_rows(x, (first ? w : w2) + (size_t) row0 * n_in, n_in, n_tok, rows, acc); + if (threadIdx.x == 0) #pragma unroll - for (int k = 0; k < NT; ++k) acc[k] = mmvf_warp_sum(partials[k][t]); - } - if (t == 0) + for (int r = 0; r < kRows; ++r) #pragma unroll - for (int k = 0; k < NT; ++k) - if (k < n_tok) y[(size_t) k * ldy + blockIdx.x] = acc[k]; + for (int k = 0; k < kMultiT; ++k) + if (r < rows && k < n_tok) (first ? y : y2)[(size_t) k * n + row0 + r] = acc[r][k]; } int mmvf_block_size(int64_t n_in) { @@ -133,26 +100,6 @@ int mmvf_block_size(int64_t n_in) { } // namespace -void bf16_gemv_fp32_mmvf_multi(const float* x, int64_t ldx, const uint16_t* w, float* y, int64_t ldy, - int64_t n_in, int64_t n_out, int n_tok, void* stream) { - if (n_tok == 1 && ldy >= n_out) { bf16_gemv_fp32_mmvf(x, w, y, n_in, n_out, stream); return; } - if (n_tok < 1 || n_tok > 8 || n_in <= 0 || (n_in & 1) != 0 || n_out <= 0 || (ldx & 1) != 0 || x == nullptr || - w == nullptr || y == nullptr || (reinterpret_cast(x) & 7u) != 0) - throw std::invalid_argument("bf16_gemv_fp32_mmvf_multi: 1..8 rows, even n_in/ldx, aligned pointers"); - const cudaStream_t st = (cudaStream_t) stream; -#define STRATA_MMVF_M(N) case N: \ - if (n_tok <= 4) bf16_f32_mmvf_multi_kernel<<<(unsigned) n_out, N, 0, st>>>(x, ldx, w, y, ldy, (int) n_in, n_tok); \ - else bf16_f32_mmvf_multi_kernel<<<(unsigned) n_out, N, 0, st>>>(x, ldx, w, y, ldy, (int) n_in, n_tok); break - switch (mmvf_block_size(n_in)) { - STRATA_MMVF_M(32); STRATA_MMVF_M(64); STRATA_MMVF_M(96); STRATA_MMVF_M(128); - STRATA_MMVF_M(160); STRATA_MMVF_M(192); STRATA_MMVF_M(224); STRATA_MMVF_M(256); - } -#undef STRATA_MMVF_M - const cudaError_t result = cudaGetLastError(); - if (result != cudaSuccess) - throw std::runtime_error(std::string("bf16_gemv_fp32_mmvf_multi launch: ") + cudaGetErrorString(result)); -} - void bf16_gemv_fp32_mmvf(const float* x, const uint16_t* w, float* y, int64_t n_in, int64_t n_out, void* stream) { if (n_in <= 0 || (n_in & 1) != 0 || n_in > std::numeric_limits::max() || @@ -183,4 +130,41 @@ void bf16_gemv_fp32_mmvf(const float* x, const uint16_t* w, float* y, } +void bf16_gemv_fp32_mmvf_multi(const float* x, const uint16_t* w, float* y, int64_t n_in, int64_t n_out, int n_tok, + void* stream, const uint16_t* w2, float* y2, int64_t n_out2) { + if (n_tok < 1 || n_tok > kMultiT) + throw std::invalid_argument("bf16_gemv_fp32_mmvf_multi: 1..8 activation rows"); + if (n_in <= 0 || (n_in & 1) != 0 || n_in > std::numeric_limits::max() || + n_out <= 0 || n_out2 < 0 || n_out + n_out2 > std::numeric_limits::max()) + throw std::invalid_argument("bf16_gemv_fp32_mmvf_multi: require positive even n_in and positive n_out <= INT_MAX"); + if (x == nullptr || w == nullptr || y == nullptr || (n_out2 > 0 && (w2 == nullptr || y2 == nullptr)) || + (reinterpret_cast(x) & 7u) != 0 || + (reinterpret_cast(w) & 3u) != 0 || (reinterpret_cast(w2) & 3u) != 0 || + (reinterpret_cast(y) & 3u) != 0 || (reinterpret_cast(y2) & 3u) != 0) + throw std::invalid_argument("bf16_gemv_fp32_mmvf_multi: null or misaligned pointer"); + if (n_tok == 1 && n_out2 == 0) { // the same sums: a row a block streams the weights faster + bf16_gemv_fp32_mmvf(x, w, y, n_in, n_out, stream); + return; + } + const cudaStream_t st = (cudaStream_t) stream; + const unsigned blocks = (unsigned) ((n_out + kRows - 1) / kRows + (n_out2 + kRows - 1) / kRows); +#define STRATA_MMVF_CASE(N) case N: \ + bf16_f32_mmvf_multi_kernel<<>>(x, w, y, (int) n_in, (int) n_out, n_tok, w2, y2, (int) n_out2); \ + break + switch (mmvf_block_size(n_in)) { + STRATA_MMVF_CASE(32); + STRATA_MMVF_CASE(64); + STRATA_MMVF_CASE(96); + STRATA_MMVF_CASE(128); + STRATA_MMVF_CASE(160); + STRATA_MMVF_CASE(192); + STRATA_MMVF_CASE(224); + STRATA_MMVF_CASE(256); + } +#undef STRATA_MMVF_CASE + const cudaError_t result = cudaGetLastError(); + if (result != cudaSuccess) + throw std::runtime_error(std::string("bf16_gemv_fp32_mmvf_multi launch: ") + cudaGetErrorString(result)); +} + } // namespace strata::kernels diff --git a/src/kernels/cuda/native_gr_norm.cu b/src/kernels/cuda/native_gr_norm.cu index 5abe013ecf3..9a1d1593dc1 100644 --- a/src/kernels/cuda/native_gr_norm.cu +++ b/src/kernels/cuda/native_gr_norm.cu @@ -45,11 +45,11 @@ __device__ __forceinline__ float norm_warp_sum(float value) { template __global__ void weighted_rms_norm(const float* __restrict__ input, const float* __restrict__ gamma, - float* __restrict__ output, int n_cols, float epsilon) { + float* __restrict__ output, int n_cols, int gamma_rows, float epsilon) { const int tid = threadIdx.x; const std::size_t row_offset = std::size_t(blockIdx.x) * n_cols; input += row_offset; - gamma += row_offset; + gamma += std::size_t(blockIdx.x % gamma_rows) * n_cols; output += row_offset; float partial = 0.0f; for (int col = tid; col < n_cols; col += BlockSize) { @@ -84,16 +84,23 @@ void check_pointer(const void* p) { void native_gr_rms_norm_weighted(const float* input, const float* gamma, float* output, int n_cols, int n_rows, float epsilon, void* stream) { - if (n_cols <= 0 || n_rows <= 0 || !std::isfinite(epsilon) || epsilon < 0.0f) + native_gr_rms_norm_weighted_repeat(input, gamma, output, n_cols, n_rows, n_rows, epsilon, stream); +} + +void native_gr_rms_norm_weighted_repeat(const float* input, const float* gamma, float* output, + int n_cols, int n_rows, int gamma_rows, float epsilon, void* stream) { + if (n_cols <= 0 || n_rows <= 0 || gamma_rows <= 0 || !std::isfinite(epsilon) || epsilon < 0.0f) throw std::invalid_argument("native GR RMSNorm requires positive dimensions and finite nonnegative epsilon"); check_pointer(input); check_pointer(gamma); check_pointer(output); const auto cuda_stream = static_cast(stream); if (n_cols < 1024) - weighted_rms_norm<256><<>>(input, gamma, output, n_cols, epsilon); + weighted_rms_norm<256><<>>(input, gamma, output, n_cols, gamma_rows, + epsilon); else - weighted_rms_norm<1024><<>>(input, gamma, output, n_cols, epsilon); + weighted_rms_norm<1024><<>>(input, gamma, output, n_cols, gamma_rows, + epsilon); const auto error = cudaGetLastError(); if (error != cudaSuccess) throw std::runtime_error(std::string("native GR RMSNorm launch: ") + cudaGetErrorString(error)); diff --git a/src/kernels/cuda/native_mmvq.cu b/src/kernels/cuda/native_mmvq.cu index a7801fa7250..ae37c545ceb 100644 --- a/src/kernels/cuda/native_mmvq.cu +++ b/src/kernels/cuda/native_mmvq.cu @@ -25,10 +25,12 @@ #include "strata/kernels/native_mmvq.hpp" #include "strata/kernels/iq_kernels.hpp" +#include "q8_1_il.cuh" #include #include +#include #include #include #include @@ -1076,6 +1078,292 @@ void launch_multi(const void* weights, const void* x_q8_1, float* y, int n_in, i } } +// ============================ plan v0.3 P6: ncols = 2..8 from interleaved activations ============================ +// +// `native_quantize_q8_1_il` also writes the columns interleaved (native_mmvq.hpp), so one load reads the same int of +// every column. A warp takes R rows and walks their blocks in the chunks of the exact layout above - its lane l's chunk +// m is what thread 32 (m % 4) + l does in iteration m / 4 - with a sum per virtual warp, added in warp order before the +// xor tree; each weight block is decoded once (the traits' `load`) and every column takes the single-column dot's +// arithmetic, so each value is bitwise the multi-column kernel's. Formats whose lanes read consecutive ints of one q8_1 +// block read the block-major copy; IQ4_XS, whose lanes read the same int of consecutive blocks, the position-major one. +// +// Per format: the traits' blocks and `load`, and `apply` for NC columns (kby = the columns' first q8_1 block of the +// weight block): the single-column `apply`'s expression per column. +struct IlIQ4XS { + using F = IQ4XSTraits; + static constexpr bool PM = true; + template + __device__ static void apply(const F::W& r, const X& x, int kby, int iqs, float (&out)[NC]) { + const int b = kby + iqs / 4; + int sumi[NC]; +#pragma unroll + for (int c = 0; c < NC; ++c) sumi[c] = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + int u0[NC], u1[NC]; + x.u(b, j, u0); + x.u(b, j + 4, u1); +#pragma unroll + for (int c = 0; c < NC; ++c) { + sumi[c] = __dp4a(r.v[j].x, u0[c], sumi[c]); + sumi[c] = __dp4a(r.v[j].y, u1[c], sumi[c]); + } + } + float d8[NC]; + x.scales(b, d8); +#pragma unroll + for (int c = 0; c < NC; ++c) { + const int s = sumi[c] * (r.ls - 32); + const float d = r.dw * d8[c]; + out[c] = d * s; + } + } +}; +struct IlQ4K { + using F = Q4KTraits; + static constexpr bool PM = false; + template + __device__ static void apply(const F::W& r, const X& x, int kby, int iqs, float (&out)[NC]) { + int u[4][NC]; + float d8[2][NC]; +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int b = kby + r.bq8_offset + i; + x.u(b, (iqs / 2) % 4, u[2 * i]); + x.u(b, (iqs / 2) % 4 + 4, u[2 * i + 1]); + x.scales(b, d8[i]); + } + const uint8_t* sc = reinterpret_cast(r.aux); + const uint8_t* m = sc + 2; + const float2 dm4f = __half22float2(r.dm); +#pragma unroll + for (int c = 0; c < NC; ++c) { + float sumf_d = 0.0f; + float sumf_m = 0.0f; +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int v0i = (r.v[0] >> (4 * i)) & 0x0f0f0f0f; + const int v1i = (r.v[1] >> (4 * i)) & 0x0f0f0f0f; + const int dot1 = __dp4a(v1i, u[2 * i + 1][c], __dp4a(v0i, u[2 * i][c], 0)); + const int dot2 = __dp4a(0x01010101, u[2 * i + 1][c], __dp4a(0x01010101, u[2 * i][c], 0)); + sumf_d += d8[i][c] * (dot1 * sc[i]); + sumf_m += d8[i][c] * (dot2 * m[i]); + } + out[c] = dm4f.x * sumf_d - dm4f.y * sumf_m; + } + } +}; +struct IlQ5K { + using F = Q5KTraits; + static constexpr bool PM = false; + template + __device__ static void apply(const F::W& r, const X& x, int kby, int iqs, float (&out)[NC]) { + int u[4][NC]; + float d8[2][NC]; +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int b = kby + r.bq8_offset + i; + x.u(b, (iqs / 2) % 4, u[2 * i]); + x.u(b, (iqs / 2) % 4 + 4, u[2 * i + 1]); + x.scales(b, d8[i]); + } + const uint8_t* sc = reinterpret_cast(r.aux); + const uint8_t* m = sc + 2; + const float2 dm5f = __half22float2(r.dm); +#pragma unroll + for (int c = 0; c < NC; ++c) { + float sumf_d = 0.0f; + float sumf_m = 0.0f; +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int vl0i = (r.vl[0] >> (4 * i)) & 0x0f0f0f0f; + const int vl1i = (r.vl[1] >> (4 * i)) & 0x0f0f0f0f; + const int vh0i = ((r.vh[0] >> i) << 4) & 0x10101010; + const int vh1i = ((r.vh[1] >> i) << 4) & 0x10101010; + const int v0i = vl0i | vh0i; + const int v1i = vl1i | vh1i; + const int dot1 = __dp4a(v0i, u[2 * i][c], __dp4a(v1i, u[2 * i + 1][c], 0)); + const int dot2 = __dp4a(0x01010101, u[2 * i][c], __dp4a(0x01010101, u[2 * i + 1][c], 0)); + sumf_d += d8[i][c] * (dot1 * sc[i]); + sumf_m += d8[i][c] * (dot2 * m[i]); + } + out[c] = dm5f.x * sumf_d - dm5f.y * sumf_m; + } + } +}; +struct IlQ6K { + using F = Q6KTraits; + static constexpr bool PM = false; + template + __device__ static void apply(const F::W& r, const X& x, int kby, int iqs, float (&out)[NC]) { + int u[2][NC]; + float d8[2][NC]; +#pragma unroll + for (int i = 0; i < 2; ++i) { + x.u(kby + r.bq8_offset + 2 * i, iqs % 8, u[i]); + x.scales(kby + r.bq8_offset + 2 * i, d8[i]); + } + int vi[2], sc[2]; +#pragma unroll + for (int i = 0; i < 2; ++i) { + sc[i] = r.scales[4 * i]; + const int vil = (r.vl >> (4 * i)) & 0x0f0f0f0f; + const int vih = ((r.vh >> (4 * i)) << 4) & 0x30303030; + vi[i] = __vsubss4(vil | vih, 0x20202020); + } +#pragma unroll + for (int c = 0; c < NC; ++c) { + float sumf = 0.0f; +#pragma unroll + for (int i = 0; i < 2; ++i) sumf += d8[i][c] * (__dp4a(vi[i], u[i][c], 0) * sc[i]); + out[c] = r.d * sumf; + } + } +}; +// Q8_0: its traits keep the block pointer (`small_q8_dot` loads it per column); here the codes and scale load once. +struct IlQ80 { + using F = SmallTraits; + static constexpr bool PM = false; + struct W { int v[2]; float d0; }; + __device__ static W load(const Q80Block* __restrict__ w, int iqs) { + W r; + r.v[0] = load_int_b2(w->qs, iqs); + r.v[1] = load_int_b2(w->qs, iqs + 1); + r.d0 = w->d; + return r; + } + template + __device__ static void apply(const W& r, const X& x, int kby, int iqs, float (&out)[NC]) { + int u0[NC], u1[NC]; + x.u(kby, iqs, u0); + x.u(kby, iqs + 1, u1); + float d8[NC]; + x.scales(kby, d8); +#pragma unroll + for (int c = 0; c < NC; ++c) { + int sumi = 0; + sumi = __dp4a(r.v[0], u0[c], sumi); + sumi = __dp4a(r.v[1], u1[c], sumi); + const float d0 = r.d0; + const float d1 = d8[c]; + out[c] = d0 * d1 * float(sumi); + } + } +}; +template struct IlLoad { + using W = typename I::F::W; + __device__ static W load(const typename I::F::Block* w, int kqs) { return I::F::load(w, kqs); } +}; +template <> struct IlLoad { + using W = IlQ80::W; + __device__ static W load(const Q80Block* w, int kqs) { return IlQ80::load(w, kqs); } +}; + +template +__launch_bounds__(4 * WARP) +__global__ void native_mmvq_il_kernel(const typename I::F::Block* __restrict__ w, const void* __restrict__ xq, + const float* __restrict__ xd, float* __restrict__ y, int n_in, int n_out) { + using F = typename I::F; + constexpr int SUB = WARP / F::T; // weight blocks a chunk covers + const int lane = int(threadIdx.x) & (WARP - 1); + const auto x = q8_1_cols(xq, xd, n_in / Q8K); + const int bpr = n_in / F::DIV; + const int kqs = F::kqs(lane); + const int nm = (bpr + SUB - 1) / SUB; + // A grid-stride loop over the row groups, though the grid gives each warp one: without it ptxas gives several + // instances fewer registers, and Q6_K's at 3 columns run 10-50% slower (the rows-a-warp table is measured on it). + for (int grp = int(blockIdx.x) * 4 + (int(threadIdx.x) >> 5); grp * R < n_out; grp += int(gridDim.x) * 4) { + const int row0 = grp * R; + float acc[4][R][NC]; +#pragma unroll + for (int v = 0; v < 4; ++v) +#pragma unroll + for (int i = 0; i < R; ++i) +#pragma unroll + for (int c = 0; c < NC; ++c) acc[v][i][c] = 0.0f; + for (int m0 = 0; m0 < nm; m0 += 4) { +#pragma unroll + for (int v = 0; v < 4; ++v) { + const int kbx = SUB * (m0 + v) + lane / F::T; + if (kbx < bpr) { +#pragma unroll + for (int i = 0; i < R; ++i) { + const int row = min(row0 + i, n_out - 1); // a partial last group recomputes its last row + const typename IlLoad::W wv = IlLoad::load(w + std::size_t(row) * bpr + kbx, kqs); + float o[NC]; + I::template apply(wv, x, kbx * F::KBY, kqs, o); +#pragma unroll + for (int c = 0; c < NC; ++c) acc[v][i][c] += o[c]; + } + } + } + } +#pragma unroll + for (int i = 0; i < R; ++i) +#pragma unroll + for (int c = 0; c < NC; ++c) { + float s = acc[0][i][c]; + s += acc[1][i][c]; + s += acc[2][i][c]; + s += acc[3][i][c]; + s = warp_sum(s); + if (lane == 0 && row0 + i < n_out) y[std::size_t(c) * n_out + row0 + i] = s; + } + } +} + +template +void launch_il(const void* weights, const void* x_q8_1, const void* x_il, float* y, int n_in, int n_out, int ncols, + cudaStream_t s) { + const auto* w = static_cast(weights); + const unsigned blocks = unsigned(((n_out + R - 1) / R + 3) / 4); + if (ncols == 1) { + native_mmvq_il_kernel<<>>(w, x_q8_1, nullptr, y, n_in, n_out); + return; + } + const Q81IlParts parts = q8_1_il_parts(x_il, n_in, ncols); + const int* xq = I::PM ? parts.pm : parts.bm; + const float* xd = parts.d; + switch (ncols) { +#define STRATA_IL_CASE(N) case N: native_mmvq_il_kernel<<>>(w, xq, xd, y, n_in, n_out); break; + STRATA_IL_CASE(2) STRATA_IL_CASE(3) STRATA_IL_CASE(4) STRATA_IL_CASE(5) STRATA_IL_CASE(6) STRATA_IL_CASE(7) + STRATA_IL_CASE(8) +#undef STRATA_IL_CASE + default: throw std::invalid_argument("native_mmvq_il requires 1 <= ncols <= 8"); + } +} + +template +void launch_il_rows(int r, const void* weights, const void* x_q8_1, const void* x_il, float* y, int n_in, int n_out, + int ncols, cudaStream_t s) { + switch (r) { + case 1: launch_il(weights, x_q8_1, x_il, y, n_in, n_out, ncols, s); break; + case 4: launch_il(weights, x_q8_1, x_il, y, n_in, n_out, ncols, s); break; + default: launch_il(weights, x_q8_1, x_il, y, n_in, n_out, ncols, s); break; + } +} + +// native_quantize_q8_1_kernel plus the interleaved copy (native_mmvq.hpp): the same values in the same bytes +__launch_bounds__(QUANT_THREADS, 1) +__global__ void native_quantize_q8_1_il_kernel(const float* __restrict__ x, Q81Block* __restrict__ y, + int8_t* __restrict__ bm, int8_t* __restrict__ pm, float* __restrict__ dl, + int n_in, int n_total, int cp) { + const int i = int(blockIdx.x) * QUANT_THREADS + int(threadIdx.x); + if (i >= n_total) return; // n_in is a multiple of 32: only whole warps return. + const float xi = x[i]; + const float amax = warp_max(fabsf(xi)); + const float sum = warp_sum(xi); + const float d = amax / 127.0f; + const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); + const half2 ds = make_half2(d, sum); + y[i / Q8K].qs[i % Q8K] = q; + if (i % Q8K == 0) y[i / Q8K].ds = ds; + const int c = i / n_in, e = i - c * n_in, b = e / Q8K, p = (e % Q8K) / 4, nb = n_in / Q8K; + bm[(std::size_t(b * 8 + p) * cp + c) * 4 + e % 4] = q; + pm[(std::size_t(p * nb + b) * cp + c) * 4 + e % 4] = q; + if (e % Q8K == 0) dl[std::size_t(b) * cp + c] = __low2float(ds); +} + void validate_shape(int n_in, int ncols, int block_elems = Q8K) { if (n_in <= 0 || n_in % block_elems != 0) { throw std::invalid_argument("native MMVQ requires n_in > 0 and divisible by its block element count"); @@ -1489,4 +1777,85 @@ void native_mmvq(int ggml_type, const void* weights, const void* x_q8_1, float* } } +std::size_t native_q8_1_il_bytes(int n_in, int ncols) { + validate_shape(n_in, ncols); + return std::size_t(n_in / Q8K) * std::size_t(native_q8_1_il_cp(ncols)) * (16 * sizeof(int) + sizeof(float)); +} + +void native_quantize_q8_1_il(const float* x, void* x_q8_1, void* x_il, int n_in, int ncols, void* stream) { + validate_shape(n_in, ncols); + if (ncols < 2) throw std::invalid_argument("native_quantize_q8_1_il requires 2 <= ncols <= 8"); + validate_pointer(x); + validate_pointer(x_q8_1); + validate_pointer(x_il); + validate_stream(stream); + const int nb = n_in / Q8K, cp = native_q8_1_il_cp(ncols); + auto* bm = static_cast(x_il); + auto* pm = bm + std::size_t(nb) * 8 * cp * 4; + auto* dl = reinterpret_cast(pm + std::size_t(nb) * 8 * cp * 4); + const int n_total = n_in * ncols; + const unsigned blocks = unsigned((std::size_t(n_total) + QUANT_THREADS - 1) / QUANT_THREADS); + native_quantize_q8_1_il_kernel<<(stream)>>>( + x, static_cast(x_q8_1), bm, pm, dl, n_in, n_total, cp); + launch_check(); +} + +namespace { +// Rows a warp at 1-4 columns ([ncols - 1]) by the row count, in classes around the model's shapes: below 2048 (the +// shared expert's 640, the keys' 512), 4096 (the output projections' 2560), 8192 (the gate projection's 6144), 12288 +// (qkv's 10240), and more (the queries' 12288, the head's); 0: native_mmvq's kernels, which spread a row over four +// warps (IQ3_S's decode the weights once a column). The fastest in the verify window on the RTX 3090 +// (scratch/mm1_bench.cu alone, CUPTI traces beside the other branches); 5-8 columns take 2 rows from 2048. +struct IlRows { int type; uint8_t r[4][5]; }; +constexpr IlRows kIlRows[] = { + {23, {{0, 0, 0, 0, 0}, {0, 2, 2, 2, 2}, {0, 2, 2, 2, 2}, {0, 2, 2, 1, 2}}}, // IQ4_XS + {12, {{0, 0, 0, 0, 0}, {0, 2, 2, 2, 2}, {0, 2, 2, 2, 2}, {0, 2, 2, 2, 2}}}, // Q4_K + {13, {{0, 0, 0, 0, 0}, {0, 2, 1, 1, 2}, {0, 2, 2, 2, 2}, {0, 2, 2, 2, 2}}}, // Q5_K + {14, {{0, 0, 0, 0, 0}, {0, 1, 1, 1, 1}, {0, 1, 1, 1, 1}, {0, 2, 2, 1, 2}}}, // Q6_K + {8, {{0, 0, 0, 0, 0}, {0, 0, 4, 2, 1}, {0, 2, 2, 2, 2}, {0, 2, 2, 2, 2}}}, // Q8_0 + {21, {{1, 1, 1, 2, 2}, {1, 2, 4, 2, 2}, {1, 1, 1, 1, 1}, {1, 2, 2, 1, 1}}}, // IQ3_S +}; +int il_rows(int type, int ncols, int n_out) { + const int cls = n_out < 2048 ? 0 : n_out < 4096 ? 1 : n_out < 8192 ? 2 : n_out < 12288 ? 3 : 4; + for (const IlRows& e : kIlRows) + if (e.type == type) return ncols <= 4 ? e.r[ncols - 1][cls] : cls ? 2 : 0; + return 0; +} + +int g_tune_rows = 0; // native_mmvq_il_tune (tests, benchmarks): rows a warp for every shape (0: the table), +bool g_tune_one = false; // and 1 column through these kernels too +} // namespace +void native_mmvq_il_tune(int rows, bool one) { + g_tune_rows = rows; + g_tune_one = one; +} + +void native_mmvq_il(int ggml_type, const void* weights, const void* x_q8_1, const void* x_il, float* y, int n_in, + int n_out, int ncols, void* stream) { + const bool il = ggml_type == 23 || ggml_type == 12 || ggml_type == 13 || ggml_type == 14 || ggml_type == 8; + const int r = g_tune_rows && (ncols > 1 || g_tune_one) ? g_tune_rows : il_rows(ggml_type, ncols, n_out); + if (r == 0 || (!il && !iq_mmvq_il_supported(ggml_type))) { + native_mmvq(ggml_type, weights, x_q8_1, y, n_in, n_out, ncols, stream); + return; + } + if (!il) { + iq_mmvq_il(ggml_type, weights, x_q8_1, x_il, y, n_in, n_out, ncols, stream, r); + return; + } + validate_shape(n_in, ncols, ggml_type == 8 ? Q8K : QK); + validate_pointer(weights); + validate_pointer(ncols == 1 ? x_q8_1 : x_il); + validate_pointer(y); + validate_stream(stream); + const auto s = static_cast(stream); + switch (ggml_type) { + case 23: launch_il_rows(r, weights, x_q8_1, x_il, y, n_in, n_out, ncols, s); break; + case 12: launch_il_rows(r, weights, x_q8_1, x_il, y, n_in, n_out, ncols, s); break; + case 13: launch_il_rows(r, weights, x_q8_1, x_il, y, n_in, n_out, ncols, s); break; + case 14: launch_il_rows(r, weights, x_q8_1, x_il, y, n_in, n_out, ncols, s); break; + case 8: launch_il_rows(r, weights, x_q8_1, x_il, y, n_in, n_out, ncols, s); break; + } + launch_check(); +} + } // namespace strata::kernels diff --git a/src/kernels/cuda/native_moe.cu b/src/kernels/cuda/native_moe.cu index 7f1f91ce6c8..fec14c0fdf8 100644 --- a/src/kernels/cuda/native_moe.cu +++ b/src/kernels/cuda/native_moe.cu @@ -34,9 +34,6 @@ std::atomic enabled{false}; __global__ void combine(const float* __restrict__ parts, const float* __restrict__ weights, const float* __restrict__ shared, float* __restrict__ output, int64_t n_embd, int k) { - // blockIdx.y = the token of a multi-token launch (0 for the single one) - const int64_t tk = blockIdx.y; - parts += tk * k * n_embd; weights += tk * k; if (shared) shared += tk * n_embd; output += tk * n_embd; const int64_t col = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; if (col >= n_embd) return; float sum = parts[col] * weights[0]; @@ -46,6 +43,70 @@ __global__ void combine(const float* __restrict__ parts, const float* __restrict if (shared) sum += shared[col]; output[col] = sum; } +constexpr int kMaxK = 15; +// A shared expert's row times its scalar gate, in shared_expert_gate_rows' instructions (shared_expert.cu builds +// without fast math: nothing flushes to zero): sigmoid(logit) as __fdividef(1, 1 + __expf(-logit)), then the product. +// The explicit rounding keeps ptxas from fusing the product into the sum after it. +__device__ __forceinline__ float gated_shared(float row, float logit) { + float e, d, g, y; + const float one = 1.0f; + asm("mul.rn.f32 %0, %1, 0fBFB8AA3B;" : "=f"(e) : "f"(logit)); + asm("ex2.approx.f32 %0, %1;" : "=f"(d) : "f"(e)); + asm("add.rn.f32 %0, %1, %2;" : "=f"(e) : "f"(d), "f"(one)); + asm("div.approx.f32 %0, %1, %2;" : "=f"(g) : "f"(one), "f"(e)); + asm("mul.rn.f32 %0, %1, %2;" : "=f"(y) : "f"(row), "f"(g)); + return y; +} +// One block row per token. The rows a GPU computed are marked from the plan's entries first; each column then +// loads its k values (the host's over PCIe, all in flight at once) and sums them as `combine` does, then adds the +// shared expert's row, times its gate when `shared_gate` holds the gates' logits. +__global__ void gather_combine(const float* __restrict__ gpu_rows, const float* host_rows, + const int32_t* __restrict__ dst, const int32_t* __restrict__ count, + const int32_t* __restrict__ dst2, const int32_t* __restrict__ count2, + const int32_t* __restrict__ dst3, const int32_t* __restrict__ count3, + const float* __restrict__ weights, const float* __restrict__ shared, + const float* __restrict__ shared_gate, float* __restrict__ output, int64_t n_embd, + int k) { + __shared__ unsigned on_gpu, as_is; // bit j: row t*k + j is in gpu_rows, a hit / a row taken as it is + const int t = blockIdx.y; + if (threadIdx.x == 0) on_gpu = as_is = 0u; + __syncthreads(); + const int c = *count, c2 = dst2 != nullptr ? *count2 : 0, c3 = dst3 != nullptr ? *count3 : 0; + for (int i = threadIdx.x; i < c + c2 + c3; i += blockDim.x) { + const int r = (i < c ? dst[i] : i < c + c2 ? dst2[i - c] : dst3[i - c - c2]) - t * k; + if (r >= 0 && r < k) atomicOr(i < c + c2 ? &on_gpu : &as_is, 1u << r); + } + __syncthreads(); + const int64_t col = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (col >= n_embd) return; + const unsigned mask = on_gpu, same = as_is; + const float* w = weights + int64_t(t) * k; + float v[kMaxK]; +#pragma unroll + for (int j = 0; j < kMaxK; ++j) { + if (j >= k) break; + const int64_t at = (int64_t(t) * k + j) * n_embd + col; + if ((same >> j) & 1u) { + v[j] = gpu_rows[at]; + } else if ((mask >> j) & 1u) { + const float h = gpu_rows[at]; + v[j] = __float_as_uint(h) == 0x80000000u ? 0.0f : h; // 0 + h: the host's zeroed row plus the hit + } else { + v[j] = host_rows[at]; + } + } + float sum = v[0] * w[0]; +#pragma unroll + for (int j = 1; j < kMaxK; ++j) { + if (j >= k) break; + sum += v[j] * w[j]; + } + if (shared) { + const float s = shared[int64_t(t) * n_embd + col]; + sum += shared_gate ? gated_shared(s, shared_gate[t]) : s; + } + output[int64_t(t) * n_embd + col] = sum; +} bool valid_span(const void* p, size_t bytes) { const auto address = reinterpret_cast(p); return p && address % alignof(float) == 0 && bytes <= UINTPTR_MAX - address; @@ -74,12 +135,20 @@ void native_moe_combine(const float* parts, const float* weights, const float* s const auto error = cudaGetLastError(); if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error)); } -void native_moe_combine_multi(const float* parts, const float* weights, const float* shared, float* output, - int64_t n_embd, int64_t k, int n_tok, void* stream) { - if (!stream || n_embd <= 0 || k < 1 || k > 15 || n_tok < 1) - throw std::invalid_argument("native MoE combine (multi) requires a stream, width, 1..15 experts, tokens"); - combine<<(stream)>>>( - parts, weights, shared, output, n_embd, int(k)); +void native_moe_gather_combine(const float* gpu_rows, const float* host_rows, const int32_t* dst, + const int32_t* count, const int32_t* dst2, const int32_t* count2, const int32_t* dst3, + const int32_t* count3, const float* weights, const float* shared, + const float* shared_gate, float* output, int64_t n_embd, int64_t k, int n_tok, + void* stream) { + if (!stream || n_embd <= 0 || n_embd > std::numeric_limits::max() || k < 1 || k > kMaxK || n_tok < 1 || + !gpu_rows || !host_rows || !dst || !count || (dst2 != nullptr && count2 == nullptr) || + (dst3 != nullptr && count3 == nullptr) || !weights || !output || (shared_gate != nullptr && !shared)) + throw std::invalid_argument("native MoE gather-combine requires a stream, positive width, 1..15 experts " + "and its buffers"); + const dim3 grid(unsigned((n_embd + 255) / 256), unsigned(n_tok)); + gather_combine<<(stream)>>>(gpu_rows, host_rows, dst, count, dst2, count2, + dst3, count3, weights, shared, shared_gate, + output, n_embd, int(k)); const auto error = cudaGetLastError(); if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error)); } diff --git a/src/kernels/cuda/native_ple_postops.cu b/src/kernels/cuda/native_ple_postops.cu index 922ae0a812f..c80b371749c 100644 --- a/src/kernels/cuda/native_ple_postops.cu +++ b/src/kernels/cuda/native_ple_postops.cu @@ -86,76 +86,52 @@ __global__ void conv_residual_kernel(const float* history, const float* normaliz // Exact hidden/result alias is safe: each thread owns one element. result[c] = __fadd_rn(hidden[c], __fadd_rn(gated[c], activation)); } -// ---- the batch: T tokens, the same arithmetic per element as the kernels above -// weighted_rms_norm (native_gr_norm.cu) with the gamma row repeating every H rows (one token's H groups) -__device__ float norm_warp_sum(float value) { -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) value += __shfl_xor_sync(0xffffffffu, value, offset, 32); - return value; -} -__global__ void rms_rep_kernel(const float* __restrict__ input, const float* __restrict__ gamma, - float* __restrict__ output) { - constexpr int BlockSize = 1024; // native_gr_rms_norm_weighted's choice for 2560 columns - const int tid = threadIdx.x; - const size_t row_offset = size_t(blockIdx.x) * N; - input += row_offset; - output += row_offset; - gamma += size_t(blockIdx.x % H) * N; - float partial = 0.0f; - for (int col = tid; col < N; col += BlockSize) { - const float value = input[col]; - partial += value * value; - } - __shared__ float sums[32]; - partial = norm_warp_sum(partial); - const int lane = tid % 32; - if (lane == 0) sums[tid / 32] = partial; - __syncthreads(); - partial = 0.0f; - if (lane < BlockSize / 32) partial = sums[lane]; - partial = norm_warp_sum(partial); - const float mean = partial / N; - const float scale = rsqrtf(mean + NG_RMS_EPS); - for (int col = tid; col < N; col += BlockSize) output[col] = scale * input[col] * gamma[col]; -} -__global__ void broadcast_batch_kernel(const float* value, const float* gate, float* gated, int T) { +// n_tok tokens: token t's taps read the normalized rows of tokens t-9, t-6 and t-3 (history rows t, t+3 and t+6 +// before the chunk while those tokens precede it) and its own. +__global__ void broadcast_tokens_kernel(const float* value, const float* gate, float* gated, int n_tok) { const size_t i = size_t(blockIdx.x) * blockDim.x + threadIdx.x; - if (i >= size_t(T) * D) return; - const size_t t = i / D, d = i % D; - gated[i] = __fmul_rn(value[t * N + d % N], gate[t * H + d / N]); + if (i < size_t(n_tok) * D) gated[i] = __fmul_rn(value[i / D * N + i % N], gate[i / N]); } -// the dilated conv (taps 9, 6, 3 tokens back and this one) and the residual; a tap before the chunk reads the history -__global__ void conv_residual_batch_kernel(const float* history, const float* normalized, const uint16_t* weights, - float* hidden, const float* gated, int T) { +__global__ void conv_residual_tokens_kernel(const float* history, const float* normalized, + const uint16_t* weights, const float* hidden, + const float* gated, float* result, int n_tok) { const size_t i = size_t(blockIdx.x) * blockDim.x + threadIdx.x; - if (i >= size_t(T) * D) return; + if (i >= size_t(n_tok) * D) return; const int t = int(i / D), c = int(i % D); float sum = 0; #pragma unroll for (int k = 0; k < 4; ++k) { - const int p = t - 9 + 3 * k; // the token this tap reads (k == 3: this one) - const float x = p >= 0 ? normalized[size_t(p) * D + c] : history[size_t(c) * HISTORY + (9 + p)]; - const float wk = __half2float(__ushort_as_half(weights[c * 4 + k])); - const float term = __fmul_rn(x, wk); + const int j = t - HISTORY + 3 * k; + const float x = j >= 0 ? normalized[size_t(j) * D + c] : history[c * HISTORY + j + HISTORY]; + const float w = __half2float(__ushort_as_half(weights[c * 4 + k])); + const float term = __fmul_rn(x, w); sum = k == 0 ? term : __fadd_rn(sum, term); } const float activation = sum / (1.0f + expf(-sum)); - hidden[i] = __fadd_rn(hidden[i], __fadd_rn(gated[i], activation)); + result[i] = __fadd_rn(hidden[i], __fadd_rn(gated[i], activation)); } -// the history after the chunk: the last nine normalized rows (older ones from the history when T < 9) -__global__ void history_batch_kernel(float* history, const float* normalized, int T) { +// The history after the chunk: its last nine normalized rows, older rows moved up when the chunk is shorter; with +// `snaps`, the history after each of its tokens too. +__global__ void history_tokens_kernel(float* history, const float* normalized, int n_tok, float* snaps) { const int c = blockIdx.x * blockDim.x + threadIdx.x; if (c >= D) return; - float h[HISTORY]; + float old[HISTORY]; +#pragma unroll + for (int r = 0; r < HISTORY; ++r) old[r] = history[c * HISTORY + r]; + if (snaps != nullptr) + for (int t = 1; t <= n_tok; ++t) +#pragma unroll + for (int r = 0; r < HISTORY; ++r) { + const int j = t - HISTORY + r; + snaps[size_t(t - 1) * HISTORY * D + c * HISTORY + r] = + j >= 0 ? normalized[size_t(j) * D + c] : old[r + t]; + } #pragma unroll for (int r = 0; r < HISTORY; ++r) { - const int p = T - HISTORY + r; - h[r] = p >= 0 ? normalized[size_t(p) * D + c] : history[size_t(c) * HISTORY + (T + r)]; + const int j = n_tok - HISTORY + r; + history[c * HISTORY + r] = j >= 0 ? normalized[size_t(j) * D + c] : old[r + n_tok]; } -#pragma unroll - for (int r = 0; r < HISTORY; ++r) history[size_t(c) * HISTORY + r] = h[r]; } - struct Span { const void* p; size_t bytes; size_t alignment; }; bool overlaps(Span a, Span b) { const auto x = reinterpret_cast(a.p), y = reinterpret_cast(b.p); @@ -202,20 +178,45 @@ void native_ple_postops(const float* projected_key, const float* hidden, launch_check(); } -void native_ple_postops_batch(float* key, float* hidden, const float* value, float* history, const PleWeights& w, - float* query_norm, float* gated, float* gate, int T, void* stream) { - if (!stream || T <= 0 || !key || !hidden || !value || !history || !query_norm || !gated || !gate) - throw std::invalid_argument("native PLE postops batch: null input or empty batch"); +void native_ple_postops_tokens(const float* projected_key, const float* hidden, const float* value, float* history, + const PleWeights& w, const NativePleTokensBuffers& b, int n_tok, void* stream, + float* snaps) { + if (!stream || n_tok < 1) throw std::invalid_argument("native PLE postops: tokens need a stream and n_tok >= 1"); + const size_t rows = size_t(n_tok) * D * 4; + const Span inputs[] = {{projected_key,rows,4}, {hidden,rows,4}, {value,size_t(n_tok)*N*4,4}, + {history,HISTORY*D*4,4}, {w.norm_key,D*4,4}, {w.norm_query,D*4,4}, + {w.norm_conv,D*4,4}, {w.conv1d_f16,4*D*2,2}}; + const Span outputs[] = {{b.key,rows,4}, {b.query,rows,4}, {b.gate,size_t(n_tok)*H*4,4}, + {b.gated,rows,4}, {b.normalized,rows,4}, {b.result,rows,4}}; + for (const auto& span : inputs) validate(span); + for (const auto& span : outputs) validate(span); + for (size_t i = 0; i < 6; ++i) { + for (size_t j = 0; j < 8; ++j) + if (!(i == 5 && j == 1 && b.result == hidden) && overlaps(outputs[i], inputs[j])) + throw std::invalid_argument("native PLE postops output overlaps an input, the history or a weight"); + for (size_t j = i + 1; j < 6; ++j) + if (!(i == 1 && j == 4 && b.query == b.normalized) && overlaps(outputs[i], outputs[j])) + throw std::invalid_argument("native PLE postops writable spans overlap"); + } + if (snaps != nullptr) { + const Span s{snaps, size_t(n_tok) * HISTORY * D * 4, 4}; + validate(s); + for (const auto& span : inputs) + if (overlaps(s, span)) throw std::invalid_argument("native PLE postops: the snapshots overlap an input"); + for (const auto& span : outputs) + if (overlaps(s, span)) throw std::invalid_argument("native PLE postops: the snapshots overlap an output"); + } + const int R = H * n_tok; + const unsigned blocks = unsigned((size_t(n_tok) * D + 255) / 256); auto st = static_cast(stream); - const unsigned rows = unsigned(T) * H; - const unsigned blocks = unsigned((size_t(T) * D + 255) / 256); - rms_rep_kernel<<>>(key, w.norm_key, key); - rms_rep_kernel<<>>(hidden, w.norm_query, query_norm); - gate_kernel<<>>(key, query_norm, gate, 1.0f / std::sqrt(float(N))); - broadcast_batch_kernel<<>>(value, gate, gated, T); - rms_rep_kernel<<>>(gated, w.norm_conv, query_norm); - conv_residual_batch_kernel<<>>(history, query_norm, w.conv1d_f16, hidden, gated, T); - history_batch_kernel<<>>(history, query_norm, T); + native_gr_rms_norm_weighted_repeat(projected_key,w.norm_key,b.key,N,R,H,NG_RMS_EPS,stream); + native_gr_rms_norm_weighted_repeat(hidden,w.norm_query,b.query,N,R,H,NG_RMS_EPS,stream); + gate_kernel<<>>(b.key,b.query,b.gate,1.0f / std::sqrt(float(N))); + broadcast_tokens_kernel<<>>(value,b.gate,b.gated,n_tok); + launch_check(); + native_gr_rms_norm_weighted_repeat(b.gated,w.norm_conv,b.normalized,N,R,H,NG_RMS_EPS,stream); + conv_residual_tokens_kernel<<>>(history,b.normalized,w.conv1d_f16,hidden,b.gated,b.result,n_tok); + history_tokens_kernel<<>>(history,b.normalized,n_tok,snaps); launch_check(); } } // namespace strata::kernels diff --git a/src/kernels/cuda/native_qsa.cu b/src/kernels/cuda/native_qsa.cu index 32e2511c213..14a29f08641 100644 --- a/src/kernels/cuda/native_qsa.cu +++ b/src/kernels/cuda/native_qsa.cu @@ -23,6 +23,7 @@ // SOFTWARE. #include "strata/kernels/native_qsa.hpp" +#include "strata/kernels/mrope.hpp" #include #include #include @@ -66,6 +67,47 @@ __global__ void norm(const float* input, const float* __restrict__ gamma, float* for (std::size_t col = tid; col < std::size_t(n_cols); col += BlockSize) output[col] = scale * input[col] * gamma[col]; } +// `norm<256>`, then native_rope.cu's `apply` on its output, one block a row of at most 256 columns: the normalized +// row waits in shared memory. Row r is read at input + r * in_stride. +__global__ void norm_rope(const float* input, int in_stride, const float* __restrict__ gamma, float* output, + int n_cols, float epsilon, int n_rot, float theta_scale, const int* __restrict__ positions, + const int32_t* __restrict__ mtab, int heads, int pos_stride) { + constexpr int BlockSize = 256; + const int tid = threadIdx.x, row = blockIdx.x; + input += std::size_t(row) * in_stride; + output += std::size_t(row) * n_cols; + float partial = 0.0f; + for (std::size_t col = tid; col < std::size_t(n_cols); col += BlockSize) { + const float x = input[col]; + partial += x * x; + } + __shared__ float sums[32]; + __shared__ float values[BlockSize]; + partial = warp_sum(partial); + const int lane = tid % 32; + if (lane == 0) sums[tid / 32] = partial; + __syncthreads(); + partial = lane < BlockSize / 32 ? sums[lane] : 0.0f; + partial = warp_sum(partial); + const float mean = partial / n_cols; + const float scale = rsqrtf(mean + epsilon); + for (std::size_t col = tid; col < std::size_t(n_cols); col += BlockSize) + values[col] = scale * input[col] * gamma[col]; + __syncthreads(); // every read of input precedes this: output may be input + const int pair = tid; + if (pair >= n_cols / 2) return; + if (pair >= n_rot / 2) { + output[2 * pair] = values[2 * pair]; + output[2 * pair + 1] = values[2 * pair + 1]; + return; + } + const int position = positions[(row / heads) * pos_stride + row % heads]; + const float theta = mrope_pos(mtab, position, pair) * powf(theta_scale, float(pair)); + const float c = cosf(theta), s = sinf(theta); + const float a = values[pair], b = values[pair + n_rot / 2]; + output[pair] = a * c - b * s; + output[pair + n_rot / 2] = a * s + b * c; +} __global__ void gate(const float* attn, const float* __restrict__ q_full, float* output, int n_head, int head_dim) { const std::size_t i = std::size_t(blockIdx.x) * blockDim.x + threadIdx.x; @@ -117,6 +159,27 @@ void native_qsa_rms_norm_weighted(const float* input, const float* gamma, float* norm<1024><<(stream)>>>(input, gamma, output, n_cols, epsilon); check_launch(); } +void native_qsa_norm_rope_tokens(const float* input, int in_stride, const float* gamma, float* output, int n_cols, + int n_rows, float epsilon, int n_rot, float freq_base, const int* positions, int heads, + int pos_stride, void* stream) { + const auto count = elements(n_cols, n_rows); + if ((n_cols != 128 && n_cols != 256) || n_rot != 64 || in_stride < n_cols || heads < 1 || pos_stride < heads || + !std::isfinite(epsilon) || epsilon < 0.0f || !std::isfinite(freq_base) || freq_base <= 1.0f || !positions) + throw std::invalid_argument("native QSA norm-rope requires 128 or 256 columns, 64 rotated, valid strides, " + "epsilon and base"); + const std::size_t in_bytes = (std::size_t(n_rows - 1) * in_stride + n_cols) * 4; + if (!stream || !valid(input, in_bytes) || !valid(gamma, std::size_t(n_cols) * 4) || !valid(output, count * 4) || + (input != output && overlap(input, in_bytes, output, count * 4)) || + (input == output && in_stride != n_cols) || overlap(output, count * 4, gamma, std::size_t(n_cols) * 4)) + throw std::invalid_argument("native QSA norm-rope requires a stream, aligned spans, and disjoint buffers or an " + "exact input/output alias"); + // pinned host-side float powf, as native_rope_apply_tokens + const float theta_scale = powf(freq_base, -2.0f / n_rot); + norm_rope<<(stream)>>>( + input, in_stride, gamma, output, n_cols, epsilon, n_rot, theta_scale, positions, mrope_table(), heads, + pos_stride); + check_launch(); +} void native_qsa_gate_apply(const float* attn, const float* q_full, float* output, int n_head, int head_dim, void* stream) { const auto count = elements(head_dim, n_head); diff --git a/src/kernels/cuda/native_qsa_indexer.cu b/src/kernels/cuda/native_qsa_indexer.cu index 8a5ecbc3c9e..607cebd84be 100644 --- a/src/kernels/cuda/native_qsa_indexer.cu +++ b/src/kernels/cuda/native_qsa_indexer.cu @@ -39,12 +39,14 @@ __device__ float warp_sum(float x) { x += __shfl_xor_sync(0xffffffffu, x, offset); return x; } -__global__ void append(const float* __restrict__ raw, const int32_t* __restrict__ pos_dev, - int pos_base, const float* __restrict__ gamma, float epsilon, - float* __restrict__ tail, float* __restrict__ dead, - float* __restrict__ pooled, int32_t* __restrict__ block_pos, - int max_cells, float theta_scale, const int32_t* __restrict__ mtab) { - const int pos = *pos_dev, d = threadIdx.x; +// One token's append; every thread of the block calls it (it synchronizes the block), a later call sees this one's +// tail writes. +__device__ void append_one(const float* __restrict__ raw, const int pos, + int pos_base, const float* __restrict__ gamma, float epsilon, + float* __restrict__ tail, float* __restrict__ dead, + float* __restrict__ pooled, int32_t* __restrict__ block_pos, + int max_cells, float theta_scale, const int32_t* __restrict__ mtab) { + const int d = threadIdx.x; if (pos < 0 || pos >= max_cells) return; const int slot = pos % R; float incoming = 0.0f; @@ -77,21 +79,40 @@ __global__ void append(const float* __restrict__ raw, const int32_t* __restrict_ const float scale = rsqrtf(square_sum / D + epsilon); if (d < D) values[d] = scale * mean * gamma[d]; __syncthreads(); - if (d >= D) return; - const int b = pos / R; - const int rope_pos = pos == 0 ? 0 : pos_base + R * b; - float y = values[d]; - if (d < ROT) { - const int pair = d % (ROT / 2); - const float theta = (pos == 0 ? 0 : mrope_pos(mtab, rope_pos, pair)) * powf(theta_scale, float(pair)); - const float c = cosf(theta), s = sinf(theta); - const float a = values[pair], z = values[pair + ROT / 2]; - y = d < ROT / 2 ? a * c - z * s : a * s + z * c; + if (d < D) { + const int b = pos / R; + const int rope_pos = pos == 0 ? 0 : pos_base + R * b; + float y = values[d]; + if (d < ROT) { + const int pair = d % (ROT / 2); + const float theta = (pos == 0 ? 0 : mrope_pos(mtab, rope_pos, pair)) * powf(theta_scale, float(pair)); + const float c = cosf(theta), s = sinf(theta); + const float a = values[pair], z = values[pair + ROT / 2]; + y = d < ROT / 2 ? a * c - z * s : a * s + z * c; + } + pooled[std::size_t(b) * D + d] = y; + if (pos == 0) dead[d] = y; + else pooled[std::size_t(b + 1) * D + d] = dead[d]; + if (d == 0 && pos != 0) *block_pos = rope_pos; + } + __syncthreads(); // values and partials are reused by the next call +} +// With `tail_snap`, the tail as it was is copied there first. +__global__ void append(const float* __restrict__ raw, const int32_t* __restrict__ pos_dev, int pos_stride, + int n_tok, int pos_base, const float* __restrict__ gamma, float epsilon, + float* __restrict__ tail, float* __restrict__ dead, + float* __restrict__ pooled, int32_t* __restrict__ block_pos, + int max_cells, float theta_scale, const int32_t* __restrict__ mtab, + float* __restrict__ tail_snap) { + if (tail_snap != nullptr) { + for (int i = threadIdx.x; i < (R - 1) * D; i += blockDim.x) tail_snap[i] = tail[i]; + __syncthreads(); + } + for (int t = 0; t < n_tok; ++t) { + append_one(raw + std::size_t(t) * D, pos_dev[std::size_t(t) * pos_stride], pos_base, gamma, epsilon, tail, + dead, pooled, block_pos, max_cells, theta_scale, mtab); + __syncthreads(); // this token's tail and pooled writes before the next token reads them } - pooled[std::size_t(b) * D + d] = y; - if (pos == 0) dead[d] = y; - else pooled[std::size_t(b + 1) * D + d] = dead[d]; - if (d == 0 && pos != 0) *block_pos = rope_pos; } // ---- C-2: the batched append. The pooled key of a completed block b (its last cell pos = 4b+3), computed with // the single append's arithmetic: keys rounded through F16, summed tail[0]+tail[1]+tail[2]+incoming in that order, @@ -201,19 +222,35 @@ bool native_qsa_indexer_enabled() { return enabled.load(std::memory_order_relaxe void native_qsa_indexer_append(const float* raw, const int32_t* relative_pos_device, int32_t pos_base, const float* gamma, float epsilon, const QsaIndexerBuffers& b, const QsaShapes& s, int64_t max_cells, float freq_base, void* stream) { + native_qsa_indexer_append_multi(raw, relative_pos_device, 0, 1, pos_base, gamma, epsilon, b, s, max_cells, + freq_base, stream); +} +void native_qsa_indexer_append_multi(const float* raw, const int32_t* relative_pos_device, int pos_stride, + int n_tok, int32_t pos_base, const float* gamma, float epsilon, + const QsaIndexerBuffers& b, const QsaShapes& s, int64_t max_cells, + float freq_base, void* stream, float* tail_snap) { + if (n_tok < 1 || pos_stride < 0 || (n_tok > 1 && pos_stride < 1)) + throw std::invalid_argument("native QSA indexer: n_tok >= 1 appends with a positive position stride"); if (!stream || s.idx_dim != D || s.idx_block != R || s.n_rot != ROT || max_cells < 1 || max_cells > INT32_MAX || pos_base < 0 || pos_base % R || int64_t(pos_base) + max_cells > INT32_MAX || !std::isfinite(epsilon) || epsilon <= 0.0f || !std::isfinite(freq_base) || freq_base <= 1.0f) throw std::invalid_argument("native QSA indexer requires fixed geometry, aligned position base, positive capacity/epsilon, valid frequency and explicit stream"); - const Span spans[] = {{raw,D*4},{relative_pos_device,4},{gamma,D*4},{b.tail,(R-1)*D*4}, + const Span spans[] = {{raw,std::size_t(n_tok)*D*4},{relative_pos_device,(std::size_t(n_tok-1)*pos_stride+1)*4}, + {gamma,D*4},{b.tail,(R-1)*D*4}, {b.dead,D*4},{b.pooled,std::size_t(max_cells/R+1)*D*4},{b.block_pos,4}}; for (const auto& span : spans) validate(span); for (int i = 0; i < 7; ++i) for (int j = i + 1; j < 7; ++j) if (overlaps(spans[i], spans[j])) throw std::invalid_argument("native QSA indexer buffers overlap"); + if (tail_snap != nullptr) { + const Span snap{tail_snap, (R - 1) * D * 4}; + validate(snap); + for (const auto& span : spans) + if (overlaps(snap, span)) throw std::invalid_argument("native QSA indexer: the tail snapshot overlaps"); + } const float theta_scale = powf(freq_base, -2.0f / ROT); - append<<<1,THREADS,0,static_cast(stream)>>>(raw,relative_pos_device,pos_base,gamma,epsilon, - b.tail,b.dead,b.pooled,b.block_pos,int(max_cells),theta_scale,mrope_table()); + append<<<1,THREADS,0,static_cast(stream)>>>(raw,relative_pos_device,pos_stride,n_tok,pos_base,gamma, + epsilon,b.tail,b.dead,b.pooled,b.block_pos,int(max_cells),theta_scale,mrope_table(),tail_snap); const auto error = cudaGetLastError(); if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error)); } diff --git a/src/kernels/cuda/native_rope.cu b/src/kernels/cuda/native_rope.cu index 6dd3f6b590a..176765beeaf 100644 --- a/src/kernels/cuda/native_rope.cu +++ b/src/kernels/cuda/native_rope.cu @@ -39,10 +39,12 @@ bool overlaps(const void* a, size_t an, const void* b, size_t bn) { return x <= y ? y - x < an : x - y < bn; } __global__ void apply(const float* x, float* out, int rows, int width, - int n_rot, float theta_scale, const int* positions, const int32_t* mtab) { + int n_rot, float theta_scale, const int* positions, const int32_t* mtab, int heads, + int pos_stride) { const int row = blockIdx.y; const int pair = blockIdx.x * blockDim.x + threadIdx.x; if (row >= rows || pair >= width / 2) return; + const int position = positions[(row / heads) * pos_stride + row % heads]; const size_t start = size_t(row) * width; if (pair >= n_rot / 2) { if (x != out) { @@ -51,7 +53,7 @@ __global__ void apply(const float* x, float* out, int rows, int width, } return; } - const float theta = mrope_pos(mtab, positions[row], pair) * powf(theta_scale, float(pair)); + const float theta = mrope_pos(mtab, position, pair) * powf(theta_scale, float(pair)); const float c = cosf(theta), s = sinf(theta); const float a = x[start + pair], b = x[start + pair + n_rot / 2]; out[start + pair] = a * c - b * s; @@ -74,7 +76,11 @@ void native_rope_set_enabled(bool value) { enabled.store(value, std::memory_orde bool native_rope_enabled() { return enabled.load(std::memory_order_relaxed); } void native_rope_apply(const float* x, float* out, int rows, int head_dim, int n_rot, float freq_base, const int* positions, void* stream) { - if (!x || !out || !positions || !stream || rows < 1 || rows > 65535 || + native_rope_apply_tokens(x, out, rows, head_dim, n_rot, freq_base, positions, rows, rows, stream); +} +void native_rope_apply_tokens(const float* x, float* out, int rows, int head_dim, int n_rot, float freq_base, + const int* positions, int heads, int pos_stride, void* stream) { + if (!x || !out || !positions || !stream || rows < 1 || rows > 65535 || heads < 1 || pos_stride < heads || (head_dim != 128 && head_dim != 256) || n_rot != 64 || !std::isfinite(freq_base) || freq_base <= 1.0f || reinterpret_cast(x) % 4 || reinterpret_cast(out) % 4 || @@ -82,15 +88,17 @@ void native_rope_apply(const float* x, float* out, int rows, int head_dim, throw std::invalid_argument("native RoPE requires aligned F32 rows, width 128/256, rotation 64, valid base and explicit stream"); } const size_t bytes = size_t(rows) * head_dim * sizeof(float); + const size_t pos_bytes = (size_t((rows - 1) / heads) * pos_stride + heads) * sizeof(int); if ((x != out && overlaps(x, bytes, out, bytes)) || - overlaps(positions, size_t(rows) * sizeof(int), out, bytes) || - overlaps(positions, size_t(rows) * sizeof(int), x, bytes)) { + overlaps(positions, pos_bytes, out, bytes) || + overlaps(positions, pos_bytes, x, bytes)) { throw std::invalid_argument("native RoPE buffers partially overlap"); } // Match pinned host-side float powf before device fast powf/trigonometry. const float theta_scale = powf(freq_base, -2.0f / n_rot); apply<<(stream)>>>(x, out, rows, head_dim, n_rot, theta_scale, positions, mrope_table()); + static_cast(stream)>>>(x, out, rows, head_dim, n_rot, theta_scale, positions, mrope_table(), + heads, pos_stride); const auto error = cudaGetLastError(); if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error)); } diff --git a/src/kernels/cuda/native_router.cu b/src/kernels/cuda/native_router.cu index 3f8b69ce21d..64784dc6dec 100644 --- a/src/kernels/cuda/native_router.cu +++ b/src/kernels/cuda/native_router.cu @@ -21,6 +21,10 @@ // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. #include "strata/kernels/native_router.hpp" +#include "strata/kernels/bf16_bits.hpp" +#include "hit_plan.cuh" +#include "mmvf_multi.cuh" +#include #include #include #include @@ -41,18 +45,22 @@ __device__ __forceinline__ float warp_max(float value) { for (int mask = 16; mask; mask >>= 1) value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, mask, 32)); return value; } -__launch_bounds__(256, 1) -__global__ void route(const float* __restrict__ logits, int32_t* __restrict__ ids, - float* __restrict__ weights) { - // Preserve the pinned 32x8 block geometry; only row zero is active here. - // blockIdx.x = the token (a multi-token launch; 0 for the single one) - logits += (size_t) blockIdx.x * 512; ids += (size_t) blockIdx.x * 10; weights += (size_t) blockIdx.x * 10; - if (threadIdx.y != 0) return; - const int lane = threadIdx.x; +// A float as a uint32 in the same order (sign flipped for positives, all bits for negatives), and back. +__device__ __forceinline__ unsigned order_key(float f) { + const unsigned b = __float_as_uint(f); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} +__device__ __forceinline__ float order_key_float(unsigned k) { + return __uint_as_float((k & 0x80000000u) ? (k & 0x7fffffffu) : ~k); +} +// One token, one warp: 512 logits -> the top 10 ids and weights; ids2 and weights2 take copies when given. The logits +// load through the L2 (another block may have written them). Each round takes the largest probability and, among +// equal ones, the lowest expert (a lane's own best, then two warp reductions); only the winning lane rescans. +__device__ __forceinline__ void route_token(const float* __restrict__ logits, int32_t* __restrict__ ids, + float* __restrict__ weights, int32_t* ids2, float* weights2, int lane) { float values[16]; #pragma unroll - for (int i = 0; i < 16; ++i) values[i] = logits[lane + i * 32]; - __syncthreads(); + for (int i = 0; i < 16; ++i) values[i] = __ldcg(logits + lane + i * 32); float maximum = -INFINITY; #pragma unroll for (int i = 0; i < 16; ++i) maximum = max(maximum, values[i]); @@ -69,23 +77,29 @@ __global__ void route(const float* __restrict__ logits, int32_t* __restrict__ id values[i] *= reciprocal; if (__isnanf(values[i])) values[i] = -FLT_MAX; } + float lbest = values[0]; // the lane's largest value, the lowest index among equal ones + int li = 0; +#pragma unroll + for (int i = 1; i < 16; ++i) + if (values[i] > lbest) { lbest = values[i]; li = i; } float selected = 0.0f, selected_sum = 0.0f; +#pragma unroll 1 for (int rank = 0; rank < 10; ++rank) { - float best = values[0]; - int expert = lane; + const unsigned mine = order_key(lbest); + const unsigned top = __reduce_max_sync(0xffffffffu, mine); + const int expert = __reduce_min_sync(0xffffffffu, mine == top ? lane + li * 32 : 0x7fffffff); + const float best = order_key_float(top); + if ((expert & 31) == lane) { #pragma unroll - for (int i = 1; i < 16; ++i) { - if (values[i] > best) { best = values[i]; expert = lane + i * 32; } - } + for (int i = 0; i < 16; ++i) + if (i == li) values[i] = -INFINITY; + lbest = values[0]; + li = 0; #pragma unroll - for (int mask = 16; mask; mask >>= 1) { - const float other = __shfl_xor_sync(0xffffffffu, best, mask, 32); - const int other_id = __shfl_xor_sync(0xffffffffu, expert, mask, 32); - if (other > best || (other == best && other_id < expert)) { best = other; expert = other_id; } - } - if ((expert & 31) == lane) { - values[expert / 32] = -INFINITY; + for (int i = 1; i < 16; ++i) + if (values[i] > lbest) { lbest = values[i]; li = i; } ids[rank] = expert; + if (ids2 != nullptr) ids2[rank] = expert; // Deliberately accumulate by WINNING EXPERT lane, not output rank. // Multiple selected experts in one lane add in selection order. selected_sum += best; @@ -94,7 +108,182 @@ __global__ void route(const float* __restrict__ logits, int32_t* __restrict__ id } selected_sum = max(warp_sum(selected_sum), 6.103515625e-5f); const float inverse_selected_sum = 1.0f / selected_sum; - if (lane < 10) weights[lane] = selected * inverse_selected_sum; + if (lane < 10) { + const float wv = selected * inverse_selected_sum; + weights[lane] = wv; + if (weights2 != nullptr) weights2[lane] = wv; + } +} +__launch_bounds__(256, 1) +__global__ void route(const float* __restrict__ logits, int32_t* __restrict__ ids, + float* __restrict__ weights) { + // Preserve the pinned 32x8 block geometry; only row zero is active here. Block b routes token b. + if (threadIdx.y != 0) return; + route_token(logits + blockIdx.x * 512, ids + blockIdx.x * 10, weights + blockIdx.x * 10, nullptr, nullptr, + threadIdx.x); +} + +// ================================ plan v0.3 P6: the verify window's router in one kernel ================================ +// +// For T tokens: the first blocks (scheduled first) write the input rows' other forms, and fence; then a block a row +// of the logits (bf16_gemv_fp32_mmvf_multi's 256-thread block, mmvf_multi_row). The last block to finish routes the +// tokens (`route_token`, a warp each); the routing warps, which wrote the mapped ids and weights, fence; then the +// ring's number, and the plan of the main GPU's hits (hit_plan_block). In place of the gemv, the top 10, +// doorbell_publish and verify_hit_plan. +// +// The rows' forms: a mapped copy; or their q8_1 blocks for the GPUs' experts (quantize_q8_1_rows) and ggml's Q8_K +// blocks for the CPU pool (quantize_row_q8_K_ref), bitwise: the quantizers' float operations are written in PTX, which +// this file's fast-math flags (no denormals, approximate division) do not change. A block takes four Q8_K blocks (292 +// bytes each), so its rows go out in 16-byte stores, full PCIe transactions: a block each, written 4 bytes at a time, +// cost twice as much a byte. +constexpr int VR_THREADS = 256; // mmvf_block_size of n_embd a multiple of 512 +constexpr int kQ81Bytes = 36, kQ8KBytes = 292; +__device__ __forceinline__ float add_rn(float x, float y) { + float r; + asm("add.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(x), "f"(y)); + return r; +} +__device__ __forceinline__ float mul_rn(float x, float y) { + float r; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(x), "f"(y)); + return r; +} +__device__ __forceinline__ float div_rn(float x, float y) { + float r; + asm("div.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(x), "f"(y)); + return r; +} +// four superblocks of 256 values a block (their 4 x 292 bytes of Q8_K and 4 x 288 of q8_1 are whole 16-byte stores; +// a warp a q8_1 block of each): block b's q8_1 blocks to xq1 (and xq1_out), its Q8_K blocks to xk_out, its floats to +// x_out, each when given +constexpr int kGroup = 4; +__device__ __forceinline__ void quantize_rows(const VerifyRouterArgs& a, int b) { + __shared__ uint4 s_q1[kGroup * 8 * kQ81Bytes / 16]; + __shared__ uint4 s_qk[(kGroup * kQ8KBytes + 15) / 16]; + __shared__ unsigned s_amax[kGroup][8]; + __shared__ float s_max[kGroup][8]; + const int t = (int) threadIdx.x, lane = t & 31, warp = t >> 5, per = a.n_embd / 256; + const int groups = (per + kGroup - 1) / kGroup, tok = b / groups, sb0 = (b % groups) * kGroup; + const int ns = min(kGroup, per - sb0); + const float* xr = a.x + (size_t) tok * a.n_embd + (size_t) sb0 * 256; + float x[kGroup]; +#pragma unroll + for (int j = 0; j < kGroup; ++j) x[j] = j < ns ? xr[j * 256 + t] : 0.0f; +#pragma unroll + for (int j = 0; j < kGroup; ++j) { + if (j >= ns) break; + const unsigned ax = __float_as_uint(x[j]) & 0x7fffffffu; // |x|, ordered as an integer + // q8_1 (quantize_q8_1_kernel): d = amax / 127, q = round(x / d), the warp's sum in xor-shuffle order + const unsigned amax = __reduce_max_sync(0xffffffffu, ax); + float sum = x[j]; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) sum = add_rn(sum, __shfl_xor_sync(0xffffffffu, sum, o)); + const float d = div_rn(__uint_as_float(amax), 127.0f); + float q = 0.0f; + if (amax != 0u) { // roundf: half away from zero + const float v = div_rn(x[j], d), r = truncf(v); + q = fabsf(v - r) >= 0.5f ? r + copysignf(1.0f, v) : r; + } + uint8_t* q1 = reinterpret_cast(s_q1) + (j * 8 + warp) * kQ81Bytes; + q1[4 + lane] = (uint8_t) (int8_t) q; + if (lane == 0) *reinterpret_cast(q1) = __halves2half2(__float2half(d), __float2half(sum)); + // Q8_K: the superblock's first value of the largest magnitude, from each warp's + const int first = __reduce_min_sync(0xffffffffu, ax == amax ? lane : 32); + const float wmax = __shfl_sync(0xffffffffu, x[j], first & 31); + if (lane == 0) { + s_amax[j][warp] = amax; + s_max[j][warp] = wmax; + } + } + __syncthreads(); + // Q8_K: iscale = -127 / that value, q = min(127, nearest_int(iscale * x)), d = 1 / iscale, the sums of 16 quants + uint8_t* qk = reinterpret_cast(s_qk); +#pragma unroll + for (int j = 0; j < kGroup; ++j) { + if (j >= ns) break; + unsigned best = 0u; + float mx = 0.0f; + for (int w = 0; w < 8; ++w) + if (s_amax[j][w] > best) { + best = s_amax[j][w]; + mx = s_max[j][w]; + } + uint8_t* blk = qk + j * kQ8KBytes; + int qi = 0; + if (best != 0u) { + const float iscale = div_rn(-127.0f, mx); + qi = min(127, (int) (__float_as_uint(add_rn(mul_rn(iscale, x[j]), 12582912.0f)) & 0x007fffffu) - + 0x00400000); + if (t == 0) *reinterpret_cast(blk) = div_rn(1.0f, iscale); + } else if (t == 0) { + *reinterpret_cast(blk) = 0.0f; + } + blk[4 + t] = (uint8_t) (int8_t) qi; + int bsum = qi; +#pragma unroll + for (int o = 8; o > 0; o >>= 1) bsum += __shfl_xor_sync(0xffffffffu, bsum, o); // within the 16-lane halves + if ((lane & 15) == 0) reinterpret_cast(blk + 4 + 256)[t >> 4] = (int16_t) bsum; + } + __syncthreads(); + const size_t row1 = (size_t) per * 8 * kQ81Bytes, at1 = (size_t) sb0 * 8 * kQ81Bytes; + const int n1 = ns * 8 * kQ81Bytes / 16, nk = (ns * kQ8KBytes + 15) / 16; + if (t < n1) { + const uint4 v = s_q1[t]; + reinterpret_cast(a.xq1 + tok * row1 + at1)[t] = v; + if (a.xq1_out != nullptr) reinterpret_cast(a.xq1_out + tok * row1 + at1)[t] = v; + } + if (a.xk_out != nullptr && t < nk) + reinterpret_cast(a.xk_out + (size_t) tok * a.xk_stride + (size_t) sb0 * kQ8KBytes)[t] = s_qk[t]; + if (a.x_out != nullptr && t < ns * 64) + reinterpret_cast(a.x_out + (size_t) tok * a.n_embd + (size_t) sb0 * 256)[t] = + reinterpret_cast(xr)[t]; +} +template +__global__ void __launch_bounds__(VR_THREADS) verify_router_kernel(VerifyRouterArgs a, int n_copy) { + __shared__ bool s_last; + const int t = (int) threadIdx.x, lane = t & 31, warp = t >> 5, b = (int) blockIdx.x; + if (b < n_copy) { + if (a.xq1 != nullptr) { + quantize_rows(a, b); + } else { + for (int i = b * VR_THREADS + t; i < TT * a.n_embd / 4; i += n_copy * VR_THREADS) + reinterpret_cast(a.x_out)[i] = reinterpret_cast(a.x)[i]; + } + __threadfence_system(); // this thread's writes before the block counts as done + } else { + const int row = b - n_copy; + float acc[kMmvfMaxRows]; + mmvf_multi_row(a.x, a.w + (size_t) row * a.n_embd, a.n_embd, TT, acc); + if (t == 0) + for (int k = 0; k < TT; ++k) + a.logits[(size_t) k * a.n_expert + row] = a.bias ? acc[k] + a.bias[row] : acc[k]; + __threadfence(); + } + __syncthreads(); + if (t == 0) s_last = atomicAdd(a.counter, 1u) == gridDim.x - 1; + __syncthreads(); + if (!s_last) return; + if (warp < TT) { + route_token(a.logits + (size_t) warp * a.n_expert, a.ids + warp * 10, a.weights + warp * 10, + a.ids_out != nullptr ? a.ids_out + warp * 10 : nullptr, + a.w_out != nullptr ? a.w_out + warp * 10 : nullptr, lane); + __threadfence_system(); + } + __syncthreads(); + if (t == 0) { // warp 0 routed and fenced + *(volatile uint32_t*) a.seq = a.ring; + *a.counter = 0u; // for the next launch + } + if (a.plan != nullptr) + hit_plan_block(a.ids, TT * 10, 10, a.res, a.n_expert, a.slot_ptr, a.plan, a.cap, a.ptr_off); +} +__global__ void __launch_bounds__(512) bias_update(const float* __restrict__ logits, + const float* __restrict__ predicted, int n_tok, float scale, + float* __restrict__ bias) { + const int e = (int) threadIdx.x; + float d = 0.0f; + for (int t = 0; t < n_tok; ++t) d += logits[t * 512 + e] - predicted[t * 512 + e]; + bias[e] = fmaf(scale, d, bias[e]); } bool valid(const void* p, size_t bytes) { const auto address = reinterpret_cast(p); @@ -116,11 +305,42 @@ void native_router_top10(const float* logits, int32_t* ids, float* weights, void const auto error = cudaGetLastError(); if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error)); } +void verify_router(const VerifyRouterArgs& a, void* stream) { + if (!stream || a.n_tok < 1 || a.n_tok > 8 || a.n_expert != 512 || a.n_embd <= 0 || a.n_embd % 512 != 0 || + !a.x || !a.w || !a.logits || !a.ids || !a.weights || !a.seq || !a.counter || (a.plan && (!a.res || !a.slot_ptr))) + throw std::invalid_argument("verify_router: 1..8 tokens, 512 experts, n_embd a multiple of 512, and buffers"); + if ((a.xk_out != nullptr && + (a.xq1 == nullptr || a.xk_stride < (a.n_embd / 256 * kQ8KBytes + 15) / 16 * 16 || a.xk_stride % 16)) || + (a.xq1_out != nullptr && a.xq1 == nullptr)) + throw std::invalid_argument("verify_router: the Q8_K and mapped q8_1 rows come with the q8_1 rows, the Q8_K " + "rows a multiple of 16 bytes apart and at least a row's"); + const auto st = static_cast(stream); + const int n_copy = a.xq1 != nullptr ? a.n_tok * ((a.n_embd / 256 + kGroup - 1) / kGroup) + : a.x_out != nullptr ? (a.n_tok * a.n_embd / 4 + VR_THREADS - 1) / VR_THREADS + : 0; + switch (a.n_tok) { +#define STRATA_VR(T) \ + case T: verify_router_kernel<<>>(a, n_copy); break; + STRATA_VR(1) STRATA_VR(2) STRATA_VR(3) STRATA_VR(4) STRATA_VR(5) STRATA_VR(6) STRATA_VR(7) STRATA_VR(8) +#undef STRATA_VR + } + const auto error = cudaGetLastError(); + if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error)); +} +void verify_router_bias_update(const float* logits, const float* predicted, int n_tok, float* bias, void* stream) { + if (!stream || n_tok < 1 || n_tok > 8 || !logits || !predicted || !bias) + throw std::invalid_argument("verify_router_bias_update: 1..8 tokens and buffers"); + bias_update<<<1, 512, 0, static_cast(stream)>>>(logits, predicted, n_tok, 0.125f / (float) n_tok, + bias); + const auto error = cudaGetLastError(); + if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error)); +} void native_router_top10_multi(const float* logits, int32_t* ids, float* weights, int n_tok, void* stream) { - if (!stream || n_tok < 1 || !valid(logits, (size_t) n_tok * 512 * 4) || !valid(ids, (size_t) n_tok * 10 * 4) || - !valid(weights, (size_t) n_tok * 10 * 4)) - throw std::invalid_argument("native router (multi) requires a stream and aligned [n,512]/[n,10] buffers"); - route<<<(unsigned) n_tok, dim3(32, 8), 0, static_cast(stream)>>>(logits, ids, weights); + const size_t nl = size_t(n_tok) * 512 * 4, no = size_t(n_tok) * 10 * 4; + if (!stream || n_tok < 1 || !valid(logits, nl) || !valid(ids, no) || !valid(weights, no) + || overlap(logits, nl, ids, no) || overlap(logits, nl, weights, no) || overlap(ids, no, weights, no)) + throw std::invalid_argument("native router requires a stream, aligned spans, and disjoint outputs"); + route<<(stream)>>>(logits, ids, weights); const auto error = cudaGetLastError(); if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error)); } diff --git a/src/kernels/cuda/ple.cu b/src/kernels/cuda/ple.cu index 66bb4220618..092ca68b4c0 100644 --- a/src/kernels/cuda/ple.cu +++ b/src/kernels/cuda/ple.cu @@ -197,6 +197,7 @@ void ck(cudaError_t e, const char* what) { } // namespace void ple_set_native_bf16(bool enabled) { native_bf16 = enabled; } +bool ple_native_bf16() { return native_bf16; } void ple_set_native_postops(bool enabled) { native_postops = enabled; } bool ple_native_postops_enabled() { return native_postops; } @@ -228,9 +229,8 @@ void ple_block(const float* emb, const float* hidden, const float* hist_rows, co PleOut& out, void* scratch, void* stream) { const bool native_key = w.key_native_data != nullptr && w.key_bf16 == nullptr; if (native_key && (!emb || !hidden || !hist_rows || !out.result || !scratch || !stream || - !w.key_native_q8_1 || - (w.key_native_type != 42 && w.key_native_type != 18 && w.key_native_type != 23))) - throw std::invalid_argument("ple_block: native key requires Q2_0, IQ3_XXS or IQ4_XS weights, input/output, private scratch and explicit stream"); + !w.key_native_q8_1 || !native_mmvq_supported(w.key_native_type))) + throw std::invalid_argument("ple_block: native key requires native MMVQ weights, input/output, private scratch and explicit stream"); if (emb == nullptr || hidden == nullptr || hist_rows == nullptr || out.result == nullptr) return; const int n_embd = NG_N_EMBD, hc = NG_HC, hc_dim = NG_HC_DIM; static_assert(NG_N_EMBD == 2560 && NG_HC_DIM == 10240, "native PLE key geometry changed"); diff --git a/src/kernels/cuda/q8_1_il.cuh b/src/kernels/cuda/q8_1_il.cuh new file mode 100644 index 00000000000..49c7a5a1fb3 --- /dev/null +++ b/src/kernels/cuda/q8_1_il.cuh @@ -0,0 +1,99 @@ +// src/kernels/cuda/q8_1_il.cuh - the interleaved copy of 2..8 q8_1 activation columns (native_mmvq.hpp) as the +// multi-column dot products read it, for native_mmvq.cu and iq_kernels.cu. +#pragma once + +#include "strata/kernels/native_mmvq.hpp" + +#include +#include + +#include +#include + +namespace strata::kernels { +namespace { + +// NC columns of the block-major (PM false) or position-major copy: u(b, p, o) = int p of q8_1 block b of every column, +// scales(b, o) = their scales +template struct Q81Il { + static constexpr int CP = native_q8_1_il_cp(NC); + const int* __restrict__ q; + const float* __restrict__ d; + int nb; // blocks a column + __device__ __forceinline__ void u(int b, int p, int (&o)[NC]) const { + const int* s = q + std::size_t(PM ? p * nb + b : b * 8 + p) * CP; + if constexpr (CP == 2) { + const int2 a = *reinterpret_cast(s); + o[0] = a.x; + o[1] = a.y; + } else { + const int4 a = *reinterpret_cast(s); + o[0] = a.x; + if constexpr (NC > 1) o[1] = a.y; + if constexpr (NC > 2) o[2] = a.z; + if constexpr (NC > 3) o[3] = a.w; + if constexpr (NC > 4) { + const int4 e = *reinterpret_cast(s + 4); + o[4] = e.x; + if constexpr (NC > 5) o[5] = e.y; + if constexpr (NC > 6) o[6] = e.z; + if constexpr (NC > 7) o[7] = e.w; + } + } + } + __device__ __forceinline__ void scales(int b, float (&o)[NC]) const { + const float* s = d + std::size_t(b) * CP; + if constexpr (CP == 2) { + const float2 a = *reinterpret_cast(s); + o[0] = a.x; + o[1] = a.y; + } else { + const float4 a = *reinterpret_cast(s); + o[0] = a.x; + if constexpr (NC > 1) o[1] = a.y; + if constexpr (NC > 2) o[2] = a.z; + if constexpr (NC > 3) o[3] = a.w; + if constexpr (NC > 4) { + const float4 e = *reinterpret_cast(s + 4); + o[4] = e.x; + if constexpr (NC > 5) o[5] = e.y; + if constexpr (NC > 6) o[6] = e.z; + if constexpr (NC > 7) o[7] = e.w; + } + } + } +}; + +// One column's plain q8_1 blocks (36 bytes: the half2 scale and sum, then 32 codes) behind Q81Il's interface +struct Q81One { + const uint8_t* __restrict__ x; + __device__ __forceinline__ void u(int b, int p, int (&o)[1]) const { + o[0] = reinterpret_cast(x + std::size_t(b) * 36 + 4)[p]; + } + __device__ __forceinline__ void scales(int b, float (&o)[1]) const { + o[0] = __low2float(*reinterpret_cast(x + std::size_t(b) * 36)); + } +}; + +// The activations a kernel for NC columns reads: one column's q8_1 blocks, 2+ columns' interleaved copy (nb blocks a +// column) +template +__device__ __forceinline__ auto q8_1_cols(const void* xq, const float* xd, int nb) { + if constexpr (NC == 1) return Q81One{static_cast(xq)}; + else return Q81Il{static_cast(xq), xd, nb}; +} + +// The copy's parts inside the buffer native_q8_1_il_bytes sizes: bm, pm (int32) and d (float). +struct Q81IlParts { + const int* bm; + const int* pm; + const float* d; +}; +inline Q81IlParts q8_1_il_parts(const void* x_il, int n_in, int ncols) { + const std::size_t words = std::size_t(n_in / 32) * 8 * std::size_t(native_q8_1_il_cp(ncols)); + const int* bm = static_cast(x_il); + return {bm, bm + words, reinterpret_cast(bm + 2 * words)}; +} + +} // namespace +} // namespace strata::kernels diff --git a/src/kernels/cuda/qsa_decode_attn.cu b/src/kernels/cuda/qsa_decode_attn.cu index 9f7fb8842ad..3ef91c3c372 100644 --- a/src/kernels/cuda/qsa_decode_attn.cu +++ b/src/kernels/cuda/qsa_decode_attn.cu @@ -1,11 +1,11 @@ // src/kernels/cuda/qsa_decode_attn.cu - see include/strata/kernels/qsa_decode_attn.hpp. #include "strata/kernels/qsa_decode_attn.hpp" #include "strata/kernels/kv_q8.hpp" -#include "strata/kernels/kv_q4.hpp" #include #include +#include #include #include #include @@ -15,9 +15,10 @@ namespace { constexpr int HD = 256; // head_dim constexpr int G = 12; // query heads per KV head (24 / 2) -constexpr int CHUNK = 64; // cells per block constexpr int THREADS = 256; constexpr int WARPS = THREADS / 32; +constexpr int PT = 64; // cells per tile +constexpr int SPLIT_MAX = 20; // a decode query's splits at most __device__ __forceinline__ float warp_sum(float v) { #pragma unroll @@ -30,261 +31,783 @@ __device__ __forceinline__ float warp_max(float v) { return v; } -// 8 consecutive values of one cell's key or value row for KV head `kvh`, dimensions [d0, d0+8). -// Per-format bodies; `load8` below is the KV_MODE dispatcher. value=false is the K side, true the V side. -__device__ __forceinline__ void load8_f16(const QsaAttnPools& p, bool value, long long row, int d0, float* out) { - const uint16_t* base = (value ? p.v_pool : p.k_pool) + row * HD + d0; - const uint4 raw = *reinterpret_cast(base); - const __half2* h2 = reinterpret_cast(&raw); -#pragma unroll - for (int j = 0; j < 4; ++j) { - const float2 f = __half22float2(h2[j]); - out[2 * j] = f.x; - out[2 * j + 1] = f.y; - } -} -__device__ __forceinline__ void load8_q8(const QsaAttnPools& p, bool value, long long row, int d0, float* out) { - const int8_t* codes = (value ? p.v_q : p.k_q) + row * HD + d0; - const uint16_t sbits = (value ? p.v_scale : p.k_scale)[row * (HD / KV_Q8_GROUP) + d0 / KV_Q8_GROUP]; - const float sc = __half2float(__ushort_as_half(sbits)); - const uint2 raw = *reinterpret_cast(codes); - const int8_t* c = reinterpret_cast(&raw); -#pragma unroll - for (int j = 0; j < 8; ++j) out[j] = (float) c[j] * sc; -} -__device__ __forceinline__ void load8_q4(const QsaAttnPools& p, bool value, long long row, int d0, float* out) { - constexpr int bytes_per_head = (HD / QK4_0) * sizeof(block_q4_0); - const int b = d0 / QK4_0; - const int rem = d0 % QK4_0; - const block_q4_0* blk = reinterpret_cast((value ? p.v_q4 : p.k_q4) + row * bytes_per_head) + b; - const float d = __half2float(__ushort_as_half(blk->d)); - const int j = (rem == 0 || rem == 16) ? 0 : 8; - const uint8_t* bytes = blk->qs + j; - if (rem < 16) { -#pragma unroll - for (int k = 0; k < 8; ++k) out[k] = (float) ((int)(bytes[k] & 0x0F) - 8) * d; - } else { -#pragma unroll - for (int k = 0; k < 8; ++k) out[k] = (float) ((int)(bytes[k] >> 4) - 8) * d; - } -} -template -__device__ __forceinline__ void load8(const QsaAttnPools& p, bool value, long long row, int d0, float* out) { - if constexpr (KV_MODE == 0) load8_f16(p, value, row, d0, out); - else if constexpr (KV_MODE == 1) load8_q8(p, value, row, d0, out); - else if constexpr (KV_MODE == 3) { - // hybrid K8V4: both sides are defined - K unrotated INT8, V rotated Q4_0 - so a value=true call - // reads the Q4_0 pool instead of dereferencing the null v_q (no call site does today; PR review) - if (value) load8_q4(p, value, row, d0, out); - else load8_q8(p, value, row, d0, out); - } else load8_q4(p, value, row, d0, out); +// The decode's partial results, per query `stride` floats: [split][head][HD] sums of p * v (p relative to the split's +// maximum), then at `m` and `l` (offsets from `acc`) [split][head] maxima and sums of p. +struct Partials { + float* acc = nullptr; + long long m = 0, l = 0, stride = 0; + int splits = 0; +}; +// the cells of each of a query's splits: split k takes [k c, min((k + 1) c, n)), whole tiles while the cells last +__host__ __device__ __forceinline__ int split_cells(int n, int splits) { + const int c = (n + splits - 1) / splits; + return c > PT ? c : PT; } -template -__global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __restrict__ q, QsaAttnPools p, - const int32_t* __restrict__ ids, - const int32_t* __restrict__ step, int n_kv_heads, - int page_size, float scale, float* __restrict__ part_acc, - float* __restrict__ part_m, float* __restrict__ part_l, - int n_chunks, int cap = 0, long long scratch_stride = 0) { - // batched form: query blockIdx.z, with its own q row, selection, step and scratch - q += (size_t) blockIdx.z * (size_t) (n_kv_heads * G) * HD; - ids += (size_t) blockIdx.z * (size_t) cap; - step += (size_t) blockIdx.z * kStepCount; - part_acc += (size_t) blockIdx.z * (size_t) scratch_stride; - part_m += (size_t) blockIdx.z * (size_t) scratch_stride; - part_l += (size_t) blockIdx.z * (size_t) scratch_stride; - __shared__ __align__(16) float sq[G][HD]; // 12 KB: this KV head's query heads - __shared__ float sp[G][CHUNK]; // scores, then probabilities - __shared__ long long srow[CHUNK]; // pool row of each cell (page, kv head, slot) - const int n_ids = __ldg(step + kStepWidth); - const int chunk = blockIdx.x, kvh = blockIdx.y; +// ---- FP16 pools: a block per (KV head, query) walks all of the query's cells in tiles, online softmax; SPLIT: a +// block per (KV head, query, split) its split's cells, writing partials. Scores: thread (cell, quarter) dots its +// cell's 64 dimensions of that quarter with the 12 query heads, the quarters summed through shared memory. Values: +// thread (dimension pair, cell parity) accumulates its two dimensions of the 12 heads over every other cell; the +// parities are added at the end. +template +__global__ void __launch_bounds__(THREADS) attn_f16_kernel(const float* __restrict__ q, QsaAttnPools p, + const int32_t* __restrict__ ids, + const int32_t* __restrict__ steps, int n_kv_heads, + int page_size, float scale, float* __restrict__ attn, + int cap, Partials part) { + const int kvh = blockIdx.x; + const size_t qi = blockIdx.y; + const size_t head0 = qi * (size_t) (n_kv_heads * G) + (size_t) kvh * G; + q += head0 * HD; + if constexpr (!SPLIT) attn += head0 * HD; + ids += qi * (size_t) cap; + int n_ids = __ldg(steps + qi * kStepCount + kStepWidth); + if constexpr (SPLIT) { + const int c = split_cells(n_ids, part.splits), c0 = (int) blockIdx.z * c; + if (c0 >= n_ids) return; + ids += c0; + n_ids = min(c, n_ids - c0); + } + __shared__ __align__(16) float sq[G][HD]; // 12 KB: this KV head's query heads + __shared__ __align__(16) float spart[4][G][PT]; // 12 KB: the quarters' partial scores; the parity sums at the end + __shared__ __align__(16) float sp[PT][G]; // probabilities, cell-major + __shared__ long long srow[PT]; + __shared__ float s_alpha[G], s_m[G], s_l[G]; const int t = threadIdx.x, lane = t & 31, warp = t >> 5; - const int c0 = chunk * CHUNK; - const int n_here = min(CHUNK, n_ids - c0); - const int slot = kvh * n_chunks + chunk; - if (n_here <= 0) { - if (t < G) { part_m[slot * G + t] = -FLT_MAX; part_l[slot * G + t] = 0.0f; } - return; + for (int i = t; i < G * HD; i += THREADS) sq[i / HD][i % HD] = q[i]; + if (t < G) { s_m[t] = -FLT_MAX; s_l[t] = 0.0f; } + const int quarter = warp >> 1, cell = ((warp & 1) << 5) | lane; // scores + const int parity = t >> 7, d0 = 2 * (t & 127); // values + float acc0[G], acc1[G]; +#pragma unroll + for (int h = 0; h < G; ++h) acc0[h] = acc1[h] = 0.0f; + for (int c0 = 0; c0 < n_ids; c0 += PT) { + const int n_here = min(PT, n_ids - c0); + if (t < PT) { + long long r = -1; + if (t < n_here) { + const int id = ids[c0 + t]; + r = ((long long) p.page_table[id / page_size] * n_kv_heads + kvh) * page_size + id % page_size; + } + srow[t] = r; + } + __syncthreads(); + float s[G]; +#pragma unroll + for (int h = 0; h < G; ++h) s[h] = 0.0f; + if (cell < n_here) { + const long long row = srow[cell]; +#pragma unroll + for (int j = 0; j < 64; j += 16) { + const int d = quarter * 64 + j; + float kf[16]; + const uint4* src = reinterpret_cast(p.k_pool + row * HD + d); +#pragma unroll + for (int u = 0; u < 2; ++u) { + const uint4 raw = src[u]; + const __half2* h2 = reinterpret_cast(&raw); +#pragma unroll + for (int v = 0; v < 4; ++v) { + const float2 f = __half22float2(h2[v]); + kf[8 * u + 2 * v] = f.x; + kf[8 * u + 2 * v + 1] = f.y; + } + } +#pragma unroll + for (int h = 0; h < G; ++h) { + const float4* qh = reinterpret_cast(&sq[h][d]); + float a = s[h]; +#pragma unroll + for (int v = 0; v < 4; ++v) { + const float4 q4 = qh[v]; + a = fmaf(kf[4 * v], q4.x, a); + a = fmaf(kf[4 * v + 1], q4.y, a); + a = fmaf(kf[4 * v + 2], q4.z, a); + a = fmaf(kf[4 * v + 3], q4.w, a); + } + s[h] = a; + } + } + } +#pragma unroll + for (int h = 0; h < G; ++h) spart[quarter][h][cell] = s[h]; + __syncthreads(); + // online softmax: warp w takes heads w and w + 8, a lane cells lane and lane + 32 + for (int h = warp; h < G; h += WARPS) { + float a = -FLT_MAX, b = -FLT_MAX; + if (lane < n_here) a = (spart[0][h][lane] + spart[1][h][lane] + spart[2][h][lane] + spart[3][h][lane]) * scale; + if (lane + 32 < n_here) + b = (spart[0][h][lane + 32] + spart[1][h][lane + 32] + spart[2][h][lane + 32] + spart[3][h][lane + 32]) * scale; + const float m_old = s_m[h]; + const float m_new = fmaxf(m_old, warp_max(fmaxf(a, b))); + const float ea = lane < n_here ? __expf(a - m_new) : 0.0f; + const float eb = lane + 32 < n_here ? __expf(b - m_new) : 0.0f; + sp[lane][h] = ea; + sp[lane + 32][h] = eb; + const float l = warp_sum(ea + eb); + if (lane == 0) { + const float alpha = __expf(m_old - m_new); + s_alpha[h] = alpha; + s_l[h] = fmaf(s_l[h], alpha, l); + s_m[h] = m_new; + } + } + __syncthreads(); +#pragma unroll + for (int h = 0; h < G; ++h) { + acc0[h] *= s_alpha[h]; + acc1[h] *= s_alpha[h]; + } + for (int c = parity; c < n_here; c += 2) { + const long long r = srow[c]; + const float2 f = __half22float2(*reinterpret_cast(p.v_pool + r * HD + d0)); + const float v0 = f.x, v1 = f.y; + const float4* pc = reinterpret_cast(sp[c]); + const float4 pa = pc[0], pb = pc[1], pd = pc[2]; + const float pr[G] = {pa.x, pa.y, pa.z, pa.w, pb.x, pb.y, pb.z, pb.w, pd.x, pd.y, pd.z, pd.w}; +#pragma unroll + for (int h = 0; h < G; ++h) { + acc0[h] = fmaf(pr[h], v0, acc0[h]); + acc1[h] = fmaf(pr[h], v1, acc1[h]); + } + } + __syncthreads(); // srow, sp and s_alpha are rewritten by the next tile } - for (int i = t; i < G * HD; i += THREADS) sq[i / HD][i % HD] = q[(size_t) (kvh * G) * HD + i]; - if (t < CHUNK) { - long long r = -1; - if (t < n_here) { - const int cell = ids[c0 + t]; - const long long page = (long long) p.page_table[cell / page_size]; - r = (page * n_kv_heads + kvh) * page_size + (cell % page_size); + // the odd cells' sums into shared memory, added to the even ones' + float* odd = &spart[0][0][0]; // G x HD floats + if (parity == 1) { +#pragma unroll + for (int h = 0; h < G; ++h) { + odd[h * HD + d0] = acc0[h]; + odd[h * HD + d0 + 1] = acc1[h]; } - srow[t] = r; } __syncthreads(); - // scores: each warp takes cells warp, warp+8, ...; each lane holds 8 of the 256 dimensions. - for (int c = warp; c < CHUNK; c += WARPS) { - if (c >= n_here) { - if (lane < G) sp[lane][c] = -FLT_MAX; - continue; + if constexpr (SPLIT) { + float* base = part.acc + qi * part.stride; + const size_t hq = (size_t) blockIdx.z * n_kv_heads * G + (size_t) kvh * G; // the block's first head + if (parity == 0) { +#pragma unroll + for (int h = 0; h < G; ++h) { + base[(hq + h) * HD + d0] = acc0[h] + odd[h * HD + d0]; + base[(hq + h) * HD + d0 + 1] = acc1[h] + odd[h * HD + d0 + 1]; + } } - float k8[8]; - load8(p, false, srow[c], lane * 8, k8); + if (t < G) { + base[part.m + hq + t] = s_m[t]; + base[part.l + hq + t] = s_l[t]; + } + } else if (parity == 0) { #pragma unroll for (int h = 0; h < G; ++h) { - const float4 qa = *reinterpret_cast(&sq[h][lane * 8]); - const float4 qb = *reinterpret_cast(&sq[h][lane * 8 + 4]); - float s = k8[0] * qa.x + k8[1] * qa.y + k8[2] * qa.z + k8[3] * qa.w + - k8[4] * qb.x + k8[5] * qb.y + k8[6] * qb.z + k8[7] * qb.w; - s = warp_sum(s); - if (lane == 0) sp[h][c] = s * scale; + const float l = s_l[h]; + attn[(size_t) h * HD + d0] = l > 0.0f ? (acc0[h] + odd[h * HD + d0]) / l : 0.0f; + attn[(size_t) h * HD + d0 + 1] = l > 0.0f ? (acc1[h] + odd[h * HD + d0 + 1]) / l : 0.0f; } } - __syncthreads(); - // per-head chunk max and exp-sum: warp w handles heads w and w+8. - for (int h = warp; h < G; h += WARPS) { - const float a = sp[h][lane], b = sp[h][lane + 32]; - const float m = warp_max(fmaxf(a, b)); - const float ea = (lane < n_here) ? __expf(a - m) : 0.0f; - const float eb = (lane + 32 < n_here) ? __expf(b - m) : 0.0f; - sp[h][lane] = ea; - sp[h][lane + 32] = eb; - const float l = warp_sum(ea + eb); - if (lane == 0) { part_m[slot * G + h] = m; part_l[slot * G + h] = l; } +} + +// ---- INT8 pools: the codes on the tensor cores (mma m16n8k32, int8). A block per (KV head, query), or per (KV head, +// query, split) with SPLIT, its 12 query heads the rows of the 16-row fragments, 64 cells a tile. q and the +// probabilities enter as 24-bit fixed point, three 8-bit limbs each; the integer sums are exact, so the error is the +// fixed point's: q to 2^-23 of its row's power-of-two bound, a probability times its V scale to 2^-23 of its tile's. +// Warp-specialized, one block an SM: +// warps 0-3 score 16 cells of each tile over all 256 dims (q's limbs in registers) and run the online softmax; +// warps 4-7 multiply a 64-dim group of V each: the probabilities times the V scales as limbs against V through +// ldmatrix.trans (cells and dims permuted inside the fragments, the same way on both sides); +// warps 8-11 fetch by cp.async, 16 cells of each tile a warp: their K and V codes and scales together (with +// K and V in warps of their own, V's copies queued behind K's, which ran ahead), the pool rows two +// tiles ahead. +// Rings in shared memory between them (K 2 slots, V 3, probabilities 2), synchronized by named barriers. +namespace pf { +constexpr int NT = 384; +constexpr int KS = 2, VS = 3; // slots of the K and V rings +enum : int { + END = 0, FULL_K = 1, EMPTY_K = FULL_K + KS, FULL_V = EMPTY_K + KS, EMPTY_V = FULL_V + VS, FULL_P = EMPTY_V + VS, + EMPTY_P = FULL_P + 2, SRED = EMPTY_P + 2 +}; +static_assert(SRED < 16, "sixteen named barriers"); +constexpr int NKB = 128 + 128, NVB = 128 + 128, NPB = 128 + 128; // threads at the K, V, P barriers (END as P) +static_assert(KV_Q8_GROUP == 64 && HD == 256, "four 64-dim scale groups a row"); +constexpr int MAGIC = 0x4B400000; // 1.5 * 2^23 as float bits: an int x (|x| < 2^22) added gives 1.5 * 2^23 + x +constexpr float MAGIC_F = 12582912.0f; + +struct __align__(16) Smem { + uint8_t k[KS][PT * HD]; // K codes, 16-byte chunk j of cell c at j ^ ksw(c) + uint8_t v[VS][PT * HD]; // V codes, chunk j of cell c at j ^ (c & 7) + uint2 ks[KS][PT]; // FP16 scales of a cell's four dim groups + uint2 vs[VS][PT]; + float p[2][16][PT]; // probabilities, float4 i of row r at i ^ ((r & 1) << 2) + float alpha[2][16]; // a row's rescale of the running sums, and the tile's largest probability + float pmax[2][16]; + float red[2][4][16]; // the score warps' tile maxima + float lsum[4][16]; + float qmax[16], qsc[16]; + float mfin[16]; // SPLIT: the rows' final maxima +}; + +// Score warp w's B fragments: column j of its n-tile n is cell 16w + 4 (j >> 1) + 2n + (j & 1), so a lane's C +// fragments hold four consecutive cells; this swizzle puts the eight cells an ldmatrix reads on distinct banks. +__device__ __forceinline__ int ksw(int c) { return ((c >> 1) & 6) | (c & 1); } +__device__ __forceinline__ void bar_sync(int id, int n) { + asm volatile("bar.sync %0, %1;\n" ::"r"(id), "r"(n) : "memory"); +} +__device__ __forceinline__ void bar_arrive(int id, int n) { + asm volatile("bar.arrive %0, %1;\n" ::"r"(id), "r"(n) : "memory"); +} +// where cell c of probability row r lives in Smem::p +__device__ __forceinline__ int p_at(int r, int c) { return (((c >> 2) ^ ((r & 1) << 2)) << 2) | (c & 3); } +__device__ __forceinline__ void mma_u8s8(int (&d)[4], const unsigned (&a)[4], unsigned b0, unsigned b1) { + asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.u8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" + : "+r"(d[0]), "+r"(d[1]), "+r"(d[2]), "+r"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} +__device__ __forceinline__ void mma_s8s8(int (&d)[4], const unsigned (&a)[4], unsigned b0, unsigned b1) { + asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" + : "+r"(d[0]), "+r"(d[1]), "+r"(d[2]), "+r"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} +__device__ __forceinline__ void ldsm4(unsigned (&r)[4], const void* s) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"((unsigned) __cvta_generic_to_shared(s))); +} +__device__ __forceinline__ void ldsm4t(unsigned (&r)[4], const void* s) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"((unsigned) __cvta_generic_to_shared(s))); +} +__device__ __forceinline__ void cp_async16(void* s, const void* g) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"((unsigned) __cvta_generic_to_shared(s)), "l"(g)); +} +__device__ __forceinline__ void cp_async8(void* s, const void* g) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n" ::"r"((unsigned) __cvta_generic_to_shared(s)), "l"(g)); +} +__device__ __forceinline__ int exp_of(float x) { return (__float_as_int(x) >> 23) & 255; } // x < 2^(E - 126) +__device__ __forceinline__ float pow2f(int e) { return __int_as_float((e + 127) << 23); } // -126 <= e <= 127 +// w (65536 r2 + 256 r1 + r0) from three limbs' sums (|r| < 2^21), w a power of two and wk = -65793 C w: the first +// multiply-add is exact, the other two round +__device__ __forceinline__ float limbs3(int r0, int r1, int r2, float w65536, float w256, float w, float wk) { + const float u0 = __int_as_float(r0 + MAGIC), u1 = __int_as_float(r1 + MAGIC), u2 = __int_as_float(r2 + MAGIC); + return fmaf(u0, w, fmaf(u1, w256, fmaf(u2, w65536, wk))); +} +// bytes 0-2 of four values as three words: byte j of word i is byte i of value j +__device__ __forceinline__ void pack3(unsigned f0, unsigned f1, unsigned f2, unsigned f3, unsigned& w0, unsigned& w1, + unsigned& w2) { + const unsigned t01 = __byte_perm(f0, f1, 0x5140), t23 = __byte_perm(f2, f3, 0x5140); + w0 = __byte_perm(t01, t23, 0x5410); + w1 = __byte_perm(t01, t23, 0x7632); + w2 = __byte_perm(__byte_perm(f0, f1, 0x0062), __byte_perm(f2, f3, 0x0062), 0x5410); +} +// x >= 0 rounded to an integer below 2^23, in the low 23 bits +__device__ __forceinline__ unsigned fx23(float x) { return __float_as_uint(fminf(x, 8388607.0f) + 8388608.0f); } +// dim group g's FP16 scale of a cell's four +__device__ __forceinline__ float scale_of(uint2 s, int g) { + return __half2float(__ushort_as_half((unsigned short) ((g < 2 ? s.x : s.y) >> (16 * (g & 1))))); +} +} // namespace pf + +template +__global__ void __launch_bounds__(pf::NT, 1) attn_mma_kernel(const float* __restrict__ q, QsaAttnPools p, + const int32_t* __restrict__ ids, + const int32_t* __restrict__ steps, int n_kv_heads, + int page_size, float scale, float* __restrict__ attn, + int cap, Partials part) { + using namespace pf; + extern __shared__ __align__(16) uint8_t smem_raw[]; + Smem& sm = *reinterpret_cast(smem_raw); + const int kvh = blockIdx.x; + const size_t qi = blockIdx.y; + const size_t head0 = qi * (size_t) (n_kv_heads * G) + (size_t) kvh * G; + q += head0 * HD; + if constexpr (!SPLIT) attn += head0 * HD; + ids += qi * (size_t) cap; + int n_ids = __ldg(steps + qi * kStepCount + kStepWidth); + if constexpr (SPLIT) { + const int c = split_cells(n_ids, part.splits), c0 = (int) blockIdx.z * c; + if (c0 >= n_ids) return; + ids += c0; + n_ids = min(c, n_ids - c0); + } + const int n_tiles = (n_ids + PT - 1) / PT; + const int t = threadIdx.x, lane = t & 31, warp = t >> 5, gid = lane >> 2, tig = lane & 3; + const float kS = -65793.0f * MAGIC_F; + + // each query row's largest magnitude, and its scale 2^(e - 23) (|q| < 2^e) times the softmax scale + if (t < 256) { + const int r = t >> 4, ch = t & 15; + float mx = 0.0f; + if (r < G) +#pragma unroll + for (int i = 0; i < 16; i += 4) { + const float4 f = *reinterpret_cast(q + (size_t) r * HD + ch * 16 + i); + mx = fmaxf(mx, fmaxf(fmaxf(fabsf(f.x), fabsf(f.y)), fmaxf(fabsf(f.z), fabsf(f.w)))); + } +#pragma unroll + for (int o = 8; o > 0; o >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, o)); + if (ch == 0) { + sm.qmax[r] = mx; + sm.qsc[r] = mx > 0.0f ? pow2f(exp_of(mx) - 126 - 23) * scale : 0.0f; + } } __syncthreads(); - // values: thread t owns dimension t for all 12 heads. - float acc[G]; -#pragma unroll - for (int h = 0; h < G; ++h) acc[h] = 0.0f; - for (int c = 0; c < n_here; ++c) { - float v; - if constexpr (KV_MODE == 0) { - v = __half2float(__ushort_as_half(p.v_pool[srow[c] * HD + t])); - } else if constexpr (KV_MODE == 1) { - const float sc = __half2float(__ushort_as_half(p.v_scale[srow[c] * (HD / KV_Q8_GROUP) + t / KV_Q8_GROUP])); - v = (float) p.v_q[srow[c] * HD + t] * sc; - } else { // modes 2 and 3: V is rotated Q4_0 (kv_q4.hpp); the caller rotates the output back - constexpr int bytes_per_head = (HD / QK4_0) * sizeof(block_q4_0); - const int b = t / QK4_0; - const int rem = t % QK4_0; - const block_q4_0* blk = reinterpret_cast(p.v_q4 + srow[c] * bytes_per_head) + b; - const float d = __half2float(__ushort_as_half(blk->d)); - const int j = rem < 16 ? rem : (rem - 16); - const uint8_t byte = blk->qs[j]; - const int nibble = (rem < 16) ? ((byte & 0x0F) - 8) : ((byte >> 4) - 8); - v = (float) nibble * d; + if (!SPLIT && n_tiles == 0) { + for (int i = t; i < G * HD; i += NT) attn[i] = 0.0f; + return; + } + + if (warp >= 8) { + // ======== fetch, K and V together: warp 8 + f takes cells 16f .. 16f+15 of each tile (lane l & 15 looks up + // a cell's row); tile i is issued, then tile i-1 signalled once its copies have landed + const int f = warp - 8, cl = 16 * f + (lane & 15); + auto row_of = [&](int id, int pg) { return (pg * n_kv_heads + kvh) * page_size + id % page_size; }; + int row = -1, id1 = 0, pg1 = 0, id2 = 0; + if (cl < n_ids) { + const int id = ids[cl]; + row = row_of(id, p.page_table[id / page_size]); + } + if (PT + cl < n_ids) { + id1 = ids[PT + cl]; + pg1 = p.page_table[id1 / page_size]; } + if (2 * PT + cl < n_ids) id2 = ids[2 * PT + cl]; + for (int i = 0; i < n_tiles; ++i) { + const int sk = i % KS, sv = i % VS, c0 = i * PT, n_here = min(PT, n_ids - c0); + if (i >= KS) bar_sync(EMPTY_K + sk, NKB); + if (i >= VS) bar_sync(EMPTY_V + sv, NVB); #pragma unroll - for (int h = 0; h < G; ++h) acc[h] = fmaf(sp[h][c], v, acc[h]); - } + for (int u = 0; u < 8; ++u) { + const int e = lane + 32 * u, cw = e >> 4, j = e & 15, c = 16 * f + cw; + const int r = __shfl_sync(0xffffffffu, row, cw); + if (c < n_here) { + cp_async16(&sm.k[sk][c * HD + ((j ^ ksw(c)) << 4)], p.k_q + (size_t) r * HD + j * 16); + cp_async16(&sm.v[sv][c * HD + ((j ^ (c & 7)) << 4)], p.v_q + (size_t) r * HD + j * 16); + } + } + const bool vl = lane >= 16; // lanes 0-15 a cell's K scales, 16-31 its V scales + uint2* sd = vl ? &sm.vs[sv][cl] : &sm.ks[sk][cl]; + if (cl < n_here) + cp_async8(sd, (vl ? p.v_scale : p.k_scale) + (size_t) row * (HD / KV_Q8_GROUP)); + else *sd = make_uint2(0u, 0u); // no scale for the multiply of a missing cell's stale codes + asm volatile("cp.async.commit_group;\n" ::); + if (i >= 1) { + asm volatile("cp.async.wait_group 1;\n" ::); + bar_arrive(FULL_K + (i - 1) % KS, NKB); + bar_arrive(FULL_V + (i - 1) % VS, NVB); + } + row = cl < n_ids - c0 - PT ? row_of(id1, pg1) : -1; + id1 = id2; + pg1 = c0 + 2 * PT + cl < n_ids ? p.page_table[id2 / page_size] : 0; + id2 = c0 + 3 * PT + cl < n_ids ? ids[c0 + 3 * PT + cl] : 0; + } + asm volatile("cp.async.wait_group 0;\n" ::); + bar_arrive(FULL_K + (n_tiles - 1) % KS, NKB); + bar_arrive(FULL_V + (n_tiles - 1) % VS, NVB); + } else if (warp < 4) { + // ======== scores: cells 16w + 4tig .. +3 of each tile in this lane's C fragments (rows gid, gid + 8) + const int w = warp; + unsigned qa[3][8][4]; // q's limbs as A fragments: [limb][k-step][register] + { + float inv[2]; +#pragma unroll + for (int r = 0; r < 2; ++r) { + const float mx = sm.qmax[gid + 8 * r]; + inv[r] = mx > 0.0f ? pow2f(23 - (exp_of(mx) - 126)) : 0.0f; + } +#pragma unroll + for (int ks = 0; ks < 8; ++ks) +#pragma unroll + for (int r = 0; r < 2; ++r) +#pragma unroll + for (int hf = 0; hf < 2; ++hf) { + const int row = gid + 8 * r; + float4 f = make_float4(0.f, 0.f, 0.f, 0.f); + if (row < G) + f = *reinterpret_cast(q + (size_t) row * HD + 32 * ks + 16 * hf + 4 * tig); + const float fv[4] = {f.x, f.y, f.z, f.w}; + unsigned u[4]; +#pragma unroll + for (int e = 0; e < 4; ++e) + u[e] = (unsigned) max(-8388607, min(8388607, __float2int_rn(fv[e] * inv[r]))); + pack3(u[0], u[1], u[2], u[3], qa[0][ks][r + 2 * hf], qa[1][ks][r + 2 * hf], qa[2][ks][r + 2 * hf]); + } + } + const float qs[2] = {sm.qsc[gid], sm.qsc[gid + 8]}; + float m_run[2] = {-FLT_MAX, -FLT_MAX}, lp[2] = {0.0f, 0.0f}; + for (int i = 0; i < n_tiles; ++i) { + const int s = i % KS, sp = i & 1, n_here = min(PT, n_ids - i * PT); + bar_sync(FULL_K + s, NKB); + uint2 kr[4]; +#pragma unroll + for (int e = 0; e < 4; ++e) kr[e] = sm.ks[s][16 * w + 4 * tig + e]; + float sv[2][4] = {{0.f, 0.f, 0.f, 0.f}, {0.f, 0.f, 0.f, 0.f}}; +#pragma unroll + for (int g = 0; g < 4; ++g) { + unsigned b[2][4]; +#pragma unroll + for (int n = 0; n < 2; ++n) { + const int cb = 16 * w + 4 * ((lane & 7) >> 1) + 2 * n + (lane & 1); + ldsm4(b[n], &sm.k[s][cb * HD + (((4 * g + (lane >> 3)) ^ ksw(cb)) << 4)]); + } + int acc[3][2][4]; +#pragma unroll + for (int j = 0; j < 3; ++j) +#pragma unroll + for (int n = 0; n < 2; ++n) +#pragma unroll + for (int ii = 0; ii < 4; ++ii) acc[j][n][ii] = 0; +#pragma unroll + for (int kk = 0; kk < 2; ++kk) +#pragma unroll + for (int j = 0; j < 3; ++j) +#pragma unroll + for (int n = 0; n < 2; ++n) { + if (j < 2) mma_u8s8(acc[j][n], qa[j][2 * g + kk], b[n][2 * kk], b[n][2 * kk + 1]); + else mma_s8s8(acc[j][n], qa[j][2 * g + kk], b[n][2 * kk], b[n][2 * kk + 1]); + } +#pragma unroll + for (int n = 0; n < 2; ++n) #pragma unroll - for (int h = 0; h < G; ++h) part_acc[((size_t) slot * G + h) * HD + t] = acc[h]; + for (int ii = 0; ii < 4; ++ii) + sv[n][ii] = fmaf(limbs3(acc[0][n][ii], acc[1][n][ii], acc[2][n][ii], 65536.0f, 256.0f, 1.0f, kS), + scale_of(kr[2 * n + (ii & 1)], g), sv[n][ii]); + } + if (i + KS < n_tiles) bar_arrive(EMPTY_K + s, NKB); + // online softmax: the tile's maximum over the four warps + float mx[2] = {-FLT_MAX, -FLT_MAX}; +#pragma unroll + for (int n = 0; n < 2; ++n) +#pragma unroll + for (int ii = 0; ii < 4; ++ii) { + const int r = ii >> 1, cell = 16 * w + 4 * tig + 2 * n + (ii & 1); + const float v = cell < n_here ? sv[n][ii] * qs[r] : -FLT_MAX; + sv[n][ii] = v; + mx[r] = fmaxf(mx[r], v); + } +#pragma unroll + for (int r = 0; r < 2; ++r) { + mx[r] = fmaxf(mx[r], __shfl_xor_sync(0xffffffffu, mx[r], 1)); + mx[r] = fmaxf(mx[r], __shfl_xor_sync(0xffffffffu, mx[r], 2)); + } + if (tig == 0) { sm.red[sp][w][gid] = mx[0]; sm.red[sp][w][gid + 8] = mx[1]; } + bar_sync(SRED, 128); + float al[2], pm[2], pv[2][4]; +#pragma unroll + for (int r = 0; r < 2; ++r) { + const int row = gid + 8 * r; + const float tmax = fmaxf(fmaxf(sm.red[sp][0][row], sm.red[sp][1][row]), + fmaxf(sm.red[sp][2][row], sm.red[sp][3][row])); + const float m_new = fmaxf(m_run[r], tmax); + al[r] = __expf(m_run[r] - m_new); + pm[r] = __expf(tmax - m_new); + m_run[r] = m_new; + } +#pragma unroll + for (int n = 0; n < 2; ++n) +#pragma unroll + for (int ii = 0; ii < 4; ++ii) { + const int r = ii >> 1, cell = 16 * w + 4 * tig + 2 * n + (ii & 1); + pv[r][2 * n + (ii & 1)] = cell < n_here ? __expf(sv[n][ii] - m_run[r]) : 0.0f; + } +#pragma unroll + for (int r = 0; r < 2; ++r) lp[r] = fmaf(lp[r], al[r], (pv[r][0] + pv[r][1]) + (pv[r][2] + pv[r][3])); + if (i >= 2) bar_sync(EMPTY_P + sp, NPB); +#pragma unroll + for (int r = 0; r < 2; ++r) { + const int row = gid + 8 * r; + *reinterpret_cast(&sm.p[sp][row][p_at(row, 16 * w + 4 * tig)]) = + make_float4(pv[r][0], pv[r][1], pv[r][2], pv[r][3]); + } + if (w == 0 && tig == 0) { + sm.alpha[sp][gid] = al[0]; sm.alpha[sp][gid + 8] = al[1]; + sm.pmax[sp][gid] = pm[0]; sm.pmax[sp][gid + 8] = pm[1]; + } + bar_arrive(FULL_P + sp, NPB); + } +#pragma unroll + for (int r = 0; r < 2; ++r) { + lp[r] += __shfl_xor_sync(0xffffffffu, lp[r], 1); + lp[r] += __shfl_xor_sync(0xffffffffu, lp[r], 2); + } + if (tig == 0) { sm.lsum[w][gid] = lp[0]; sm.lsum[w][gid + 8] = lp[1]; } + if (SPLIT && w == 0 && tig == 0) { sm.mfin[gid] = m_run[0]; sm.mfin[gid + 8] = m_run[1]; } + bar_sync(END, NPB); + } else { + // ======== values: dims 64g .. 64g+63 as 16-dim blocks db; n-tile (db, par) column j is dim + // 64g + 16db + 2j + par, and k index 4tig + e of a k-step's half hf is cell 16hf + {2tig, 2tig+1, 2tig+8, + // 2tig+9}[e] (as ldmatrix.trans hands out V and two byte permutes rearrange it) + const int g = warp - 4; + float o[4][2][4]; +#pragma unroll + for (int a = 0; a < 4; ++a) +#pragma unroll + for (int b = 0; b < 2; ++b) +#pragma unroll + for (int ii = 0; ii < 4; ++ii) o[a][b][ii] = 0.0f; + for (int i = 0; i < n_tiles; ++i) { + const int sp = i & 1, s = i % VS; + bar_sync(FULL_P + sp, NPB); + bar_sync(FULL_V + s, NVB); + // the fixed point of p * vs: 2^-23 of the tile's largest p times its largest V scale in this group + float vmax = fmaxf(scale_of(sm.vs[s][lane], g), scale_of(sm.vs[s][lane + 32], g)); +#pragma unroll + for (int o_ = 16; o_ > 0; o_ >>= 1) vmax = fmaxf(vmax, __shfl_xor_sync(0xffffffffu, vmax, o_)); + float inv[2], ps[2], al[2]; +#pragma unroll + for (int r = 0; r < 2; ++r) { + const int row = gid + 8 * r; + const int eb = exp_of(sm.pmax[sp][row] * vmax) - 126; + const bool ok = eb >= -103; + inv[r] = ok ? pow2f(23 - eb) : 0.0f; + ps[r] = ok ? pow2f(eb - 23) : 0.0f; + al[r] = sm.alpha[sp][row]; + } + unsigned pa[3][2][4]; // [limb][k-step][register] +#pragma unroll + for (int kk = 0; kk < 2; ++kk) +#pragma unroll + for (int hf = 0; hf < 2; ++hf) { + const int c = 32 * kk + 16 * hf + 2 * tig; // cells c, c+1, c+8, c+9 + const float v0 = scale_of(sm.vs[s][c], g), v1 = scale_of(sm.vs[s][c + 1], g); + const float v8 = scale_of(sm.vs[s][c + 8], g), v9 = scale_of(sm.vs[s][c + 9], g); +#pragma unroll + for (int r = 0; r < 2; ++r) { + const int row = gid + 8 * r; + const float2 pa2 = *reinterpret_cast(&sm.p[sp][row][p_at(row, c)]); + const float2 pb2 = *reinterpret_cast(&sm.p[sp][row][p_at(row, c + 8)]); + const int x = r + 2 * hf; + pack3(fx23(pa2.x * v0 * inv[r]), fx23(pa2.y * v1 * inv[r]), fx23(pb2.x * v8 * inv[r]), + fx23(pb2.y * v9 * inv[r]), pa[0][kk][x], pa[1][kk][x], pa[2][kk][x]); + } + } + if (i + 2 < n_tiles) bar_arrive(EMPTY_P + sp, NPB); +#pragma unroll + for (int db = 0; db < 4; ++db) { + unsigned bv[2][2][2]; // [k-step][parity][b0, b1] +#pragma unroll + for (int kk = 0; kk < 2; ++kk) { + const int cell = 32 * kk + 8 * (lane >> 3) + (lane & 7); + unsigned r4[4]; + ldsm4t(r4, &sm.v[s][cell * HD + (((4 * g + db) ^ (cell & 7)) << 4)]); + bv[kk][0][0] = __byte_perm(r4[0], r4[1], 0x6420); + bv[kk][1][0] = __byte_perm(r4[0], r4[1], 0x7531); + bv[kk][0][1] = __byte_perm(r4[2], r4[3], 0x6420); + bv[kk][1][1] = __byte_perm(r4[2], r4[3], 0x7531); + } + int acc[3][2][4]; +#pragma unroll + for (int j = 0; j < 3; ++j) +#pragma unroll + for (int par = 0; par < 2; ++par) +#pragma unroll + for (int ii = 0; ii < 4; ++ii) acc[j][par][ii] = 0; +#pragma unroll + for (int kk = 0; kk < 2; ++kk) +#pragma unroll + for (int j = 0; j < 3; ++j) +#pragma unroll + for (int par = 0; par < 2; ++par) mma_u8s8(acc[j][par], pa[j][kk], bv[kk][par][0], bv[kk][par][1]); +#pragma unroll + for (int par = 0; par < 2; ++par) +#pragma unroll + for (int ii = 0; ii < 4; ++ii) { + const int r = ii >> 1; + o[db][par][ii] = fmaf(o[db][par][ii], al[r], + limbs3(acc[0][par][ii], acc[1][par][ii], acc[2][par][ii], 65536.0f * ps[r], + 256.0f * ps[r], ps[r], kS * ps[r])); + } + } + if (i + VS < n_tiles) bar_arrive(EMPTY_V + s, NVB); + } + bar_sync(END, NPB); + if constexpr (SPLIT) { // o as it is, with each row's maximum and sum + float* base = part.acc + qi * part.stride; + const size_t hq = (size_t) blockIdx.z * n_kv_heads * G + (size_t) kvh * G; // the block's first head +#pragma unroll + for (int r = 0; r < 2; ++r) { + const int row = gid + 8 * r; + if (row >= G) continue; +#pragma unroll + for (int db = 0; db < 4; ++db) + *reinterpret_cast(&base[(hq + row) * HD + 64 * g + 16 * db + 4 * tig]) = + make_float4(o[db][0][2 * r], o[db][1][2 * r], o[db][0][2 * r + 1], o[db][1][2 * r + 1]); + } + if (g == 0 && lane < G) { + base[part.m + hq + lane] = sm.mfin[lane]; + base[part.l + hq + lane] = + (sm.lsum[0][lane] + sm.lsum[1][lane]) + (sm.lsum[2][lane] + sm.lsum[3][lane]); + } + } else { +#pragma unroll + for (int r = 0; r < 2; ++r) { + const int row = gid + 8 * r; + if (row >= G) continue; + const float l = (sm.lsum[0][row] + sm.lsum[1][row]) + (sm.lsum[2][row] + sm.lsum[3][row]); +#pragma unroll + for (int db = 0; db < 4; ++db) { // dims 64g + 16db + 4tig + {0, 1, 2, 3} + float4 out = make_float4(0.f, 0.f, 0.f, 0.f); + if (l > 0.0f) out = make_float4(o[db][0][2 * r] / l, o[db][1][2 * r] / l, o[db][0][2 * r + 1] / l, + o[db][1][2 * r + 1] / l); + *reinterpret_cast(&attn[(size_t) row * HD + 64 * g + 16 * db + 4 * tig]) = out; + } + } + } + } } -__global__ void __launch_bounds__(HD) attn_merge_kernel(const float* __restrict__ part_acc, - const float* __restrict__ part_m, - const float* __restrict__ part_l, int n_chunks, - float* __restrict__ attn, long long scratch_stride = 0) { - part_acc += (size_t) blockIdx.y * (size_t) scratch_stride; - part_m += (size_t) blockIdx.y * (size_t) scratch_stride; - part_l += (size_t) blockIdx.y * (size_t) scratch_stride; - attn += (size_t) blockIdx.y * (size_t) gridDim.x * HD; - const int h = blockIdx.x; // global query head - const int kvh = h / G, hl = h % G; - const int d = threadIdx.x; - float M = -FLT_MAX; - for (int c = 0; c < n_chunks; ++c) M = fmaxf(M, part_m[(kvh * n_chunks + c) * G + hl]); - float L = 0.0f, acc = 0.0f; - for (int c = 0; c < n_chunks; ++c) { - const int slot = kvh * n_chunks + c; - const float m = part_m[slot * G + hl]; - if (m == -FLT_MAX) continue; - const float w = __expf(m - M); - L = fmaf(part_l[slot * G + hl], w, L); - acc = fmaf(part_acc[((size_t) slot * G + hl) * HD + d], w, acc); +// x * sigmoid(raw) as native_qsa_gate_apply computes it: its TU's fast-math instructions +__device__ __forceinline__ float gated(float x, float raw) { + float r; + asm("{\n\t.reg .f32 e;\n\tmul.ftz.f32 e, %2, 0fBFB8AA3B;\n\tex2.approx.ftz.f32 e, e;\n\t" + "add.ftz.f32 e, e, 0f3F800000;\n\trcp.approx.ftz.f32 e, e;\n\tmul.ftz.f32 %0, %1, e;\n\t}" + : "=f"(r) + : "f"(x), "f"(raw)); + return r; +} + +// The decode's merge: a block per (head, query), 4 groups of 64 threads adding every 4th split's sums, a thread 4 +// dims. Every thread reads all the splits' maxima and sums, so all of its loads go out at once. +__global__ void __launch_bounds__(THREADS) attn_merge_kernel(Partials part, const int32_t* __restrict__ steps, + const float* __restrict__ gate, float* __restrict__ attn) { + const int h = blockIdx.x, n_head = gridDim.x; + const size_t qi = blockIdx.y; + const int n_ids = __ldg(steps + qi * kStepCount + kStepWidth), c = split_cells(n_ids, part.splits); + const int n_sp = (n_ids + c - 1) / c; + const float* base = part.acc + qi * part.stride; + const int t = threadIdx.x, grp = t >> 6, d4 = 4 * (t & 63); + float4 o[SPLIT_MAX / 4]; + float m[SPLIT_MAX], l[SPLIT_MAX]; +#pragma unroll + for (int k = 0; k < SPLIT_MAX / 4; ++k) + if (grp + 4 * k < n_sp) + o[k] = *reinterpret_cast(base + ((size_t) (grp + 4 * k) * n_head + h) * HD + d4); +#pragma unroll + for (int k = 0; k < SPLIT_MAX; ++k) + if (k < n_sp) { + m[k] = base[part.m + k * n_head + h]; + l[k] = base[part.l + k * n_head + h]; + } + float4 raw = make_float4(0.f, 0.f, 0.f, 0.f); + if (gate && grp == 0) raw = *reinterpret_cast(gate + (qi * n_head + h) * 2 * HD + HD + d4); + float M = -FLT_MAX, L = 0.0f; +#pragma unroll + for (int k = 0; k < SPLIT_MAX; ++k) + if (k < n_sp) M = fmaxf(M, m[k]); + float4 a = make_float4(0.f, 0.f, 0.f, 0.f); +#pragma unroll + for (int k = 0; k < SPLIT_MAX; ++k) + if (k < n_sp) { + const float w = __expf(m[k] - M); + L = fmaf(l[k], w, L); + if ((k & 3) == grp) { + const float4 x = o[k >> 2]; + a = make_float4(fmaf(x.x, w, a.x), fmaf(x.y, w, a.y), fmaf(x.z, w, a.z), fmaf(x.w, w, a.w)); + } + } + __shared__ float4 red[3][64]; + if (grp > 0) red[grp - 1][t & 63] = a; + __syncthreads(); + if (grp > 0) return; +#pragma unroll + for (int k = 0; k < 3; ++k) { + const float4 b = red[k][t]; + a = make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); } - attn[(size_t) h * HD + d] = L > 0.0f ? acc / L : 0.0f; + float4 r = make_float4(0.f, 0.f, 0.f, 0.f); + if (L > 0.0f) r = make_float4(a.x / L, a.y / L, a.z / L, a.w / L); + if (gate) r = make_float4(gated(r.x, raw.x), gated(r.y, raw.y), gated(r.z, raw.z), gated(r.w, raw.w)); + *reinterpret_cast(attn + (qi * n_head + h) * HD + d4) = r; +} + +// A decode query's splits at most: SPLIT_MAX (more cost the merge and each block's setup more than they gain), whole +// tiles. +int64_t max_splits(int64_t cap) { return std::max(1, std::min(SPLIT_MAX, (cap + PT - 1) / PT)); } +// n_q queries' splits: a block of the INT8 kernel takes an SM, so all queries' blocks in one wave +int decode_splits(int64_t cap, int64_t n_q, int64_t n_kv) { + int dev = 0, sms = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + return (int) std::max(1, std::min(max_splits(cap), sms / (n_kv * n_q))); +} + +// the per-query scratch's layout for `splits` +Partials partials(float* scratch, int splits, int64_t n_head, int64_t stride) { + Partials part; + part.acc = scratch; + part.m = (long long) splits * n_head * HD; + part.l = part.m + (long long) splits * n_head; + part.stride = stride; + part.splits = splits; + return part; +} + +// the INT8 kernel's shared memory, above the default limit (once for each form) +template +void allow_mma_smem() { + static const cudaError_t e = cudaFuncSetAttribute( + attn_mma_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) sizeof(pf::Smem)); + (void) e; } } // namespace -void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, - int64_t cap, const QsaShapes& s, float* scratch, float* attn, int64_t n_q, void* stream) { +void qsa_prefill_attn(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, int64_t cap, + const QsaShapes& s, float* attn, int64_t n_q, void* stream) { if (n_q <= 0) return; - if (s.head_dim != HD || s.n_head != (int64_t) G * s.n_head_kv || cap <= 0 || !scratch || !ids || !steps || - !pools.page_table || n_q > 65535) { - std::fprintf(stderr, "qsa_decode_attn_batch: unsupported geometry or missing buffers\n"); + const bool int8 = pools.k_q != nullptr; + if (s.head_dim != HD || s.n_head != (int64_t) G * s.n_head_kv || cap <= 0 || !ids || !steps || !pools.page_table || + (int8 ? (!pools.v_q || !pools.k_scale || !pools.v_scale) : (!pools.k_pool || !pools.v_pool))) { + std::fprintf(stderr, "qsa_prefill_attn: unsupported geometry or missing buffers\n"); std::exit(1); } - const int kv_mode = pools.k_q4 != nullptr ? 2 : (pools.k_q != nullptr && pools.v_q4 != nullptr ? 3 - : (pools.k_q != nullptr ? 1 : 0)); - const int n_chunks = (int) ((cap + CHUNK - 1) / CHUNK); - // per query: [acc: n_chunks*n_head*HD][m: n_chunks*n_head][l: n_chunks*n_head], all offsets from one stride - const long long stride = (long long) qsa_decode_attn_scratch_floats(cap, s); - float* part_acc = scratch; - float* part_m = scratch + (size_t) n_chunks * s.n_head * HD; - float* part_l = part_m + (size_t) n_chunks * s.n_head; const float scale = 1.0f / sqrtf((float) HD); - const dim3 grid((unsigned) n_chunks, (unsigned) s.n_head_kv, (unsigned) n_q); cudaStream_t st = (cudaStream_t) stream; - if (kv_mode == 3) - attn_chunk_kernel<3><<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); - else if (kv_mode == 2) - attn_chunk_kernel<2><<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); - else if (kv_mode == 1) - attn_chunk_kernel<1><<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); - else - attn_chunk_kernel<0><<>>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride); - attn_merge_kernel<<>>(part_acc, part_m, part_l, n_chunks, - attn, stride); + if (int8) allow_mma_smem(); + for (int64_t q0 = 0; q0 < n_q; q0 += 65535) { + const int64_t nb = n_q - q0 < 65535 ? n_q - q0 : 65535; + const dim3 grid((unsigned) s.n_head_kv, (unsigned) nb); + const size_t qo = (size_t) q0 * (size_t) s.n_head * HD; + const int32_t* ids_q = ids + (size_t) q0 * (size_t) cap; + const int32_t* steps_q = steps + (size_t) q0 * kStepCount; + if (int8) + attn_mma_kernel<<>>( + q + qo, pools, ids_q, steps_q, (int) s.n_head_kv, (int) s.page_size, scale, attn + qo, (int) cap, {}); + else + attn_f16_kernel<<>>( + q + qo, pools, ids_q, steps_q, (int) s.n_head_kv, (int) s.page_size, scale, attn + qo, (int) cap, {}); + } const cudaError_t e = cudaGetLastError(); if (e != cudaSuccess) { - std::fprintf(stderr, "qsa_decode_attn_batch: %s\n", cudaGetErrorString(e)); + std::fprintf(stderr, "qsa_prefill_attn: %s\n", cudaGetErrorString(e)); std::exit(1); } } -uint64_t qsa_decode_attn_scratch_floats(int64_t cap, const QsaShapes& s) { - const int64_t chunks = (cap + CHUNK - 1) / CHUNK; - return (uint64_t) chunks * (uint64_t) s.n_head * (HD + 2) + 64; -} - -void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* step, - int64_t cap, const QsaShapes& s, float* scratch, float* attn, void* stream) { - if (s.head_dim != HD || s.n_head != (int64_t) G * s.n_head_kv || cap <= 0 || !scratch || !ids || !step || - !pools.page_table) { +void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps, + int64_t cap, const QsaShapes& s, float* scratch, float* attn, int64_t n_q, void* stream, + const float* gate) { + if (n_q <= 0) return; + const bool int8 = pools.k_q != nullptr; + if (s.head_dim != HD || s.n_head != (int64_t) G * s.n_head_kv || cap <= 0 || !scratch || !ids || !steps || + !pools.page_table || n_q > 65535 || + (int8 ? (!pools.v_q || !pools.k_scale || !pools.v_scale) : (!pools.k_pool || !pools.v_pool))) { std::fprintf(stderr, "qsa_decode_attn: unsupported geometry or missing buffers\n"); std::exit(1); } - const int kv_mode = pools.k_q4 != nullptr ? 2 : (pools.k_q != nullptr && pools.v_q4 != nullptr ? 3 - : (pools.k_q != nullptr ? 1 : 0)); - if (kv_mode == 3 ? (!pools.k_scale || !pools.v_q4) - : (kv_mode == 2 ? (!pools.v_q4) : (kv_mode == 1 ? (!pools.v_q || !pools.k_scale || !pools.v_scale) - : (!pools.k_pool || !pools.v_pool)))) { - std::fprintf(stderr, "qsa_decode_attn: incomplete KV pools\n"); - std::exit(1); - } - const int n_chunks = (int) ((cap + CHUNK - 1) / CHUNK); - float* part_acc = scratch; - float* part_m = scratch + (size_t) n_chunks * s.n_head * HD; - float* part_l = part_m + (size_t) n_chunks * s.n_head; + const int splits = decode_splits(cap, n_q, s.n_head_kv); + const Partials part = partials(scratch, splits, s.n_head, (long long) qsa_decode_attn_scratch_floats(cap, s)); const float scale = 1.0f / sqrtf((float) HD); - const dim3 grid((unsigned) n_chunks, (unsigned) s.n_head_kv); + const dim3 grid((unsigned) s.n_head_kv, (unsigned) n_q, (unsigned) splits); cudaStream_t st = (cudaStream_t) stream; - if (kv_mode == 3) - attn_chunk_kernel<3><<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks); - else if (kv_mode == 2) - attn_chunk_kernel<2><<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks); - else if (kv_mode == 1) - attn_chunk_kernel<1><<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks); - else - attn_chunk_kernel<0><<>>(q, pools, ids, step, (int) s.n_head_kv, (int) s.page_size, - scale, part_acc, part_m, part_l, n_chunks); - attn_merge_kernel<<<(unsigned) s.n_head, HD, 0, st>>>(part_acc, part_m, part_l, n_chunks, attn); + if (int8) { + allow_mma_smem(); + attn_mma_kernel<<>>( + q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, scale, nullptr, (int) cap, part); + } else { + attn_f16_kernel<<>>( + q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size, scale, nullptr, (int) cap, part); + } + attn_merge_kernel<<>>(part, steps, gate, attn); const cudaError_t e = cudaGetLastError(); if (e != cudaSuccess) { std::fprintf(stderr, "qsa_decode_attn: %s\n", cudaGetErrorString(e)); @@ -292,4 +815,13 @@ void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32 } } +uint64_t qsa_decode_attn_scratch_floats(int64_t cap, const QsaShapes& s) { + return (uint64_t) max_splits(cap) * (uint64_t) s.n_head * (HD + 2) + 64; +} + +void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* step, + int64_t cap, const QsaShapes& s, float* scratch, float* attn, void* stream) { + qsa_decode_attn_batch(q, pools, ids, step, cap, s, scratch, attn, 1, stream); +} + } // namespace strata::kernels diff --git a/src/kernels/cuda/qsa_select.cu b/src/kernels/cuda/qsa_select.cu index 85b7342c1ee..39ececf2929 100644 --- a/src/kernels/cuda/qsa_select.cu +++ b/src/kernels/cuda/qsa_select.cu @@ -6,13 +6,26 @@ #include #include #include +#include namespace strata::kernels { namespace { constexpr int IDX_DIM = 128, IDX_HEADS = 4, R = 4; +constexpr unsigned FULL = 0xffffffffu; constexpr int SCORE_WARPS = 8; -constexpr int TOPK_T = 256; +constexpr int SCORE_KB = 2; // key blocks per 8-lane group +constexpr int SCORE_TILE = SCORE_WARPS * 4 * SCORE_KB; // key blocks per thread block and round +constexpr int SCORE_QC = 16; // queries staged in shared memory at a time (32 KB) +constexpr int BATCH_T = 256, BATCH_W = BATCH_T / 32; // batch_scores_kernel: one block an SM +constexpr int BATCH_QT = 2, BATCH_BT = 4; // a lane: 2 queries x 4 key blocks, all heads +constexpr int BATCH_TILE = 32 * BATCH_BT; // key blocks in shared memory (64 KB) +constexpr int BATCH_QC = BATCH_W * BATCH_QT; // queries a unit (a warp's slots: 32 KB) +constexpr int QD = IDX_HEADS * IDX_DIM; // a query's floats +constexpr int kBatchMinQueries = 12; // fewer (decode windows): block_scores_kernel +constexpr int TOPK_T = 1024, TOPK_W = TOPK_T / 32; +constexpr int TOPK_ROWS = 32; // rows of TOPK_T blocks held in registers +constexpr int TOPK_MAX_ROWS = 64; // 65536 blocks: 262144 cells __device__ __forceinline__ uint32_t order_key(float s) { const float v = s + 0.0f; @@ -21,503 +34,520 @@ __device__ __forceinline__ uint32_t order_key(float s) { return (b & 0x80000000u) ? ~b : (b | 0x80000000u); } -__global__ void __launch_bounds__(SCORE_WARPS * 32) block_scores_kernel(const float* __restrict__ pooled, - const float* __restrict__ dead, - const float* __restrict__ q_idx, - const int32_t* __restrict__ steps, - int64_t max_blocks, float* __restrict__ out) { - const int64_t qi = blockIdx.y; - const int32_t* st = steps + qi * kStepCount; - const int64_t n_kv = st[kStepNKv], n_bid = st[kStepNBid]; - const int64_t b = (int64_t) blockIdx.x * SCORE_WARPS + (threadIdx.x >> 5); - if (b > n_bid || b >= max_blocks) return; - const int lane = threadIdx.x & 31; - const float* key = (b == n_bid) ? dead : pooled + b * IDX_DIM; - const float4 k4 = *reinterpret_cast(key + lane * 4); - const float* q = q_idx + qi * IDX_HEADS * IDX_DIM + lane * 4; - float score = 0.0f; +// One leaf of a score (dims 4l..4l+3 of a head): k.x*q.x + k.y*q.y + k.z*q.z + k.w*q.w as the original warp-per-pair +// kernel compiled it (qsa_select_parity keeps it) +__device__ __forceinline__ float dot4(float4 k, float4 q) { + return __fmaf_rn(k.w, q.w, __fmaf_rn(k.z, q.z, __fmaf_rn(k.x, q.x, __fmul_rn(k.y, q.y)))); +} + +// One query's scores for the group's SCORE_KB key blocks. Lane j of the 8-lane group holds each key's dims 4l..4l+3 +// for l = j, j+8, j+16, j+24; per head its four 4-term dot products d_l are summed in the lane as +// (d_j + d_j+16) + (d_j+8 + d_j+24), then across the group by xor 4, 2 and 1: the tree in which a warp-wide xor +// butterfly (16, 8, 4, 2, 1) over lanes l sums the same d_l, so each score is bit for bit that of one warp per +// (query, block). TAIL: key `tail_kb` (-1: none) is the query's tail block and takes the `dead` key instead. +template +__device__ __forceinline__ void group_scores(const float4 (&k)[SCORE_KB][4], const float* __restrict__ dead, + int tail_kb, const float* q, int j, float (&score)[SCORE_KB]) { +#pragma unroll + for (int kb = 0; kb < SCORE_KB; ++kb) score[kb] = 0.0f; #pragma unroll for (int h = 0; h < IDX_HEADS; ++h) { - const float4 q4 = *reinterpret_cast(q + h * IDX_DIM); - float d = k4.x * q4.x + k4.y * q4.y + k4.z * q4.z + k4.w * q4.w; + float4 q4[4]; #pragma unroll - for (int o = 16; o > 0; o >>= 1) d += __shfl_xor_sync(0xffffffffu, d, o); - score += d > 0.0f ? d : 0.0f; - } - if (lane == 0) { - if (b == n_bid && n_kv % R != 0) score += 1e9f; - out[qi * max_blocks + b] = score; + for (int i = 0; i < 4; ++i) q4[i] = *reinterpret_cast(q + h * IDX_DIM + 4 * (j + 8 * i)); +#pragma unroll + for (int kb = 0; kb < SCORE_KB; ++kb) { + float d[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + float4 k4 = k[kb][i]; + if (TAIL && kb == tail_kb) k4 = *reinterpret_cast(dead + 4 * (j + 8 * i)); + d[i] = dot4(k4, q4[i]); + } + float c = (d[0] + d[2]) + (d[1] + d[3]); + c += __shfl_xor_sync(FULL, c, 4); + c += __shfl_xor_sync(FULL, c, 2); + c += __shfl_xor_sync(FULL, c, 1); + score[kb] += c > 0.0f ? c : 0.0f; + } } } -__global__ void __launch_bounds__(TOPK_T) block_topk_kernel(const float* __restrict__ scores, - const int32_t* __restrict__ steps, int64_t max_blocks, - int64_t cap, int32_t* __restrict__ ids) { - __shared__ int hist[256]; - __shared__ int s_a[TOPK_T], s_b[TOPK_T]; - __shared__ int s_digit, s_above; - const int64_t qi = blockIdx.x; - const int32_t* st = steps + qi * kStepCount; - const int64_t n_kv = st[kStepNKv], n_bid = st[kStepNBid], width = st[kStepWidth]; - int32_t* out = ids + qi * cap; - const int t = threadIdx.x; - if (n_kv <= width) { // everything is selected: the identity, ascending - for (int64_t j = t; j < n_kv; j += TOPK_T) out[j] = (int32_t) j; - return; +// 8-lane groups of SCORE_KB consecutive key blocks each, read once for all the queries, which go through shared +// memory SCORE_QC at a time. Block n_bid of a query is its incomplete tail: the `dead` key, +1e9 when it has cells. +__global__ void __launch_bounds__(SCORE_WARPS * 32) block_scores_kernel(const float* __restrict__ pooled, + const float* __restrict__ dead, + const float* __restrict__ q_idx, + const int32_t* __restrict__ steps, int nq, + int64_t max_blocks, float* __restrict__ out) { + __shared__ __align__(16) float qs[SCORE_QC * IDX_HEADS * IDX_DIM]; + __shared__ int s_nkv[SCORE_QC], s_nbid[SCORE_QC]; + const int t = threadIdx.x, lane = t & 31, j = lane & 7; + const int slot = ((t >> 5) * 4 + (lane >> 3)) * SCORE_KB; // the group's first block within a tile + int64_t top = -1; // the queries' last block: the rounds end there + for (int qi = 0; qi < nq; ++qi) { + const int64_t n_bid = steps[(int64_t) qi * kStepCount + kStepNBid]; + const int64_t last = n_bid < max_blocks - 1 ? n_bid : max_blocks - 1; + top = last > top ? last : top; } - const float* sc = scores + qi * max_blocks; - const int64_t nb = n_bid + 1; // blocks 0..n_bid, the last possibly empty - const int64_t per = (nb + TOPK_T - 1) / TOPK_T; - const int64_t b0 = (int64_t) t * per, b1 = (b0 + per < nb) ? b0 + per : nb; - auto weight = [&](int64_t b) -> int { return b < n_bid ? R : (int) (n_kv - n_bid * R); }; - // ---- radix select: the largest key thr with (cells with key >= thr) >= width, 8 bits at a time - uint32_t prefix = 0; - int above = 0; // cells strictly above the digits fixed so far - for (int shift = 24; shift >= 0; shift -= 8) { - for (int i = t; i < 256; i += TOPK_T) hist[i] = 0; - __syncthreads(); - const uint32_t hi_mask = shift == 24 ? 0u : (0xffffffffu << (shift + 8)); - for (int64_t b = b0; b < b1; ++b) { - const int w = weight(b); - if (w == 0) continue; - const uint32_t k = order_key(sc[b]); - if ((k & hi_mask) == (prefix & hi_mask)) atomicAdd(&hist[(k >> shift) & 255], w); + for (int64_t b0 = (int64_t) blockIdx.x * SCORE_TILE; b0 <= top; b0 += (int64_t) gridDim.x * SCORE_TILE) { + const int64_t bg = b0 + slot; + float4 k[SCORE_KB][4]; +#pragma unroll + for (int kb = 0; kb < SCORE_KB; ++kb) { + const float* key = pooled + (bg + kb <= top ? bg + kb : 0) * IDX_DIM; +#pragma unroll + for (int i = 0; i < 4; ++i) k[kb][i] = __ldg(reinterpret_cast(key + 4 * (j + 8 * i))); } - __syncthreads(); - if (t == 0) { - int cum = above, d = 255; - for (; d > 0; --d) { - if (cum + hist[d] >= width) break; - cum += hist[d]; + for (int q0 = 0; q0 < nq; q0 += SCORE_QC) { + const int qn = nq - q0 < SCORE_QC ? nq - q0 : SCORE_QC; + __syncthreads(); + const float4* src = reinterpret_cast(q_idx + (int64_t) q0 * IDX_HEADS * IDX_DIM); + for (int i = t; i < qn * IDX_HEADS * IDX_DIM / 4; i += SCORE_WARPS * 32) + reinterpret_cast(qs)[i] = src[i]; + if (t < qn) { + s_nkv[t] = steps[(int64_t) (q0 + t) * kStepCount + kStepNKv]; + s_nbid[t] = steps[(int64_t) (q0 + t) * kStepCount + kStepNBid]; + } + __syncthreads(); + for (int qq = 0; qq < qn; ++qq) { + const int64_t n_kv = s_nkv[qq], n_bid = s_nbid[qq]; + const int64_t last = n_bid < max_blocks - 1 ? n_bid : max_blocks - 1; + const int tail_kb = n_bid >= bg && n_bid < bg + SCORE_KB ? (int) (n_bid - bg) : -1; + const float* q = qs + qq * IDX_HEADS * IDX_DIM; + float score[SCORE_KB]; + // warp-uniform branches: every lane of the warp runs the group shuffles + if (__any_sync(FULL, tail_kb >= 0)) group_scores(k, dead, tail_kb, q, j, score); + else group_scores(k, dead, -1, q, j, score); + if (j == 0) { + float* o = out + (int64_t) (q0 + qq) * max_blocks; +#pragma unroll + for (int kb = 0; kb < SCORE_KB; ++kb) { + if (bg + kb > last) break; + o[bg + kb] = kb == tail_kb && n_kv % R != 0 ? score[kb] + 1e9f : score[kb]; + } + } } - s_digit = d; - s_above = cum; - } - __syncthreads(); - prefix |= (uint32_t) s_digit << shift; - above = s_above; - __syncthreads(); - } - const uint32_t thr = prefix; - const int64_t eq_budget = width - above; // cells equal to thr that fit, lowest index first - // ---- per-thread counts of cells above and at the threshold, then their exclusive prefixes - int gt = 0, eq = 0; - for (int64_t b = b0; b < b1; ++b) { - const int w = weight(b); - if (w == 0) continue; - const uint32_t k = order_key(sc[b]); - if (k > thr) gt += w; - else if (k == thr) eq += w; - } - s_a[t] = gt; - s_b[t] = eq; - __syncthreads(); - if (t == 0) { - int ag = 0, ae = 0; - for (int i = 0; i < TOPK_T; ++i) { - const int g = s_a[i], e = s_b[i]; - s_a[i] = ag; s_b[i] = ae; - ag += g; ae += e; - } - } - __syncthreads(); - const int64_t eq_before = s_b[t]; - int64_t my_eq = eq_budget - eq_before; - if (my_eq < 0) my_eq = 0; - if (my_eq > eq) my_eq = eq; - const int sel = gt + (int) my_eq; - __syncthreads(); - s_a[t] = sel; - __syncthreads(); - if (t == 0) { - int a = 0; - for (int i = 0; i < TOPK_T; ++i) { const int c = s_a[i]; s_a[i] = a; a += c; } - } - __syncthreads(); - int64_t wpos = s_a[t]; - int64_t eq_left = my_eq; - for (int64_t b = b0; b < b1; ++b) { - const int w = weight(b); - if (w == 0) continue; - const uint32_t k = order_key(sc[b]); - if (k > thr) { - for (int c = 0; c < w; ++c) out[wpos++] = (int32_t) (b * R + c); - } else if (k == thr) { - for (int c = 0; c < w && eq_left > 0; ++c, --eq_left) out[wpos++] = (int32_t) (b * R + c); } } } - -// ---- QSA select on tensor cores (perf-review, after D-1): the block scores of many queries are one GEMM, -// rows (query, indexer head) x columns (blocks), K = 128, with relu per head summed. 3xTF32 (each operand split -// into a TF32 hi and lo part, hi*hi + hi*lo + lo*hi) keeps FP32-level accuracy; the summation order differs from -// the warp kernel, so a score can move in its last bits and a near-tie can select differently (not bitwise). -// Blocks < n_bid only; the tail block n_bid (the `dead` key, +1e9) is scored by the warp kernel's own code. -constexpr int TC_QT = 16; // queries per CTA (one m16 tile per indexer head) -constexpr int TC_NB = 32; // blocks per tile (4 warps x n8) -constexpr int TC_ITER = 4; // tiles per CTA (the query tile is loaded once) -constexpr int TC_QS = IDX_HEADS * IDX_DIM + 4; // query row stride in floats -constexpr int TC_KS = IDX_DIM + 4; // key row stride - -// TF32 conversion and MMA need sm_80: below it they compile to a trap and qsa_block_scores_tc refuses the device -#if defined(__HIPCC__) // AMD: no mma.sync / cp.async; the host keeps the warp kernel (below) -#define STRATA_SEL_SM80 0 -#elif !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 -#define STRATA_SEL_SM80 1 -#else -#define STRATA_SEL_SM80 0 -#endif -__device__ __forceinline__ uint32_t tf32_hi(float x) { -#if STRATA_SEL_SM80 - uint32_t r; - asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(r) : "f"(x)); - return r; -#else - return __float_as_uint(x); -#endif +__host__ __device__ constexpr int rev5(int i) { + return ((i & 1) << 4) | ((i & 2) << 2) | (i & 4) | ((i & 8) >> 2) | ((i & 16) >> 4); } -__device__ __forceinline__ void mma_tf32(float* c, const uint32_t* a, const uint32_t* b) { -#if !STRATA_SEL_SM80 - __trap(); -#else - asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " - "{%0,%1,%2,%3};\n" - : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); -#endif +__host__ __device__ constexpr int popc5(int i) { + return (i & 1) + ((i >> 1) & 1) + ((i >> 2) & 1) + ((i >> 3) & 1) + ((i >> 4) & 1); +} +// key rows in shared memory, 16-byte chunk c of row r at c ^ (r & 7): a quarter warp's rows on distinct bank quads +__device__ __forceinline__ int key_at(int r, int c) { return r * IDX_DIM + 4 * (c ^ (r & 7)); } +__device__ __forceinline__ void cp_async16(void* s, const void* g) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"((unsigned) __cvta_generic_to_shared(s)), "l"(g)); } -__global__ void __launch_bounds__(128) block_scores_tc_kernel(const float* __restrict__ pooled, - const float* __restrict__ q_idx, - const int32_t* __restrict__ steps, int64_t nq, - int64_t max_blocks, int64_t reach, - float* __restrict__ out) { - extern __shared__ __align__(16) float sm[]; - float* sQ = sm; // [TC_QT][TC_QS] - float* sK = sm + TC_QT * TC_QS; // [TC_NB][TC_KS] - __shared__ int s_nbid[TC_QT]; - const int t = threadIdx.x, lane = t & 31, warp = t >> 5, gid = lane >> 2, tig = lane & 3; - const int64_t q0 = (int64_t) blockIdx.y * TC_QT; - for (int i = t; i < TC_QT * IDX_HEADS * IDX_DIM / 4; i += 128) { - const int r = i / (IDX_HEADS * IDX_DIM / 4), c = i % (IDX_HEADS * IDX_DIM / 4); - float4 v = make_float4(0.f, 0.f, 0.f, 0.f); - if (q0 + r < nq) v = reinterpret_cast(q_idx + (q0 + r) * IDX_HEADS * IDX_DIM)[c]; - *reinterpret_cast(sQ + r * TC_QS + c * 4) = v; - } - if (t < TC_QT) s_nbid[t] = q0 + t < nq ? steps[(q0 + t) * kStepCount + kStepNBid] : 0; - int lo_nbid = 0x7fffffff, hi_nbid = 0; - __syncthreads(); - for (int i = 0; i < TC_QT; ++i) { - if (q0 + i >= nq) break; - lo_nbid = min(lo_nbid, s_nbid[i]); - hi_nbid = max(hi_nbid, s_nbid[i]); +// The same scores for batches of queries, in a per-thread GEMM form. Lane l of a warp scores key blocks l + 32 j +// (j < BATCH_BT) of the tile against the warp's BATCH_QT queries, all heads: a dot4 per leaf (dims 4l..4l+3), the 32 +// leaves summed on a stack in the tree of the warp form's butterfly (bit-reversed order: leaves l and l + 16 first), +// so each score is bit for bit the warp form's. The warp reads its queries' values all at once, a shared-memory +// broadcast at about half the cost of distinct reads. Persistent and balanced: block g takes units +// [g U / G, (g + 1) U / G) of the U = tiles x chunks of BATCH_QC queries, tile-major. The tile's keys are loaded +// between barriers when the tile changes; each warp stages its queries into a slot of its own, the next unit's even +// leaves once this unit is done with them, then the odd ones, so the copies run while it computes. A query's tail +// block scores the `dead` key in the warp form. +__global__ void __launch_bounds__(BATCH_T, 1) batch_scores_kernel(const float* __restrict__ pooled, + const float* __restrict__ dead, + const float* __restrict__ q_idx, + const int32_t* __restrict__ steps, int nq, + int64_t max_blocks, float* __restrict__ out) { + extern __shared__ __align__(16) float smem[]; + __shared__ int s_top; + const int t = threadIdx.x, lane = t & 31, warp = t >> 5; + float* ks = smem; // BATCH_TILE x IDX_DIM + float* qs = smem + BATCH_TILE * IDX_DIM + warp * BATCH_QT * QD; + { + int top = -1; // the batch's last block + for (int i = t; i < nq; i += BATCH_T) { + const int n_bid = steps[(int64_t) i * kStepCount + kStepNBid]; + top = max(top, (int64_t) n_bid < max_blocks - 1 ? n_bid : (int) (max_blocks - 1)); + } + top = __reduce_max_sync(FULL, top); + if (t == 0) s_top = -1; + __syncthreads(); + if (lane == 0) atomicMax(&s_top, top); + __syncthreads(); } - for (int it = 0; it < TC_ITER; ++it) { - const int64_t b0 = ((int64_t) blockIdx.x * TC_ITER + it) * TC_NB; - if (b0 >= reach || b0 >= hi_nbid) break; // blocks >= every query's n_bid: nothing to score - __syncthreads(); // the previous tile's reads are done - for (int i = t; i < TC_NB * IDX_DIM / 4; i += 128) { - const int r = i / (IDX_DIM / 4), c = i % (IDX_DIM / 4); - float4 v = make_float4(0.f, 0.f, 0.f, 0.f); - if (b0 + r < hi_nbid) v = reinterpret_cast(pooled + (b0 + r) * IDX_DIM)[c]; - *reinterpret_cast(sK + r * TC_KS + c * 4) = v; + const int top = s_top; + if (top < 0) return; + const int n_chunks = (nq + BATCH_QC - 1) / BATCH_QC, units = (top / BATCH_TILE + 1) * n_chunks; + const int u1 = (int) ((int64_t) (blockIdx.x + 1) * units / gridDim.x); + // leaves of parity `half` of queries q0, q0 + 1 into the warp's slot, one cp.async group + auto stage = [&](int q0, int half) { +#pragma unroll + for (int e = lane; e < BATCH_QT * IDX_HEADS * 16; e += 32) { + const int a = e / (IDX_HEADS * 16), off = (e / 16) % IDX_HEADS * IDX_DIM + 4 * (2 * (e % 16) + half); + if (q0 + a < nq) cp_async16(qs + a * QD + off, q_idx + (int64_t) (q0 + a) * QD + off); } + asm volatile("cp.async.commit_group;\n" ::); + }; + int have = -1; // the first query of the slot's rows + for (int u = (int) ((int64_t) blockIdx.x * units / gridDim.x); u < u1;) { + const int b0 = u / n_chunks * BATCH_TILE, ue = min(u1, (u / n_chunks + 1) * n_chunks); + __syncthreads(); // every warp is done with the previous tile + for (int e = t; e < BATCH_TILE * 32; e += BATCH_T) { + const int r = e >> 5, c = e & 31; + cp_async16(ks + key_at(r, c), pooled + (int64_t) min(b0 + r, top) * IDX_DIM + 4 * c); + } + asm volatile("cp.async.commit_group;\n" ::); + asm volatile("cp.async.wait_group 0;\n" ::); __syncthreads(); - float acc[IDX_HEADS][4]; -#pragma unroll - for (int h = 0; h < IDX_HEADS; ++h) acc[h][0] = acc[h][1] = acc[h][2] = acc[h][3] = 0.f; - const float* kr = sK + (warp * 8 + gid) * TC_KS; -#pragma unroll 4 - for (int k0 = 0; k0 < IDX_DIM; k0 += 8) { - const float kx0 = kr[k0 + tig], kx1 = kr[k0 + tig + 4]; - uint32_t bh[2], bl[2]; - bh[0] = tf32_hi(kx0); - bh[1] = tf32_hi(kx1); - bl[0] = tf32_hi(kx0 - __uint_as_float(bh[0])); - bl[1] = tf32_hi(kx1 - __uint_as_float(bh[1])); -#pragma unroll - for (int h = 0; h < IDX_HEADS; ++h) { - const float* qa = sQ + h * IDX_DIM + k0 + tig; - const float x0 = qa[gid * TC_QS], x1 = qa[(gid + 8) * TC_QS]; - const float x2 = qa[gid * TC_QS + 4], x3 = qa[(gid + 8) * TC_QS + 4]; - uint32_t ah[4], al[4]; - ah[0] = tf32_hi(x0); ah[1] = tf32_hi(x1); ah[2] = tf32_hi(x2); ah[3] = tf32_hi(x3); - al[0] = tf32_hi(x0 - __uint_as_float(ah[0])); - al[1] = tf32_hi(x1 - __uint_as_float(ah[1])); - al[2] = tf32_hi(x2 - __uint_as_float(ah[2])); - al[3] = tf32_hi(x3 - __uint_as_float(ah[3])); - mma_tf32(acc[h], al, bh); - mma_tf32(acc[h], ah, bl); - mma_tf32(acc[h], ah, bh); + for (; u < ue; ++u) { + const int q0 = u % n_chunks * BATCH_QC + warp * BATCH_QT; + int n_bid[BATCH_QT], n_kv[BATCH_QT], wlast = -1; +#pragma unroll + for (int a = 0; a < BATCH_QT; ++a) { + const bool in = q0 + a < nq; + n_bid[a] = in ? steps[(int64_t) (q0 + a) * kStepCount + kStepNBid] : -1; + n_kv[a] = in ? steps[(int64_t) (q0 + a) * kStepCount + kStepNKv] : 0; + wlast = max(wlast, n_bid[a]); } - } - // relu per head, heads added in order (as the warp kernel), written where the block completed for the query - const int64_t bc = b0 + warp * 8 + 2 * tig; -#pragma unroll - for (int half = 0; half < 2; ++half) { - const int qr = gid + half * 8; - const int64_t qi = q0 + qr; - if (qi >= nq) continue; - const int nb = s_nbid[qr]; -#pragma unroll - for (int j = 0; j < 2; ++j) { - const int64_t b = bc + j; - if (b >= nb || b >= max_blocks) continue; - float score = 0.0f; -#pragma unroll - for (int h = 0; h < IDX_HEADS; ++h) { - const float d = acc[h][half * 2 + j]; - score += d > 0.0f ? d : 0.0f; + if (wlast < b0) continue; // the warp's queries end before the tile + if (have != q0) { + // a skipped unit's prefetch may still be landing in the slot: copies to one address are unordered + asm volatile("cp.async.wait_group 0;\n" ::); + __syncwarp(); + stage(q0, 0); + stage(q0, 1); + have = q0; + } + const int q0n = u + 1 < u1 ? (u + 1) % n_chunks * BATCH_QC + warp * BATCH_QT : -1; + const bool ahead = q0n >= 0 && q0n != q0; + asm volatile("cp.async.wait_group 1;\n" ::); // the even leaves + __syncwarp(); + float stk[BATCH_QT][IDX_HEADS][BATCH_BT][6]; + auto leaves = [&](auto first) { +#pragma unroll + for (int ii = 0; ii < 16; ++ii) { + const int i = decltype(first)::value + ii, l = rev5(i), depth = popc5(i); + float4 kf[BATCH_BT], qf[BATCH_QT][IDX_HEADS]; +#pragma unroll + for (int j = 0; j < BATCH_BT; ++j) + kf[j] = *reinterpret_cast(ks + key_at(lane + 32 * j, l)); +#pragma unroll + for (int a = 0; a < BATCH_QT; ++a) +#pragma unroll + for (int h = 0; h < IDX_HEADS; ++h) + qf[a][h] = *reinterpret_cast(qs + a * QD + h * IDX_DIM + 4 * l); +#pragma unroll + for (int a = 0; a < BATCH_QT; ++a) +#pragma unroll + for (int h = 0; h < IDX_HEADS; ++h) +#pragma unroll + for (int j = 0; j < BATCH_BT; ++j) { + float* s = stk[a][h][j]; + s[depth] = dot4(kf[j], qf[a][h]); +#pragma unroll + for (int lv = 0; lv < 5; ++lv) { // complete subtrees combine + if (((i + 1) >> lv) & 1) break; + s[depth - lv - 1] = s[depth - lv - 1] + s[depth - lv]; + } + } + } + }; + leaves(std::integral_constant()); // rev5(0..15): the even leaves + __syncwarp(); + if (ahead) { + stage(q0n, 0); + asm volatile("cp.async.wait_group 1;\n" ::); // this unit's odd leaves + } else { + asm volatile("cp.async.wait_group 0;\n" ::); + } + __syncwarp(); + leaves(std::integral_constant()); + __syncwarp(); + if (ahead) { + stage(q0n, 1); + have = q0n; + } +#pragma unroll + for (int a = 0; a < BATCH_QT; ++a) { + const int qi = q0 + a; + const int64_t last = (int64_t) n_bid[a] < max_blocks - 1 ? n_bid[a] : max_blocks - 1; +#pragma unroll + for (int j = 0; j < BATCH_BT; ++j) { + float score = 0.0f; +#pragma unroll + for (int h = 0; h < IDX_HEADS; ++h) { + const float d = stk[a][h][j][0]; + score += d > 0.0f ? d : 0.0f; + } + const int b = b0 + lane + 32 * j; + if (qi < nq && b <= last && b != n_bid[a]) out[(int64_t) qi * max_blocks + b] = score; + } + if (qi < nq && n_bid[a] >= b0 && n_bid[a] < b0 + BATCH_TILE && (int64_t) n_bid[a] <= max_blocks - 1) { + const float4 k4 = *reinterpret_cast(dead + 4 * lane); + float score = 0.0f; +#pragma unroll + for (int h = 0; h < IDX_HEADS; ++h) { + float d = dot4(k4, *reinterpret_cast(q_idx + (int64_t) qi * QD + h * IDX_DIM + + 4 * lane)); +#pragma unroll + for (int o = 16; o > 0; o >>= 1) d += __shfl_xor_sync(FULL, d, o); + score += d > 0.0f ? d : 0.0f; + } + if (lane == 0) out[(int64_t) qi * max_blocks + n_bid[a]] = n_kv[a] % R != 0 ? score + 1e9f : score; } - out[qi * max_blocks + b] = score; } } } } -// the tail block n_bid of each query: exactly block_scores_kernel's arithmetic for that block -__global__ void __launch_bounds__(32) block_scores_tail_kernel(const float* __restrict__ dead, - const float* __restrict__ q_idx, - const int32_t* __restrict__ steps, - int64_t max_blocks, float* __restrict__ out) { - const int64_t qi = blockIdx.x; - const int32_t* st = steps + qi * kStepCount; - const int64_t n_kv = st[kStepNKv], n_bid = st[kStepNBid]; - if (n_bid >= max_blocks) return; - const int lane = threadIdx.x & 31; - const float4 k4 = *reinterpret_cast(dead + lane * 4); - const float* q = q_idx + qi * IDX_HEADS * IDX_DIM + lane * 4; - float score = 0.0f; -#pragma unroll - for (int h = 0; h < IDX_HEADS; ++h) { - const float4 q4 = *reinterpret_cast(q + h * IDX_DIM); - float d = k4.x * q4.x + k4.y * q4.y + k4.z * q4.z + k4.w * q4.w; +__device__ __forceinline__ int warp_sum(int v) { #pragma unroll - for (int o = 16; o > 0; o >>= 1) d += __shfl_xor_sync(0xffffffffu, d, o); - score += d > 0.0f ? d : 0.0f; - } - if (lane == 0) { - if (n_kv % R != 0) score += 1e9f; - out[qi * max_blocks + n_bid] = score; - } + for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(FULL, v, o); + return v; } -// ---- the same top-k with each query's keys read once: 1,024 threads hold up to TK_PER consecutive blocks' keys in -// registers (contexts up to 4 * 1024 * TK_PER cells), per-warp histograms, block-wide scans. The selection rule is -// block_topk_kernel's (radix threshold, ties to the lowest index, cells ascending): identical ids. -constexpr int TK_T = 1024; -constexpr int TK_PER = 33; - -__device__ __forceinline__ int block_excl_scan(int v, int* s_warp, int& total) { +// The exclusive prefixes of a and b over the block's threads in thread order, and their totals. +__device__ __forceinline__ void block_exclusive2(int a, int b, int& ea, int& eb, int& ta, int& tb, int* s_w) { const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; - int x = v; + int ia = a, ib = b; #pragma unroll for (int o = 1; o < 32; o <<= 1) { - const int y = __shfl_up_sync(0xffffffffu, x, o); - if (lane >= o) x += y; + const int ya = __shfl_up_sync(FULL, ia, o), yb = __shfl_up_sync(FULL, ib, o); + if (lane >= o) { ia += ya; ib += yb; } } - if (lane == 31) s_warp[warp] = x; + if (lane == 31) { s_w[warp] = ia; s_w[TOPK_W + warp] = ib; } __syncthreads(); if (warp == 0) { - int w = s_warp[lane]; - int z = w; + const int wa = s_w[lane], wb = s_w[TOPK_W + lane]; + int xa = wa, xb = wb; #pragma unroll for (int o = 1; o < 32; o <<= 1) { - const int y = __shfl_up_sync(0xffffffffu, z, o); - if (lane >= o) z += y; + const int ya = __shfl_up_sync(FULL, xa, o), yb = __shfl_up_sync(FULL, xb, o); + if (lane >= o) { xa += ya; xb += yb; } } - s_warp[lane] = z - w; // exclusive per warp - if (lane == 31) s_warp[32] = z; // total + s_w[lane] = xa - wa; + s_w[TOPK_W + lane] = xb - wb; + if (lane == 31) { s_w[2 * TOPK_W] = xa; s_w[2 * TOPK_W + 1] = xb; } } __syncthreads(); - const int r = s_warp[warp] + x - v; - total = s_warp[32]; + ea = s_w[warp] + ia - a; + eb = s_w[TOPK_W + warp] + ib - b; + ta = s_w[2 * TOPK_W]; + tb = s_w[2 * TOPK_W + 1]; __syncthreads(); - return r; } -__global__ void __launch_bounds__(TK_T) block_topk_reg_kernel(const float* __restrict__ scores, - const int32_t* __restrict__ steps, int64_t max_blocks, - int64_t cap, int32_t* __restrict__ ids) { - __shared__ int hist[TK_T / 32][256]; - __shared__ int s_warp[33]; +// The same selection as the per-cell topk_kernel: the `width` cells of the largest block scores (a block's score +// counts for each of its cells), ties to the lowest cell, emitted ascending. A radix select, 8 bits at a time, over +// the order keys, with a histogram per warp. The complete blocks go in rows of TOPK_T: thread t has block +// row * TOPK_T + t, the first TOPK_ROWS rows in registers (read once, 131K cells), the others from memory in each +// pass. A block's cells then land at (cells above the threshold before it) + min(cells at it before it, budget): +// ballots within a (row, warp), one scan over the (row, warp) totals. The tail block (n_bid) is thread 0's. +__global__ void __launch_bounds__(TOPK_T) block_topk_kernel(const float* __restrict__ scores, + const int32_t* __restrict__ steps, int64_t max_blocks, + int64_t cap, int32_t* __restrict__ ids) { + __shared__ int sm[TOPK_W * 256]; // the per-warp histograms; then the scan table and the output + __shared__ int hist[256]; + __shared__ int s_w[2 * TOPK_W + 2]; __shared__ int s_digit, s_above; const int64_t qi = blockIdx.x; const int32_t* st = steps + qi * kStepCount; const int64_t n_kv = st[kStepNKv], n_bid = st[kStepNBid], width = st[kStepWidth]; int32_t* out = ids + qi * cap; const int t = threadIdx.x, lane = t & 31, warp = t >> 5; - if (n_kv <= width) { - for (int64_t j = t; j < n_kv; j += TK_T) out[j] = (int32_t) j; + if (n_kv <= width) { // everything is selected: the identity, ascending + for (int64_t j = t; j < n_kv; j += TOPK_T) out[j] = (int32_t) j; return; } const float* sc = scores + qi * max_blocks; - const int64_t nb = n_bid + 1; - const int64_t per = (nb + TK_T - 1) / TK_T; // <= TK_PER (the caller checks) - const int64_t b0 = (int64_t) t * per, b1 = (b0 + per < nb) ? b0 + per : nb; - uint32_t key[TK_PER]; + const int tail_w = (int) (n_kv - n_bid * R); // the cells of block n_bid, 0..3 + const int nrows = (int) ((n_bid + TOPK_T - 1) / TOPK_T); + uint32_t key[TOPK_ROWS]; #pragma unroll - for (int j = 0; j < TK_PER; ++j) key[j] = (b0 + j < b1) ? order_key(sc[b0 + j]) : 0u; - auto weight = [&](int64_t b) -> int { return b < n_bid ? R : (int) (n_kv - n_bid * R); }; + for (int i = 0; i < TOPK_ROWS; ++i) { + const int64_t b = (int64_t) i * TOPK_T + t; + key[i] = b < n_bid ? order_key(sc[b]) : 0u; + } + // ---- radix select: the largest key thr with (cells with key >= thr) >= width uint32_t prefix = 0; - int above = 0; + int above = 0; // cells strictly above the digits fixed so far + int* whist = sm + warp * 256; for (int shift = 24; shift >= 0; shift -= 8) { - for (int i = lane; i < 256; i += 32) hist[warp][i] = 0; + for (int i = lane; i < 256; i += 32) whist[i] = 0; __syncwarp(); const uint32_t hi_mask = shift == 24 ? 0u : (0xffffffffu << (shift + 8)); + int cd = -1, cc = 0; // a run of one digit is added at once + auto count = [&](uint32_t k) { + if ((k & hi_mask) != (prefix & hi_mask)) return; + const int d = (int) ((k >> shift) & 255); + if (d != cd) { + if (cc != 0) atomicAdd(&whist[cd], cc); + cd = d; + cc = 0; + } + cc += R; + }; #pragma unroll - for (int j = 0; j < TK_PER; ++j) { - const int64_t b = b0 + j; - if (b >= b1) break; - const int w = weight(b); - if (w == 0) continue; - if ((key[j] & hi_mask) == (prefix & hi_mask)) atomicAdd(&hist[warp][(key[j] >> shift) & 255], w); + for (int i = 0; i < TOPK_ROWS; ++i) + if ((int64_t) i * TOPK_T + t < n_bid) count(key[i]); + for (int i = TOPK_ROWS; i < nrows; ++i) { + const int64_t b = (int64_t) i * TOPK_T + t; + if (b < n_bid) count(order_key(sc[b])); + } + if (cc != 0) atomicAdd(&whist[cd], cc); + if (t == 0 && tail_w > 0) { + const uint32_t k = order_key(sc[n_bid]); + if ((k & hi_mask) == (prefix & hi_mask)) atomicAdd(&whist[(k >> shift) & 255], tail_w); } __syncthreads(); - if (t < 256) { // fold the warps' histograms into warp 0's + if (t < 256) { int s = 0; - for (int w2 = 0; w2 < TK_T / 32; ++w2) s += hist[w2][t]; - hist[0][t] = s; +#pragma unroll 8 + for (int w = 0; w < TOPK_W; ++w) s += sm[w * 256 + t]; + hist[t] = s; } __syncthreads(); - if (t == 0) { - int cum = above, d = 255; - for (; d > 0; --d) { - if (cum + hist[0][d] >= width) break; - cum += hist[0][d]; + if (warp == 0) { + // the digit: the largest d >= 1 with above + (cells with a digit >= d) >= width, else 0; lane l holds the + // bins 255 - 8l down to 248 - 8l + int v[8], s = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) { v[i] = hist[255 - (lane * 8 + i)]; s += v[i]; } + int incl = s; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int y = __shfl_up_sync(FULL, incl, o); + if (lane >= o) incl += y; + } + int run = incl - s, below = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) { + run += v[i]; + if (255 - (lane * 8 + i) >= 1 && above + run < width) ++below; + } + const int digit = 255 - warp_sum(below); + int gt = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) + if (255 - (lane * 8 + i) > digit) gt += v[i]; + gt = warp_sum(gt); + if (lane == 0) { + s_digit = digit; + s_above = above + gt; } - s_digit = d; - s_above = cum; } __syncthreads(); prefix |= (uint32_t) s_digit << shift; above = s_above; - __syncthreads(); } const uint32_t thr = prefix; - const int64_t eq_budget = width - above; - int gt = 0, eq = 0; -#pragma unroll - for (int j = 0; j < TK_PER; ++j) { - const int64_t b = b0 + j; - if (b >= b1) break; - const int w = weight(b); - if (w == 0) continue; - if (key[j] > thr) gt += w; - else if (key[j] == thr) eq += w; + const int budget = (int) (width - above); // cells equal to thr that fit, lowest index first + // ---- the cells above and at thr per (row, warp), their exclusive prefixes in block order + int* tab_gt = sm; + int* tab_eq = sm + TOPK_MAX_ROWS * 32; + int* obuf = sm + 2 * TOPK_MAX_ROWS * 32; // the selection, copied out coalesced + __syncthreads(); // the histograms are no longer read + auto row_key = [&](int i) { return i < TOPK_ROWS ? 0u : order_key(sc[(int64_t) i * TOPK_T + t]); }; +#pragma unroll + for (int i = 0; i < TOPK_ROWS; ++i) { + if (i >= nrows) break; + const bool in = (int64_t) i * TOPK_T + t < n_bid; + const unsigned mg = __ballot_sync(FULL, in && key[i] > thr), me = __ballot_sync(FULL, in && key[i] == thr); + if (lane == 0) { tab_gt[i * 32 + warp] = R * __popc(mg); tab_eq[i * 32 + warp] = R * __popc(me); } } - int tot; - const int eq_before = block_excl_scan(eq, s_warp, tot); - int64_t my_eq = eq_budget - eq_before; - if (my_eq < 0) my_eq = 0; - if (my_eq > eq) my_eq = eq; - const int sel = gt + (int) my_eq; - int64_t wpos = block_excl_scan(sel, s_warp, tot); - int64_t eq_left = my_eq; -#pragma unroll - for (int j = 0; j < TK_PER; ++j) { - const int64_t b = b0 + j; - if (b >= b1) break; - const int w = weight(b); - if (w == 0) continue; - if (key[j] > thr) { - for (int c = 0; c < w; ++c) out[wpos++] = (int32_t) (b * R + c); - } else if (key[j] == thr) { - for (int c = 0; c < w && eq_left > 0; ++c, --eq_left) out[wpos++] = (int32_t) (b * R + c); - } + for (int i = TOPK_ROWS; i < nrows; ++i) { + const bool in = (int64_t) i * TOPK_T + t < n_bid; + const uint32_t k = in ? row_key(i) : 0u; + const unsigned mg = __ballot_sync(FULL, in && k > thr), me = __ballot_sync(FULL, in && k == thr); + if (lane == 0) { tab_gt[i * 32 + warp] = R * __popc(mg); tab_eq[i * 32 + warp] = R * __popc(me); } + } + __syncthreads(); + int gt_all, eq_all; + { + const int e = nrows * 32, k0 = 2 * t; // entries (row, warp) in block order, two per thread + const int g0 = k0 < e ? tab_gt[k0] : 0, e0 = k0 < e ? tab_eq[k0] : 0; + const int g1 = k0 + 1 < e ? tab_gt[k0 + 1] : 0, e1 = k0 + 1 < e ? tab_eq[k0 + 1] : 0; + int gx, ex; + block_exclusive2(g0 + g1, e0 + e1, gx, ex, gt_all, eq_all, s_w); + if (k0 < e) { tab_gt[k0] = gx; tab_eq[k0] = ex; } + if (k0 + 1 < e) { tab_gt[k0 + 1] = gx + g0; tab_eq[k0 + 1] = ex + e0; } } + __syncthreads(); + // ---- the cells, into shared memory + const unsigned lt = (1u << lane) - 1u; + auto emit = [&](int i, uint32_t k, bool in) { + const bool g = in && k > thr, q = in && k == thr; + const unsigned mg = __ballot_sync(FULL, g), me = __ballot_sync(FULL, q); + if (!g && !q) return; + const int gb = tab_gt[i * 32 + warp] + R * __popc(mg & lt); + const int eb = tab_eq[i * 32 + warp] + R * __popc(me & lt); + const int pos = gb + (eb < budget ? eb : budget); + const int take = g ? R : (budget - eb < 0 ? 0 : (budget - eb > R ? R : budget - eb)); + const int64_t b = (int64_t) i * TOPK_T + t; + for (int c = 0; c < take; ++c) obuf[pos + c] = (int32_t) (b * R + c); + }; +#pragma unroll + for (int i = 0; i < TOPK_ROWS; ++i) { + if (i >= nrows) break; + emit(i, key[i], (int64_t) i * TOPK_T + t < n_bid); + } + for (int i = TOPK_ROWS; i < nrows; ++i) { + const bool in = (int64_t) i * TOPK_T + t < n_bid; + emit(i, in ? row_key(i) : 0u, in); + } + if (t == 0 && tail_w > 0) { // the tail block comes last + const uint32_t k = order_key(sc[n_bid]); + const int pos = gt_all + (eq_all < budget ? eq_all : budget); + const int take = k > thr ? tail_w : (k == thr ? (budget - eq_all < 0 ? 0 : (budget - eq_all > tail_w ? tail_w + : budget - eq_all)) : 0); + for (int c = 0; c < take; ++c) obuf[pos + c] = (int32_t) (n_bid * R + c); + } + __syncthreads(); + for (int64_t j = t; j < width; j += TOPK_T) out[j] = obuf[j]; } } // namespace void qsa_block_scores(const float* pooled, const float* dead, const float* q_idx, const int32_t* steps, int64_t nq, - int64_t max_blocks, const QsaShapes& s, float* scores, void* stream, int64_t active_blocks) { + int64_t max_blocks, const QsaShapes& s, float* scores, void* stream, int64_t grid_blocks) { if (nq <= 0) return; - if (s.idx_dim != IDX_DIM || s.idx_n_head != IDX_HEADS || s.idx_block != R || nq > 65535) { + if (s.idx_dim != IDX_DIM || s.idx_n_head != IDX_HEADS || s.idx_block != R || nq > 65535 || grid_blocks < 1) { std::fprintf(stderr, "qsa_block_scores: unsupported indexer geometry\n"); std::exit(1); } - // a block past a query's n_bid returns at once: the grid need only reach the batch's largest n_bid (C-1) - const int64_t reach = active_blocks > 0 && active_blocks < max_blocks ? active_blocks : max_blocks; - const dim3 grid((unsigned) ((reach + SCORE_WARPS - 1) / SCORE_WARPS), (unsigned) nq); - block_scores_kernel<<>>(pooled, dead, q_idx, steps, max_blocks, - scores); - const cudaError_t e = cudaGetLastError(); - if (e != cudaSuccess) { std::fprintf(stderr, "qsa_block_scores: %s\n", cudaGetErrorString(e)); std::exit(1); } -} - -bool qsa_block_scores_tc(const float* pooled, const float* dead, const float* q_idx, const int32_t* steps, int64_t nq, - int64_t max_blocks, const QsaShapes& s, float* scores, void* stream, int64_t active_blocks) { - if (nq <= 0) return true; - if (s.idx_dim != IDX_DIM || s.idx_n_head != IDX_HEADS || s.idx_block != R || nq > 65535 * TC_QT) return false; - { // sm_80 or newer (TF32 MMA); an older card keeps the warp kernel - static int cc_major[64] = {}; - int dev = 0; - if (cudaGetDevice(&dev) != cudaSuccess || dev < 0 || dev >= 64) { cudaGetLastError(); return false; } - if (cc_major[dev] == 0) { - int major = 0; - if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev) != cudaSuccess) { - cudaGetLastError(); - return false; - } - cc_major[dev] = major; + const int64_t covered = grid_blocks < max_blocks ? grid_blocks : max_blocks; + if (nq >= kBatchMinQueries) { + // one block an SM (the prompt path's GPU), each with the same share of the units + constexpr int smem = (BATCH_TILE * IDX_DIM + BATCH_QC * QD) * (int) sizeof(float); + static int sms = 0; + if (sms == 0) { + int dev = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + cudaFuncSetAttribute(batch_scores_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); } - if (cc_major[dev] < 8) return false; + const int64_t units = ((covered + BATCH_TILE - 1) / BATCH_TILE) * ((nq + BATCH_QC - 1) / BATCH_QC); + const unsigned grid = (unsigned) (units < sms ? units : sms); + batch_scores_kernel<<>>(pooled, dead, q_idx, steps, (int) nq, + max_blocks, scores); + } else { + const unsigned grid = (unsigned) ((covered + SCORE_TILE - 1) / SCORE_TILE); + block_scores_kernel<<>>(pooled, dead, q_idx, steps, (int) nq, + max_blocks, scores); } -#if defined(__HIPCC__) - return false; // the tensor-core kernel is compiled out on AMD (its major version is not a CUDA sm) -#endif - static bool attr[64] = {}; // the shared-memory opt-in is per device (a layer split runs it on several) - int adev = 0; - cudaGetDevice(&adev); - const int bytes = (TC_QT * TC_QS + TC_NB * TC_KS) * (int) sizeof(float); - if (!attr[adev]) { - if (cudaFuncSetAttribute(block_scores_tc_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes) != - cudaSuccess) { - cudaGetLastError(); - return false; - } - attr[adev] = true; - } - const int64_t reach = active_blocks > 0 && active_blocks < max_blocks ? active_blocks : max_blocks; - const int64_t per = (int64_t) TC_NB * TC_ITER; - const dim3 grid((unsigned) ((reach + per - 1) / per), (unsigned) ((nq + TC_QT - 1) / TC_QT)); - block_scores_tc_kernel<<>>(pooled, q_idx, steps, nq, max_blocks, reach, - scores); - block_scores_tail_kernel<<<(unsigned) nq, 32, 0, (cudaStream_t) stream>>>(dead, q_idx, steps, max_blocks, scores); const cudaError_t e = cudaGetLastError(); - if (e != cudaSuccess) { std::fprintf(stderr, "qsa_block_scores_tc: %s\n", cudaGetErrorString(e)); std::exit(1); } - return true; -} - -void qsa_block_topk_ref(const float* scores, const int32_t* steps, int64_t nq, int64_t max_blocks, int64_t cap, - const QsaShapes& s, int32_t* ids, void* stream) { - if (nq <= 0) return; - if (s.idx_block != R || cap < qsa_selection_width(kTopkMaxCells, s)) { - std::fprintf(stderr, "qsa_block_topk: unsupported geometry or cap\n"); - std::exit(1); - } - block_topk_kernel<<<(unsigned) nq, TOPK_T, 0, (cudaStream_t) stream>>>(scores, steps, max_blocks, cap, ids); - const cudaError_t e = cudaGetLastError(); - if (e != cudaSuccess) { std::fprintf(stderr, "qsa_block_topk: %s\n", cudaGetErrorString(e)); std::exit(1); } + if (e != cudaSuccess) { std::fprintf(stderr, "qsa_block_scores: %s\n", cudaGetErrorString(e)); std::exit(1); } } void qsa_block_topk(const float* scores, const int32_t* steps, int64_t nq, int64_t max_blocks, int64_t cap, const QsaShapes& s, int32_t* ids, void* stream) { - // keys in registers when every query's blocks fit (contexts up to ~135K cells); the same ids. STRATA_TOPK_OLD=1: - // the kernel that reads them from memory on every pass - static const bool old = std::getenv("STRATA_TOPK_OLD") != nullptr; if (nq <= 0) return; - if (old || max_blocks > (int64_t) TK_T * TK_PER) { - qsa_block_topk_ref(scores, steps, nq, max_blocks, cap, s, ids, stream); - return; - } - if (s.idx_block != R || cap < qsa_selection_width(kTopkMaxCells, s)) { + if (s.idx_block != R || cap < qsa_selection_width(kTopkMaxCells, s) || + max_blocks > (int64_t) TOPK_MAX_ROWS * TOPK_T) { std::fprintf(stderr, "qsa_block_topk: unsupported geometry or cap\n"); std::exit(1); } - block_topk_reg_kernel<<<(unsigned) nq, TK_T, 0, (cudaStream_t) stream>>>(scores, steps, max_blocks, cap, ids); + block_topk_kernel<<<(unsigned) nq, TOPK_T, 0, (cudaStream_t) stream>>>(scores, steps, max_blocks, cap, ids); const cudaError_t e = cudaGetLastError(); if (e != cudaSuccess) { std::fprintf(stderr, "qsa_block_topk: %s\n", cudaGetErrorString(e)); std::exit(1); } } diff --git a/src/kernels/cuda/quantize_act.cu b/src/kernels/cuda/quantize_act.cu index ddd14dfbb0a..0e05a6c3546 100644 --- a/src/kernels/cuda/quantize_act.cu +++ b/src/kernels/cuda/quantize_act.cu @@ -97,33 +97,35 @@ __global__ void quantize_q8_0_kernel(const float* __restrict__ x, uint8_t* __res /// `moe_hit_parity`. This function's job is different: it must match **this engine's own CPU reference**, /// because a hit and a miss for the same expert on the same layer have to produce the same number. The CPU /// path is the reference - C1 passes on it - so the hit path is brought to it, not the reverse. +// A warp a block, lane i on value i: the block's |max| is a max over the same values (0 for none above it, as the +// serial scan from 0 through fmaxf), so every byte is the per-thread loop's. __global__ void quantize_q8_0_scaled_kernel(const float* __restrict__ x, uint8_t* __restrict__ blocks, float* __restrict__ scales, long long n_blocks) { - const long long b = (long long) blockIdx.x * blockDim.x + threadIdx.x; + const long long b = ((long long) blockIdx.x * blockDim.x + threadIdx.x) >> 5; + const int i = (int) (threadIdx.x & 31); if (b >= n_blocks) return; - const float* xb = x + b * QK8_0; + const float xi = x[b * QK8_0 + i]; uint8_t* out = blocks + b * 34; - float amax = 0.0f; - for (int i = 0; i < QK8_0; ++i) amax = fmaxf(amax, fabsf(xb[i])); + float amax = fmaxf(0.0f, fabsf(xi)); + for (int o = 16; o > 0; o >>= 1) amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, o)); // VERBATIM from `cpu/expert.cpp:144-145`, including the `amax > 0` guard, so the fp32 value written here // is bit-identical to the `s` the CPU path used. const float s = amax > 0.f ? amax / 127.f : 0.f; const float inv = s > 0.f ? 1.f / s : 0.f; - scales[b] = s; - - const uint16_t d16bits = f16_from_f32(s); - out[0] = (uint8_t) (d16bits & 0xFF); - out[1] = (uint8_t) (d16bits >> 8); - for (int i = 0; i < QK8_0; ++i) { - // VERBATIM from `cpu/expert.cpp:159-162`: reciprocal multiply, then `t + copysign(0.5, t)` truncated - // toward zero, which is `lround`'s rule - round half away from zero. - const float t = xb[i] * inv; - const float r = t + (t >= 0.f ? 0.5f : -0.5f); - int v = (int) r; - v = v < -127 ? -127 : (v > 127 ? 127 : v); - out[2 + i] = (uint8_t) (int8_t) v; + if (i == 0) { + scales[b] = s; + const uint16_t d16bits = f16_from_f32(s); + out[0] = (uint8_t) (d16bits & 0xFF); + out[1] = (uint8_t) (d16bits >> 8); } + // VERBATIM from `cpu/expert.cpp:159-162`: reciprocal multiply, then `t + copysign(0.5, t)` truncated + // toward zero, which is `lround`'s rule - round half away from zero. + const float t = xi * inv; + const float r = t + (t >= 0.f ? 0.5f : -0.5f); + int v = (int) r; + v = v < -127 ? -127 : (v > 127 ? 127 : v); + out[2 + i] = (uint8_t) (int8_t) v; } __global__ void dequant_q8_0_kernel(const uint8_t* __restrict__ blocks, float* __restrict__ x, @@ -254,7 +256,7 @@ void quantize_q8_0_scaled(const float* x, uint8_t* blocks, float* scales, int64_ } const long long nb = n / QK8_0; const int threads = 128; - const unsigned grid = (unsigned) ((nb + threads - 1) / threads); + const unsigned grid = (unsigned) ((nb + threads / 32 - 1) / (threads / 32)); // a warp a block quantize_q8_0_scaled_kernel<<>>(x, blocks, scales, nb); const cudaError_t e = cudaGetLastError(); if (e != cudaSuccess) { diff --git a/src/kernels/cuda/s2_gemv_quads.cu b/src/kernels/cuda/s2_gemv_quads.cu index 7cb7cb4422c..90973ab70a8 100644 --- a/src/kernels/cuda/s2_gemv_quads.cu +++ b/src/kernels/cuda/s2_gemv_quads.cu @@ -28,7 +28,6 @@ namespace { // QK = 64 elements per group; a quad of 4 elements is 1/16 of a group, so // group = (quad * 4) >> 6 = quad >> 4 constexpr int QK_S2 = 64; -constexpr int CODES_PER_BYTE_S2 = 4; __global__ void s2_gemv_quads_kernel(const uint16_t* __restrict__ x, const uint8_t* __restrict__ codes, const float* __restrict__ scales, float* __restrict__ y, long long n_in, diff --git a/src/kernels/cuda/shared_expert.cu b/src/kernels/cuda/shared_expert.cu index d795bf84001..16e025e0ca3 100644 --- a/src/kernels/cuda/shared_expert.cu +++ b/src/kernels/cuda/shared_expert.cu @@ -68,11 +68,6 @@ __global__ void native_swiglu_kernel(const float* gate, const float* up, float* out[i] = __fdividef(gate[i], 1.0f + __expf(-gate[i])) * up[i]; } -__global__ void to_f16_kernel(const float* __restrict__ in, uint16_t* __restrict__ out, int n) { - const int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) out[i] = f16_from_f32(in[i]); -} - // The per-token scalar gate: sigmoid(dot(x, w)) with w = `ffn_gate_inp_shexp` (n_embd,). // // **ONE THREAD WAS THE LARGEST SINGLE KERNEL IN THE MODEL, AND THE COMMENT DEFENDING IT SAID WHY IT SHOULD @@ -133,9 +128,7 @@ __global__ void scale_kernel(float* __restrict__ out, const float* __restrict__ __global__ void native_scalar_sigmoid_kernel(float* gate) { // Match the single-token CUDA sigmoid's FP32 fast-math operations without changing legacy kernels' // compilation flags. The dot product was already reduced by the pinned native MMVF implementation. - gate[0] = __fdividef(1.0f, 1.0f + __expf(-gate[0])); -} -__global__ void native_scalar_sigmoid_multi_kernel(float* gate) { // thread t = token t, same expression + // One thread per token. gate[threadIdx.x] = __fdividef(1.0f, 1.0f + __expf(-gate[threadIdx.x])); } @@ -157,6 +150,7 @@ __global__ void moe_combine_kernel(const float* __restrict__ parts, const float* } // namespace void shared_expert_set_native_bf16(bool enabled) { native_bf16 = enabled; } +bool shared_expert_native_bf16() { return native_bf16; } namespace { __global__ void scale_rows_kernel(float* __restrict__ out, const float* __restrict__ g, int n) { @@ -166,34 +160,36 @@ __global__ void scale_rows_kernel(float* __restrict__ out, const float* __restri } } // namespace -void shared_expert_multi(int n_tok, const float* x, const uint16_t* x_bf16, const NativeSharedWeights& nw, - const uint16_t* gate_inp_bf16, float* gate, float* up, float* g, float* out, int64_t n_embd, - int64_t n_ff, void* stream) { - if (n_tok < 1 || n_tok > 8 || !nw.q8_1 || !nw.gate_data || !nw.up_data || !nw.down_data || !stream) - throw std::invalid_argument("shared_expert_multi: needs 1..8 tokens, native weights, scratch and a stream"); +void shared_expert_gate_rows(float* out, float* g, int64_t n_embd, int n_tok, void* stream) { cudaStream_t cs = (cudaStream_t) stream; - native_quantize_q8_1(x, nw.q8_1, (int) n_embd, n_tok, stream); - native_mmvq(nw.gate_type, nw.gate_data, nw.q8_1, gate, (int) n_embd, (int) n_ff, n_tok, stream); - native_mmvq(nw.up_type, nw.up_data, nw.q8_1, up, (int) n_embd, (int) n_ff, n_tok, stream); + native_scalar_sigmoid_kernel<<<1, (unsigned) n_tok, 0, cs>>>(g); + scale_rows_kernel<<>>( + out, g, (int) n_embd); +} + +void shared_expert_multi(int n_tok, const float* x, const NativeSharedWeights& nw, const uint16_t* gate_inp_bf16, + float* gate, float* up, float* g, float* out, int64_t n_embd, int64_t n_ff, void* stream, + void* gate_stream) { + if (n_tok < 1 || n_tok > 8 || !nw.q8_1 || !nw.gate_data || !nw.up_data || !nw.down_data || !stream || + !gate_stream || !native_bf16) + throw std::invalid_argument("shared_expert_multi: needs 1..8 tokens, native weights and gate, scratch and " + "streams"); + cudaStream_t cs = (cudaStream_t) stream; + bf16_gemv_fp32_mmvf_multi(x, gate_inp_bf16, g, n_embd, 1, n_tok, gate_stream); + // native_mmvq_il's kernels where they are the faster (1 token, or 2+ with the interleaved copy) + const bool il = n_tok == 1 || nw.q8_1_il; + if (n_tok >= 2 && nw.q8_1_il) native_quantize_q8_1_il(x, nw.q8_1, nw.q8_1_il, (int) n_embd, n_tok, stream); + else native_quantize_q8_1(x, nw.q8_1, (int) n_embd, n_tok, stream); + auto proj = [&](int type, const void* w, float* y) { + if (il) native_mmvq_il(type, w, nw.q8_1, nw.q8_1_il, y, (int) n_embd, (int) n_ff, n_tok, stream); + else native_mmvq(type, w, nw.q8_1, y, (int) n_embd, (int) n_ff, n_tok, stream); + }; + proj(nw.gate_type, nw.gate_data, gate); + proj(nw.up_type, nw.up_data, up); const int n = (int) (n_ff * n_tok); native_swiglu_kernel<<<(unsigned) ((n + THREADS - 1) / THREADS), THREADS, 0, cs>>>(gate, up, gate, n); native_quantize_q8_1(gate, nw.q8_1, (int) n_ff, n_tok, stream); native_mmvq(nw.down_type, nw.down_data, nw.q8_1, out, (int) n_ff, (int) n_embd, n_tok, stream); - static const bool batch = [] { const char* v = std::getenv("STRATA_DEC_BATCH"); return v == nullptr || std::atoi(v) != 0; }(); - if (native_bf16 && batch && n_tok > 1) { // one gemv for all rows (outputs identical), one sigmoid launch - bf16_gemv_fp32_mmvf_multi(x, n_embd, gate_inp_bf16, g, 1, n_embd, 1, n_tok, stream); - native_scalar_sigmoid_multi_kernel<<<1, n_tok, 0, cs>>>(g); - } else - for (int t = 0; t < n_tok; ++t) { - if (native_bf16) { - bf16_gemv_fp32_mmvf(x + (size_t) t * n_embd, gate_inp_bf16, g + t, n_embd, 1, stream); - native_scalar_sigmoid_kernel<<<1, 1, 0, cs>>>(g + t); - } else { - scalar_gate_kernel<<<1, 256, 0, cs>>>(x_bf16 + (size_t) t * n_embd, gate_inp_bf16, g + t, (int) n_embd); - } - } - scale_rows_kernel<<>>( - out, g, (int) n_embd); const cudaError_t e = cudaGetLastError(); if (e != cudaSuccess) throw std::runtime_error(std::string("shared_expert_multi: ") + cudaGetErrorString(e)); } @@ -307,11 +303,11 @@ void shared_expert(const uint8_t* x_q8_0, const uint8_t* x_q8k, const uint16_t* // round trip. 256 threads is the reduction's width, not the problem's size. if (use_native) { bf16_gemv_fp32_mmvf(x_f32, gate_inp_bf16, g, n_embd, 1, stream); - native_scalar_sigmoid_kernel<<<1, 1, 0, (cudaStream_t) stream>>>(g); + shared_expert_gate_rows(out, g, n_embd, 1, stream); } else { scalar_gate_kernel<<<1, 256, 0, (cudaStream_t) stream>>>(x_bf16, gate_inp_bf16, g, (int) n_embd); + scale_kernel<<>>(out, g, (int) n_embd); } - scale_kernel<<>>(out, g, (int) n_embd); if (stream == nullptr) { const cudaError_t e = cudaDeviceSynchronize(); diff --git a/src/kernels/cuda/verify_kernels.cu b/src/kernels/cuda/verify_kernels.cu index b8a3a619f79..786937f07f3 100644 --- a/src/kernels/cuda/verify_kernels.cu +++ b/src/kernels/cuda/verify_kernels.cu @@ -3,9 +3,11 @@ // The per-token arithmetic of every kernel here is transcribed from its single-token original (fused_gdn.cu, // elementwise.cu) with the same operation order, so a verify window reproduces plain decode bit for bit. #include "strata/kernels/verify_kernels.hpp" +#include "hit_plan.cuh" #include +#include #include #include @@ -21,10 +23,10 @@ void check(const char* what) { if (e != cudaSuccess) { std::fprintf(stderr, "%s: %s\n", what, cudaGetErrorString(e)); std::exit(1); } } -__global__ void __launch_bounds__(S) gdn_conv_l2_multi_kernel(const float* __restrict__ hist, - const float* __restrict__ qkv, +// `commit` (one token): the history then keeps it - [hist1, hist2, x_0], each thread its channel's, after reading them +__global__ void __launch_bounds__(S) gdn_conv_l2_multi_kernel(float* hist, const float* __restrict__ qkv, const float* __restrict__ w, float* __restrict__ h, - int C, int qk_heads, float eps, int t_begin) { + int C, int qk_heads, float eps, int t_begin, bool commit) { __shared__ float part[S / 32]; const int t = t_begin + blockIdx.y; const int c = blockIdx.x * S + threadIdx.x; @@ -36,6 +38,11 @@ __global__ void __launch_bounds__(S) gdn_conv_l2_multi_kernel(const float* __res win[j] = src < 3 ? hist[c * 3 + src] : qkv[(size_t) (src - 3) * C + c]; } const float v0 = win[0], v1 = win[1], v2 = win[2], x = qkv[(size_t) t * C + c]; + if (commit) { + hist[c * 3] = v1; + hist[c * 3 + 1] = v2; + hist[c * 3 + 2] = x; + } float sum = v0 * w[c * 4] + v1 * w[c * 4 + 1] + v2 * w[c * 4 + 2] + x * w[c * 4 + 3]; float y = sum / (1.0f + __expf(-sum)); if ((int) blockIdx.x < qk_heads) { @@ -224,11 +231,37 @@ __global__ void fetch_blobs_kernel(const unsigned long long* __restrict__ src, c } } +// block (x, i): part x of listed row i, block (0, i) also its entry of the device list +__global__ void fetch_listed_rows_kernel(const int32_t* list, const float4* src, float4* __restrict__ dst, + int32_t* __restrict__ list_dev, int cap, long long n4) { + const int n = min(list[0], cap), i = (int) blockIdx.y; + const bool first = blockIdx.x == 0 && threadIdx.x == 0; + if (i >= n) { + if (i == 0 && first) list_dev[0] = 0; + return; + } + const int r = list[4 + i]; + if (first) { + list_dev[4 + i] = r; + if (i == 0) list_dev[0] = n; + } + if (r < 0 || r >= cap) return; + const long long step = (long long) gridDim.x * blockDim.x; + for (long long c = (long long) blockIdx.x * blockDim.x + threadIdx.x; c < n4; c += step) + dst[r * n4 + c] = src[r * n4 + c]; +} + __global__ void rebase_ptrs_kernel(unsigned long long* ptr, const int32_t* n, unsigned long long base, long long bytes) { const int k = threadIdx.x; if (k < *n) ptr[k] = base + (unsigned long long) k * (unsigned long long) bytes; } +__global__ void hit_plan_kernel(const int32_t* __restrict__ ids, int n, int k, const int32_t* __restrict__ res, + int n_expert, const unsigned long long* __restrict__ slot_ptr, int32_t* __restrict__ plan, + int cap, int ptr_off) { + hit_plan_block(ids, n, k, res, n_expert, slot_ptr, plan, cap, ptr_off); +} + __global__ void add_streams_broadcast_kernel(const float* __restrict__ h, const float* __restrict__ e, float* __restrict__ R, int64_t n, int hc) { const int t = blockIdx.y; @@ -244,14 +277,13 @@ __global__ void ident_hits_kernel(const int32_t* __restrict__ ids, int n, int32_ if (i == 0) *count = n; } -// E = the widest element the row size divides into (16, 4 or 1 bytes): a Q6_K head row of 2560 values is 2100 bytes -template -__global__ void gather_rows_kernel(const E* __restrict__ src, long long row_e, const int32_t* __restrict__ ids, - long long n, E* __restrict__ dst) { - const long long total = n * row_e; +template +__global__ void gather_rows_kernel(const U* __restrict__ src, long long row_units, const int32_t* __restrict__ ids, + long long n, U* __restrict__ dst) { + const long long total = n * row_units; for (long long i = (long long) blockIdx.x * blockDim.x + threadIdx.x; i < total; i += (long long) gridDim.x * blockDim.x) { - const long long r = i / row_e, o = i - r * row_e; - dst[i] = src[(long long) ids[r] * row_e + o]; + const long long r = i / row_units, o = i - r * row_units; + dst[i] = src[(long long) ids[r] * row_units + o]; } } @@ -311,6 +343,15 @@ void fetch_blobs(const unsigned long long* src, const int32_t* n, uint8_t* dst, check("fetch_blobs"); } +void fetch_listed_rows(const int32_t* list, const float* src, float* dst, int32_t* list_dev, int cap, int64_t n_embd, + void* stream) { + if (cap <= 0) return; + const long long n4 = n_embd / 4; + fetch_listed_rows_kernel<<>>( + list, (const float4*) src, (float4*) dst, list_dev, cap, n4); + check("fetch_listed_rows"); +} + void rebase_ptrs(unsigned long long* ptr, const int32_t* n, uint8_t* base, int64_t blob_bytes, void* stream) { rebase_ptrs_kernel<<<1, 128, 0, (cudaStream_t) stream>>>(ptr, n, (unsigned long long) base, (long long) blob_bytes); check("rebase_ptrs"); @@ -322,12 +363,41 @@ void add_streams_broadcast(const float* h, const float* e, float* R, int64_t n_e check("add_streams_broadcast"); } +void verify_hit_plan(const int32_t* ids, int n, int k, const int32_t* res, int n_expert, + const unsigned long long* slot_ptr, int32_t* plan, int64_t cap, int64_t ptr_off, void* stream) { + if (n < 1 || n > kHitPlanMax || n > cap || k < 1 || (ptr_off & 1) != 0) { + std::fprintf(stderr, "verify_hit_plan: invalid arguments\n"); + std::exit(1); + } + hit_plan_kernel<<<1, kHitPlanMax, 0, (cudaStream_t) stream>>>(ids, n, k, res, n_expert, slot_ptr, plan, (int) cap, + (int) ptr_off); + check("verify_hit_plan"); +} + void ident_hits(const int32_t* ids, int n, int32_t* slot, int32_t* dst, int32_t* count, void* stream) { if (n < 1 || n > 1024) { std::fprintf(stderr, "ident_hits: n out of range\n"); std::exit(1); } ident_hits_kernel<<<1, 1024, 0, (cudaStream_t) stream>>>(ids, n, slot, dst, count); check("ident_hits"); } +namespace { +__global__ void copy_row_to_first_kernel(const int32_t* __restrict__ row_dev, float* a, int64_t a_n, float* b, + int64_t b_n, float* c, int64_t c_n) { + const int64_t row = *row_dev; + if (row == 0) return; + const int64_t i0 = (int64_t) blockIdx.x * blockDim.x + threadIdx.x, st = (int64_t) gridDim.x * blockDim.x; + for (int64_t i = i0; i < a_n; i += st) a[i] = a[row * a_n + i]; + for (int64_t i = i0; i < b_n; i += st) b[i] = b[row * b_n + i]; + for (int64_t i = i0; i < c_n; i += st) c[i] = c[row * c_n + i]; +} +} // namespace + +void copy_row_to_first(const int32_t* row_dev, float* a, int64_t a_n, float* b, int64_t b_n, float* c, int64_t c_n, + void* stream) { + copy_row_to_first_kernel<<<16, 256, 0, (cudaStream_t) stream>>>(row_dev, a, a_n, b, b_n, c, c_n); + check("copy_row_to_first"); +} + void mtp_select(const float* R_src, int64_t R_stride, const int32_t* ids, const int32_t* row_dev, float* R_dst, int32_t* tok_dst, int32_t* out, int j, void* stream, const float* probs, float* out_p) { mtp_select_kernel<<<16, 256, 0, (cudaStream_t) stream>>>(R_src, R_stride, ids, row_dev, R_dst, tok_dst, out, j, @@ -336,13 +406,13 @@ void mtp_select(const float* R_src, int64_t R_stride, const int32_t* ids, const } void gather_rows(const uint8_t* src, int64_t row_bytes, const int32_t* ids, int64_t n, uint8_t* dst, void* stream) { - cudaStream_t s = (cudaStream_t) stream; + if (n <= 0) return; if (row_bytes % 16 == 0) - gather_rows_kernel<<<48 * 8, 256, 0, s>>>((const uint4*) src, row_bytes / 16, ids, n, (uint4*) dst); - else if (row_bytes % 4 == 0) - gather_rows_kernel<<<48 * 8, 256, 0, s>>>((const uint32_t*) src, row_bytes / 4, ids, n, (uint32_t*) dst); - else - gather_rows_kernel<<<48 * 8, 256, 0, s>>>(src, row_bytes, ids, n, dst); + gather_rows_kernel<<<48 * 8, 256, 0, (cudaStream_t) stream>>>((const uint4*) src, row_bytes / 16, ids, n, (uint4*) dst); + else if (row_bytes % 4 == 0) // e.g. Q6_K rows of 2,560 values: 2,100 bytes + gather_rows_kernel<<<48 * 8, 256, 0, (cudaStream_t) stream>>>((const uint32_t*) src, row_bytes / 4, ids, n, + (uint32_t*) dst); + else { std::fprintf(stderr, "gather_rows: row size must be a multiple of 4\n"); std::exit(1); } check("gather_rows"); } @@ -356,6 +426,166 @@ void row_top_prob(const float* logits, int n_rows, int n_vocab, const int32_t* i check("row_top_prob"); } +namespace { + +constexpr int kArgMaxBlocks = 128, kArgThreads = 256, kTopBlocks = 8; + +// the larger value, on equality the lower index: sampler_greedy_kernel's order +__device__ __forceinline__ void arg_take(float ov, int oi, float& bv, int& bi) { + if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; } +} + +// a block's (value, index) pick, in thread 0; `none` is the index of no candidate +__device__ __forceinline__ void arg_block(float& bv, int& bi, int none) { + __shared__ float sv[32]; + __shared__ int si[32]; + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_down_sync(0xffffffffu, bv, off); + const int oi = __shfl_down_sync(0xffffffffu, bi, off); + arg_take(ov, oi, bv, bi); + } + const int warp = (int) (threadIdx.x >> 5), lane = (int) (threadIdx.x & 31); + if (lane == 0) { sv[warp] = bv; si[warp] = bi; } + __syncthreads(); + if (warp == 0) { + const int nw = (int) (blockDim.x >> 5); + bv = lane < nw ? sv[lane] : __int_as_float(0xff800000); + bi = lane < nw ? si[lane] : none; + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_down_sync(0xffffffffu, bv, off); + const int oi = __shfl_down_sync(0xffffffffu, bi, off); + arg_take(ov, oi, bv, bi); + } + } +} + +// A row's scratch, whatever the launch's row count (one scratch serves windows of every size): its counter, then +// the blocks' values and indices. +struct ArgRow { + unsigned counter; + unsigned pad[3]; + float v[kArgMaxBlocks]; + int i[kArgMaxBlocks]; +}; + +// grid (blocks, rows): block b scans [b * per_block, +per_block) of its row +__global__ void argmax_rows_kernel(const float* __restrict__ logits, int n, int per_block, ArgRow* __restrict__ rows, + int32_t* out, float* out_val) { + const int row = (int) blockIdx.y, b = (int) blockIdx.x, nb = (int) gridDim.x; + unsigned* counter = &rows[row].counter; + float* pv = rows[row].v; + int* pi = rows[row].i; + const float* l = logits + (size_t) row * n; + const int lo = b * per_block, hi = min(n, lo + per_block); + float bv = __int_as_float(0xff800000); // -inf + int bi = n; + for (int v = lo + (int) threadIdx.x; v < hi; v += (int) blockDim.x) { + const float s = l[v]; + if (s > bv) { bv = s; bi = v; } + } + arg_block(bv, bi, n); + __shared__ bool last; + if (threadIdx.x == 0) { + pv[b] = bv; + pi[b] = bi; + __threadfence(); + last = atomicAdd(counter, 1u) == (unsigned) (nb - 1); + } + __syncthreads(); + if (!last) return; + __threadfence(); + bv = __int_as_float(0xff800000); + bi = n; + if ((int) threadIdx.x < nb) { + bv = ((volatile const float*) pv)[threadIdx.x]; + bi = ((volatile const int*) pi)[threadIdx.x]; + } + arg_block(bv, bi, n); + if (threadIdx.x == 0) { + out[row] = bi < n ? bi : 0; // no value above -inf: 0, as the one-block kernel + if (out_val != nullptr) out_val[row] = bv; + *counter = 0; + } +} + +__global__ void mapped_bump_kernel(uint32_t* flag) { + __threadfence_system(); + *(volatile uint32_t*) flag = *(volatile uint32_t*) flag + 1u; +} + +__global__ void join_rows_kernel(float* __restrict__ dst, int64_t n, const float* __restrict__ a, int64_t na, + const float* b) { + const int64_t t = blockIdx.y; + const int64_t nb = n - na; + for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x * blockDim.x) + dst[t * n + i] = i < na ? a[t * na + i] : ((const volatile float*) b)[t * nb + i - na]; +} + +// a row's top-probability scratch: its counter, then its 32 warps' sums +struct TopRow { + unsigned counter; + unsigned pad[3]; + float parts[32]; +}; + +// grid (kTopBlocks, rows), 1024 / kTopBlocks threads: thread j of block b is row_top_prob's thread b * blockDim + j +__global__ void row_top_prob_split_kernel(const float* __restrict__ logits, int n_vocab, const int32_t* __restrict__ ids, + float* __restrict__ probs, TopRow* __restrict__ rows) { + const int row = (int) blockIdx.y; + float* parts = rows[row].parts; + const float* l = logits + (size_t) row * n_vocab; + const float m = l[ids[row]]; + const int vt = (int) (blockIdx.x * blockDim.x + threadIdx.x); + float s = 0.0f; + for (int i = vt; i < n_vocab; i += 1024) s += __expf(l[i] - m); + for (int o = 16; o > 0; o >>= 1) s += __shfl_xor_sync(0xffffffffu, s, o); + if ((threadIdx.x & 31) == 0) parts[vt >> 5] = s; + __threadfence(); + __syncthreads(); + __shared__ bool last; + if (threadIdx.x == 0) last = atomicAdd(&rows[row].counter, 1u) == (unsigned) (gridDim.x - 1); + __syncthreads(); + if (!last || threadIdx.x != 0) return; + __threadfence(); + float tot = 0.0f; + for (int w = 0; w < 32; ++w) tot += ((volatile const float*) parts)[w]; + probs[row] = 1.0f / tot; + rows[row].counter = 0; +} + +} // namespace + +uint64_t argmax_rows_scratch_bytes(int n_rows) { return (uint64_t) n_rows * sizeof(ArgRow); } + +void argmax_rows(const float* logits, int n_rows, int n, void* scratch, int32_t* out, void* stream, float* out_val) { + if (n_rows <= 0) return; + const int nb = std::min(kArgMaxBlocks, std::max(1, (n + 4095) / 4096)); + const int per_block = (n + nb - 1) / nb; + argmax_rows_kernel<<>>( + logits, n, per_block, (ArgRow*) scratch, out, out_val); + check("argmax_rows"); +} + +uint64_t row_top_prob_scratch_bytes(int n_rows) { return (uint64_t) n_rows * sizeof(TopRow); } + +void mapped_bump(uint32_t* flag, void* stream) { + mapped_bump_kernel<<<1, 1, 0, (cudaStream_t) stream>>>(flag); + check("mapped_bump"); +} + +void join_rows(float* dst, int64_t n, const float* a, int64_t na, const float* b, int rows, void* stream) { + join_rows_kernel<<>>(dst, n, a, na, b); + check("join_rows"); +} + +void row_top_prob_split(const float* logits, int n_rows, int n_vocab, const int32_t* ids, float* probs, void* scratch, + void* stream) { + if (n_rows <= 0) return; + row_top_prob_split_kernel<<>>( + logits, n_vocab, ids, probs, (TopRow*) scratch); + check("row_top_prob_split"); +} + namespace { __global__ void window_ids_kernel(int32_t* steps, int window, int32_t* ids, long long stride) { const int q = blockIdx.y; @@ -381,13 +611,14 @@ void dense_steps(const int32_t* cells, int n, int32_t* steps, void* stream) { } void gdn_conv_l2_multi(const float* history, const float* qkv, const float* conv_w, float* h, int channels, - int qk_heads, float eps, int n_tok, void* stream, int t_begin) { - if (!history || !qkv || !conv_w || !h || channels % S != 0 || n_tok < 1 || n_tok > kVerifyMaxT) { + int qk_heads, float eps, int n_tok, void* stream, int t_begin, bool commit) { + if (!history || !qkv || !conv_w || !h || channels % S != 0 || n_tok < 1 || n_tok > 65535 || + (commit && (n_tok != 1 || t_begin != 0))) { std::fprintf(stderr, "gdn_conv_l2_multi: invalid arguments\n"); std::exit(1); } gdn_conv_l2_multi_kernel<<>>( - history, qkv, conv_w, h, channels, qk_heads, eps, t_begin); + const_cast(history), qkv, conv_w, h, channels, qk_heads, eps, t_begin, commit); check("gdn_conv_l2_multi"); } @@ -428,96 +659,22 @@ __global__ void wait_flag_ge_kernel(const volatile uint32_t* flag, uint32_t valu } } // namespace -namespace { -__global__ void resident_plan_kernel(const int32_t* __restrict__ ids, int n, int k, const int32_t* __restrict__ res, - int n_expert, const uint8_t* cache_base, const unsigned long long* slot_off, - long long blob, int32_t* __restrict__ pl, long long capx, uint32_t* skip, - uint32_t ring) { - // one thread: at most kVerifyMaxT * 10 entries, the host's exact loop - for (int i = 0; i < n; ++i) { - const int32_t e = ids[i]; - if (e < 0 || e >= n_expert || res[e] < 0) { *skip = 0; return; } - } - int32_t* counts = pl; - int32_t* start = pl + 4; - int32_t* dst = start + capx + 1; - int32_t* tok = dst + capx; - const long long ptr_off = ((4 + (capx + 1) + 2 * capx) + 1) & ~1ll; - unsigned long long* ptr = (unsigned long long*) (pl + ptr_off); - int32_t* start2 = pl + ptr_off + 4 * capx; - int groups = 0, entries = 0; - for (int i0 = 0; i0 < n; ++i0) { - bool first = true; - for (int j = 0; j < i0; ++j) if (ids[j] == ids[i0]) { first = false; break; } - if (!first) continue; - const int32_t slot = res[ids[i0]]; - ptr[groups] = (unsigned long long) (cache_base + (slot_off ? (size_t) slot_off[slot] : (size_t) slot * (size_t) blob)); - start[groups] = entries; - for (int i = i0; i < n; ++i) - if (ids[i] == ids[i0]) { - // an entry belongs to i0's group when its first occurrence is i0: the same expert id - dst[entries] = i; - tok[entries] = i / k; - ++entries; - } - ++groups; - } - start[groups] = entries; - start2[0] = entries; - counts[0] = groups; - counts[1] = entries; - counts[2] = 0; - __threadfence(); - *skip = ring; -} -__global__ void wait_flag_ge_or_kernel(const volatile uint32_t* flag, uint32_t value, const volatile uint32_t* skip) { - if (*skip == value) return; - while (*flag < value) __nanosleep(100); - __threadfence_system(); -} -__global__ void copy_i32_unless_kernel(int32_t* __restrict__ dst, const volatile int32_t* src, int n, - const uint32_t* skip, uint32_t value) { - if (*skip == value) return; - for (int i = threadIdx.x; i < n; i += blockDim.x) dst[i] = src[i]; -} -__global__ void copy_or_zero_kernel(float4* __restrict__ dst, const volatile float4* src, long long n4, - const uint32_t* skip, uint32_t value) { - const bool zero = *skip == value; - for (long long i = (long long) blockIdx.x * blockDim.x + threadIdx.x; i < n4; i += (long long) gridDim.x * blockDim.x) - dst[i] = zero ? make_float4(0.f, 0.f, 0.f, 0.f) : const_cast(src)[i]; +void wait_flag_ge(const uint32_t* flag, uint32_t value, void* stream) { + wait_flag_ge_kernel<<<1, 1, 0, (cudaStream_t) stream>>>(flag, value); + check("wait_flag_ge"); } -} // namespace -void resident_plan(const int32_t* ids, int n_entries, int k, const int32_t* res_layer, int n_expert, - const uint8_t* cache_base, const unsigned long long* slot_off, long long blob, int32_t* plan, - long long capx, uint32_t* skip, uint32_t ring, void* stream) { - resident_plan_kernel<<<1, 1, 0, (cudaStream_t) stream>>>(ids, n_entries, k, res_layer, n_expert, cache_base, slot_off, - blob, plan, capx, skip, ring); - check("resident_plan"); -} -void wait_flag_ge_or(const uint32_t* flag, uint32_t value, const uint32_t* skip, void* stream) { - wait_flag_ge_or_kernel<<<1, 1, 0, (cudaStream_t) stream>>>(flag, value, skip); - check("wait_flag_ge_or"); -} -void copy_i32_from_mapped_unless(int32_t* dst, const int32_t* src, long long n, const uint32_t* skip, uint32_t value, - void* stream) { - if (n <= 0) return; - copy_i32_unless_kernel<<<1, 128, 0, (cudaStream_t) stream>>>(dst, (const volatile int32_t*) src, (int) n, skip, value); - check("copy_i32_from_mapped_unless"); -} -void copy_or_zero_from_mapped(float* dst, const float* src, long long n, const uint32_t* skip, uint32_t value, - void* stream) { - if (n <= 0) return; - const long long n4 = n / 4; - const int blocks = (int) ((n4 + 255) / 256 < 64 ? (n4 + 255) / 256 : 64); - copy_or_zero_kernel<<>>((float4*) dst, (const volatile float4*) src, n4, skip, - value); - check("copy_or_zero_from_mapped"); +namespace { +__global__ void gpu_stamp_kernel(unsigned long long* dst) { + unsigned long long t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + *dst = t; } +} // namespace -void wait_flag_ge(const uint32_t* flag, uint32_t value, void* stream) { - wait_flag_ge_kernel<<<1, 1, 0, (cudaStream_t) stream>>>(flag, value); - check("wait_flag_ge"); +void gpu_stamp(unsigned long long* dst, void* stream) { + gpu_stamp_kernel<<<1, 1, 0, (cudaStream_t) stream>>>(dst); + check("gpu_stamp"); } void embedding_gather_dev(const uint8_t* codes, const float* scales, const float* offsets, const int32_t* tokens, @@ -541,18 +698,4 @@ void copy_indexed(float* dst, const float* src, int64_t stride, const int32_t* i check("copy_indexed"); } -// a GPU timestamp (ns, %globaltimer) into buf[i] - the verify window's stage profiler -namespace { __global__ void gpu_stamp_kernel(unsigned long long* buf, int i) { - unsigned long long t; -#if defined(__HIPCC__) - t = wall_clock64() * 10ull; // gfx11: a constant 100 MHz counter, in ns -#else - asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); -#endif - buf[i] = t; -} } -void gpu_stamp(unsigned long long* buf, int i, void* stream) { - gpu_stamp_kernel<<<1, 1, 0, (cudaStream_t) stream>>>(buf, i); -} - } // namespace strata::kernels diff --git a/src/kernels/cuda/xor_scatter.cuh b/src/kernels/cuda/xor_scatter.cuh new file mode 100644 index 00000000000..9db9496b80c --- /dev/null +++ b/src/kernels/cuda/xor_scatter.cuh @@ -0,0 +1,73 @@ +// src/kernels/cuda/xor_scatter.cuh - several warp sums in one xor tree, each bitwise the butterfly's. +// +// `v += __shfl_xor_sync(v, o)` for o = 16, 8, 4, 2, 1 leaves every lane with the same sum. For NV sums a lane does not +// need all of them in every lane: at distance o, while it holds more than one, it keeps the half its bit o selects and +// adds its partner's copy of them - the butterfly's own additions in the same order - then the last steps as the +// butterfly. NV + NV/2 + ... shuffles instead of 5 NV. +#pragma once + +#include + +namespace strata::kernels { +namespace { + +// c ? a : b as a select of values: the compiler turns `c ? v[i] : v[j]` into a load from a selected address, which +// puts the array in local memory +__device__ __forceinline__ float select_f(bool c, float a, float b) { + float r; + asm("{ .reg .pred p; setp.ne.u32 p, %3, 0; selp.f32 %0, %1, %2, p; }" + : "=f"(r) : "f"(a), "f"(b), "r"((unsigned) c)); + return r; +} + +// step S of xor_scatter (distance 16 >> S); the steps are template recursion so that every loop has a constant trip +// count and unrolls (a rolled loop indexes v at run time, in local memory) +template +__device__ __forceinline__ void xor_scatter_step(float (&v)[NV], int lane) { + if constexpr (S < 5) { + constexpr int o = 16 >> S, n = NV >> S; // n: the values left before this step + if constexpr (n > 1) { + const bool hi = (lane & o) != 0; +#pragma unroll + for (int e = 0; e < n / 2; ++e) { + const float mine = select_f(hi, v[n / 2 + e], v[e]); + const float other = select_f(hi, v[e], v[n / 2 + e]); + v[e] = mine + __shfl_xor_sync(0xffffffffu, other, o); + } + } else { + v[0] += __shfl_xor_sync(0xffffffffu, v[0], o); + } + xor_scatter_step(v, lane); + } +} + +// NV (a power of two, at most 32) values a lane; afterwards v[0] is the sum of value xor_scatter_index(lane) +template +__device__ __forceinline__ void xor_scatter(float (&v)[NV], int lane) { + static_assert(NV >= 1 && NV <= 32 && (NV & (NV - 1)) == 0, "NV must be a power of two <= 32"); + xor_scatter_step(v, lane); +} + +// the value a lane's v[0] sums after xor_scatter: the halves its bits 16, 8, ... selected +template +__device__ __forceinline__ int xor_scatter_index(int lane) { + int vi = 0; +#pragma unroll + for (int s = 0; (NV >> s) > 1; ++s) vi += (lane & (16 >> s)) ? NV >> (s + 1) : 0; + return vi; +} + +// the lowest lane holding value vi +template +__device__ __forceinline__ int xor_scatter_lane(int vi) { + int lane = 0; +#pragma unroll + for (int s = 0; (NV >> s) > 1; ++s) lane += (vi & (NV >> (s + 1))) ? 16 >> s : 0; + return lane; +} + +// the smallest power of two >= n +__host__ __device__ constexpr int pow2_at_least(int n) { return n <= 1 ? 1 : 2 * pow2_at_least((n + 1) / 2); } + +} // namespace +} // namespace strata::kernels diff --git a/src/kernels/expert_grouped_parity.cu b/src/kernels/expert_grouped_parity.cu new file mode 100644 index 00000000000..110185b7cce --- /dev/null +++ b/src/kernels/expert_grouped_parity.cu @@ -0,0 +1,522 @@ +// src/kernels/expert_grouped_parity.cu - `native_expert_grouped` (the verify window's and the second GPU's grouped +// experts) against the kernels it replaced and the i-quant dot products as they were, kept here verbatim: bitwise +// over every GPU gate/up and down format, on synthetic experts (random weights quantized by ggml) and plans of 1-8 +// tokens. +// +// expert_grouped_parity [--bench] +// +// --bench times both inside CUDA graphs (a kernel launched outside one pays WDDM's ~30-40 us) on the model files' +// format pairs, with routings shaped like the verify windows' (tokens share about a quarter of their experts). +#include "strata/kernels/iq_kernels.hpp" +#include "iq_dot.cuh" + +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int64_t H = 2560, FF = 640, K = 10, MAXT = 8, SLOTS = 48, DISTINCT = 4; + +void check(cudaError_t e, const char* what) { + if (e != cudaSuccess) { + std::fprintf(stderr, "%s: %s\n", what, cudaGetErrorString(e)); + std::exit(1); + } +} + +} // namespace + +// ---- the kernels native_expert_grouped launched before, verbatim, and the i-quant dot products as they were +// (sign masks through __vcmpne4 / __vsub4) +namespace strata::kernels::ref { + +__device__ __forceinline__ float ref_vec_dot_iq2_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq2_xxs* bq2 = (const block_iq2_xxs*) vbq + kbx; + const int q2 = get_int_b2(bq2->qs, iqs); + const uint8_t* aux8 = (const uint8_t*) &q2; + const uint32_t aux32 = get_int_b2(bq2->qs, iqs + 1); + int sumi = 0; +#pragma unroll + for (int k0 = 0; k0 < 8; k0 += 2) { + const uint2 grid_pos = ((const uint2*) iq2xxs_grid)[aux8[k0 / 2]]; + const uint32_t signs = unpack_ksigns(aux32 >> (7 * k0 / 2)); + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid0 = __vsub4(grid_pos.x ^ signs0, signs0); + const int u0 = get_int_b4(bq8_1[iqs / 2].qs, k0 + 0); + sumi = ggml_cuda_dp4a(grid0, u0, sumi); + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid1 = __vsub4(grid_pos.y ^ signs1, signs1); + const int u1 = get_int_b4(bq8_1[iqs / 2].qs, k0 + 1); + sumi = ggml_cuda_dp4a(grid1, u1, sumi); + } + const int ls = aux32 >> 27 | 1; + sumi = sumi * ls / 8; + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs / 2].ds); + return d * sumi; +} + +__device__ __forceinline__ float ref_vec_dot_iq2_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq2_xs* bq2 = (const block_iq2_xs*) vbq + kbx; + const int2 q2_packed = make_int2(get_int_b2(bq2->qs, iqs + 0), get_int_b2(bq2->qs, iqs + 1)); + const uint16_t* q2 = (const uint16_t*) &q2_packed; + const int ls0 = bq2->scales[iqs / 2] & 0x0F; + const int ls1 = bq2->scales[iqs / 2] >> 4; + int sumi0 = 0, sumi1 = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const uint2 grid_pos = ((const uint2*) iq2xs_grid)[q2[l0 / 2] & 0x1FF]; + const uint32_t signs = unpack_ksigns(q2[l0 / 2] >> 9); + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); + if (l0 < 4) { + sumi0 = ggml_cuda_dp4a(grid_l, u0, sumi0); + sumi0 = ggml_cuda_dp4a(grid_h, u1, sumi0); + } else { + sumi1 = ggml_cuda_dp4a(grid_l, u0, sumi1); + sumi1 = ggml_cuda_dp4a(grid_h, u1, sumi1); + } + } + const int sumi = (sumi0 * ls0 + sumi1 * ls1 + (sumi0 + sumi1) / 2) / 4; + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs / 2].ds); + return d * sumi; +} + +__device__ __forceinline__ float ref_vec_dot_iq2_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq2_s* bq2 = (const block_iq2_s*) vbq + kbx; + const int qs_packed = get_int_b2(bq2->qs, iqs / 2); + const uint8_t* qs = (const uint8_t*) &qs_packed; + const int qh = bq2->qh[iqs / 2]; + const int signs_packed_32 = get_int_b2(bq2->qs, QK_K / 32 + iqs / 2); + const uint8_t* signs_packed_8 = (const uint8_t*) &signs_packed_32; + const int ls0 = bq2->scales[iqs / 2] & 0x0F; + const int ls1 = bq2->scales[iqs / 2] >> 4; + int sumi0 = 0, sumi1 = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int* grid_pos = (const int*) (iq2s_grid + (qs[l0 / 2] | ((qh << (8 - l0)) & 0x300))); + const int signs0 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x03) << 7) | ((signs_packed_8[l0 / 2] & 0x0C) << 21), 0x00000000); + const int signs1 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x30) << 3) | ((signs_packed_8[l0 / 2] & 0xC0) << 17), 0x00000000); + const int grid_l = __vsub4(grid_pos[0] ^ signs0, signs0); + const int grid_h = __vsub4(grid_pos[1] ^ signs1, signs1); + const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); + if (l0 < 4) { + sumi0 = ggml_cuda_dp4a(grid_l, u0, sumi0); + sumi0 = ggml_cuda_dp4a(grid_h, u1, sumi0); + } else { + sumi1 = ggml_cuda_dp4a(grid_l, u0, sumi1); + sumi1 = ggml_cuda_dp4a(grid_h, u1, sumi1); + } + } + const int sumi = (sumi0 * ls0 + sumi1 * ls1 + (sumi0 + sumi1) / 2) / 4; + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs / 2].ds); + return d * sumi; +} + +__device__ __forceinline__ float ref_vec_dot_iq3_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq3_xxs* bq3 = (const block_iq3_xxs*) vbq + kbx; + const int2 q3_packed = make_int2(get_int_b2(bq3->qs, iqs), get_int_b2(bq3->qs, iqs + 1)); + const uint8_t* q3 = (const uint8_t*) &q3_packed; + const uint32_t aux32 = get_int_b2(bq3->qs, QK_K / 16 + iqs / 2); + int sumi = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int2 grid_pos = make_int2(iq3xxs_grid[q3[l0 + 0]], iq3xxs_grid[q3[l0 + 1]]); + const uint32_t signs = unpack_ksigns(aux32 >> (7 * l0 / 2)); + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); + sumi = ggml_cuda_dp4a(grid_l, u0, sumi); + sumi = ggml_cuda_dp4a(grid_h, u1, sumi); + } + const int ls = aux32 >> 28; + sumi = (ls * sumi + sumi / 2) / 2; + const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs / 2].ds); + return d * sumi; +} + +__device__ __forceinline__ float ref_vec_dot_iq3_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, + const int& kbx, const int& iqs) { + const block_iq3_s* bq3 = (const block_iq3_s*) vbq + kbx; + const int2 qs_packed = make_int2(get_int_b2(bq3->qs, iqs + 0), get_int_b2(bq3->qs, iqs + 1)); + const uint8_t* qs = (const uint8_t*) &qs_packed; + const int qh = bq3->qh[iqs / 2]; + const int signs_packed_32 = get_int_b2(bq3->signs, iqs / 2); + const uint8_t* signs_packed_8 = (const uint8_t*) &signs_packed_32; + int sumi = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int2 grid_pos = make_int2(iq3s_grid[qs[l0 + 0] | ((qh << (8 - l0)) & 0x100)], + iq3s_grid[qs[l0 + 1] | ((qh << (7 - l0)) & 0x100)]); + const int signs0 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x03) << 7) | ((signs_packed_8[l0 / 2] & 0x0C) << 21), 0x00000000); + const int signs1 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x30) << 3) | ((signs_packed_8[l0 / 2] & 0xC0) << 17), 0x00000000); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + const int u0 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs / 2].qs, l0 + 1); + sumi = ggml_cuda_dp4a(grid_l, u0, sumi); + sumi = ggml_cuda_dp4a(grid_h, u1, sumi); + } + sumi *= 1 + 2 * ((bq3->scales[iqs / 4] >> ((iqs << 1) & 0x04)) & 0x0F); + const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs / 2].ds); + return d * sumi; +} + +template struct RefFmt : Fmt {}; +template<> struct RefFmt<16> : Fmt<16> { + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return ref_vec_dot_iq2_xxs_q8_1(v, y, kbx, iqs); } }; +template<> struct RefFmt<17> : Fmt<17> { + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return ref_vec_dot_iq2_xs_q8_1(v, y, kbx, iqs); } }; +template<> struct RefFmt<18> : Fmt<18> { + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return ref_vec_dot_iq3_xxs_q8_1(v, y, kbx, iqs); } }; +template<> struct RefFmt<21> : Fmt<21> { + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return ref_vec_dot_iq3_s_q8_1(v, y, kbx, iqs); } }; +template<> struct RefFmt<22> : Fmt<22> { + __device__ static float dot(const void* v, const block_q8_1* y, int kbx, int iqs) { return ref_vec_dot_iq2_s_q8_1(v, y, kbx, iqs); } }; + +template +__device__ __forceinline__ float ref_row_dot(const uint8_t* row, const block_q8_1* x, int nb, int lane) { + using F = RefFmt; + float s = 0.0f; + for (int k = lane; k < nb * F::ipb; k += 32) { + const int kbx = k / F::ipb, iqs = F::step * (k % F::ipb); + s += F::dot(row, x + kbx * (F::qk / 32), kbx, iqs); + } + return warp_sum(s); +} + +constexpr int GU_ROWS = 8; // rows per block (one warp each) + +template +__global__ void __launch_bounds__(256) native_gu_kernel(const unsigned long long* __restrict__ grp_ptr, + const int32_t* __restrict__ grp_start, + const int32_t* __restrict__ n_groups, + const int32_t* __restrict__ ent_tok, + const block_q8_1* __restrict__ xq, NativeExpertLayout L, + float* __restrict__ gate, float* __restrict__ up) { + const int g = blockIdx.y; + if (g >= *n_groups) return; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int row = blockIdx.x * GU_ROWS + warp; // 0 .. 2*n_ff + if (row >= 2 * L.n_ff) return; + const bool is_up = row >= L.n_ff; + const int r = is_up ? row - (int) L.n_ff : row; + const uint8_t* blob = (const uint8_t*) grp_ptr[g]; + const uint8_t* wr = blob + (is_up ? L.up_off : 0) + (size_t) r * L.gu_row; + const int nb = (int) (L.n_embd / RefFmt::qk), xb = (int) (L.n_embd / 32); + const int e0 = grp_start[g], e1 = grp_start[g + 1]; + for (int e = e0; e < e1; ++e) { + const float s = ref_row_dot(wr, xq + (size_t) ent_tok[e] * xb, nb, lane); + if (lane == 0) (is_up ? up : gate)[(size_t) e * L.n_ff + r] = s; + } +} + +__global__ void swiglu_entries_kernel(const float* __restrict__ gate, const float* __restrict__ up, float* __restrict__ h, + long long n) { + const long long i = (long long) blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + const float g = gate[i]; + h[i] = (g / (1.0f + __expf(-g))) * up[i]; +} + +template +__global__ void __launch_bounds__(256) native_down_kernel(const unsigned long long* __restrict__ grp_ptr, + const int32_t* __restrict__ grp_start, + const int32_t* __restrict__ n_groups, + const int32_t* __restrict__ ent_dst, + const block_q8_1* __restrict__ hq, NativeExpertLayout L, + float* __restrict__ out) { + const int g = blockIdx.y; + if (g >= *n_groups) return; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int r = blockIdx.x * 8 + warp; + if (r >= L.n_embd) return; + const uint8_t* blob = (const uint8_t*) grp_ptr[g]; + const uint8_t* wr = blob + L.down_off + (size_t) r * L.d_row; + const int nb = (int) (L.n_ff / RefFmt::qk), hb = (int) (L.n_ff / 32); + const int e0 = grp_start[g], e1 = grp_start[g + 1]; + for (int e = e0; e < e1; ++e) { + const float s = ref_row_dot(wr, hq + (size_t) e * hb, nb, lane); + if (lane == 0) out[(size_t) ent_dst[e] * L.n_embd + r] = s; + } +} + +__global__ void quantize_q8_1_kernel(const float* __restrict__ x, block_q8_1* __restrict__ y, long long n) { + const long long i = (long long) blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + const float xi = x[i]; + float amax = fabsf(xi), sum = xi; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, o)); + sum += __shfl_xor_sync(0xffffffffu, sum, o); + } + const float d = amax / 127.0f; + const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); + const long long ib = i / 32, iqs = i % 32; + y[ib].qs[iqs] = q; + if (iqs == 0) y[ib].ds = make_half2(d, sum); +} + +size_t scratch_bytes(int64_t cap, int64_t n_ff) { + const size_t f = (size_t) cap * (size_t) n_ff * sizeof(float); + return 3 * ((f + 255) & ~(size_t) 255) + (((size_t) cap * (size_t) (n_ff / 32) * sizeof(block_q8_1) + 255) & ~(size_t) 255); +} + +void grouped(const NativeExpertLayout& L, const unsigned long long* grp_ptr, const int32_t* grp_start, + const int32_t* n_groups, const int32_t* ent_dst, const int32_t* ent_tok, int64_t cap_groups, + int64_t cap_entries, const void* x_q8_1, void* scratch, float* out, cudaStream_t s) { + const size_t f = (size_t) cap_entries * (size_t) L.n_ff * sizeof(float), fa = (f + 255) & ~(size_t) 255; + float* gate = (float*) scratch; + float* up = (float*) ((uint8_t*) scratch + fa); + float* h = (float*) ((uint8_t*) scratch + 2 * fa); + block_q8_1* hq = (block_q8_1*) ((uint8_t*) scratch + 3 * fa); + const auto* X = (const block_q8_1*) x_q8_1; + const dim3 ggu((unsigned) ((2 * L.n_ff + GU_ROWS - 1) / GU_ROWS), (unsigned) cap_groups); + switch (L.gu_type) { +#define STRATA_GU(T) case T: native_gu_kernel<<>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; + STRATA_FMTS(STRATA_GU) +#undef STRATA_GU + default: std::fprintf(stderr, "ref: gate/up type %d\n", L.gu_type); std::exit(1); + } + const long long nh = (long long) cap_entries * L.n_ff; + swiglu_entries_kernel<<<(unsigned) ((nh + 255) / 256), 256, 0, s>>>(gate, up, h, nh); + quantize_q8_1_kernel<<<(unsigned) ((nh + 255) / 256), 256, 0, s>>>(h, hq, nh); + const dim3 gd((unsigned) ((L.n_embd + 7) / 8), (unsigned) cap_groups); + switch (L.d_type) { +#define STRATA_DOWN(T) case T: native_down_kernel<<>>(grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break; + STRATA_FMTS(STRATA_DOWN) +#undef STRATA_DOWN + default: std::fprintf(stderr, "ref: down type %d\n", L.d_type); std::exit(1); + } +} + +} // namespace strata::kernels::ref + +namespace { + +// DISTINCT experts of random weights (rows of different magnitude) quantized by ggml, repeated over SLOTS slots of one +// device arena so the kernels stream them from memory as they do the VRAM tier +uint8_t* make_arena(const strata::kernels::NativeExpertLayout& L, size_t stride, int seed) { + std::vector blobs((size_t) DISTINCT * L.bytes); + std::mt19937 rng(seed); + std::normal_distribution nd(0.f, 1.f); + auto quant = [&](int type, int64_t rows, int64_t cols, uint8_t* dst) { + std::vector w((size_t) (rows * cols)); + for (int64_t r = 0; r < rows; ++r) { + const float scale = 0.02f * (0.5f + (float) (r % 7) / 7.0f); + for (int64_t c = 0; c < cols; ++c) w[(size_t) (r * cols + c)] = scale * nd(rng); + } + const std::vector imatrix((size_t) cols, 1.0f); + ggml_quantize_chunk((ggml_type) type, w.data(), dst, 0, rows, cols, + ggml_quantize_requires_imatrix((ggml_type) type) ? imatrix.data() : nullptr); + }; + for (int b = 0; b < DISTINCT; ++b) { + uint8_t* blob = blobs.data() + (size_t) b * L.bytes; + quant(L.gu_type, FF, H, blob); + quant(L.gu_type, FF, H, blob + L.up_off); + quant(L.d_type, H, FF, blob + L.down_off); + } + uint8_t* arena = nullptr; + check(cudaMalloc(&arena, (size_t) SLOTS * stride), "arena"); + for (int i = 0; i < SLOTS; ++i) + check(cudaMemcpy(arena + (size_t) i * stride, blobs.data() + (size_t) (i % DISTINCT) * L.bytes, L.bytes, + cudaMemcpyHostToDevice), "arena upload"); + return arena; +} + +// A window's routing: n_tok tokens of K distinct experts each, a quarter drawn from the earlier tokens' ones (the +// verify windows' overlap); `hit` of the experts resident. Grouped as the plan is: distinct resident experts in +// routing order, each with its entries in routing order. +struct Plan { + std::vector ptr; + std::vector start, dst, tok; +}; +Plan make_plan(std::mt19937& rng, int n_tok, double hit, const uint8_t* arena, size_t stride) { + std::vector ids; + std::vector resident(SLOTS); + for (auto& r : resident) r = std::uniform_real_distribution(0, 1)(rng) < hit; + for (int t = 0; t < n_tok; ++t) + for (int j = 0; j < K; ++j) { + int32_t e; + bool again; + do { + const bool reuse = t > 0 && rng() % 4 == 0; + e = reuse ? ids[rng() % (size_t) (t * K)] : (int32_t) (rng() % SLOTS); + again = false; + for (int m = 0; m < j; ++m) again = again || ids[(size_t) (t * K + m)] == e; + } while (again); + ids.push_back(e); + } + Plan p; + const int n = n_tok * (int) K; + for (int i = 0; i < n; ++i) { + bool first = resident[(size_t) ids[(size_t) i]] != 0; + for (int j = 0; j < i && first; ++j) first = ids[(size_t) j] != ids[(size_t) i]; + if (!first) continue; + p.ptr.push_back((unsigned long long) (arena + (size_t) ids[(size_t) i] * stride)); + p.start.push_back((int32_t) p.dst.size()); + for (int j = i; j < n; ++j) + if (ids[(size_t) j] == ids[(size_t) i]) { p.dst.push_back(j); p.tok.push_back(j / (int) K); } + } + p.start.push_back((int32_t) p.dst.size()); + return p; +} + +struct Dev { + unsigned long long* ptr = nullptr; + int32_t *start = nullptr, *n = nullptr, *dst = nullptr, *tok = nullptr; + void *xq = nullptr, *scratch_ref = nullptr, *scratch_new = nullptr; + float *out_ref = nullptr, *out_new = nullptr; +}; + +void upload(const Plan& p, const Dev& d) { + const int32_t groups = (int32_t) p.ptr.size(); + if (groups > 0) { + check(cudaMemcpy(d.ptr, p.ptr.data(), p.ptr.size() * 8, cudaMemcpyHostToDevice), "plan"); + check(cudaMemcpy(d.dst, p.dst.data(), p.dst.size() * 4, cudaMemcpyHostToDevice), "plan"); + check(cudaMemcpy(d.tok, p.tok.data(), p.tok.size() * 4, cudaMemcpyHostToDevice), "plan"); + } + check(cudaMemcpy(d.start, p.start.data(), p.start.size() * 4, cudaMemcpyHostToDevice), "plan"); + check(cudaMemcpy(d.n, &groups, 4, cudaMemcpyHostToDevice), "plan"); +} + +void run(bool reference, const strata::kernels::NativeExpertLayout& L, const Dev& d, int64_t cap, cudaStream_t s) { + if (reference) + strata::kernels::ref::grouped(L, d.ptr, d.start, d.n, d.dst, d.tok, cap, cap, d.xq, d.scratch_ref, d.out_ref, s); + else + strata::kernels::native_expert_grouped(L, d.ptr, d.start, d.n, d.dst, d.tok, cap, cap, d.xq, d.scratch_new, + d.out_new, s); +} + +// microseconds per call, inside a graph of `reps` calls +double time_graph(bool reference, const strata::kernels::NativeExpertLayout& L, const Dev& d, int64_t cap, + cudaStream_t s) { + const int reps = 50; + cudaGraph_t graph = nullptr; + cudaGraphExec_t exec = nullptr; + check(cudaStreamBeginCapture(s, cudaStreamCaptureModeThreadLocal), "capture"); + for (int i = 0; i < reps; ++i) run(reference, L, d, cap, s); + check(cudaStreamEndCapture(s, &graph), "capture"); + check(cudaGraphInstantiate(&exec, graph, 0), "instantiate"); + const auto w0 = std::chrono::steady_clock::now(); // until the GPU has clocked up + while (std::chrono::duration(std::chrono::steady_clock::now() - w0).count() < 0.3) { + check(cudaGraphLaunch(exec, s), "launch"); + check(cudaStreamSynchronize(s), "warm"); + } + double best = 1e30; + for (int k = 0; k < 10; ++k) { + const auto t0 = std::chrono::steady_clock::now(); + check(cudaGraphLaunch(exec, s), "launch"); + check(cudaStreamSynchronize(s), "sync"); + best = std::min(best, std::chrono::duration(std::chrono::steady_clock::now() - t0).count() / reps); + } + cudaGraphExecDestroy(exec); + cudaGraphDestroy(graph); + return best; +} + +} // namespace + +int main(int argc, char** argv) { + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + if (argc > 2 || (argc == 2 && !bench)) { + std::fprintf(stderr, "usage: expert_grouped_parity [--bench]\n"); + return 2; + } + cudaStream_t s; + check(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking), "stream"); + const int64_t cap = MAXT * K; + Dev d; + check(cudaMalloc(&d.ptr, (size_t) cap * 8), "dev"); + check(cudaMalloc(&d.start, (size_t) (cap + 1) * 4), "dev"); + check(cudaMalloc(&d.n, 4), "dev"); + check(cudaMalloc(&d.dst, (size_t) cap * 4), "dev"); + check(cudaMalloc(&d.tok, (size_t) cap * 4), "dev"); + check(cudaMalloc(&d.xq, (size_t) MAXT * (H / 32) * sizeof(block_q8_1)), "dev"); + check(cudaMalloc(&d.scratch_ref, strata::kernels::ref::scratch_bytes(cap, FF)), "dev"); + check(cudaMalloc(&d.scratch_new, strata::kernels::native_expert_scratch_bytes(cap, FF)), "dev"); + check(cudaMalloc(&d.out_ref, (size_t) cap * H * 4), "dev"); + check(cudaMalloc(&d.out_new, (size_t) cap * H * 4), "dev"); + { // the tokens' activations + std::mt19937 rng(7); + std::normal_distribution nd(0.f, 1.f); + std::vector x((size_t) (MAXT * H)); + for (auto& v : x) v = nd(rng); + float* dx = nullptr; + check(cudaMalloc(&dx, x.size() * 4), "dev"); + check(cudaMemcpy(dx, x.data(), x.size() * 4, cudaMemcpyHostToDevice), "x"); + strata::kernels::quantize_q8_1_rows(dx, MAXT, H, d.xq, s); + check(cudaStreamSynchronize(s), "x"); + cudaFree(dx); + } + // parity: every gate/up format with a down format in turn (down rows are n_ff = 640 long: 32- and 64-blocks); + // bench: the model files' pairs (IQ2_XXS .. IQ3_S with IQ4_NL / Q2_0: the IQ3_XXS file; Q4_K, Q5_K with Q5_1, + // Q8_0: UD-Q4_K_XL) + const int gus[] = {7, 8, 12, 13, 14, 16, 17, 18, 19, 20, 21, 22, 23, 29, 42}, downs[] = {7, 8, 20, 42}; + std::vector> pairs; + if (bench) pairs = {{16, 42}, {17, 20}, {18, 20}, {21, 42}, {22, 20}, {12, 7}, {13, 8}}; + else + for (size_t i = 0; i < sizeof gus / sizeof gus[0]; ++i) pairs.emplace_back(gus[i], downs[i % 4]); + int bad = 0; + std::mt19937 rng(11); + for (const auto& [gu, dn] : pairs) { + const strata::kernels::NativeExpertLayout L = strata::kernels::native_expert_layout(gu, dn, H, FF); + const size_t stride = (L.bytes + 255) & ~(size_t) 255; + uint8_t* arena = make_arena(L, stride, gu * 100 + dn); + const char* names = ggml_type_name((ggml_type) gu); + if (!bench) { + int cases = 0, diff = 0; + for (int n_tok : {1, 2, 3, 4, 8}) + for (int rep = 0; rep < 4; ++rep) { + const Plan p = make_plan(rng, n_tok, rep == 0 ? 1.0 : 0.8, arena, stride); + upload(p, d); + check(cudaMemset(d.out_ref, 0xff, (size_t) cap * H * 4), "memset"); + check(cudaMemset(d.out_new, 0xff, (size_t) cap * H * 4), "memset"); + run(true, L, d, n_tok * K, s); + run(false, L, d, n_tok * K, s); + check(cudaStreamSynchronize(s), "run"); + std::vector a((size_t) cap * H), b((size_t) cap * H); + check(cudaMemcpy(a.data(), d.out_ref, a.size() * 4, cudaMemcpyDeviceToHost), "down"); + check(cudaMemcpy(b.data(), d.out_new, b.size() * 4, cudaMemcpyDeviceToHost), "down"); + diff += a != b; + ++cases; + } + std::printf("%-8s/%-7s %d plans: %s\n", names, ggml_type_name((ggml_type) dn), cases, + diff ? "MISMATCH" : "bitwise equal"); + bad += diff; + } else { + std::printf("%-8s/%-7s", names, ggml_type_name((ggml_type) dn)); + for (int n_tok : {1, 2, 4}) { + const Plan p = make_plan(rng, n_tok, 0.92, arena, stride); + upload(p, d); + const double tr = time_graph(true, L, d, n_tok * K, s), tn = time_graph(false, L, d, n_tok * K, s); + std::printf(" T%d %2zu experts %6.1f -> %5.1f us (%4.0f GB/s)", n_tok, p.ptr.size(), tr, tn, + (double) (p.ptr.size() * L.bytes) / tn * 1e-3); + } + std::printf("\n"); + } + cudaFree(arena); + } + if (!bench) std::printf("expert_grouped_parity: %s\n", bad ? "FAIL" : "PASS"); + return bad ? 1 : 0; +} diff --git a/src/kernels/mmvq_il_parity.cu b/src/kernels/mmvq_il_parity.cu new file mode 100644 index 00000000000..61d3a531d8f --- /dev/null +++ b/src/kernels/mmvq_il_parity.cu @@ -0,0 +1,217 @@ +// src/kernels/mmvq_il_parity.cu - native_mmvq_il (1 column from its q8_1 blocks, 2..8 from the interleaved copy) +// against native_mmvq, bitwise, on synthetic matrices of the model files' dense formats and shapes, with every rows a +// warp; the interleaving quantizer's q8_1 blocks against native_quantize_q8_1's. --bench times both inside CUDA graphs +// (weights cycled past the L2). +#include "strata/kernels/native_mmvq.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace strata::kernels; + +#define CK(x) do { const cudaError_t e_ = (x); if (e_ != cudaSuccess) { std::printf("%s: %s\n", #x, cudaGetErrorString(e_)); std::exit(1); } } while (0) + +namespace { + +struct Fmt { int type; const char* name; int block_elems, block_bytes, d_off, dmin_off; }; +// the dense formats of the IQ3_XXS and UD-Q4_K_XL files; d_off: the fp16 block scale (random bytes elsewhere) +const Fmt kFmts[] = { + {23, "IQ4_XS", 256, 136, 0, -1}, {12, "Q4_K", 256, 144, 0, 2}, {13, "Q5_K", 256, 176, 0, 2}, + {14, "Q6_K", 256, 210, 208, -1}, {8, "Q8_0", 32, 34, 0, -1}, {21, "IQ3_S", 256, 110, 0, -1}, +}; + +cudaStream_t g_s; +std::mt19937 g_rng(29); + +void* dev_weights(const Fmt& f, int n_in, int n_out) { + const size_t bytes = native_mmvq_weight_bytes(f.type, n_in, n_out); + std::vector h(bytes); + for (auto& b : h) b = (uint8_t) g_rng(); + std::uniform_int_distribution e(0x1c00, 0x2400); // fp16 scales in ~[0.004, 0.016] + for (size_t o = 0; o < bytes; o += (size_t) f.block_bytes) { + const uint16_t d = (uint16_t) e(g_rng); + std::memcpy(&h[o + f.d_off], &d, 2); + if (f.dmin_off >= 0) { const uint16_t m = (uint16_t) e(g_rng); std::memcpy(&h[o + f.dmin_off], &m, 2); } + } + void* p; + CK(cudaMalloc(&p, bytes)); + CK(cudaMemcpy(p, h.data(), bytes, cudaMemcpyHostToDevice)); + CK(cudaDeviceSynchronize()); // a pageable copy may return before its DMA lands + return p; +} + +struct Act { + float* x = nullptr; + void *q_ref = nullptr, *q = nullptr, *il = nullptr; + int n_in = 0; +}; +Act make_act(int n_in) { + Act a; + a.n_in = n_in; + std::vector h((size_t) 8 * n_in); + std::normal_distribution nd(0.0f, 1.0f); + for (auto& v : h) v = nd(g_rng); + for (size_t i = 0; i < h.size(); i += 97) h[i] *= 20.0f; // outliers set some blocks' scales + for (int j = 0; j < 32; ++j) h[(size_t) 3 * n_in + 64 + j] = 0.0f; // an all-zero block + CK(cudaMalloc(&a.x, h.size() * 4)); + CK(cudaMemcpy(a.x, h.data(), h.size() * 4, cudaMemcpyHostToDevice)); + CK(cudaMalloc(&a.q_ref, native_q8_1_bytes(n_in, 8))); + CK(cudaMalloc(&a.q, native_q8_1_bytes(n_in, 8))); + CK(cudaMalloc(&a.il, native_q8_1_il_bytes(n_in, 8))); + CK(cudaDeviceSynchronize()); + return a; +} + +template +std::vector fetch(const void* d, size_t n) { + std::vector h(n); + CK(cudaMemcpy(h.data(), d, n * sizeof(T), cudaMemcpyDeviceToHost)); + return h; +} + +// every column count 1..8 (1: the one-column kernels), with the table's rows a warp and each of 1, 2, 4 +int check_shape(const Fmt& f, const Act& a, int n_out) { + const int n_in = a.n_in; + void* w = dev_weights(f, n_in, n_out); + float *y0, *y1; + CK(cudaMalloc(&y0, (size_t) 8 * n_out * 4)); + CK(cudaMalloc(&y1, (size_t) 8 * n_out * 4)); + int bad = 0; + for (int nc = 1; nc <= 8; ++nc) { + native_quantize_q8_1(a.x, a.q_ref, n_in, nc, g_s); + if (nc >= 2) native_quantize_q8_1_il(a.x, a.q, a.il, n_in, nc, g_s); + else native_quantize_q8_1(a.x, a.q, n_in, nc, g_s); + CK(cudaMemsetAsync(y0, 0xff, (size_t) 8 * n_out * 4, g_s)); + native_mmvq(f.type, w, a.q_ref, y0, n_in, n_out, nc, g_s); + CK(cudaStreamSynchronize(g_s)); + const size_t qb = native_q8_1_bytes(n_in, nc); + if (fetch(a.q_ref, qb) != fetch(a.q, qb)) { + ++bad; + std::printf(" %s %d->%d, %d columns: the q8_1 blocks differ\n", f.name, n_in, n_out, nc); + } + const auto r0 = fetch(y0, (size_t) nc * n_out); + for (int rows : {0, 1, 2, 4}) { + CK(cudaMemsetAsync(y1, 0xee, (size_t) 8 * n_out * 4, g_s)); + native_mmvq_il_tune(rows, nc == 1); + native_mmvq_il(f.type, w, a.q, a.il, y1, n_in, n_out, nc, g_s); + native_mmvq_il_tune(0, false); + CK(cudaStreamSynchronize(g_s)); + const auto r1 = fetch(y1, (size_t) nc * n_out); + size_t diff = 0; + for (size_t i = 0; i < r0.size(); ++i) diff += r0[i] != r1[i]; + if (diff) { + ++bad; + std::printf(" %s %d->%d, %d columns, rows %d: %zu of %zu values differ\n", f.name, n_in, n_out, nc, + rows, diff, r0.size()); + } + } + } + cudaFree(w); + cudaFree(y0); + cudaFree(y1); + return bad; +} + +__global__ void spin_kernel(long long cycles) { + const long long t0 = clock64(); + while (clock64() - t0 < cycles) {} +} + +// The median of 20 launches, each after ~3 ms of a one-thread spin, as the engine's load: the GPU active at little +// power (a continuous load reaches the 3090's 260 W cap and lowers the clocks, an idle GPU lowers them too). +double time_calls(const std::function& fn, int reps = 48) { + cudaGraph_t g; + CK(cudaStreamBeginCapture(g_s, cudaStreamCaptureModeThreadLocal)); + for (int r = 0; r < reps; ++r) fn(r); + CK(cudaStreamEndCapture(g_s, &g)); + cudaGraphExec_t ex; + CK(cudaGraphInstantiate(&ex, g, 0)); + cudaGraphDestroy(g); + cudaEvent_t a, b; + CK(cudaEventCreate(&a)); + CK(cudaEventCreate(&b)); + std::vector ts; + for (int k = 0; k < 24; ++k) { + spin_kernel<<<1, 1, 0, g_s>>>(5'500'000); + CK(cudaEventRecord(a, g_s)); + CK(cudaGraphLaunch(ex, g_s)); + CK(cudaEventRecord(b, g_s)); + CK(cudaEventSynchronize(b)); + float ms = 0; + CK(cudaEventElapsedTime(&ms, a, b)); + if (k >= 4) ts.push_back(ms * 1000.0 / reps); + } + std::sort(ts.begin(), ts.end()); + cudaGraphExecDestroy(ex); + cudaEventDestroy(a); + cudaEventDestroy(b); + return ts[ts.size() / 2]; +} + +void bench_shape(const Fmt& f, const Act& a, int n_out) { + const int n_in = a.n_in; + const double mb = (double) native_mmvq_weight_bytes(f.type, n_in, n_out) / 1e6; + const int copies = std::max(4, (int) (24.0 / mb) + 1); + std::vector w((size_t) copies); + for (auto& p : w) p = dev_weights(f, n_in, n_out); + float* y; + CK(cudaMalloc(&y, (size_t) 8 * n_out * 4)); + std::printf("%-6s %5d -> %5d (%5.2f MB), us (GB/s):", f.name, n_in, n_out, mb); + for (int nc : {2, 3, 4}) { + native_quantize_q8_1(a.x, a.q_ref, n_in, nc, g_s); + native_quantize_q8_1_il(a.x, a.q, a.il, n_in, nc, g_s); + const double t0 = time_calls([&](int r) { native_mmvq(f.type, w[(size_t) (r % copies)], a.q_ref, y, n_in, n_out, nc, g_s); }); + const double t1 = time_calls([&](int r) { + native_mmvq_il(f.type, w[(size_t) (r % copies)], a.q, a.il, y, n_in, n_out, nc, g_s); + }); + std::printf(" %d: %5.1f -> %5.1f (%3.0f)", nc, t0, t1, mb * 1e3 / t1); + } + std::printf("\n"); + for (auto p : w) cudaFree(p); + cudaFree(y); +} + +} // namespace + +int main(int argc, char** argv) { + const bool bench = argc > 1 && std::strcmp(argv[1], "--bench") == 0; + CK(cudaStreamCreateWithFlags(&g_s, cudaStreamNonBlocking)); + Act a2560 = make_act(2560), a6144 = make_act(6144), a640 = make_act(640); + if (bench) { + for (const Fmt& f : kFmts) { + bench_shape(f, a2560, 12288); + bench_shape(f, a2560, 10240); + bench_shape(f, a2560, 6144); + bench_shape(f, a6144, 2560); + bench_shape(f, a2560, 640); + } + return 0; + } + int bad = 0, cases = 0; + for (const Fmt& f : kFmts) { + const std::pair shapes[] = {{&a2560, 10240}, {&a2560, 6144}, {&a6144, 2560}, {&a2560, 12288}, + {&a2560, 2050}, {&a2560, 4097}, {&a2560, 8192}, {&a2560, 512}, + {&a2560, 640}}; + for (const auto& sh : shapes) { + bad += check_shape(f, *sh.first, sh.second); + ++cases; + } + if (f.block_elems == 32) { // Q8_0 also at the shared expert's down shape + bad += check_shape(f, a640, 2560); + ++cases; + } + } + std::printf("mmvq_il_parity: %d shape/format cases x 1..8 columns x rows a warp: %s\n", cases, + bad ? "FAILED" : "bitwise equal"); + return bad ? 1 : 0; +} diff --git a/src/kernels/native_expert_parity.cpp b/src/kernels/native_expert_parity.cpp index 496cf5ea7a0..ed8dc262cc8 100644 --- a/src/kernels/native_expert_parity.cpp +++ b/src/kernels/native_expert_parity.cpp @@ -1,16 +1,16 @@ -// src/kernels/native_expert_parity.cpp - plan v0.3 P6: one native expert three ways, on real GGUF rows. +// src/kernels/native_expert_parity.cpp - plan v0.3 P6: one native expert three ways. // -// build/native_expert_parity [layer ...] +// build/native_expert_parity [layer ...] real rows (a split model's shards are found by name) +// build/native_expert_parity --synthetic GU/DOWN ... random weights quantized by ggml, e.g. iq1_s/iq4_nl // // (a) float reference: ggml's own dequantizer (`to_float`) and a float SwiGLU expert, (b) the CPU path // (ggml-cpu vec_dot with its quantized activations), (c) the GPU path (`native_expert_grouped`, q8_1 -// activations). (b) and (c) each differ from (a) by their activation rounding only (a few 1e-3 relative). +// activations). (b) and (c) each differ from (a) by their activation rounding only (about 1e-2 relative). +// (d) the GPU dequantizers of the prompt path and the embedding, against `to_float`: exact. #include "strata/artifact/gguf_reader.hpp" #include "strata/kernels/cpu/native_expert.hpp" #include "strata/kernels/cpu/expert.hpp" #include "strata/kernels/cpu/iq_avx512.hpp" -#include "strata/kernels/cpu/iq_avx2.hpp" -#include "strata/kernels/cpu/expert_layout.hpp" #include "ggml-cpu.h" #include "strata/kernels/iq_kernels.hpp" @@ -18,6 +18,7 @@ #include +#include #include #include #include @@ -28,233 +29,266 @@ namespace cpu = strata::kernels::cpu; -static double rel(const std::vector& a, const std::vector& b) { +namespace { + +constexpr int NT = 3; +constexpr int64_t H = 2560, FF = 640; + +double rel(const std::vector& a, const std::vector& b) { double n = 0, d = 0; for (size_t i = 0; i < a.size(); ++i) { n += std::fabs((double) a[i] - b[i]); d += std::fabs((double) b[i]); } return n / (d + 1e-30); } -int main(int argc, char** argv) { - setvbuf(stdout, nullptr, _IONBF, 0); // keep the trail on a crash - if (argc < 2) { std::fprintf(stderr, "usage: native_expert_parity [layer ...]\n"); return 2; } - strata::GgufFile gguf(argv[1]); - std::vector layers; - for (int i = 2; i < argc; ++i) layers.push_back(std::atoi(argv[i])); - if (layers.empty()) layers = {0, 1, 2, 3, 20, 47}; - const int NT = 3, E = 7; - const int64_t H = 2560, FF = 640; +double us_since(std::chrono::steady_clock::time_point t0, int it) { + return std::chrono::duration(std::chrono::steady_clock::now() - t0).count() / it; +} + +// The three ways and the dequantizers for one expert blob; returns the number of failed checks. +int check_blob(const cpu::NativeFmt& f, const std::vector& blob, int seed, const std::string& label, + cudaStream_t s) { int failures = 0; - cudaStream_t s; - cudaStreamCreate(&s); - for (int l : layers) { - const strata::TensorInfo* t[3] = {}; - const char* roles[3] = {"gate", "up", "down"}; - for (const auto& ti : gguf.tensors()) - for (int r = 0; r < 3; ++r) - if (ti.name == "blk." + std::to_string(l) + ".ffn_" + roles[r] + "_exps.weight") t[r] = &ti; - if (!t[0] || !t[1] || !t[2]) { std::printf("layer %d: no expert tensors\n", l); ++failures; continue; } - cpu::NativeFmt f; - std::string err; - if (!cpu::native_fmt((int) t[0]->type, (int) t[2]->type, H, FF, f, err)) { - std::printf("layer %d: %s\n", l, err.c_str()); ++failures; continue; - } - std::vector blob(f.bytes); - std::memcpy(blob.data(), gguf.tensor_data(*t[0]) + (size_t) E * f.up_off, f.up_off); - std::memcpy(blob.data() + f.up_off, gguf.tensor_data(*t[1]) + (size_t) E * f.up_off, f.up_off); - std::memcpy(blob.data() + f.down_off, gguf.tensor_data(*t[2]) + (size_t) E * (f.bytes - f.down_off), - f.bytes - f.down_off); - // (a) the float reference - const auto* tg = ggml_get_type_traits((ggml_type) f.gu_type); - const auto* td = ggml_get_type_traits((ggml_type) f.d_type); - std::vector G((size_t) FF * H), U((size_t) FF * H), D((size_t) H * FF); + // (a) the float reference + const auto* tg = ggml_get_type_traits((ggml_type) f.gu_type); + const auto* td = ggml_get_type_traits((ggml_type) f.d_type); + std::vector G((size_t) FF * H), U((size_t) FF * H), D((size_t) H * FF); + for (int64_t r = 0; r < FF; ++r) { + tg->to_float(blob.data() + r * f.gu_row, G.data() + r * H, H); + tg->to_float(blob.data() + f.up_off + r * f.gu_row, U.data() + r * H, H); + } + for (int64_t r = 0; r < H; ++r) td->to_float(blob.data() + f.down_off + r * f.d_row, D.data() + r * FF, FF); + std::mt19937 rng(11 + seed); + std::normal_distribution nd(0.f, 1.f); + std::vector x((size_t) NT * H); + for (auto& v : x) v = nd(rng); + std::vector ref((size_t) NT * H), got_c((size_t) NT * H), got_g((size_t) NT * H); + for (int k = 0; k < NT; ++k) { + std::vector h(FF); for (int64_t r = 0; r < FF; ++r) { - tg->to_float(blob.data() + r * f.gu_row, G.data() + r * H, H); - tg->to_float(blob.data() + f.up_off + r * f.gu_row, U.data() + r * H, H); + double g = 0, u = 0; + for (int64_t i = 0; i < H; ++i) { g += (double) G[r * H + i] * x[k * H + i]; u += (double) U[r * H + i] * x[k * H + i]; } + h[r] = (float) (g / (1.0 + std::exp(-g)) * u); } - for (int64_t r = 0; r < H; ++r) td->to_float(blob.data() + f.down_off + r * f.d_row, D.data() + r * FF, FF); - std::mt19937 rng(11 + l); - std::normal_distribution nd(0.f, 1.f); - std::vector x((size_t) NT * H); - for (auto& v : x) v = nd(rng); - std::vector ref((size_t) NT * H), got_c((size_t) NT * H), got_g((size_t) NT * H); + for (int64_t r = 0; r < H; ++r) { + double o = 0; + for (int64_t i = 0; i < FF; ++i) o += (double) D[r * FF + i] * h[i]; + ref[k * H + r] = (float) o; + } + } + // (b) the CPU + { + std::vector> act(NT, std::vector(cpu::kNativeActBytes)); + std::vector> hq(NT, std::vector(cpu::kNativeHBytes)); + std::vector> ff(NT, std::vector(FF)); + const void* a[NT]; + float* ffp[NT]; + const void* hp[NT]; + float* op[NT]; for (int k = 0; k < NT; ++k) { - std::vector h(FF); - for (int64_t r = 0; r < FF; ++r) { - double g = 0, u = 0; - for (int64_t i = 0; i < H; ++i) { g += (double) G[r * H + i] * x[k * H + i]; u += (double) U[r * H + i] * x[k * H + i]; } - h[r] = (float) (g / (1.0 + std::exp(-g)) * u); - } - for (int64_t r = 0; r < H; ++r) { - double o = 0; - for (int64_t i = 0; i < FF; ++i) o += (double) D[r * FF + i] * h[i]; - ref[k * H + r] = (float) o; - } + cpu::native_quant_act(f, x.data() + k * H, act[k].data()); + a[k] = act[k].data(); + ffp[k] = ff[k].data(); } - // (b) the CPU - { - std::vector> act(NT, std::vector(cpu::kNativeActBytes)); - std::vector> hq(NT, std::vector(cpu::kNativeHBytes)); - std::vector> ff(NT, std::vector(FF)); - const void* a[NT]; - float* ffp[NT]; - const void* hp[NT]; - float* op[NT]; - for (int k = 0; k < NT; ++k) { - cpu::native_quant_act(f, x.data() + k * H, act[k].data()); - a[k] = act[k].data(); - ffp[k] = ff[k].data(); - } - cpu::native_gu_rows(f, blob.data(), a, NT, ffp, 0, (int) FF); - if (cpu::iq512_supported(f.gu_type)) { - // ggml's own vec_dot, same Q8_K activations: the reference for both multi-token kernels - // (float-order differences only) - const auto* tc = ggml_get_type_traits_cpu((ggml_type) f.gu_type); - std::vector gref((size_t) NT * FF); - for (int k = 0; k < NT; ++k) - for (int64_t r = 0; r < FF; ++r) - tc->vec_dot((int) H, &gref[k * FF + r], 0, blob.data() + r * f.gu_row, 0, a[k], 0, 1); - double usg = 0.0; - { // ggml single-token throughput, the weights cache-resident (compute bound) - const int it = 50; - auto t0 = std::chrono::steady_clock::now(); - for (int i = 0; i < it; ++i) - for (int64_t r = 0; r < FF; ++r) { - float gg, uu; - tc->vec_dot((int) H, &gg, 0, blob.data() + r * f.gu_row, 0, a[0], 0, 1); - tc->vec_dot((int) H, &uu, 0, blob.data() + f.up_off + r * f.gu_row, 0, a[0], 0, 1); - } - auto t1 = std::chrono::steady_clock::now(); - usg = std::chrono::duration(t1 - t0).count() / it; + cpu::native_gu_rows(f, blob.data(), a, NT, ffp, 0, (int) FF); + const auto* tc = ggml_get_type_traits_cpu((ggml_type) f.gu_type); + const int it = 50; + if (cpu::iq512_supported(f.gu_type)) { + // the AVX-512 rows against ggml's own vec_dot, same Q8_K activations: float-order differences only + std::vector g512((size_t) NT * FF), gref((size_t) NT * FF); + float* gp[NT]; + for (int k = 0; k < NT; ++k) gp[k] = g512.data() + k * FF; + cpu::iq512_rows(f.gu_type, blob.data(), f.gu_row, (int) H, a, NT, gp, 0, (int) FF); + for (int k = 0; k < NT; ++k) + for (int64_t r = 0; r < FF; ++r) + tc->vec_dot((int) H, &gref[k * FF + r], 0, blob.data() + r * f.gu_row, 0, a[k], 0, 1); + std::printf(" %s AVX-512 gate rows vs ggml vec_dot: rel %.2e\n", + ggml_type_name((ggml_type) f.gu_type), rel(g512, gref)); + float* f1[1] = {ffp[0]}; + auto t0 = std::chrono::steady_clock::now(); + for (int i = 0; i < it; ++i) cpu::iq512_gu_rows(f.gu_type, blob.data(), f.gu_row, f.up_off, (int) H, a, 1, f1, 0, (int) FF); + const double us1 = us_since(t0, it); + t0 = std::chrono::steady_clock::now(); + for (int i = 0; i < it; ++i) cpu::iq512_gu_rows(f.gu_type, blob.data(), f.gu_row, f.up_off, (int) H, a, NT, ffp, 0, (int) FF); + std::printf(" gate+up one thread, weights cache-resident: avx512 %.0f us one token (%.2f GB/s), " + "%.0f us %d tokens\n", us1, 2.0 * f.up_off / us1 / 1e3, us_since(t0, it), NT); + } + { // ggml-cpu's own rows, one thread (what the pool runs for formats without an AVX-512 kernel) + auto t0 = std::chrono::steady_clock::now(); + for (int i = 0; i < it; ++i) + for (int64_t r = 0; r < FF; ++r) { + float gg, uu; + tc->vec_dot((int) H, &gg, 0, blob.data() + r * f.gu_row, 0, a[0], 0, 1); + tc->vec_dot((int) H, &uu, 0, blob.data() + f.up_off + r * f.gu_row, 0, a[0], 0, 1); } - // one multi-token kernel: parity against ggml's vec_dot, then one-token and NT-token throughput - auto check = [&](const char* tag, bool is512) { - std::vector g((size_t) NT * FF); - float* gp[NT]; - for (int k = 0; k < NT; ++k) gp[k] = g.data() + k * FF; - (is512 ? cpu::iq512_rows : cpu::iq256_rows)(f.gu_type, blob.data(), f.gu_row, (int) H, a, NT, gp, 0, (int) FF); - const double rg = rel(g, gref); - std::printf(" %s %s gate rows vs ggml vec_dot: rel %.2e\n", - ggml_type_name((ggml_type) f.gu_type), tag, rg); - if (rg > 1e-5) { // ggml's own vec_dot is the reference: above 1e-5 it is a real mismatch, - std::printf(" %s %s gate rows MISMATCH (rel %.2e > 1e-5)\n", - ggml_type_name((ggml_type) f.gu_type), tag, rg); // not rounding (measured 3e-8) - ++failures; - } - float* f1[1] = {ffp[0]}; - const int it = 50; - const auto gu = is512 ? cpu::iq512_gu_rows : cpu::iq256_gu_rows; - auto t0 = std::chrono::steady_clock::now(); - for (int i = 0; i < it; ++i) gu(f.gu_type, blob.data(), f.gu_row, f.up_off, (int) H, a, 1, f1, 0, (int) FF); - auto t1 = std::chrono::steady_clock::now(); - for (int i = 0; i < it; ++i) gu(f.gu_type, blob.data(), f.gu_row, f.up_off, (int) H, a, NT, ffp, 0, (int) FF); - auto t2 = std::chrono::steady_clock::now(); - const double us1 = std::chrono::duration(t1 - t0).count() / it; - const double usn = std::chrono::duration(t2 - t1).count() / it; - std::printf(" gate+up one token one thread: %s %.0f us (%.2f GB/s), ggml %.0f us (%.2f GB/s)\n", - tag, us1, 2.0 * f.up_off / us1 / 1e3, usg, 2.0 * f.up_off / usg / 1e3); - std::printf(" gate+up %d tokens one thread: %s %.0f us vs ggml %d x %.0f us\n", NT, tag, usn, NT, usg); - }; - if (cpu::cpu_avx512_ok()) check("avx512", true); // guarded: the binary runs on AVX-2 CPUs too - check("avx2", false); - } + const double usg = us_since(t0, it); + std::printf(" gate+up one thread, weights cache-resident: ggml %.0f us one token (%.2f GB/s)\n", usg, + 2.0 * f.up_off / usg / 1e3); + } + for (int k = 0; k < NT; ++k) { + cpu::native_quant_h(f, ff[k].data(), hq[k].data()); + hp[k] = hq[k].data(); + op[k] = got_c.data() + k * H; + } + cpu::native_down_rows(f, blob.data(), hp, NT, op, 0, (int) H); + if (f.d_type == 42) { + // (b2) the AVX-512 GGUF-layout Q2_0 kernel the pool uses for Q2_0 down projections + std::vector a2(NT); + const cpu::ActQ* ap[NT]; + std::vector alt((size_t) NT * H); + float* altp[NT]; for (int k = 0; k < NT; ++k) { - cpu::native_quant_h(f, ff[k].data(), hq[k].data()); - hp[k] = hq[k].data(); - op[k] = got_c.data() + k * H; - } - cpu::native_down_rows(f, blob.data(), hp, NT, op, 0, (int) H); - if (f.d_type == 42) { - // (b2) the GGUF-layout Q2_0 kernel the pool uses for Q2_0 down projections - the AVX-512 one - // where the CPU has it, the AVX-2 one (q2_avx2.cpp) where it does not. Calling the AVX-512 - // kernel unconditionally faults on a Zen 2/3 CPU. - std::vector a2(NT); - const cpu::ActQ* ap[NT]; - std::vector alt((size_t) NT * H); - float* altp[NT]; - for (int k = 0; k < NT; ++k) { - cpu::act_quant_any(ff[k].data(), (int) FF, a2[k]); // gated: AVX-512 or AVX-2 per CPU - ap[k] = &a2[k]; - altp[k] = alt.data() + k * H; - } - cpu::q2_rows_any(blob.data() + f.down_off, f.d_row, (int) (FF / 64), ap, NT, altp, 0, (int) H); - std::printf(" q2_0 %s down vs ggml down: rel %.2e\n", - cpu::cpu_avx512_ok() ? "AVX-512" : "AVX-2", rel(alt, got_c)); - } - if (f.d_type == 20) { - // (b3) the IQ4_NL multi-token AVX-2 kernel the pool now uses for IQ4_NL down projections, - // against ggml-cpu's single-token vec_dot on the SAME Q8_0 activations (h), plus timing. - std::vector alt((size_t) NT * H), refd((size_t) NT * H); - float* altp[NT]; - for (int k = 0; k < NT; ++k) altp[k] = alt.data() + k * H; - cpu::iq4nl256_down_rows(blob.data() + f.down_off, f.d_row, (int) FF, hp, NT, altp, 0, (int) H); - const auto* tdc = ggml_get_type_traits_cpu((ggml_type) f.d_type); - for (int k = 0; k < NT; ++k) - for (int64_t r = 0; r < H; ++r) - tdc->vec_dot((int) FF, refd.data() + k * H + r, 0, - blob.data() + f.down_off + (size_t) r * f.d_row, 0, hp[k], 0, 1); - const double rdn = rel(alt, refd); - std::printf(" iq4_nl AVX-2 multi-token down vs ggml down: rel %.2e\n", rdn); - if (rdn > 1e-5) { // measured 1e-7; same reasoning as the gate/up rows above - std::printf(" iq4_nl down MISMATCH (rel %.2e > 1e-5)\n", rdn); - ++failures; - } - const int it = 50; - const auto t0 = std::chrono::steady_clock::now(); - for (int i = 0; i < it; ++i) - cpu::iq4nl256_down_rows(blob.data() + f.down_off, f.d_row, (int) FF, hp, NT, altp, 0, (int) H); - const auto t1 = std::chrono::steady_clock::now(); - for (int i = 0; i < it; ++i) - for (int k = 0; k < NT; ++k) - for (int64_t r = 0; r < H; ++r) - tdc->vec_dot((int) FF, refd.data() + k * H + r, 0, - blob.data() + f.down_off + (size_t) r * f.d_row, 0, hp[k], 0, 1); - const auto t2 = std::chrono::steady_clock::now(); - const double usn = std::chrono::duration(t1 - t0).count() / it; - const double usg = std::chrono::duration(t2 - t1).count() / it; - std::printf(" down %d tokens one thread: avx2 %.0f us (%.2f GB/s) vs ggml %.0f us (%.2f GB/s)\n", - NT, usn, (double) (f.bytes - f.down_off) / usn / 1e3, usg, - (double) (f.bytes - f.down_off) / usg / 1e3); + cpu::act_quant_q8_1(ff[k].data(), (int) FF, a2[k]); + ap[k] = &a2[k]; + altp[k] = alt.data() + k * H; } + cpu::q2_0_gguf_rows_multi(blob.data() + f.down_off, f.d_row, (int) (FF / 64), ap, NT, altp, 0, (int) H); + std::printf(" q2_0 AVX-512 down vs ggml down: rel %.2e\n", rel(alt, got_c)); } - // (c) the GPU: one group holding the NT entries - { - const auto L = strata::kernels::native_expert_layout(f.gu_type, f.d_type, H, FF); - void *dblob, *dx, *dxq, *dscr; - float* dout; - unsigned long long* dptr; - int32_t *dstart, *dn, *ddst, *dtok; - cudaMalloc(&dblob, blob.size()); - cudaMalloc(&dx, x.size() * 4); - cudaMalloc(&dxq, (size_t) NT * H / 32 * 36); - cudaMalloc(&dscr, strata::kernels::native_expert_scratch_bytes(NT, FF)); - cudaMalloc((void**) &dout, (size_t) NT * H * 4); - cudaMalloc((void**) &dptr, 8); - cudaMalloc((void**) &dstart, 8); - cudaMalloc((void**) &dn, 4); - cudaMalloc((void**) &ddst, NT * 4); - cudaMalloc((void**) &dtok, NT * 4); - cudaMemcpy(dblob, blob.data(), blob.size(), cudaMemcpyHostToDevice); - cudaMemcpy(dx, x.data(), x.size() * 4, cudaMemcpyHostToDevice); - const unsigned long long p = (unsigned long long) dblob; - const int32_t st[2] = {0, NT}, one = 1, idx[NT] = {0, 1, 2}; - cudaMemcpy(dptr, &p, 8, cudaMemcpyHostToDevice); - cudaMemcpy(dstart, st, 8, cudaMemcpyHostToDevice); - cudaMemcpy(dn, &one, 4, cudaMemcpyHostToDevice); - cudaMemcpy(ddst, idx, NT * 4, cudaMemcpyHostToDevice); - cudaMemcpy(dtok, idx, NT * 4, cudaMemcpyHostToDevice); - strata::kernels::quantize_q8_1_rows((const float*) dx, NT, H, dxq, s); - strata::kernels::native_expert_grouped(L, dptr, dstart, dn, ddst, dtok, 1, NT, dxq, dscr, dout, s); + } + // (c) the GPU: one group holding the NT entries; (d) its dequantizers + { + const auto L = strata::kernels::native_expert_layout(f.gu_type, f.d_type, H, FF); + void *dblob, *dx, *dxq, *dscr; + float* dout; + unsigned long long* dptr; + int32_t *dstart, *dn, *ddst, *dtok; + cudaMalloc(&dblob, blob.size()); + cudaMalloc(&dx, x.size() * 4); + cudaMalloc(&dxq, (size_t) NT * H / 32 * 36); + cudaMalloc(&dscr, strata::kernels::native_expert_scratch_bytes(NT, FF)); + cudaMalloc((void**) &dout, (size_t) NT * H * 4); + cudaMalloc((void**) &dptr, 8); + cudaMalloc((void**) &dstart, 8); + cudaMalloc((void**) &dn, 4); + cudaMalloc((void**) &ddst, NT * 4); + cudaMalloc((void**) &dtok, NT * 4); + cudaMemcpy(dblob, blob.data(), blob.size(), cudaMemcpyHostToDevice); + cudaMemcpy(dx, x.data(), x.size() * 4, cudaMemcpyHostToDevice); + const unsigned long long p = (unsigned long long) dblob; + const int32_t st[2] = {0, NT}, one = 1, idx[NT] = {0, 1, 2}; + cudaMemcpy(dptr, &p, 8, cudaMemcpyHostToDevice); + cudaMemcpy(dstart, st, 8, cudaMemcpyHostToDevice); + cudaMemcpy(dn, &one, 4, cudaMemcpyHostToDevice); + cudaMemcpy(ddst, idx, NT * 4, cudaMemcpyHostToDevice); + cudaMemcpy(dtok, idx, NT * 4, cudaMemcpyHostToDevice); + strata::kernels::quantize_q8_1_rows((const float*) dx, NT, H, dxq, s); + strata::kernels::native_expert_grouped(L, dptr, dstart, dn, ddst, dtok, 1, NT, dxq, dscr, dout, s); + cudaStreamSynchronize(s); + cudaMemcpy(got_g.data(), dout, got_g.size() * 4, cudaMemcpyDeviceToHost); + float* dq = nullptr; + cudaMalloc((void**) &dq, (size_t) FF * H * 4); + std::vector got_dq((size_t) FF * H); + double dq_err = 0.0; + const struct { const uint8_t* src; int type; const std::vector* want; } mats[3] = { + {(const uint8_t*) dblob, f.gu_type, &G}, {(const uint8_t*) dblob + f.up_off, f.gu_type, &U}, + {(const uint8_t*) dblob + f.down_off, f.d_type, &D}}; + for (const auto& m : mats) { + strata::kernels::iq_dequant_f32(m.type, m.src, FF * H, dq, s); cudaStreamSynchronize(s); - cudaMemcpy(got_g.data(), dout, got_g.size() * 4, cudaMemcpyDeviceToHost); - cudaFree(dblob); cudaFree(dx); cudaFree(dxq); cudaFree(dscr); cudaFree(dout); cudaFree(dptr); - cudaFree(dstart); cudaFree(dn); cudaFree(ddst); cudaFree(dtok); + cudaMemcpy(got_dq.data(), dq, got_dq.size() * 4, cudaMemcpyDeviceToHost); + dq_err = (std::max)(dq_err, rel(got_dq, *m.want)); + } + std::printf(" GPU dequant vs ggml to_float: rel %.2e\n", dq_err); + if (!(dq_err < 1e-6)) ++failures; + cudaFree(dq); + cudaFree(dblob); cudaFree(dx); cudaFree(dxq); cudaFree(dscr); cudaFree(dout); cudaFree(dptr); + cudaFree(dstart); cudaFree(dn); cudaFree(ddst); cudaFree(dtok); + } + const double ec = rel(got_c, ref), eg = rel(got_g, ref), ecg = rel(got_c, got_g); + const bool ok = ec < 3e-2 && eg < 3e-2 && std::isfinite(ec) && std::isfinite(eg); + std::printf("%-9s %-8s/%-7s blob %8zu cpu rel %.2e gpu rel %.2e cpu-gpu %.2e %s\n", label.c_str(), + ggml_type_name((ggml_type) f.gu_type), ggml_type_name((ggml_type) f.d_type), f.bytes, ec, eg, ecg, + ok ? "ok" : "FAIL"); + return failures + (ok ? 0 : 1); +} + +int type_of(std::string name) { + for (auto& c : name) c = (char) std::tolower((unsigned char) c); + for (int t = 0; t < GGML_TYPE_COUNT; ++t) { + const char* n = ggml_type_name((ggml_type) t); + if (n == nullptr) continue; + std::string m(n); + for (auto& c : m) c = (char) std::tolower((unsigned char) c); + if (m == name) return t; + } + return -1; +} + +// Random weights (rows of different magnitude) quantized by ggml into the blob layout. +std::vector synthetic_blob(const cpu::NativeFmt& f, int seed) { + std::vector blob(f.bytes); + std::mt19937 rng(seed); + std::normal_distribution nd(0.f, 1.f); + auto quant = [&](int type, int64_t rows, int64_t cols, uint8_t* dst) { + std::vector w((size_t) (rows * cols)); + for (int64_t r = 0; r < rows; ++r) { + const float scale = 0.02f * (0.5f + (float) (r % 7) / 7.0f); + for (int64_t c = 0; c < cols; ++c) w[(size_t) (r * cols + c)] = scale * nd(rng); + } + const std::vector imatrix((size_t) cols, 1.0f); + ggml_quantize_chunk((ggml_type) type, w.data(), dst, 0, rows, cols, + ggml_quantize_requires_imatrix((ggml_type) type) ? imatrix.data() : nullptr); + }; + quant(f.gu_type, FF, H, blob.data()); + quant(f.gu_type, FF, H, blob.data() + f.up_off); + quant(f.d_type, H, FF, blob.data() + f.down_off); + return blob; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc < 2) { + std::fprintf(stderr, "usage: native_expert_parity [layer ...]\n" + " native_expert_parity --synthetic GU/DOWN ... (ggml type names, e.g. q6_K/q8_0)\n"); + return 2; + } + int failures = 0; + cudaStream_t s; + cudaStreamCreate(&s); + if (std::string(argv[1]) == "--synthetic") { + for (int i = 2; i < argc; ++i) { + const std::string arg = argv[i]; + const size_t slash = arg.find('/'); + const int gu = slash == std::string::npos ? -1 : type_of(arg.substr(0, slash)); + const int dn = slash == std::string::npos ? -1 : type_of(arg.substr(slash + 1)); + cpu::NativeFmt f; + std::string err; + 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 supported pair of GGML types" : err.c_str()); + ++failures; + continue; + } + failures += check_blob(f, synthetic_blob(f, i), i, "synthetic", s); + } + } else { + const strata::GgufModel model(strata::gguf_split_paths(argv[1])); + std::vector layers; + for (int i = 2; i < argc; ++i) layers.push_back(std::atoi(argv[i])); + if (layers.empty()) layers = {0, 1, 2, 3, 20, 47}; + const int E = 7; + for (int l : layers) { + const strata::TensorInfo* t[3] = {}; + const uint8_t* data[3] = {}; + const char* roles[3] = {"gate", "up", "down"}; + for (int r = 0; r < 3; ++r) { + size_t at = 0; + t[r] = model.find("blk." + std::to_string(l) + ".ffn_" + roles[r] + "_exps.weight", &at); + if (t[r]) data[r] = model.shard(at).tensor_data(*t[r]); + } + if (!t[0] || !t[1] || !t[2]) { std::printf("layer %d: no expert tensors\n", l); ++failures; continue; } + cpu::NativeFmt f; + std::string err; + if (!cpu::native_fmt((int) t[0]->type, (int) t[2]->type, H, FF, f, err)) { + std::printf("layer %d: %s\n", l, err.c_str()); ++failures; continue; + } + std::vector blob(f.bytes); + std::memcpy(blob.data(), data[0] + (size_t) E * f.up_off, f.up_off); + std::memcpy(blob.data() + f.up_off, data[1] + (size_t) E * f.up_off, f.up_off); + std::memcpy(blob.data() + f.down_off, data[2] + (size_t) E * (f.bytes - f.down_off), f.bytes - f.down_off); + failures += check_blob(f, blob, l, "layer " + std::to_string(l), s); } - const double ec = rel(got_c, ref), eg = rel(got_g, ref), ecg = rel(got_c, got_g); - const bool ok = ec < 3e-2 && eg < 3e-2 && std::isfinite(ec) && std::isfinite(eg); - std::printf("layer %2d %-8s/%-7s blob %8zu cpu rel %.2e gpu rel %.2e cpu-gpu %.2e %s\n", l, - ggml_type_name((ggml_type) f.gu_type), ggml_type_name((ggml_type) f.d_type), f.bytes, ec, eg, ecg, - ok ? "ok" : "FAIL"); - if (!ok) ++failures; } std::printf("native_expert_parity: %d failures\n", failures); return failures ? 1 : 0; diff --git a/src/kernels/ngram.cpp b/src/kernels/ngram.cpp index b1bd98b2b4f..8db9298852a 100644 --- a/src/kernels/ngram.cpp +++ b/src/kernels/ngram.cpp @@ -1,10 +1,16 @@ -// src/kernels/ngram.cpp - P2.S4: the PLE n-gram hash and the IQ4_NL table read. +// src/kernels/ngram.cpp - P2.S4: the PLE n-gram hash and the table read (IQ4_NL or Q8_0). // // See include/strata/kernels/ngram.hpp for the semantics, the rival readings, and the note on MADV_RANDOM. #include "strata/kernels/ngram.hpp" #include "strata/artifact/gguf_reader.hpp" #include "strata/kernels/f16_bits.hpp" #include "strata/ngram/ple_reader.hpp" +#include "strata/platform/direct_file.hpp" + +#include +#include +#include +#include #include #include @@ -107,11 +113,37 @@ void iq4nl_dequant_row(const uint8_t* row, float* out160) { } } +void ple_dequant_row(int type, const uint8_t* row, float* out160) { + if (type != 8) { + iq4nl_dequant_row(row, out160); + return; + } + for (int b = 0; b < PLE_HEAD_DIM / 32; ++b) { // Q8_0: fp16 scale + 32 int8 + const uint8_t* blk = row + (size_t) b * 34; + uint16_t dbits; + std::memcpy(&dbits, blk, 2); + const float d = f32_from_f16(dbits); + for (int j = 0; j < 32; ++j) out160[b * 32 + j] = d * (float) (int8_t) blk[2 + j]; + } +} + +std::string ple_table_shard(const std::vector& shards) { + for (const auto& path : shards) { + try { + if (GgufFile(path).find("per_layer_token_embd.weight") != nullptr) return path; + } catch (const std::exception&) { + } + } + return {}; +} + // --------------------------------------------------------------------------------------------------- struct PleTable::Impl { GgufFile* file = nullptr; const uint8_t* data = nullptr; uint64_t n_rows = 0; + int type = 20; // IQ4_NL or Q8_0 + uint32_t row_bytes = PLE_ROW_BYTES; mutable uint64_t bytes_read = 0; // Direct mode (plan v0.3 P2): the mapping above is released after the header parse and every row comes // from an unbuffered SSD read into `raw`. @@ -120,9 +152,82 @@ struct PleTable::Impl { strata::ngram::PleReader::Ticket ticket; bool pending = false; uint32_t rows[PLE_N_HEADS] = {}; - uint8_t raw[PLE_N_HEADS * PLE_ROW_BYTES] = {}; + uint8_t raw[PLE_N_HEADS * PLE_ROW_BYTES_Q8_0] = {}; + // Ram mode: the table read into `ram_base` by `loader`; `ram` (the table's first row) valid once `ram_ready` + std::string path; + uint64_t table_offset = 0; + uint8_t* ram_base = nullptr; + uint64_t ram_bytes = 0; + const uint8_t* ram = nullptr; + std::atomic ram_ready{false}, cancel{false}; + std::thread loader; + std::string ram_note; + bool pending_ram = false; // the pending token's rows come from RAM + // the read-ahead slots: a token's rows in flight from the SSD (`via_reader`), else read at collect + struct Ahead { + bool pending = false, via_reader = false; + uint32_t rows[PLE_N_HEADS] = {}; + uint8_t raw[PLE_N_HEADS * PLE_ROW_BYTES_Q8_0] = {}; + strata::ngram::PleReader::Ticket ticket; + }; + Ahead aheads[PleTable::kAheadSlots]; + + bool from_ram() const { return mode == PleIo::Ram && ram_ready.load(std::memory_order_acquire); } + bool direct() const { return (mode == PleIo::Direct || mode == PleIo::Ram) && !from_ram(); } + void load(); }; +// The whole table in 8 MiB unbuffered reads, four in flight. Pageable: locking it resident first (VirtualLock of +// 28.8 GB) faults in its pages under the process's working-set lock and stalled the startup's other threads ~8 s. +void PleTable::Impl::load() { + using strata::platform::DirectFile; + const auto t0 = std::chrono::steady_clock::now(); + const uint64_t a0 = table_offset & ~(uint64_t) (DirectFile::alignment() - 1); + const uint64_t need = table_offset + n_rows * (uint64_t) row_bytes - a0; + DirectFile f; + std::string err; + if (!f.open(path, err)) { ram_note = "PLE table in RAM: " + err; std::fprintf(stderr, "strata: %s\n", ram_note.c_str()); return; } + constexpr uint32_t kChunk = 8u << 20; + constexpr int kDepth = 4; + uint64_t next = 0, got = 0; + int inflight = 0; + bool ok = true; + while (ok && !cancel.load() && (next < need || inflight > 0)) { + while (ok && inflight < kDepth && next < need) { + const uint32_t len = (uint32_t) std::min(kChunk, ram_bytes - next); + ok = f.submit(a0 + next, ram_base + next, len, next, err); + next += len; + ++inflight; + } + strata::platform::Completion c[kDepth]; + const int k = ok ? f.wait(c, kDepth, -1) : 0; + for (int i = 0; i < k; ++i) { + if (c[i].tag == DirectFile::WAKE_TAG) continue; + --inflight; + ok = ok && c[i].ok; + got += c[i].bytes; + } + } + while (inflight > 0) { // a failed or cancelled load still collects what it queued + strata::platform::Completion c[kDepth]; + const int k = f.wait(c, kDepth, -1); + for (int i = 0; i < k; ++i) if (c[i].tag != DirectFile::WAKE_TAG) --inflight; + } + if (cancel.load()) return; + if (!ok || got < need) { + ram_note = "PLE table in RAM: the load failed" + (err.empty() ? std::string() : ": " + err) + "; rows stay on the SSD"; + std::fprintf(stderr, "strata: %s\n", ram_note.c_str()); + return; + } + ram = ram_base + (table_offset - a0); + const double s = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + char buf[120]; + std::snprintf(buf, sizeof buf, "PLE table in RAM: %.1f GB in %.1f s", (double) need / 1e9, s); + ram_note = buf; + ram_ready.store(true, std::memory_order_release); + std::fprintf(stderr, "strata: %s\n", ram_note.c_str()); +} + PleTable::PleTable() : impl_(new Impl) {} PleTable::~PleTable() { close(); delete impl_; } @@ -150,11 +255,13 @@ bool PleTable::open(const std::string& gguf_path, std::string& err, const PleIoO close(); return false; } - if (std::strcmp(t->type_name(), "IQ4_NL") != 0) { - err = std::string("per_layer_token_embd.weight is ") + t->type_name() + ", not IQ4_NL"; + if (t->type != 20 && t->type != 8) { + err = std::string("per_layer_token_embd.weight is ") + t->type_name() + ", not IQ4_NL or Q8_0"; close(); return false; } + impl_->type = (int) t->type; + impl_->row_bytes = t->type == 8 ? PLE_ROW_BYTES_Q8_0 : PLE_ROW_BYTES; impl_->n_rows = t->shape[1]; impl_->data = impl_->file->tensor_data(*t); @@ -162,34 +269,38 @@ bool PleTable::open(const std::string& gguf_path, std::string& err, const PleIoO // the offset within the GGUF's DATA SECTION: the file's first 192 bytes are a header, and reading at 0 // would decode the header plus 192 bytes of shifted rows - still plausible IQ4_NL, and wrong for every // row. `GgufFile` parses the header, so `tensor_data` is already correct; this asserts the tensor - // exactly fills the file from there, which is what makes the whole arrangement checkable rather than - // assumed. A wrong data offset would leave a different remainder. - // A shard may hold other tensors too (Swift 1.5's shard 1 holds layers 0-12 and the table): the table must - // then fit inside the file at its own offset; alone in its shard (the original's shard 2) it fills it exactly. - const uint64_t need = impl_->n_rows * (uint64_t) PLE_ROW_BYTES; - const uint64_t have = impl_->file->file_size() - impl_->file->data_start(); - const bool alone = impl_->file->tensors().size() == 1; - if (alone ? need != have : t->offset + need > have) { + // exactly fills its span - up to the next tensor or the end of the file, less the alignment padding - + // which is what makes the whole arrangement checkable rather than assumed. A wrong data offset would + // leave a different remainder. (The table fills the ISTA shard 2 alone; the unsloth shard 2 holds other + // tensors too.) + const uint64_t need = impl_->n_rows * (uint64_t) impl_->row_bytes; + uint64_t end = impl_->file->file_size() - impl_->file->data_start(); + for (const auto& o : impl_->file->tensors()) + if (o.offset > t->offset && o.offset < end) end = o.offset; + const uint64_t have = end - t->offset; + if (need > have || have - need >= impl_->file->alignment()) { char buf[256]; std::snprintf(buf, sizeof buf, - "PLE table size mismatch: %llu rows x %d B = %llu at offset %llu, but the file holds %llu from " - "data_start %llu", - (unsigned long long) impl_->n_rows, PLE_ROW_BYTES, (unsigned long long) need, - (unsigned long long) t->offset, (unsigned long long) have, - (unsigned long long) impl_->file->data_start()); + "PLE table size mismatch: %llu rows x %d B = %llu, but its span in the file is %llu B " + "(data_start %llu)", + (unsigned long long) impl_->n_rows, (int) impl_->row_bytes, (unsigned long long) need, + (unsigned long long) have, (unsigned long long) impl_->file->data_start()); err = buf; close(); return false; } - if (io.mode == PleIo::Direct) { + if (io.mode == PleIo::Direct || io.mode == PleIo::Ram) { // The parse above is the validated source of the offset; the mapping itself is not kept, so no page of // the table can enter this process's working set or the file cache through it. const uint64_t table_offset = impl_->file->data_start() + t->offset; + impl_->path = gguf_path; + impl_->table_offset = table_offset; const uint64_t n_rows = impl_->n_rows; delete impl_->file; impl_->file = nullptr; impl_->data = nullptr; - if (!impl_->reader.open(gguf_path, table_offset, n_rows, io.max_inflight, io.cache_rows, err, io.io_thread)) { + if (!impl_->reader.open(gguf_path, table_offset, n_rows, io.max_inflight, io.cache_rows, err, io.io_thread, + impl_->row_bytes)) { close(); return false; } @@ -200,7 +311,19 @@ bool PleTable::open(const std::string& gguf_path, std::string& err, const PleIoO } void PleTable::close() { + if (impl_->loader.joinable()) { + impl_->cancel = true; + impl_->loader.join(); + } + impl_->ram_ready = false; + impl_->cancel = false; + impl_->ram = nullptr; + if (impl_->ram_base != nullptr) { + strata::platform::DirectFile::free_aligned(impl_->ram_base); + impl_->ram_base = nullptr; + } impl_->reader.close(); + for (auto& a : impl_->aheads) a.pending = false; impl_->pending = false; impl_->mode = PleIo::Mmap; delete impl_->file; @@ -210,36 +333,60 @@ void PleTable::close() { } bool PleTable::is_open() const { return impl_->data != nullptr || impl_->reader.is_open(); } + +bool PleTable::start_ram_load(std::string& err) { + using strata::platform::DirectFile; + if (impl_->mode != PleIo::Ram || impl_->ram_base != nullptr) return true; + const uint64_t a0 = impl_->table_offset & ~(uint64_t) (DirectFile::alignment() - 1); + const uint64_t end = impl_->table_offset + impl_->n_rows * (uint64_t) impl_->row_bytes; + impl_->ram_bytes = (end - a0 + DirectFile::alignment() - 1) / DirectFile::alignment() * DirectFile::alignment(); + impl_->ram_base = (uint8_t*) DirectFile::alloc_aligned((size_t) impl_->ram_bytes); + if (impl_->ram_base == nullptr) { + err = "PLE table in RAM: " + std::to_string(impl_->ram_bytes >> 20) + " MiB could not be allocated"; + return false; + } + impl_->loader = std::thread([m = impl_] { m->load(); }); + return true; +} + +bool PleTable::ram_ready() const { return impl_->from_ram(); } PleIo PleTable::mode() const { return impl_->mode; } uint64_t PleTable::rows() const { return impl_->n_rows; } uint64_t PleTable::bytes_read() const { return impl_->bytes_read; } void PleTable::read_row(uint32_t row, float* out160) const { - if (impl_->mode == PleIo::Direct && impl_->reader.is_open()) { - uint8_t raw[PLE_ROW_BYTES]; + if (impl_->from_ram()) { + if (row >= impl_->n_rows) std::memset(out160, 0, (size_t) PLE_HEAD_DIM * sizeof(float)); + else ple_dequant_row(impl_->type, impl_->ram + (size_t) row * impl_->row_bytes, out160); + impl_->bytes_read += impl_->row_bytes; + return; + } + if (impl_->direct() && impl_->reader.is_open()) { + uint8_t raw[PLE_ROW_BYTES_Q8_0]; std::string err; const auto t = impl_->reader.issue(&row, 1, raw); if (!impl_->reader.collect(t, err)) { std::memset(out160, 0, (size_t) PLE_HEAD_DIM * sizeof(float)); return; } - iq4nl_dequant_row(raw, out160); - impl_->bytes_read += PLE_ROW_BYTES; + ple_dequant_row(impl_->type, raw, out160); + impl_->bytes_read += impl_->row_bytes; return; } if (impl_->data == nullptr || row >= impl_->n_rows) { std::memset(out160, 0, (size_t) PLE_HEAD_DIM * sizeof(float)); return; } - iq4nl_dequant_row(impl_->data + (size_t) row * PLE_ROW_BYTES, out160); - impl_->bytes_read += PLE_ROW_BYTES; + ple_dequant_row(impl_->type, impl_->data + (size_t) row * impl_->row_bytes, out160); + impl_->bytes_read += impl_->row_bytes; } bool PleTable::issue(const uint32_t* rows16) { std::memcpy(impl_->rows, rows16, sizeof impl_->rows); - if (impl_->mode == PleIo::Direct) { + if (impl_->mode != PleIo::Mmap) { if (impl_->pending) return false; // one token in flight per table - impl_->ticket = impl_->reader.issue(impl_->rows, PLE_N_HEADS, impl_->raw); + impl_->pending_ram = impl_->from_ram(); + if (!impl_->pending_ram) impl_->ticket = impl_->reader.issue(impl_->rows, PLE_N_HEADS, impl_->raw); impl_->pending = true; return true; } @@ -249,8 +396,8 @@ bool PleTable::issue(const uint32_t* rows16) { ULONG_PTR n = 0; for (int h = 0; h < PLE_N_HEADS; ++h) { if (rows16[h] >= impl_->n_rows) continue; - ranges[n].VirtualAddress = (PVOID) (impl_->data + (size_t) rows16[h] * PLE_ROW_BYTES); - ranges[n].NumberOfBytes = PLE_ROW_BYTES; + ranges[n].VirtualAddress = (PVOID) (impl_->data + (size_t) rows16[h] * impl_->row_bytes); + ranges[n].NumberOfBytes = impl_->row_bytes; ++n; } if (n > 0) (void) PrefetchVirtualMemory(GetCurrentProcess(), n, ranges, 0); @@ -263,11 +410,15 @@ bool PleTable::issue(const uint32_t* rows16) { bool PleTable::collect(float* out2560, std::string& err) { if (!impl_->pending) { err = "PleTable::collect without issue"; return false; } impl_->pending = false; - if (impl_->mode == PleIo::Direct) { + if (impl_->pending_ram) { + for (int h = 0; h < PLE_N_HEADS; ++h) read_row(impl_->rows[h], out2560 + (size_t) h * PLE_HEAD_DIM); + return true; + } + if (impl_->mode != PleIo::Mmap) { if (!impl_->reader.collect(impl_->ticket, err)) return false; for (int h = 0; h < PLE_N_HEADS; ++h) - iq4nl_dequant_row(impl_->raw + (size_t) h * PLE_ROW_BYTES, out2560 + (size_t) h * PLE_HEAD_DIM); - impl_->bytes_read += (uint64_t) PLE_N_HEADS * PLE_ROW_BYTES; + ple_dequant_row(impl_->type, impl_->raw + (size_t) h * impl_->row_bytes, out2560 + (size_t) h * PLE_HEAD_DIM); + impl_->bytes_read += (uint64_t) PLE_N_HEADS * impl_->row_bytes; return true; } for (int h = 0; h < PLE_N_HEADS; ++h) read_row(impl_->rows[h], out2560 + (size_t) h * PLE_HEAD_DIM); @@ -277,22 +428,61 @@ bool PleTable::collect(float* out2560, std::string& err) { bool PleTable::gather_batch(const uint32_t* rows, size_t n_tokens, float* out, std::string& err) { if (impl_->pending) { err = "PleTable::gather_batch while a token is in flight"; return false; } const size_t n = n_tokens * (size_t) PLE_N_HEADS; - if (impl_->mode == PleIo::Direct) { - std::vector raw(n * PLE_ROW_BYTES); + if (impl_->direct()) { + std::vector raw(n * impl_->row_bytes); const auto ticket = impl_->reader.issue(rows, n, raw.data()); if (!impl_->reader.collect(ticket, err)) return false; - for (size_t i = 0; i < n; ++i) iq4nl_dequant_row(raw.data() + i * PLE_ROW_BYTES, out + i * PLE_HEAD_DIM); - impl_->bytes_read += (uint64_t) n * PLE_ROW_BYTES; + for (size_t i = 0; i < n; ++i) + ple_dequant_row(impl_->type, raw.data() + i * impl_->row_bytes, out + i * PLE_HEAD_DIM); + impl_->bytes_read += (uint64_t) n * impl_->row_bytes; return true; } for (size_t i = 0; i < n; ++i) read_row(rows[i], out + i * PLE_HEAD_DIM); return true; } +void PleTable::ahead(int slot, const uint32_t* rows16) { + if (slot < 0 || slot >= kAheadSlots) return; + Impl::Ahead& a = impl_->aheads[slot]; + if (a.pending && a.via_reader) { // the slot's earlier rows (a draft the window left out): long since read + std::string err; + (void) impl_->reader.collect(a.ticket, err); + } + std::memcpy(a.rows, rows16, sizeof a.rows); + a.via_reader = impl_->direct() && impl_->reader.is_open(); + if (a.via_reader) a.ticket = impl_->reader.issue(a.rows, PLE_N_HEADS, a.raw); + a.pending = true; +} + +bool PleTable::ahead_holds(int slot, const uint32_t* rows16) const { + if (slot < 0 || slot >= kAheadSlots) return false; + const Impl::Ahead& a = impl_->aheads[slot]; + return a.pending && std::memcmp(a.rows, rows16, sizeof a.rows) == 0; +} + +bool PleTable::ahead_collect(int slot, float* out2560, std::string& err) { + if (slot < 0 || slot >= kAheadSlots || !impl_->aheads[slot].pending) { + err = "PleTable::ahead_collect: slot " + std::to_string(slot) + " holds no rows"; + return false; + } + Impl::Ahead& a = impl_->aheads[slot]; + a.pending = false; + if (!a.via_reader) { + for (int h = 0; h < PLE_N_HEADS; ++h) read_row(a.rows[h], out2560 + (size_t) h * PLE_HEAD_DIM); + return true; + } + if (!impl_->reader.collect(a.ticket, err)) return false; + for (int h = 0; h < PLE_N_HEADS; ++h) + ple_dequant_row(impl_->type, a.raw + (size_t) h * impl_->row_bytes, out2560 + (size_t) h * PLE_HEAD_DIM); + impl_->bytes_read += (uint64_t) PLE_N_HEADS * impl_->row_bytes; + return true; +} + void PleTable::set_injected_delay_us(double us) { impl_->reader.set_injected_delay_us(us); } std::string PleTable::io_report() const { - if (impl_->mode != PleIo::Direct || !impl_->reader.is_open()) return {}; + if (impl_->from_ram()) return "ple io: " + impl_->ram_note; + if (impl_->mode == PleIo::Mmap || !impl_->reader.is_open()) return {}; const strata::ngram::ReaderStats& s = impl_->reader.stats(); char buf[320]; std::snprintf(buf, sizeof buf, @@ -306,7 +496,11 @@ std::string PleTable::io_report() const { } void PleTable::gather(const uint32_t* rows16, float* out2560) const { - if (impl_->mode == PleIo::Direct) { + if (impl_->from_ram()) { + for (int h = 0; h < PLE_N_HEADS; ++h) read_row(rows16[h], out2560 + (size_t) h * PLE_HEAD_DIM); + return; + } + if (impl_->direct()) { // `gather` stays const for its existing callers; the reader's state is the table's I/O state. PleTable* self = const_cast(this); std::string err; @@ -336,8 +530,8 @@ void PleTable::gather(const uint32_t* rows16, float* out2560) const { for (int h = 0; h < PLE_N_HEADS; ++h) { // Out-of-range rows are handled by `read_row` as zeros and have no address to prefetch. if (rows16[h] >= impl_->n_rows) continue; - ranges[n].VirtualAddress = (PVOID) (impl_->data + (size_t) rows16[h] * PLE_ROW_BYTES); - ranges[n].NumberOfBytes = PLE_ROW_BYTES; + ranges[n].VirtualAddress = (PVOID) (impl_->data + (size_t) rows16[h] * impl_->row_bytes); + ranges[n].NumberOfBytes = impl_->row_bytes; ++n; } if (n > 0) (void) PrefetchVirtualMemory(GetCurrentProcess(), n, ranges, 0); diff --git a/src/kernels/qsa_parity.cpp b/src/kernels/qsa_parity.cpp index 5a8186ddae1..517db948aa9 100644 --- a/src/kernels/qsa_parity.cpp +++ b/src/kernels/qsa_parity.cpp @@ -763,7 +763,7 @@ int main(int argc, char** argv) { const int64_t max_blocks = max_n / R + 1; const int64_t cap = strata::kernels::qsa_selection_width(max_n, S); Dev pooled, query, scores((size_t) max_n); - Dev ids((size_t) cap), step(strata::kernels::kStepCount); + Dev sel((size_t) cap), step(strata::kernels::kStepCount); pooled.put(std::vector((size_t) max_blocks * IDXD, 0.0f)); query.put(std::vector((size_t) IDXN * IDXD, 0.0f)); @@ -773,7 +773,7 @@ int main(int argc, char** argv) { check(cudaStreamCreate(&stream), "tail stream"); check(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal), "tail capture begin"); strata::kernels::qsa_index_step(pooled.p, query.p, nullptr, S, step.p, max_blocks, scores.p, stream); - strata::kernels::topk_512_step(scores.p, S, cap, step.p, ids.p, stream); + strata::kernels::topk_512_step(scores.p, S, cap, step.p, sel.p, stream); check(cudaStreamEndCapture(stream, &graph), "tail capture end"); check(cudaGraphInstantiate(&exec, graph, 0), "tail instantiate"); @@ -789,15 +789,15 @@ int main(int argc, char** argv) { for (int64_t j = n - tail; j < n; ++j) want.push_back((int32_t) j); strata::kernels::qsa_index(pooled.p, n / R, query.p, nullptr, S, n, scores.p, nullptr); - strata::kernels::topk_512(scores.p, n, S, cap, ids.p, nullptr); - require("tail static n=" + std::to_string(n), ids.get((size_t) width) == want); + strata::kernels::topk_512(scores.p, n, S, cap, sel.p, nullptr); + require("tail static n=" + std::to_string(n), sel.get((size_t) width) == want); std::vector values(strata::kernels::kStepCount); strata::kernels::qsa_step_fill(values.data(), n - 1, S); step.put(values); check(cudaGraphLaunch(exec, stream), "tail replay"); check(cudaStreamSynchronize(stream), "tail replay sync"); - require("tail captured n=" + std::to_string(n), ids.get((size_t) width) == want); + require("tail captured n=" + std::to_string(n), sel.get((size_t) width) == want); const auto actual_scores = scores.get((size_t) n); bool exact_scores = true; for (int64_t j = 0; j < n; ++j) { diff --git a/src/kernels/qsa_select_parity.cu b/src/kernels/qsa_select_parity.cu new file mode 100644 index 00000000000..933006af576 --- /dev/null +++ b/src/kernels/qsa_select_parity.cu @@ -0,0 +1,390 @@ +// src/kernels/qsa_select_parity.cu - the QSA block scores and top-k (qsa_select.cu) against the kernels they +// replaced, kept here as the reference: scores bitwise, selections identical, for decode windows (1-8 queries, +// the captured graph's fixed grid) and the prompt path's batches (up to 256 queries), from the identity case to 200K +// cells, with exact score ties, zero scores and every tail length. --bench times both. +#include "strata/kernels/qsa.hpp" +#include "strata/kernels/qsa_select.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ref { + +using namespace strata::kernels; +constexpr int IDX_DIM = 128, IDX_HEADS = 4, R = 4; +constexpr int SCORE_WARPS = 8; +constexpr int TOPK_T = 256; + +__device__ __forceinline__ uint32_t order_key(float s) { + const float v = s + 0.0f; + if (!(v == v)) return 0u; + const uint32_t b = __float_as_uint(v); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} + +__global__ void __launch_bounds__(SCORE_WARPS * 32) block_scores_kernel(const float* __restrict__ pooled, + const float* __restrict__ dead, + const float* __restrict__ q_idx, + const int32_t* __restrict__ steps, + int64_t max_blocks, float* __restrict__ out) { + const int64_t qi = blockIdx.y; + const int32_t* st = steps + qi * kStepCount; + const int64_t n_kv = st[kStepNKv], n_bid = st[kStepNBid]; + const int64_t last = n_bid < max_blocks - 1 ? n_bid : max_blocks - 1; + const int lane = threadIdx.x & 31; + const float* q = q_idx + qi * IDX_HEADS * IDX_DIM + lane * 4; + float4 q4[IDX_HEADS]; +#pragma unroll + for (int h = 0; h < IDX_HEADS; ++h) q4[h] = *reinterpret_cast(q + h * IDX_DIM); + for (int64_t b = (int64_t) blockIdx.x * SCORE_WARPS + (threadIdx.x >> 5); b <= last; + b += (int64_t) gridDim.x * SCORE_WARPS) { + const float* key = (b == n_bid) ? dead : pooled + b * IDX_DIM; + const float4 k4 = *reinterpret_cast(key + lane * 4); + float score = 0.0f; +#pragma unroll + for (int h = 0; h < IDX_HEADS; ++h) { + float d = k4.x * q4[h].x + k4.y * q4[h].y + k4.z * q4[h].z + k4.w * q4[h].w; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) d += __shfl_xor_sync(0xffffffffu, d, o); + score += d > 0.0f ? d : 0.0f; + } + if (lane == 0) { + if (b == n_bid && n_kv % R != 0) score += 1e9f; + out[qi * max_blocks + b] = score; + } + } +} + +__global__ void __launch_bounds__(TOPK_T) block_topk_kernel(const float* __restrict__ scores, + const int32_t* __restrict__ steps, int64_t max_blocks, + int64_t cap, int32_t* __restrict__ ids) { + __shared__ int hist[256]; + __shared__ int s_a[TOPK_T], s_b[TOPK_T]; + __shared__ int s_digit, s_above; + const int64_t qi = blockIdx.x; + const int32_t* st = steps + qi * kStepCount; + const int64_t n_kv = st[kStepNKv], n_bid = st[kStepNBid], width = st[kStepWidth]; + int32_t* out = ids + qi * cap; + const int t = threadIdx.x; + if (n_kv <= width) { + for (int64_t j = t; j < n_kv; j += TOPK_T) out[j] = (int32_t) j; + return; + } + const float* sc = scores + qi * max_blocks; + const int64_t nb = n_bid + 1; + const int64_t per = (nb + TOPK_T - 1) / TOPK_T; + const int64_t b0 = (int64_t) t * per, b1 = (b0 + per < nb) ? b0 + per : nb; + auto weight = [&](int64_t b) -> int { return b < n_bid ? R : (int) (n_kv - n_bid * R); }; + uint32_t prefix = 0; + int above = 0; + for (int shift = 24; shift >= 0; shift -= 8) { + for (int i = t; i < 256; i += TOPK_T) hist[i] = 0; + __syncthreads(); + const uint32_t hi_mask = shift == 24 ? 0u : (0xffffffffu << (shift + 8)); + for (int64_t b = b0; b < b1; ++b) { + const int w = weight(b); + if (w == 0) continue; + const uint32_t k = order_key(sc[b]); + if ((k & hi_mask) == (prefix & hi_mask)) atomicAdd(&hist[(k >> shift) & 255], w); + } + __syncthreads(); + if (t == 0) { + int cum = above, d = 255; + for (; d > 0; --d) { + if (cum + hist[d] >= width) break; + cum += hist[d]; + } + s_digit = d; + s_above = cum; + } + __syncthreads(); + prefix |= (uint32_t) s_digit << shift; + above = s_above; + __syncthreads(); + } + const uint32_t thr = prefix; + const int64_t eq_budget = width - above; + int gt = 0, eq = 0; + for (int64_t b = b0; b < b1; ++b) { + const int w = weight(b); + if (w == 0) continue; + const uint32_t k = order_key(sc[b]); + if (k > thr) gt += w; + else if (k == thr) eq += w; + } + s_a[t] = gt; + s_b[t] = eq; + __syncthreads(); + if (t == 0) { + int ag = 0, ae = 0; + for (int i = 0; i < TOPK_T; ++i) { + const int g = s_a[i], e = s_b[i]; + s_a[i] = ag; s_b[i] = ae; + ag += g; ae += e; + } + } + __syncthreads(); + const int64_t eq_before = s_b[t]; + int64_t my_eq = eq_budget - eq_before; + if (my_eq < 0) my_eq = 0; + if (my_eq > eq) my_eq = eq; + const int sel = gt + (int) my_eq; + __syncthreads(); + s_a[t] = sel; + __syncthreads(); + if (t == 0) { + int a = 0; + for (int i = 0; i < TOPK_T; ++i) { const int c = s_a[i]; s_a[i] = a; a += c; } + } + __syncthreads(); + int64_t wpos = s_a[t]; + int64_t eq_left = my_eq; + for (int64_t b = b0; b < b1; ++b) { + const int w = weight(b); + if (w == 0) continue; + const uint32_t k = order_key(sc[b]); + if (k > thr) { + for (int c = 0; c < w; ++c) out[wpos++] = (int32_t) (b * R + c); + } else if (k == thr) { + for (int c = 0; c < w && eq_left > 0; ++c, --eq_left) out[wpos++] = (int32_t) (b * R + c); + } + } +} + +void scores(const float* pooled, const float* dead, const float* q_idx, const int32_t* steps, int64_t nq, + int64_t max_blocks, float* out, cudaStream_t s, int64_t grid_blocks) { + const int64_t covered = grid_blocks < max_blocks ? grid_blocks : max_blocks; + const dim3 grid((unsigned) ((covered + SCORE_WARPS - 1) / SCORE_WARPS), (unsigned) nq); + block_scores_kernel<<>>(pooled, dead, q_idx, steps, max_blocks, out); +} + +void topk(const float* scores, const int32_t* steps, int64_t nq, int64_t max_blocks, int64_t cap, int32_t* ids, + cudaStream_t s) { + block_topk_kernel<<<(unsigned) nq, TOPK_T, 0, s>>>(scores, steps, max_blocks, cap, ids); +} + +} // namespace ref + +namespace { + +using namespace strata::kernels; + +void check(cudaError_t e, const char* what) { + if (e != cudaSuccess) { + std::fprintf(stderr, "%s: %s\n", what, cudaGetErrorString(e)); + std::exit(1); + } +} + +template T* dev(size_t n) { + T* p = nullptr; + check(cudaMalloc(&p, n * sizeof(T)), "cudaMalloc"); + return p; +} + +constexpr int64_t kMaxCells = 200192; // the configs' --max-context +constexpr int QD = 512; // a query's floats: 4 heads x 128 + +struct Data { + QsaShapes s = qsa_real_shapes(); + int64_t max_blocks = kMaxCells / 4 + 2; + int64_t cap = qsa_selection_width(kTopkMaxCells, qsa_real_shapes()); + float *pooled = nullptr, *dead = nullptr, *q = nullptr, *sc_ref = nullptr, *sc_new = nullptr; + int32_t *steps = nullptr, *ids_ref = nullptr, *ids_new = nullptr; + int64_t max_q = 256; +}; + +// Keys like the indexer's: most N(0,1); runs of copies (exact score ties, so the threshold often falls inside a +// run of equal keys and blocks are taken in part), all-zero keys (scores exactly 0) and -0.0 components. +void fill(Data& d, std::mt19937& rng) { + std::normal_distribution n01(0.0f, 1.0f); + std::vector pooled((size_t) d.max_blocks * 128); + for (int64_t b = 0; b < d.max_blocks; ++b) { + float* row = pooled.data() + (size_t) b * 128; + const int kind = (int) (rng() % 64); + if (b > 0 && kind < 6) std::memcpy(row, row - 128, 128 * sizeof(float)); // a copy of the previous block + else if (kind == 6) std::fill(row, row + 128, 0.0f); + else + for (int i = 0; i < 128; ++i) row[i] = (rng() % 97 == 0) ? -0.0f : n01(rng); + } + std::vector dead(128); + for (auto& x : dead) x = n01(rng); + std::vector q((size_t) d.max_q * QD); + for (auto& x : q) x = 0.3f * n01(rng); + d.pooled = dev(pooled.size()); + d.dead = dev(128); + d.q = dev(q.size()); + check(cudaMemcpy(d.pooled, pooled.data(), pooled.size() * 4, cudaMemcpyHostToDevice), "up pooled"); + check(cudaMemcpy(d.dead, dead.data(), 128 * 4, cudaMemcpyHostToDevice), "up dead"); + check(cudaMemcpy(d.q, q.data(), q.size() * 4, cudaMemcpyHostToDevice), "up q"); + d.sc_ref = dev((size_t) d.max_q * d.max_blocks); + d.sc_new = dev((size_t) d.max_q * d.max_blocks); + d.steps = dev((size_t) d.max_q * kStepCount); + d.ids_ref = dev((size_t) d.max_q * d.cap); + d.ids_new = dev((size_t) d.max_q * d.cap); +} + +std::vector steps_from(const Data& d, int64_t pos0, int64_t nq) { + std::vector st((size_t) nq * kStepCount); + for (int64_t i = 0; i < nq; ++i) qsa_step_fill(st.data() + i * kStepCount, pos0 + i, d.s); + check(cudaMemcpy(d.steps, st.data(), st.size() * 4, cudaMemcpyHostToDevice), "up steps"); + return st; +} + +// both implementations on queries at positions pos0.., then scores bitwise and selections equal +int compare(Data& d, int64_t pos0, int64_t nq, int64_t grid_blocks, cudaStream_t s, const char* what) { + const std::vector st = steps_from(d, pos0, nq); + check(cudaMemset(d.sc_ref, 0xff, (size_t) nq * d.max_blocks * 4), "clear"); + check(cudaMemset(d.sc_new, 0xff, (size_t) nq * d.max_blocks * 4), "clear"); + check(cudaMemset(d.ids_ref, 0xff, (size_t) nq * d.cap * 4), "clear"); + check(cudaMemset(d.ids_new, 0xff, (size_t) nq * d.cap * 4), "clear"); + ref::scores(d.pooled, d.dead, d.q, d.steps, nq, d.max_blocks, d.sc_ref, s, grid_blocks); + ref::topk(d.sc_ref, d.steps, nq, d.max_blocks, d.cap, d.ids_ref, s); + qsa_block_scores(d.pooled, d.dead, d.q, d.steps, nq, d.max_blocks, d.s, d.sc_new, s, grid_blocks); + qsa_block_topk(d.sc_new, d.steps, nq, d.max_blocks, d.cap, d.s, d.ids_new, s); + check(cudaStreamSynchronize(s), what); + std::vector a((size_t) nq * d.max_blocks), b(a.size()); + std::vector ia((size_t) nq * d.cap), ib(ia.size()); + check(cudaMemcpy(a.data(), d.sc_ref, a.size() * 4, cudaMemcpyDeviceToHost), "down"); + check(cudaMemcpy(b.data(), d.sc_new, b.size() * 4, cudaMemcpyDeviceToHost), "down"); + check(cudaMemcpy(ia.data(), d.ids_ref, ia.size() * 4, cudaMemcpyDeviceToHost), "down"); + check(cudaMemcpy(ib.data(), d.ids_new, ib.size() * 4, cudaMemcpyDeviceToHost), "down"); + int bad = 0; + for (int64_t i = 0; i < nq; ++i) { + const int64_t n_bid = st[(size_t) (i * kStepCount + kStepNBid)]; + const int64_t width = st[(size_t) (i * kStepCount + kStepWidth)]; + const int64_t last = std::min(n_bid, d.max_blocks - 1); + for (int64_t k = 0; k <= last; ++k) { + const size_t o = (size_t) (i * d.max_blocks + k); + if (a[o] != b[o]) { + if (bad < 5) { + float fa, fb; + std::memcpy(&fa, &a[o], 4); + std::memcpy(&fb, &b[o], 4); + std::fprintf(stderr, "%s: pos %lld, block %lld: score %.9g (0x%08x) vs %.9g (0x%08x)\n", what, + (long long) (pos0 + i), (long long) k, fa, a[o], fb, b[o]); + } + ++bad; + } + } + for (int64_t k = 0; k < width; ++k) { + const size_t o = (size_t) (i * d.cap + k); + if (ia[o] != ib[o]) { + if (bad < 5) + std::fprintf(stderr, "%s: pos %lld, selection entry %lld: cell %d vs %d\n", what, + (long long) (pos0 + i), (long long) k, ia[o], ib[o]); + ++bad; + } + } + } + return bad; +} + +int selftest(Data& d, cudaStream_t s) { + int bad = 0, cases = 0; + // decode windows: the captured graph's fixed grid, 1-8 queries at consecutive positions + const int64_t pos_dec[] = {0, 3, 2046, 2047, 2050, 2051, 2052, 8190, 16381, 32766, 65533, 131070, 200180}; + for (int64_t p : pos_dec) + for (int64_t nq : {1, 2, 3, 4, 8}) { + if (p + nq > kMaxCells) continue; + bad += compare(d, p, nq, qsa_score_grid_blocks, s, "decode"); + ++cases; + } + // prompt batches: 256 queries and a sub-chunk's last, shorter batch (11 go to the decode windows' kernel), the grid + // covering the batch's last block + const int64_t pos_pre[][2] = {{0, 256}, {1900, 256}, {16384, 256}, {65536 - 100, 256}, + {131072 - 256, 256}, {199936 - 1, 256}, {5000, 12}, {70001, 37}, + {131072 - 100, 100}, {199936 - 255, 255}, {9000, 11}}; + for (const auto& pq : pos_pre) { + const int64_t p = pq[0], nq = pq[1], last = p + nq - 1; + bad += compare(d, p, nq, (last + 1) / 4 + 1, s, "prompt"); + ++cases; + } + std::printf("qsa_select: %s (%d cases: scores bitwise, selections equal)\n", bad ? "MISMATCH" : "PASS", cases); + return bad; +} + +// GPU time per call: `reps` calls captured in a graph, as the engine runs them (a launch outside a graph costs more +// than these kernels on Windows) +template float time_ms(cudaStream_t s, int reps, const F& fn) { + cudaGraph_t g = nullptr; + cudaGraphExec_t ge = nullptr; + check(cudaStreamBeginCapture(s, cudaStreamCaptureModeThreadLocal), "capture"); + for (int i = 0; i < reps; ++i) fn(); + check(cudaStreamEndCapture(s, &g), "capture"); + check(cudaGraphInstantiate(&ge, g, 0), "instantiate"); + cudaEvent_t a, b; + cudaEventCreate(&a); + cudaEventCreate(&b); + check(cudaGraphLaunch(ge, s), "launch"); + cudaEventRecord(a, s); + for (int i = 0; i < 3; ++i) check(cudaGraphLaunch(ge, s), "launch"); + cudaEventRecord(b, s); + check(cudaEventSynchronize(b), "bench"); + float ms = 0; + cudaEventElapsedTime(&ms, a, b); + cudaEventDestroy(a); + cudaEventDestroy(b); + cudaGraphExecDestroy(ge); + cudaGraphDestroy(g); + return ms / (float) (3 * reps); +} + +void bench(Data& d, cudaStream_t s) { + std::printf("%-26s %12s %12s %12s %12s\n", "case", "ref scores", "new scores", "ref top-k", "new top-k"); + struct Case { const char* name; int64_t pos; int64_t nq; bool prompt; }; + const Case cs[] = {{"decode 4 q, 1K cells", 1024, 4, false}, {"decode 4 q, 2K cells", 2048, 4, false}, + {"decode 4 q, 16K cells", 16384, 4, false}, + {"decode 4 q, 64K cells", 65536, 4, false}, {"decode 4 q, 131K cells", 131072, 4, false}, + {"decode 1 q, 131K cells", 131072, 1, false}, {"prompt 256 q at 16K", 16384 - 256, 256, true}, + {"prompt 256 q at 64K", 65536 - 256, 256, true}, + {"prompt 256 q at 131K", 131072 - 256, 256, true}, + {"prompt 16 q at 131K", 131072 - 16, 16, true}, {"prompt 64 q at 131K", 131072 - 64, 64, true}}; + for (const Case& c : cs) { + steps_from(d, c.pos, c.nq); + const int64_t grid = c.prompt ? (c.pos + c.nq) / 4 + 1 : qsa_score_grid_blocks; + const int reps = c.prompt ? 5 : 50; + const float rs = time_ms(s, reps, [&] { + ref::scores(d.pooled, d.dead, d.q, d.steps, c.nq, d.max_blocks, d.sc_ref, s, grid); + }); + const float ns = time_ms(s, reps, [&] { + qsa_block_scores(d.pooled, d.dead, d.q, d.steps, c.nq, d.max_blocks, d.s, d.sc_new, s, grid); + }); + const float rt = time_ms(s, reps, [&] { + ref::topk(d.sc_ref, d.steps, c.nq, d.max_blocks, d.cap, d.ids_ref, s); + }); + const float nt = time_ms(s, reps, [&] { + qsa_block_topk(d.sc_new, d.steps, c.nq, d.max_blocks, d.cap, d.s, d.ids_new, s); + }); + std::printf("%-26s %9.1f us %9.1f us %9.1f us %9.1f us\n", c.name, rs * 1e3f, ns * 1e3f, rt * 1e3f, nt * 1e3f); + } +} + +} // namespace + +int main(int argc, char** argv) { + bool do_bench = false; + for (int i = 1; i < argc; ++i) { + if (std::string(argv[i]) == "--bench") do_bench = true; + else { + std::fprintf(stderr, "usage: qsa_select_parity [--bench]\n"); + return 2; + } + } + std::mt19937 rng(20260926); + cudaStream_t s = nullptr; + check(cudaStreamCreate(&s), "stream"); + Data d; + fill(d, rng); + const int bad = selftest(d, s); + if (do_bench) bench(d, s); + return bad ? 1 : 0; +} diff --git a/src/kernels/quantize_act_parity.cpp b/src/kernels/quantize_act_parity.cpp index b2cd7ee4149..e9c794e22c4 100644 --- a/src/kernels/quantize_act_parity.cpp +++ b/src/kernels/quantize_act_parity.cpp @@ -318,10 +318,10 @@ int main(int argc, char** argv) { {0.0f, 0x0000u, "0.0"}, }; int fp16_bad = 0; - for (const F16Case& c : cases) { - const uint16_t got = strata::kernels::f16_from_f32(c.in); - if (got != c.want) { - std::printf(" *** %s: want 0x%04X got 0x%04X\n", c.what, c.want, got); + for (const F16Case& fc : cases) { + const uint16_t got = strata::kernels::f16_from_f32(fc.in); + if (got != fc.want) { + std::printf(" *** %s: want 0x%04X got 0x%04X\n", fc.what, fc.want, got); ++fp16_bad; } } @@ -388,16 +388,16 @@ int main(int argc, char** argv) { std::vector b127, b128; reference_q8_K(ka.data(), NK, r127, b127, -127.0f); reference_q8_K(ka.data(), NK, r128, b128, -128.0f); - double d = 0, m = 0; + double diff = 0, m = 0; long long qdiffer = 0; - for (long long i = 0; i < NK; ++i) { d += std::fabs(b127[(size_t) i] - b128[(size_t) i]); m += std::fabs(b127[(size_t) i]); } + for (long long i = 0; i < NK; ++i) { diff += std::fabs(b127[(size_t) i] - b128[(size_t) i]); m += std::fabs(b127[(size_t) i]); } for (size_t i = 4; i < r127.size(); ++i) if (r127[i] != r128[i]) ++qdiffer; // The expected size of the effect is exactly (1 - 127/128) = 0.78% of the magnitude, so the floor is // 0.5% - below what the rival constant must produce and far above any noise. The byte count is the // blunter evidence and is reported alongside it: 172,017 bytes differ, i.e. a third of the quants. std::printf(" %-26s %-4s (%.3f%% of magnitude, %lld bytes differ)\n", "-127 vs -128 scale observable", - d / m > 0.005 ? "yes" : "*** NO ***", 100.0 * d / m, qdiffer); - if (!(d / m > 0.005)) ++bad_k; + diff / m > 0.005 ? "yes" : "*** NO ***", 100.0 * diff / m, qdiffer); + if (!(diff / m > 0.005)) ++bad_k; } std::printf("\nquantize_act: %d failures over 5 Q8_0 + 4 Q8_K distributions\n", bad + bad_k); diff --git a/src/kernels/router_top10_parity.cpp b/src/kernels/router_top10_parity.cpp index ac7e652eb53..bcb34f2236b 100644 --- a/src/kernels/router_top10_parity.cpp +++ b/src/kernels/router_top10_parity.cpp @@ -150,11 +150,11 @@ int main(int argc, char** argv) { // distribution). Here k/n = 10/512 = 0.01953125, which is 320x the 2**-14 = 6.1035e-05 clamp, so no input // whatsoever can make this model's renormalisation clamp - not an extreme one, not an adversarial one. // Asserting the bound is worth more than a test that cannot reach it. - const double min_possible = (double) K / (double) NE; + constexpr double min_possible = (double) K / (double) NE; std::printf(" %-26s top-%d sum >= k/n = %.5f, clamp is %.3e -> margin %.0fx %s\n", "clamp reachability", K, min_possible, RENORM_CLAMP, min_possible / RENORM_CLAMP, min_possible > RENORM_CLAMP ? "(CLAMP IS UNREACHABLE, asserted)" : "*** REACHABLE ***"); - if (min_possible <= RENORM_CLAMP) { + if constexpr (min_possible <= RENORM_CLAMP) { std::fprintf(stderr, "the clamp CAN trigger, so it needs a test\n"); ++bad; } diff --git a/src/kernels/shared_expert_parity.cpp b/src/kernels/shared_expert_parity.cpp index 92de10ad00b..e469593a0aa 100644 --- a/src/kernels/shared_expert_parity.cpp +++ b/src/kernels/shared_expert_parity.cpp @@ -436,7 +436,7 @@ int main(int argc, char** argv) { } return y; }; - const std::vector want = ref(true, false); + const std::vector want_mc = ref(true, false); float *d_p = nullptr, *d_w = nullptr, *d_s = nullptr, *d_y = nullptr; check(cudaMalloc(&d_p, parts.size() * 4), "mc p"); @@ -447,12 +447,12 @@ int main(int argc, char** argv) { check(cudaMemcpy(d_w, w.data(), w.size() * 4, cudaMemcpyHostToDevice), "mc cw"); check(cudaMemcpy(d_s, shared.data(), shared.size() * 4, cudaMemcpyHostToDevice), "mc cs"); strata::kernels::moe_combine(d_p, d_w, d_s, d_y, n_embd2, k2, nullptr); - std::vector got((size_t) n_embd2); - check(cudaMemcpy(got.data(), d_y, got.size() * 4, cudaMemcpyDeviceToHost), "mc cy"); + std::vector got_mc((size_t) n_embd2); + check(cudaMemcpy(got_mc.data(), d_y, got_mc.size() * 4, cudaMemcpyDeviceToHost), "mc cy"); for (const auto& trap : {std::make_pair(false, false), std::make_pair(true, true)}) { const std::vector wrong = ref(trap.first, trap.second); - const double rel = rel_l1(want, wrong); + const double rel = rel_l1(want_mc, wrong); const char* what = (!trap.first && !trap.second) ? "routed sum left UNWEIGHTED is observable" : "shared output router-WEIGHTED is observable"; const bool visible = rel > 0.05; @@ -460,15 +460,15 @@ int main(int argc, char** argv) { if (!visible) ++bad; } - const double rel_k = rel_l1(want, got); + const double rel_k = rel_l1(want_mc, got_mc); std::printf(" %-42s rel %.3e\n", "moe_combine vs reference", rel_k); // Both sides accumulate in double over the same k terms, so this is order-only. if (!(rel_k <= 1e-6)) { std::printf(" *** over 1e-6 ***\n"); ++bad; } // a null `shared` must mean "no shared expert", not "add nothing but leave it undefined" strata::kernels::moe_combine(d_p, d_w, nullptr, d_y, n_embd2, k2, nullptr); - check(cudaMemcpy(got.data(), d_y, got.size() * 4, cudaMemcpyDeviceToHost), "mc cy2"); - const double rel_no_shared = rel_l1(want, got); + check(cudaMemcpy(got_mc.data(), d_y, got_mc.size() * 4, cudaMemcpyDeviceToHost), "mc cy2"); + const double rel_no_shared = rel_l1(want_mc, got_mc); const bool differs = rel_no_shared > 0.05; std::printf(" %-42s %s (%.2f%% apart - null shared really omits it)\n", "a null `shared` is observable", differs ? "yes" : "*** NO ***", rel_no_shared * 100); diff --git a/src/kernels/verify_parity.cpp b/src/kernels/verify_parity.cpp new file mode 100644 index 00000000000..d233fad2c2e --- /dev/null +++ b/src/kernels/verify_parity.cpp @@ -0,0 +1,1328 @@ +// src/kernels/verify_parity.cpp - batched kernels against the per-token kernel sequences they replace. +// +// A verify window's token t must come out bit for bit as it would in a window of any other size (the drafts are +// accepted exactly when greedy decode would have produced them), so every batched kernel is checked BITWISE against +// the per-token kernels it stands in for, on random inputs that include -0.0, denormals and large values. The +// prompt path's batched PLE and GDN arithmetic is held to the same standard; its attention kernel, which sums in +// another order, is checked against the decode kernel within a tolerance. +#include "strata/kernels/bf16_gemv.hpp" +#include "strata/kernels/cpu/native_expert.hpp" +#include "strata/kernels/elementwise.hpp" +#include "strata/kernels/fused_gr.hpp" +#include "strata/kernels/iq_kernels.hpp" +#include "strata/kernels/kv_q8.hpp" +#include "strata/kernels/native_moe.hpp" +#include "strata/kernels/native_ple_postops.hpp" +#include "strata/kernels/native_qsa.hpp" +#include "strata/kernels/native_qsa_indexer.hpp" +#include "strata/kernels/native_rope.hpp" +#include "strata/kernels/qsa_decode_attn.hpp" +#include "strata/kernels/sampler.hpp" +#include "strata/kernels/native_router.hpp" +#include "strata/kernels/s2_expert_grouped.hpp" +#include "strata/kernels/shared_expert.hpp" +#include "strata/kernels/verify_kernels.hpp" +#include "strata/prefill/kernels.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +void check(cudaError_t e, const char* what) { + if (e != cudaSuccess) { + std::fprintf(stderr, "%s: %s\n", what, cudaGetErrorString(e)); + std::exit(1); + } +} + +template T* dev(size_t n) { + T* p = nullptr; + check(cudaMalloc(&p, n * sizeof(T)), "cudaMalloc"); + return p; +} +template void up(T* d, const std::vector& h) { + check(cudaMemcpy(d, h.data(), h.size() * sizeof(T), cudaMemcpyHostToDevice), "upload"); +} +template std::vector down(const T* d, size_t n) { + std::vector h(n); + check(cudaMemcpy(h.data(), d, n * sizeof(T), cudaMemcpyDeviceToHost), "download"); + return h; +} + +// values that exercise the rounding edges: ordinary, large, tiny, denormal, -0.0 +float edgy(std::mt19937& rng) { + std::uniform_real_distribution u(-1.0f, 1.0f); + const int kind = (int) (rng() % 16); + if (kind == 0) return -0.0f; + if (kind == 1) return 0.0f; + if (kind == 2) return u(rng) * 1e-39f; // denormal + if (kind == 3) return u(rng) * 1e4f; + return u(rng); +} + +int bitwise_diff(const std::vector& a, const std::vector& b, const char* what) { + int bad = 0; + for (size_t i = 0; i < a.size(); ++i) { + uint32_t x, y; + std::memcpy(&x, &a[i], 4); + std::memcpy(&y, &b[i], 4); + if (x != y && !(std::isnan(a[i]) && std::isnan(b[i]))) { + if (bad < 5) std::fprintf(stderr, " %s: element %zu: %.9g (0x%08x) vs %.9g (0x%08x)\n", what, i, a[i], x, b[i], y); + ++bad; + } + } + return bad; +} + +// ---- the MoE combine: copy all rows + add the hits + one combine per token vs the second GPU's rows taken into VRAM +// (fetch_listed_rows) and native_moe_gather_combine; from rep 2 on the shared rows gated by shared_expert_gate_rows vs +// the gather-combine gating them +int test_gather_combine(std::mt19937& rng, cudaStream_t s) { + const int64_t N = 2560, K = 10; + int bad = 0; + for (int n_tok : {1, 2, 3, 4, 8}) { + for (int rep = 0; rep < 8; ++rep) { + const int64_t rows = n_tok * K; + std::vector hit((size_t) (rows * N)), host((size_t) (rows * N)), w((size_t) rows), + shared((size_t) (n_tok * N)); + for (auto& v : hit) v = edgy(rng); + for (auto& v : host) v = edgy(rng); + for (auto& v : w) v = std::uniform_real_distribution(0.0f, 1.0f)(rng); + for (auto& v : shared) v = edgy(rng); + std::vector glog((size_t) n_tok); // the gates' logits: sigmoid near 1, near 0 (denormal products) + for (auto& v : glog) + v = rng() % 4 == 0 ? std::uniform_real_distribution(-100.0f, -80.0f)(rng) + : std::uniform_real_distribution(-8.0f, 8.0f)(rng); + const bool gated = rep >= 2; + // the GPU's entries, in a random order and a random share (none, some, all); from rep 4 on the second GPU + // takes some of the rest + std::vector dst, raw; + for (int32_t r = 0; r < rows; ++r) { + if (rep == 1 || (rep != 0 && rng() % 3 != 0)) dst.push_back(r); + else if (rep >= 4 && rng() % 2 == 0) raw.push_back(r); + } + std::shuffle(dst.begin(), dst.end(), rng); + std::shuffle(raw.begin(), raw.end(), rng); + const int32_t count = (int32_t) dst.size(); + // the old path's host rows: the pool zeroed the GPU's, the host copied the second GPU's in + std::vector host_zeroed = host; + for (int32_t r : dst) std::fill(host_zeroed.begin() + r * N, host_zeroed.begin() + (r + 1) * N, 0.0f); + for (int32_t r : raw) + std::copy(hit.begin() + r * N, hit.begin() + (r + 1) * N, host_zeroed.begin() + r * N); + // the new path's mapped rows: the GPU's rows hold garbage it must not read, the second GPU wrote its own; + // its device rows hold garbage until they are fetched + std::vector host_poison = host, hit_poison = hit; + for (int32_t r : dst) std::fill(host_poison.begin() + r * N, host_poison.begin() + (r + 1) * N, NAN); + for (int32_t r : raw) { + std::copy(hit.begin() + r * N, hit.begin() + (r + 1) * N, host_poison.begin() + r * N); + std::fill(hit_poison.begin() + r * N, hit_poison.begin() + (r + 1) * N, NAN); + } + + float *d_hit = dev(hit.size()), *d_parts = dev(hit.size()), *d_w = dev(w.size()), + *d_sh = dev(shared.size()), *d_out_old = dev(shared.size()), + *d_out_new = dev(shared.size()), *d_shg = dev(shared.size()), + *d_g = dev(glog.size()), *d_glog = dev(glog.size()); + int32_t *d_dst = dev(std::max(1, dst.size()) + 64), *d_count = dev(3), + *d_list = dev((size_t) (4 + rows)); + up(d_hit, hit); + up(d_w, w); + up(d_sh, shared); + up(d_shg, shared); + up(d_g, glog); + up(d_glog, glog); + if (!dst.empty()) up(d_dst, dst); + // the new path takes them as two lists (the window's VRAM share and the pool's PCIe share), split at random + const int32_t split = rep % 2 == 0 ? count : (int32_t) (rng() % (dst.size() + 1)); + up(d_count, std::vector{count, split, count - split}); + check(cudaMemset(d_list, 0xff, (size_t) (4 + rows) * 4), "memset"); + float *h_map = nullptr, *m_map = nullptr; + int32_t *h_list = nullptr, *m_list = nullptr; + check(cudaHostAlloc((void**) &h_map, host.size() * sizeof(float), cudaHostAllocMapped), "cudaHostAlloc"); + check(cudaHostGetDevicePointer((void**) &m_map, h_map, 0), "cudaHostGetDevicePointer"); + check(cudaHostAlloc((void**) &h_list, (size_t) (4 + rows) * 4, cudaHostAllocMapped), "cudaHostAlloc"); + check(cudaHostGetDevicePointer((void**) &m_list, h_list, 0), "cudaHostGetDevicePointer"); + h_list[0] = (int32_t) raw.size(); + std::copy(raw.begin(), raw.end(), h_list + 4); + + // old: copy all rows, add the hits, combine per token + std::memcpy(h_map, host_zeroed.data(), host.size() * sizeof(float)); + strata::kernels::copy_from_mapped(d_parts, m_map, rows * N, s); + strata::kernels::moe_hit_add(d_parts, d_hit, d_dst, d_count, rows, N, s); + if (gated) strata::kernels::shared_expert_gate_rows(d_shg, d_g, N, n_tok, s); + for (int t = 0; t < n_tok; ++t) + strata::kernels::native_moe_combine(d_parts + t * K * N, d_w + t * K, (gated ? d_shg : d_sh) + t * N, + d_out_old + t * N, N, K, s); + check(cudaStreamSynchronize(s), "old path"); + // new + std::memcpy(h_map, host_poison.data(), host.size() * sizeof(float)); + up(d_hit, hit_poison); + const bool second = rep >= 4; + if (second) strata::kernels::fetch_listed_rows(m_list, m_map, d_hit, d_list, (int) rows, N, s); + strata::kernels::native_moe_gather_combine(d_hit, m_map, d_dst, d_count + 1, + rep % 2 == 0 ? nullptr : d_dst + split, d_count + 2, + second ? d_list + 4 : nullptr, second ? d_list : nullptr, d_w, + d_sh, gated ? d_glog : nullptr, d_out_new, N, K, n_tok, s); + check(cudaStreamSynchronize(s), "new path"); + int b = bitwise_diff(down(d_out_old, shared.size()), down(d_out_new, shared.size()), "gather_combine"); + if (second) { + std::vector want((size_t) (4 + rows), -1), got = down(d_list, (size_t) (4 + rows)); + want[0] = (int32_t) raw.size(); + std::copy(raw.begin(), raw.end(), want.begin() + 4); + for (int i = 1; i < 4; ++i) got[(size_t) i] = -1; + if (got != want) { + std::fprintf(stderr, "fetch_listed_rows: n_tok %d rep %d: the device list differs\n", n_tok, rep); + ++b; + } + } + if (b) std::fprintf(stderr, "gather_combine: n_tok %d rep %d: %d of %zu differ\n", n_tok, rep, b, shared.size()); + bad += b; + cudaFreeHost(h_map); + cudaFreeHost(h_list); + for (void* p : {(void*) d_hit, (void*) d_parts, (void*) d_w, (void*) d_sh, (void*) d_out_old, + (void*) d_out_new, (void*) d_shg, (void*) d_g, (void*) d_glog, (void*) d_dst, + (void*) d_count, (void*) d_list}) + cudaFree(p); + } + } + std::printf("gather_combine: %s\n", + bad ? "MISMATCH" + : "bitwise equal (n_tok 1-4, 8; 8 plans each, one or two lists, the second GPU's rows fetched, " + "the shared rows gated)"); + return bad; +} + +// ---- the main GPU's hit plan decided on the device vs the host's rule: distinct experts in routing order, the +// resident ones as groups of their entries in routing order +int test_hit_plan(std::mt19937& rng, cudaStream_t s) { + const int NE = 512, K = 10, slots = 300; + const int64_t cap = 8 * K, ptr_off = ((4 + (cap + 1) + 2 * cap) + 1) & ~1ll, words = ptr_off + 2 * cap; + std::vector slot_ptr(slots); + for (int i = 0; i < slots; ++i) slot_ptr[(size_t) i] = 0x700000000ull + (unsigned long long) i * 1382400ull; + unsigned long long* d_sp = dev(slots); + up(d_sp, slot_ptr); + int32_t *d_ids = dev(cap), *d_res = dev(NE), *d_plan = dev(words); + int bad = 0; + for (int n_tok : {1, 2, 3, 4, 8}) { + for (int rep = 0; rep < 24; ++rep) { + const int n = n_tok * K; + // each token's k experts distinct (rep 3: repeats within a token too), drawn from a small set so tokens share + std::vector set((size_t) (4 + rng() % 40)), ids((size_t) n), res((size_t) NE, -1); + for (auto& e : set) e = (int32_t) (rng() % NE); + for (int t = 0; t < n_tok; ++t) + for (int j = 0; j < K; ++j) { + int32_t e; + bool again; + do { + e = rng() % 4 == 0 ? (int32_t) (rng() % NE) : set[rng() % set.size()]; + again = false; + for (int m = 0; m < j && rep != 3; ++m) again = again || ids[(size_t) (t * K + m)] == e; + } while (again); + ids[(size_t) (t * K + j)] = e; + } + for (int e = 0; e < NE; ++e) // resident: none, all, or at random + if (rep == 1 || (rep != 0 && rng() % 3 != 0)) res[(size_t) e] = (int32_t) (rng() % slots); + std::vector first((size_t) n), start, dst, tok; + std::vector ptr; + for (int i = 0; i < n; ++i) { + first[(size_t) i] = i; + for (int j = 0; j < i; ++j) + if (ids[(size_t) j] == ids[(size_t) i]) { first[(size_t) i] = j; break; } + } + for (int i = 0; i < n; ++i) { + if (first[(size_t) i] != i || res[(size_t) ids[(size_t) i]] < 0) continue; + ptr.push_back(slot_ptr[(size_t) res[(size_t) ids[(size_t) i]]]); + start.push_back((int32_t) dst.size()); + for (int j = i; j < n; ++j) + if (first[(size_t) j] == i) { dst.push_back(j); tok.push_back(j / K); } + } + start.push_back((int32_t) dst.size()); + up(d_ids, ids); + up(d_res, res); + check(cudaMemset(d_plan, 0xff, (size_t) words * 4), "memset"); + strata::kernels::verify_hit_plan(d_ids, n, K, d_res, NE, d_sp, d_plan, cap, ptr_off, s); + check(cudaStreamSynchronize(s), "hit plan"); + const std::vector pl = down(d_plan, (size_t) words); + const int groups = (int) ptr.size(), entries = (int) dst.size(); + bool ok = pl[0] == groups && pl[1] == entries; + for (int q = 0; ok && q <= groups; ++q) ok = pl[(size_t) (4 + q)] == start[(size_t) q]; + for (int q = 0; ok && q < groups; ++q) { + unsigned long long p; + std::memcpy(&p, &pl[(size_t) (ptr_off + 2 * q)], 8); + ok = p == ptr[(size_t) q]; + } + for (int q = 0; ok && q < entries; ++q) + ok = pl[(size_t) (4 + cap + 1 + q)] == dst[(size_t) q] && pl[(size_t) (4 + 2 * cap + 1 + q)] == tok[(size_t) q]; + if (!ok) { + std::fprintf(stderr, "hit_plan: n_tok %d rep %d: groups %d/%d entries %d/%d\n", n_tok, rep, pl[0], groups, + pl[1], entries); + ++bad; + } + } + } + for (void* p : {(void*) d_sp, (void*) d_ids, (void*) d_res, (void*) d_plan}) cudaFree(p); + std::printf("hit_plan: %s\n", bad ? "MISMATCH" : "the host's groups (n_tok 1-4, 8; 24 routings each)"); + return bad; +} + +// ---- BF16 x FP32 MMVF: one call per row vs bf16_gemv_fp32_mmvf_multi (the router, indexer and gate shapes), and +// with a second matrix of 125 rows in the same launch (the indexer's keys and queries; partial groups of 4 rows) +int test_mmvf_multi(std::mt19937& rng, cudaStream_t s) { + auto bf16_matrix = [&](int64_t n) { + std::vector w((size_t) n); + for (auto& v : w) { + const float f = edgy(rng); + uint32_t b; + std::memcpy(&b, &f, 4); + v = (uint16_t) (b >> 16); + } + uint16_t* d = dev(w.size()); + up(d, w); + return d; + }; + int bad = 0; + const int64_t n_in = 2560, n_out2 = 125; + uint16_t* d_w2 = bf16_matrix(n_in * n_out2); + for (int64_t n_out : {512, 130, 1}) { + uint16_t* d_w = bf16_matrix(n_in * n_out); + for (int n_tok = 1; n_tok <= 8; ++n_tok) { + std::vector x((size_t) (n_tok * n_in)); + for (auto& v : x) v = edgy(rng); + const size_t ny = (size_t) (n_tok * n_out), ny2 = (size_t) (n_tok * n_out2); + float *d_x = dev(x.size()), *d_a = dev(ny), *d_b = dev(ny), *d_c = dev(ny), + *d_a2 = dev(ny2), *d_c2 = dev(ny2); + up(d_x, x); + for (int t = 0; t < n_tok; ++t) { + strata::kernels::bf16_gemv_fp32_mmvf(d_x + t * n_in, d_w, d_a + t * n_out, n_in, n_out, s); + strata::kernels::bf16_gemv_fp32_mmvf(d_x + t * n_in, d_w2, d_a2 + t * n_out2, n_in, n_out2, s); + } + strata::kernels::bf16_gemv_fp32_mmvf_multi(d_x, d_w, d_b, n_in, n_out, n_tok, s); + strata::kernels::bf16_gemv_fp32_mmvf_multi(d_x, d_w, d_c, n_in, n_out, n_tok, s, d_w2, d_c2, n_out2); + check(cudaStreamSynchronize(s), "mmvf"); + const int b = bitwise_diff(down(d_a, ny), down(d_b, ny), "mmvf_multi") + + bitwise_diff(down(d_a, ny), down(d_c, ny), "mmvf_multi, two matrices") + + bitwise_diff(down(d_a2, ny2), down(d_c2, ny2), "mmvf_multi, the second matrix"); + if (b) std::fprintf(stderr, "mmvf_multi: n_out %lld n_tok %d: %d differ\n", (long long) n_out, n_tok, b); + bad += b; + for (float* p : {d_x, d_a, d_b, d_c, d_a2, d_c2}) cudaFree(p); + } + cudaFree(d_w); + } + cudaFree(d_w2); + std::printf("mmvf_multi: %s\n", + bad ? "MISMATCH" : "bitwise equal (n_out 512, 130, 1; 1-8 rows; with a second matrix)"); + return bad; +} + +// ---- the top-10 router: one launch per token vs native_router_top10_multi +int test_router_multi(std::mt19937& rng, cudaStream_t s) { + int bad = 0; + for (int n_tok = 1; n_tok <= 8; ++n_tok) { + std::vector logits((size_t) n_tok * 512); + std::normal_distribution nd(0.0f, 2.0f); + for (auto& v : logits) v = nd(rng); + for (int t = 0; t < n_tok; ++t) // exact ties, which the lower expert index must win + for (int i = 0; i < 6; ++i) logits[(size_t) t * 512 + rng() % 512] = logits[(size_t) t * 512 + 7]; + float *d_l = dev(logits.size()), *d_wa = dev((size_t) n_tok * 10), + *d_wb = dev((size_t) n_tok * 10); + int32_t *d_ia = dev((size_t) n_tok * 10), *d_ib = dev((size_t) n_tok * 10); + up(d_l, logits); + for (int t = 0; t < n_tok; ++t) + strata::kernels::native_router_top10(d_l + t * 512, d_ia + t * 10, d_wa + t * 10, s); + strata::kernels::native_router_top10_multi(d_l, d_ib, d_wb, n_tok, s); + check(cudaStreamSynchronize(s), "router"); + int b = bitwise_diff(down(d_wa, (size_t) n_tok * 10), down(d_wb, (size_t) n_tok * 10), "router weights"); + if (down(d_ia, (size_t) n_tok * 10) != down(d_ib, (size_t) n_tok * 10)) { + std::fprintf(stderr, "router_multi: n_tok %d: the ids differ\n", n_tok); + ++b; + } + bad += b; + for (void* p : {(void*) d_l, (void*) d_wa, (void*) d_wb, (void*) d_ia, (void*) d_ib}) cudaFree(p); + } + std::printf("router_multi: %s\n", bad ? "MISMATCH" : "bitwise equal (1-8 tokens, with ties)"); + return bad; +} + +// ---- a layer's router: the gemv, the top 10, doorbell_publish and verify_hit_plan vs verify_router +int test_verify_router(std::mt19937& rng, cudaStream_t s) { + const int N = 2560, NE = 512, K = 10, slots = 300; + const int64_t cap = 8 * K, ptr_off = ((4 + (cap + 1) + 2 * cap) + 1) & ~1ll, words = ptr_off + 4 * cap + 2; + std::vector w((size_t) NE * N); + std::normal_distribution nd(0.0f, 0.05f); + for (auto& v : w) { + const float f = rng() % 64 == 0 ? edgy(rng) : nd(rng); + uint32_t b; + std::memcpy(&b, &f, 4); + v = (uint16_t) (b >> 16); + } + std::vector slot_ptr(slots); + for (int i = 0; i < slots; ++i) slot_ptr[(size_t) i] = 0x700000000ull + (unsigned long long) i * 1382400ull; + uint16_t* d_w = dev(w.size()); + unsigned long long* d_sp = dev(slots); + int32_t* d_res = dev(NE); + unsigned* d_counter = dev(1); + up(d_w, w); + up(d_sp, slot_ptr); + check(cudaMemset(d_counter, 0, 4), "memset"); + // per path: logits, ids, weights, plan (device); x, ids, weights, seq (mapped) + float *d_logits[2], *d_wts[2], *h_x[2], *m_x[2], *h_w[2], *m_w[2]; + int32_t *d_ids[2], *d_plan[2], *h_ids[2], *m_ids[2]; + uint32_t *h_seq[2], *m_seq[2]; + for (int p = 0; p < 2; ++p) { + d_logits[p] = dev((size_t) 8 * NE); + d_wts[p] = dev((size_t) 8 * K); + d_ids[p] = dev((size_t) 8 * K); + d_plan[p] = dev((size_t) words); + check(cudaHostAlloc((void**) &h_x[p], (size_t) 8 * N * 4, cudaHostAllocMapped), "cudaHostAlloc"); + check(cudaHostAlloc((void**) &h_ids[p], (size_t) 8 * K * 4, cudaHostAllocMapped), "cudaHostAlloc"); + check(cudaHostAlloc((void**) &h_w[p], (size_t) 8 * K * 4, cudaHostAllocMapped), "cudaHostAlloc"); + check(cudaHostAlloc((void**) &h_seq[p], 64, cudaHostAllocMapped), "cudaHostAlloc"); + check(cudaHostGetDevicePointer((void**) &m_x[p], h_x[p], 0), "mapped"); + check(cudaHostGetDevicePointer((void**) &m_ids[p], h_ids[p], 0), "mapped"); + check(cudaHostGetDevicePointer((void**) &m_w[p], h_w[p], 0), "mapped"); + check(cudaHostGetDevicePointer((void**) &m_seq[p], h_seq[p], 0), "mapped"); + } + float *d_x = dev((size_t) 8 * N), *d_bias = dev((size_t) NE); + int bad = 0; + for (int n_tok = 1; n_tok <= 8; ++n_tok) { + for (int rep = 0; rep < 4; ++rep) { + std::vector x((size_t) n_tok * N); + for (auto& v : x) v = rng() % 32 == 0 ? edgy(rng) : std::normal_distribution(0.0f, 1.0f)(rng); + std::vector res((size_t) NE, -1); + for (int e = 0; e < NE; ++e) // resident: none, all, or at random + if (rep == 1 || (rep != 0 && rng() % 3 != 0)) res[(size_t) e] = (int32_t) (rng() % slots); + up(d_x, x); + up(d_res, res); + for (int p = 0; p < 2; ++p) { + check(cudaMemset(d_plan[p], 0xff, (size_t) words * 4), "memset"); + std::memset(h_x[p], 0xcd, (size_t) 8 * N * 4); + *h_seq[p] = 5; + } + // the kernels verify_router stands in for + strata::kernels::bf16_gemv_fp32_mmvf_multi(d_x, d_w, d_logits[0], N, NE, n_tok, s); + strata::kernels::native_router_top10_multi(d_logits[0], d_ids[0], d_wts[0], n_tok, s); + strata::kernels::doorbell_publish(d_x, d_ids[0], d_wts[0], (int64_t) n_tok * N, (int64_t) n_tok * K, m_x[0], + m_ids[0], m_w[0], m_seq[0], s); + strata::kernels::verify_hit_plan(d_ids[0], n_tok * K, K, d_res, NE, d_sp, d_plan[0], cap, ptr_off, s); + strata::kernels::VerifyRouterArgs a; + a.x = d_x; a.w = d_w; a.logits = d_logits[1]; a.ids = d_ids[1]; a.weights = d_wts[1]; + a.x_out = m_x[1]; a.ids_out = m_ids[1]; a.w_out = m_w[1]; a.seq = m_seq[1]; a.ring = 6; + a.res = d_res; a.slot_ptr = d_sp; a.plan = d_plan[1]; a.cap = (int) cap; a.ptr_off = (int) ptr_off; + a.counter = d_counter; a.n_tok = n_tok; a.n_embd = N; a.n_expert = NE; + strata::kernels::verify_router(a, s); + check(cudaStreamSynchronize(s), "router"); + const size_t nk = (size_t) n_tok * K; + int b = bitwise_diff(down(d_logits[0], (size_t) n_tok * NE), down(d_logits[1], (size_t) n_tok * NE), "logits"); + b += bitwise_diff(down(d_wts[0], nk), down(d_wts[1], nk), "weights"); + b += bitwise_diff(std::vector(h_w[0], h_w[0] + nk), std::vector(h_w[1], h_w[1] + nk), "mapped weights"); + b += bitwise_diff(std::vector(h_x[0], h_x[0] + (size_t) n_tok * N), + std::vector(h_x[1], h_x[1] + (size_t) n_tok * N), "mapped x"); + if (down(d_ids[0], nk) != down(d_ids[1], nk) || + std::vector(h_ids[0], h_ids[0] + nk) != std::vector(h_ids[1], h_ids[1] + nk)) { + std::fprintf(stderr, "verify_router: n_tok %d rep %d: the ids differ\n", n_tok, rep); + ++b; + } + if (*h_seq[0] != 6 || *h_seq[1] != 6 || down(d_counter, 1)[0] != 0u) { + std::fprintf(stderr, "verify_router: n_tok %d rep %d: seq %u / %u, counter %u\n", n_tok, rep, *h_seq[0], + *h_seq[1], down(d_counter, 1)[0]); + ++b; + } + if (down(d_plan[0], (size_t) words) != down(d_plan[1], (size_t) words)) { + std::fprintf(stderr, "verify_router: n_tok %d rep %d: the hit plans differ\n", n_tok, rep); + ++b; + } + // with a bias the logits are the plain ones plus it; the update then corrects it by the plain ones + std::vector bias((size_t) NE); + for (auto& v : bias) v = std::normal_distribution(0.0f, 0.5f)(rng); + up(d_bias, bias); + a.bias = d_bias; + strata::kernels::verify_router(a, s); + strata::kernels::verify_router_bias_update(d_logits[0], d_logits[1], n_tok, d_bias, s); + check(cudaStreamSynchronize(s), "router bias"); + const std::vector plain = down(d_logits[0], (size_t) n_tok * NE); + std::vector want((size_t) n_tok * NE), fixed((size_t) NE); + for (int t = 0; t < n_tok; ++t) + for (int e = 0; e < NE; ++e) + want[(size_t) (t * NE + e)] = plain[(size_t) (t * NE + e)] + bias[(size_t) e]; + for (int e = 0; e < NE; ++e) { + float dsum = 0.0f; + for (int t = 0; t < n_tok; ++t) dsum += plain[(size_t) (t * NE + e)] - want[(size_t) (t * NE + e)]; + fixed[(size_t) e] = std::fmaf(0.125f / (float) n_tok, dsum, bias[(size_t) e]); + } + b += bitwise_diff(want, down(d_logits[1], (size_t) n_tok * NE), "biased logits"); + b += bitwise_diff(fixed, down(d_bias, (size_t) NE), "corrected bias"); + if (b) std::fprintf(stderr, "verify_router: n_tok %d rep %d: %d mismatches\n", n_tok, rep, b); + bad += b; + } + } + for (int p = 0; p < 2; ++p) { + for (void* q : {(void*) d_logits[p], (void*) d_wts[p], (void*) d_ids[p], (void*) d_plan[p]}) cudaFree(q); + for (void* q : {(void*) h_x[p], (void*) h_ids[p], (void*) h_w[p], (void*) h_seq[p]}) cudaFreeHost(q); + } + for (void* q : {(void*) d_w, (void*) d_sp, (void*) d_res, (void*) d_counter, (void*) d_x, (void*) d_bias}) + cudaFree(q); + std::printf("verify_router: %s\n", + bad ? "MISMATCH" + : "bitwise the gemv, top 10, doorbell and hit plan (1-8 tokens, 4 routings each), and with a bias " + "the logits plus it and its update"); + return bad; +} + +// ---- the router kernel's forms of its input: the q8_1 rows vs quantize_q8_1_rows, the Q8_K rows vs ggml's +// quantize_row_q8_K (the CPU pool's native_quant_act), the floats vs the input +int test_router_rows(std::mt19937& rng, cudaStream_t s) { + const int N = 2560, NE = 512, FF = 640; + const size_t row1 = N / 32 * 36, stride = 2928; + strata::kernels::cpu::NativeFmt f; + std::string ferr; + if (!strata::kernels::cpu::native_fmt(18, 20, N, FF, f, ferr) || f.act_bytes != (size_t) (N / 256 * 292)) { + std::fprintf(stderr, "router_rows: %s\n", ferr.empty() ? "IQ3_XXS's activation is not Q8_K" : ferr.c_str()); + return 1; + } + std::vector w((size_t) NE * N); + std::normal_distribution nd(0.0f, 0.05f); + for (auto& v : w) { + const float x = nd(rng); + uint32_t b; + std::memcpy(&b, &x, 4); + v = (uint16_t) (b >> 16); + } + uint16_t* d_w = dev(w.size()); + up(d_w, w); + float *d_x = dev((size_t) 8 * N), *d_logits = dev((size_t) 8 * NE), *d_wts = dev(80); + int32_t* d_ids = dev(80); + unsigned* d_counter = dev(1); + check(cudaMemset(d_counter, 0, 4), "memset"); + uint8_t *d_q1 = dev(8 * row1), *d_ref = dev(8 * row1); + uint8_t *h_q1 = nullptr, *m_q1 = nullptr, *h_xk = nullptr, *m_xk = nullptr; + float *h_x = nullptr, *m_x = nullptr; + uint32_t *h_seq = nullptr, *m_seq = nullptr; + check(cudaHostAlloc((void**) &h_q1, 8 * row1, cudaHostAllocMapped), "cudaHostAlloc"); + check(cudaHostAlloc((void**) &h_xk, 8 * stride, cudaHostAllocMapped), "cudaHostAlloc"); + check(cudaHostAlloc((void**) &h_x, (size_t) 8 * N * 4, cudaHostAllocMapped), "cudaHostAlloc"); + check(cudaHostAlloc((void**) &h_seq, 64, cudaHostAllocMapped), "cudaHostAlloc"); + check(cudaHostGetDevicePointer((void**) &m_q1, h_q1, 0), "mapped"); + check(cudaHostGetDevicePointer((void**) &m_xk, h_xk, 0), "mapped"); + check(cudaHostGetDevicePointer((void**) &m_x, h_x, 0), "mapped"); + check(cudaHostGetDevicePointer((void**) &m_seq, h_seq, 0), "mapped"); + int bad = 0; + for (int n_tok = 1; n_tok <= 8; ++n_tok) { + for (int rep = 0; rep < 4; ++rep) { + std::vector x((size_t) n_tok * N); + for (auto& v : x) v = rng() % 32 == 0 ? edgy(rng) : std::normal_distribution(0.0f, 1.0f)(rng); + if (rep == 1) std::fill(x.begin(), x.begin() + 256, 0.0f); // a zero block (and eight q8_1 ones) + if (rep == 2) { // the largest magnitude twice, in two warps and in one: the first one counts + x[7] = 5e4f; + x[40] = -5e4f; + x[256 + 3] = -6e4f; + x[256 + 20] = 6e4f; + } + up(d_x, x); + std::memset(h_q1, 0xcd, 8 * row1); + std::memset(h_xk, 0xcd, 8 * stride); + std::memset(h_x, 0xcd, (size_t) 8 * N * 4); + strata::kernels::VerifyRouterArgs a; + a.x = d_x; a.w = d_w; a.logits = d_logits; a.ids = d_ids; a.weights = d_wts; + a.xq1 = d_q1; a.xq1_out = m_q1; + a.xk_out = rep == 3 ? nullptr : m_xk; a.xk_stride = (int) stride; + a.x_out = rep == 3 ? m_x : nullptr; // a layer the pool takes as floats + a.seq = m_seq; a.ring = 1; a.counter = d_counter; a.n_tok = n_tok; a.n_embd = N; a.n_expert = NE; + strata::kernels::verify_router(a, s); + strata::kernels::quantize_q8_1_rows(d_x, n_tok, N, d_ref, s); + check(cudaStreamSynchronize(s), "router rows"); + int b = 0; + const std::vector ref = down(d_ref, (size_t) n_tok * row1); + if (down(d_q1, ref.size()) != ref || std::memcmp(h_q1, ref.data(), ref.size()) != 0) { + std::fprintf(stderr, "router_rows: n_tok %d rep %d: the q8_1 rows differ\n", n_tok, rep); + ++b; + } + for (int t = 0; t < n_tok && rep != 3; ++t) { + std::vector want(strata::kernels::cpu::kNativeActBytes, 0); + strata::kernels::cpu::native_quant_act(f, x.data() + (size_t) t * N, want.data()); + const uint8_t* got = h_xk + (size_t) t * stride; + for (size_t i = 0; i < f.act_bytes; ++i) + if (got[i] != want[i]) { + std::fprintf(stderr, "router_rows: n_tok %d rep %d token %d: Q8_K byte %zu (block %zu) %02x " + "vs ggml's %02x\n", n_tok, rep, t, i, i / 292, got[i], want[i]); + ++b; + break; + } + } + if (rep == 3 && std::memcmp(h_x, x.data(), x.size() * 4) != 0) { + std::fprintf(stderr, "router_rows: n_tok %d: the float copy differs\n", n_tok); + ++b; + } + bad += b; + } + } + for (void* q : {(void*) h_q1, (void*) h_xk, (void*) h_x, (void*) h_seq}) cudaFreeHost(q); + for (void* q : {(void*) d_w, (void*) d_x, (void*) d_logits, (void*) d_wts, (void*) d_ids, (void*) d_counter, + (void*) d_q1, (void*) d_ref}) + cudaFree(q); + std::printf("router_rows: %s\n", bad ? "MISMATCH" + : "the q8_1 rows bitwise quantize_q8_1_rows', the Q8_K rows ggml's, the float " + "copy the input (1-8 tokens, denormals, zero blocks, tied maxima)"); + return bad; +} + +// ---- the hyper-connection read: every token of an n-token read bitwise its 1-token read, and within a rounding +// bound of an FP64 reference (|error| <= 1e-4 x the sum of the terms' magnitudes, carried through silu and sigmoid) +int test_gr_read(std::mt19937& rng, cudaStream_t s) { + const int64_t N = 2560, HC = 4, D = N * HC, LR = 320; + auto bf16 = [&](size_t n) { + std::vector v(n); + for (auto& x : v) { + const float f = edgy(rng) * 0.05f; + uint32_t b; + std::memcpy(&b, &f, 4); + x = (uint16_t) (b >> 16); + } + return v; + }; + auto bf = [](uint16_t v) { const uint32_t b = (uint32_t) v << 16; float f; std::memcpy(&f, &b, 4); return (double) f; }; + std::vector wn((size_t) D); + for (auto& v : wn) v = edgy(rng); + const std::vector wd = bf16((size_t) (LR * D)), wu = bf16((size_t) (D * LR)), wi = bf16((size_t) (HC * D)); + uint16_t *d_wd = dev(wd.size()), *d_wu = dev(wu.size()), *d_wi = dev(wi.size()); + float* d_wn = dev(wn.size()); + up(d_wd, wd); + up(d_wu, wu); + up(d_wi, wi); + up(d_wn, wn); + const size_t sb = strata::kernels::fused_gr_scratch_bytes(); + float* d_scratch = (float*) dev(sb); + check(cudaMemset(d_scratch, 0, sb), "memset"); + int bad = 0, off = 0; + double worst = 0.0; // the largest error as a fraction of its bound + for (int n_tok = 1; n_tok <= 8; ++n_tok) { + for (int variant = 0; variant < 3; ++variant) { // apply + inject, no apply, no inject + const bool apply = variant != 1, inject = variant != 2; + std::vector R((size_t) (n_tok * D)), bo((size_t) (n_tok * N)), inj((size_t) (n_tok * HC)); + for (auto& v : R) v = edgy(rng); + for (auto& v : bo) v = edgy(rng); + for (auto& v : inj) v = edgy(rng); + // path 0: one read of n_tok tokens; path 1: n_tok reads of one token. R is updated in place when apply. + float* d[2][5]; // R, lo, rs, inject, mixed + for (int p = 0; p < 2; ++p) { + d[p][0] = dev(R.size()); + up(d[p][0], R); + d[p][1] = dev((size_t) (n_tok * LR)); + d[p][2] = dev((size_t) (n_tok * HC)); + d[p][3] = dev((size_t) (n_tok * HC)); + d[p][4] = dev((size_t) (n_tok * N)); + check(cudaMemset(d[p][3], 0, (size_t) (n_tok * HC) * 4), "memset"); + } + float *d_bo = dev(bo.size()), *d_inj = dev(inj.size()); + up(d_bo, bo); + up(d_inj, inj); + // path 0 also keeps the gates' running average and estimates another read (its norm weights wn2, the + // gates eg) + std::vector ema0((size_t) D), wn2((size_t) D), eg((size_t) D); + for (auto& v : ema0) v = std::uniform_real_distribution(0.0f, 1.0f)(rng); + for (auto& v : eg) v = std::uniform_real_distribution(0.0f, 1.0f)(rng); + for (auto& v : wn2) v = edgy(rng); + float *d_ema = dev((size_t) D), *d_est = dev((size_t) (n_tok * N)); + float *d_wn2 = dev((size_t) D), *d_eg = dev((size_t) D); + up(d_ema, ema0); + up(d_wn2, wn2); + up(d_eg, eg); + std::vector args[2]; + for (int p = 0; p < 2; ++p) + for (int t = 0; t < n_tok; ++t) { + strata::kernels::FusedGrArgs a; + a.R = d[p][0] + t * D; a.R_out = d[p][0] + t * D; a.apply = apply; + a.bo_prev = d_bo + t * N; a.inj_prev = d_inj + t * HC; + a.w_norm = d_wn; a.w_down = d_wd; a.w_up = d_wu; a.w_inject = inject ? d_wi : nullptr; + a.eps = 1e-6f; + a.lo = d[p][1] + t * LR; a.rs = d[p][2] + t * HC; a.inject_out = d[p][3] + t * HC; + a.mixed = d[p][4] + t * N; + if (p == 0) { a.gate_ema = d_ema; a.est_norm = d_wn2; a.est_gates = d_eg; a.est = d_est + t * N; } + args[p].push_back(a); + } + strata::kernels::fused_gr_read_multi(args[0].data(), n_tok, d_scratch, s); + for (int t = 0; t < n_tok; ++t) strata::kernels::fused_gr_read_multi(&args[1][(size_t) t], 1, d_scratch, s); + check(cudaStreamSynchronize(s), "gr read"); + const std::vector ema1 = down(d_ema, (size_t) D), est = down(d_est, (size_t) (n_tok * N)); + const char* names[5] = {"R", "lo", "rs", "inject", "mixed"}; + const size_t sizes[5] = {R.size(), (size_t) (n_tok * LR), (size_t) (n_tok * HC), (size_t) (n_tok * HC), + (size_t) (n_tok * N)}; + std::vector got[5]; + for (int o = 0; o < 5; ++o) { + got[o] = down(d[0][o], sizes[o]); + const int b = bitwise_diff(got[o], down(d[1][o], sizes[o]), names[o]); + if (b) std::fprintf(stderr, "gr_read: n_tok %d variant %d: %s: %d differ\n", n_tok, variant, names[o], b); + bad += b; + } + // FP64 for the first token of the 1- and 4-token reads + if (n_tok == 1 || n_tok == 4) { + const float* Rt = R.data(); + // xa: the magnitude xn's rounding scales with (R' = R + gw bo may cancel) + std::vector Rp((size_t) D), Ra((size_t) D), xn((size_t) D), xa((size_t) D), rs(HC); + for (int c = 0; c < HC; ++c) { + const double gw = apply ? 2.0 / (1.0 + std::exp(-(double) inj[(size_t) c] / HC)) : 0.0; + double ss = 0.0; + for (int64_t dd = 0; dd < N; ++dd) { + const size_t i = (size_t) (c * N + dd); + Rp[i] = (double) Rt[i] + gw * bo[(size_t) dd]; + Ra[i] = std::fabs((double) Rt[i]) + std::fabs(gw * bo[(size_t) dd]); + ss += Rp[i] * Rp[i]; + } + rs[(size_t) c] = 1.0 / std::sqrt(ss / N + 1e-6); + for (int64_t dd = 0; dd < N; ++dd) { + const size_t i = (size_t) (c * N + dd); + xn[i] = Rp[i] * wn[i] * rs[(size_t) c]; + xa[i] = Ra[i] * std::fabs((double) wn[i]) * rs[(size_t) c]; + } + } + auto near = [&](double gpu, double ref, double bound, const char* what, int64_t i) { + if (!std::isfinite(ref) && std::isinf(gpu) && (gpu > 0) == (ref > 0)) return; + const double e = std::fabs(gpu - ref); + worst = std::max(worst, e / (bound + 1e-300)); + if (!(e <= bound) && off++ < 5) + std::fprintf(stderr, " gr_read FP64: n_tok %d variant %d: %s[%lld] %.9g vs %.9g (bound %.3g)\n", n_tok, + variant, what, (long long) i, gpu, ref, bound); + }; + for (int c = 0; c < HC; ++c) near(got[2][(size_t) c], rs[(size_t) c], 1e-5 * rs[(size_t) c], "rs", c); + if (apply) + for (int64_t i = 0; i < D; ++i) near(got[0][(size_t) i], Rp[(size_t) i], 1e-6 * Ra[(size_t) i] + 1e-30, "R", i); + std::vector lo((size_t) LR), lo_b((size_t) LR); + for (int64_t r = 0; r < LR + (inject ? HC : 0); ++r) { + const uint16_t* w = r < LR ? wd.data() + (size_t) (r * D) : wi.data() + (size_t) ((r - LR) * D); + double y = 0.0, mag = 0.0; + for (int64_t i = 0; i < D; ++i) { + y += bf(w[i]) * xn[(size_t) i]; + mag += std::fabs(bf(w[i])) * xa[(size_t) i]; + } + if (r < LR) { + const double x = y / HC; + lo[(size_t) r] = x / (1.0 + std::exp(-x)); + lo_b[(size_t) r] = 1.1 * 1e-4 * mag / HC + 1e-30; + near(got[1][(size_t) r], lo[(size_t) r], lo_b[(size_t) r], "lo", r); + } else { + near(got[3][(size_t) (r - LR)], y, 1e-4 * mag + 1e-30, "inject", r - LR); + } + } + for (int64_t dd = 0; dd < N; ++dd) { + double mixed = 0.0, bound = 0.0; + for (int c = 0; c < HC; ++c) { + const size_t i = (size_t) (c * N + dd); + double u = 0.0, du = 0.0; + for (int64_t k = 0; k < LR; ++k) { + const double w = bf(wu[i * LR + (size_t) k]); + u += w * lo[(size_t) k]; + du += std::fabs(w) * (1e-4 * std::fabs(lo[(size_t) k]) + lo_b[(size_t) k]); + } + const double sg = 1.0 / (1.0 + std::exp(-u)); + mixed += xn[i] * sg / HC; + bound += (1e-4 * xa[i] + 0.25 * std::fabs(xn[i]) * du) / HC; + if (n_tok == 1) + near(ema1[i], ema0[i] + 0.25 * (sg - ema0[i]), 0.25 * (0.25 * du + 1e-6) + 1e-6, + "gate average", (int64_t) i); + } + near(got[4][(size_t) dd], mixed, bound + 1e-30, "mixed", dd); + } + for (int64_t dd = 0; dd < N; ++dd) { // the other read's estimate + double ref = 0.0, mag = 0.0; + for (int c = 0; c < HC; ++c) { + const size_t i = (size_t) (c * N + dd); + ref += Rp[i] * rs[(size_t) c] * wn2[i] * eg[i] / HC; + mag += Ra[i] * rs[(size_t) c] * std::fabs((double) wn2[i]) * eg[i] / HC; + } + near(est[(size_t) dd], ref, 1e-4 * mag + 1e-30, "estimate", dd); + } + } + for (float* q : {d_ema, d_est, d_wn2, d_eg}) cudaFree(q); + for (int p = 0; p < 2; ++p) + for (float* q : d[p]) cudaFree(q); + cudaFree(d_bo); + cudaFree(d_inj); + } + } + for (void* p : {(void*) d_wd, (void*) d_wu, (void*) d_wi, (void*) d_wn, (void*) d_scratch}) cudaFree(p); + bad += off; + char msg[256]; + std::snprintf(msg, sizeof msg, + "each token bitwise its 1-token read (1-8 tokens, 3 variants, keeping the gates' average and " + "estimating another read or not); FP64, the gates' average and the estimate within bounds (largest " + "error %.2g of its bound)", + worst); + std::printf("gr_read: %s\n", bad ? "MISMATCH" : msg); + return bad; +} + +// ---- RoPE over several tokens' heads: native_rope_apply per token vs native_rope_apply_tokens +int test_rope_tokens(std::mt19937& rng, cudaStream_t s) { + const int NH = 24; + int bad = 0; + for (int heads : {2, 4, 24}) { + for (int head_dim : {128, 256}) { + for (int n_tok = 1; n_tok <= 8; ++n_tok) { + std::vector x((size_t) n_tok * heads * head_dim); + for (auto& v : x) v = edgy(rng); + std::vector pos((size_t) n_tok * NH); // a window's per-token position vectors + const int p0 = (int) (rng() % 200000); + for (int t = 0; t < n_tok; ++t) + for (int h = 0; h < NH; ++h) pos[(size_t) t * NH + h] = p0 + t; + float *d_a = dev(x.size()), *d_b = dev(x.size()); + int32_t* d_pos = dev(pos.size()); + up(d_a, x); + up(d_b, x); + up(d_pos, pos); + for (int t = 0; t < n_tok; ++t) + strata::kernels::native_rope_apply(d_a + (size_t) t * heads * head_dim, d_a + (size_t) t * heads * head_dim, + heads, head_dim, 64, 5000000.0f, d_pos + t * NH, s); + strata::kernels::native_rope_apply_tokens(d_b, d_b, n_tok * heads, head_dim, 64, 5000000.0f, d_pos, heads, + NH, s); + check(cudaStreamSynchronize(s), "rope"); + bad += bitwise_diff(down(d_a, x.size()), down(d_b, x.size()), "rope_tokens"); + cudaFree(d_a); + cudaFree(d_b); + cudaFree(d_pos); + } + } + } + std::printf("rope_tokens: %s\n", bad ? "MISMATCH" : "bitwise equal (2, 4, 24 heads; 1-8 tokens)"); + return bad; +} + +// ---- the QSA heads' norm and RoPE: native_qsa_rms_norm_weighted + native_rope_apply_tokens vs +// native_qsa_norm_rope_tokens, in place and from rows of twice the width (the queries' halves) +int test_norm_rope(std::mt19937& rng, cudaStream_t s) { + const int NH = 24; + int bad = 0; + for (int heads : {2, 4, 24}) { + for (int head_dim : {128, 256}) { + std::vector gamma((size_t) head_dim); + for (auto& g : gamma) g = 0.5f + std::uniform_real_distribution(0.0f, 1.0f)(rng); + float* d_gamma = dev(gamma.size()); + up(d_gamma, gamma); + for (int n_tok = 1; n_tok <= 8; ++n_tok) { + const int rows = n_tok * heads; + std::vector x((size_t) rows * 2 * head_dim); // rows of 2 * head_dim: the first half is read + for (auto& v : x) v = std::normal_distribution(0.0f, 3.0f)(rng); + std::vector half((size_t) rows * head_dim); + for (int r = 0; r < rows; ++r) + std::copy(x.begin() + (size_t) r * 2 * head_dim, x.begin() + (size_t) r * 2 * head_dim + head_dim, + half.begin() + (size_t) r * head_dim); + std::vector pos((size_t) n_tok * NH); + const int p0 = (int) (rng() % 200000); + for (int t = 0; t < n_tok; ++t) + for (int h = 0; h < NH; ++h) pos[(size_t) t * NH + h] = p0 + t; + float *d_a = dev(half.size()), *d_b = dev(half.size()), *d_x = dev(x.size()), + *d_c = dev(half.size()); + int32_t* d_pos = dev(pos.size()); + up(d_a, half); + up(d_b, half); + up(d_x, x); + up(d_pos, pos); + strata::kernels::native_qsa_rms_norm_weighted(d_a, d_gamma, d_a, head_dim, rows, 1e-6f, s); + strata::kernels::native_rope_apply_tokens(d_a, d_a, rows, head_dim, 64, 5000000.0f, d_pos, heads, NH, + s); + strata::kernels::native_qsa_norm_rope_tokens(d_b, head_dim, d_gamma, d_b, head_dim, rows, 1e-6f, 64, + 5000000.0f, d_pos, heads, NH, s); + strata::kernels::native_qsa_norm_rope_tokens(d_x, 2 * head_dim, d_gamma, d_c, head_dim, rows, 1e-6f, 64, + 5000000.0f, d_pos, heads, NH, s); + check(cudaStreamSynchronize(s), "norm_rope"); + const int b = bitwise_diff(down(d_a, half.size()), down(d_b, half.size()), "norm_rope") + + bitwise_diff(down(d_a, half.size()), down(d_c, half.size()), "norm_rope, strided"); + if (b) + std::fprintf(stderr, "norm_rope: %d heads of %d, n_tok %d: %d differ\n", heads, head_dim, n_tok, b); + bad += b; + for (float* p : {d_a, d_b, d_x, d_c}) cudaFree(p); + cudaFree(d_pos); + } + cudaFree(d_gamma); + } + } + std::printf("norm_rope_tokens: %s\n", bad ? "MISMATCH" : "bitwise equal (2, 4, 24 heads of 128, 256; 1-8 tokens; " + "in place and strided)"); + return bad; +} + +// ---- the int8 KV append: kv_append_q8_step per token vs kv_append_q8_steps +int test_kv_append(std::mt19937& rng, cudaStream_t s) { + const strata::kernels::QsaShapes sh = strata::kernels::qsa_real_shapes(); + const int cells = 4096, pages = cells / (int) sh.page_size, per_tok = (int) (sh.n_head_kv * sh.head_dim); + const size_t codes = (size_t) cells * per_tok, scales = codes / strata::kernels::KV_Q8_GROUP; + std::vector table((size_t) pages); + for (int i = 0; i < pages; ++i) table[(size_t) i] = pages - 1 - i; // a permuted page table + int32_t* d_table = dev(table.size()); + up(d_table, table); + int bad = 0; + for (int n_tok = 1; n_tok <= 8; ++n_tok) { + const int p0 = 500 + (int) (rng() % 2000); + std::vector k((size_t) n_tok * per_tok), v((size_t) n_tok * per_tok); + for (auto& x : k) x = edgy(rng); + for (auto& x : v) x = edgy(rng); + std::vector steps((size_t) n_tok * strata::kernels::kStepCount, 0); + for (int t = 0; t < n_tok; ++t) steps[(size_t) t * strata::kernels::kStepCount + strata::kernels::kStepPos] = p0 + t; + float *d_k = dev(k.size()), *d_v = dev(v.size()); + int32_t* d_steps = dev(steps.size()); + up(d_k, k); + up(d_v, v); + up(d_steps, steps); + int8_t* q[2][2]; + uint16_t* sc[2][2]; + for (int p = 0; p < 2; ++p) + for (int kv = 0; kv < 2; ++kv) { + q[p][kv] = dev(codes); + sc[p][kv] = dev(scales); + check(cudaMemset(q[p][kv], 0, codes), "memset"); + check(cudaMemset(sc[p][kv], 0, scales * 2), "memset"); + } + for (int t = 0; t < n_tok; ++t) + strata::kernels::kv_append_q8_step(q[0][0], q[0][1], sc[0][0], sc[0][1], d_table, + d_steps + t * strata::kernels::kStepCount, d_k + (size_t) t * per_tok, + d_v + (size_t) t * per_tok, sh, s); + strata::kernels::kv_append_q8_steps(q[1][0], q[1][1], sc[1][0], sc[1][1], d_table, d_steps, + strata::kernels::kStepCount, d_k, d_v, n_tok, sh, s); + check(cudaStreamSynchronize(s), "kv append"); + for (int kv = 0; kv < 2; ++kv) { + if (down(q[0][kv], codes) != down(q[1][kv], codes) || down(sc[0][kv], scales) != down(sc[1][kv], scales)) { + std::fprintf(stderr, "kv_append: n_tok %d: the %s cells differ\n", n_tok, kv ? "V" : "K"); + ++bad; + } + } + for (int p = 0; p < 2; ++p) + for (int kv = 0; kv < 2; ++kv) { + cudaFree(q[p][kv]); + cudaFree(sc[p][kv]); + } + cudaFree(d_k); + cudaFree(d_v); + cudaFree(d_steps); + } + cudaFree(d_table); + std::printf("kv_append_steps: %s\n", bad ? "MISMATCH" : "bitwise equal (1-8 tokens)"); + return bad; +} + +// ---- the indexer append: one call per token vs native_qsa_indexer_append_multi (blocks completed inside the +// window, the first cell, and rejected tokens' -1 positions) +int test_indexer_append(std::mt19937& rng, cudaStream_t s) { + const strata::kernels::QsaShapes sh = strata::kernels::qsa_real_shapes(); + const int D = 128, max_cells = 4096; + const size_t pooled_n = (size_t) (max_cells / 4 + 1) * D; + std::vector gamma((size_t) D); + for (auto& g : gamma) g = 0.5f + std::uniform_real_distribution(0.0f, 1.0f)(rng); + float* d_gamma = dev(gamma.size()); + up(d_gamma, gamma); + int bad = 0; + for (int n_tok = 1; n_tok <= 8; ++n_tok) { + for (int start : {0, 1, 2, 3, 5, 6, 7, 1022}) { + for (int rejected = 0; rejected < 2; ++rejected) { + std::vector raw((size_t) n_tok * D); + for (auto& x : raw) x = std::normal_distribution(0.0f, 1.0f)(rng); + std::vector pos((size_t) n_tok); + for (int t = 0; t < n_tok; ++t) pos[(size_t) t] = rejected && t >= (n_tok + 1) / 2 ? -1 : start + t; + // the state both paths start from: a random earlier tail, spare and pooled blocks + std::vector tail((size_t) 3 * D), dead((size_t) D), pooled(pooled_n); + for (auto& x : tail) x = edgy(rng); + for (auto& x : dead) x = edgy(rng); + for (auto& x : pooled) x = edgy(rng); + float *d_raw = dev(raw.size()); + int32_t* d_pos = dev(pos.size()); + up(d_raw, raw); + up(d_pos, pos); + float* bufs[2][3]; + int32_t* bp[2]; + for (int p = 0; p < 2; ++p) { + bufs[p][0] = dev(tail.size()); + bufs[p][1] = dev(dead.size()); + bufs[p][2] = dev(pooled.size()); + bp[p] = dev(1); + up(bufs[p][0], tail); + up(bufs[p][1], dead); + up(bufs[p][2], pooled); + up(bp[p], std::vector{-7}); + } + const strata::kernels::QsaIndexerBuffers ib0{bufs[0][0], bufs[0][1], bufs[0][2], bp[0]}; + const strata::kernels::QsaIndexerBuffers ib1{bufs[1][0], bufs[1][1], bufs[1][2], bp[1]}; + for (int t = 0; t < n_tok; ++t) + strata::kernels::native_qsa_indexer_append(d_raw + (size_t) t * D, d_pos + t, 0, d_gamma, 1e-6f, ib0, sh, + max_cells, 5000000.0f, s); + float* d_snap = dev(tail.size()); + strata::kernels::native_qsa_indexer_append_multi(d_raw, d_pos, 1, n_tok, 0, d_gamma, 1e-6f, ib1, sh, + max_cells, 5000000.0f, s, d_snap); + check(cudaStreamSynchronize(s), "indexer"); + if (down(d_snap, tail.size()) != tail) { + std::fprintf(stderr, "indexer: n_tok %d start %d: the tail snapshot differs\n", n_tok, start); + ++bad; + } + cudaFree(d_snap); + const size_t sizes[3] = {tail.size(), dead.size(), pooled.size()}; + const char* names[3] = {"tail", "dead", "pooled"}; + for (int o = 0; o < 3; ++o) { + const int b = bitwise_diff(down(bufs[0][o], sizes[o]), down(bufs[1][o], sizes[o]), names[o]); + if (b) std::fprintf(stderr, "indexer: n_tok %d start %d rejected %d: %s differ\n", n_tok, start, rejected, names[o]); + bad += b; + } + if (down(bp[0], 1) != down(bp[1], 1)) { + std::fprintf(stderr, "indexer: n_tok %d start %d: block_pos differs\n", n_tok, start); + ++bad; + } + for (int p = 0; p < 2; ++p) { + for (float* q : bufs[p]) cudaFree(q); + cudaFree(bp[p]); + } + cudaFree(d_raw); + cudaFree(d_pos); + } + } + } + cudaFree(d_gamma); + std::printf("indexer_append_multi: %s\n", + bad ? "MISMATCH" : "bitwise equal (1-8 tokens, 8 starts, rejected cells), the tail snapshot the tail"); + return bad; +} + +// ---- the PLE block after its projections: native_ple_postops + ple_history_advance per token vs +// native_ple_postops_tokens (chunks shorter and longer than the nine-row history) +int test_ple_tokens(std::mt19937& rng, cudaStream_t s) { + const int N = 2560, D = 10240, H = 4, HIST = 9; + std::normal_distribution nd(0.0f, 1.0f); + std::vector nk((size_t) D), nq((size_t) D), nc((size_t) D); + for (auto* v : {&nk, &nq, &nc}) + for (auto& x : *v) x = 0.5f + std::uniform_real_distribution(0.0f, 1.0f)(rng); + std::vector taps((size_t) 4 * D); + for (auto& t : taps) { // F16 in [-1, 1): sign, exponent 0..14, random mantissa + const uint16_t sign = (uint16_t) ((rng() & 1u) << 15), e = (uint16_t) (rng() % 15), m = (uint16_t) (rng() & 0x3FFu); + t = (uint16_t) (sign | (e << 10) | m); + } + float *d_nk = dev(D), *d_nq = dev(D), *d_nc = dev(D); + uint16_t* d_taps = dev(taps.size()); + up(d_nk, nk); up(d_nq, nq); up(d_nc, nc); up(d_taps, taps); + strata::kernels::PleWeights w; + w.norm_key = d_nk; w.norm_query = d_nq; w.norm_conv = d_nc; w.conv1d_f16 = d_taps; + int bad = 0; + for (int n_tok : {1, 2, 5, 9, 10, 23}) { + std::vector key((size_t) n_tok * D), hidden((size_t) n_tok * D), value((size_t) n_tok * N), + hist((size_t) HIST * D); + for (auto& x : key) x = edgy(rng); + for (auto& x : hidden) x = edgy(rng); + for (auto& x : value) x = nd(rng); + for (auto& x : hist) x = nd(rng); + float *d_key = dev(key.size()), *d_hidden = dev(hidden.size()), *d_value = dev(value.size()); + up(d_key, key); up(d_hidden, hidden); up(d_value, value); + // per token: the decode path's calls, the history after each token kept + float *h0 = dev(hist.size()), *r0 = dev(hidden.size()), + *s0 = dev((size_t) n_tok * hist.size()); + float *k1 = dev(D), *q1 = dev(D), *g1 = dev(H), *gd1 = dev(D), *c1 = dev(D); + up(h0, hist); + for (int t = 0; t < n_tok; ++t) { + const strata::kernels::NativePlePostopsBuffers b{k1, q1, g1, gd1, q1, c1, r0 + (size_t) t * D}; + strata::kernels::native_ple_postops(d_key + (size_t) t * D, d_hidden + (size_t) t * D, + d_value + (size_t) t * N, h0, w, b, s); + strata::kernels::ple_history_advance(h0, q1, s); + check(cudaMemcpyAsync(s0 + (size_t) t * hist.size(), h0, hist.size() * 4, cudaMemcpyDeviceToDevice, s), + "ple snapshot"); + } + // the chunk at once + float *h2 = dev(hist.size()), *r2 = dev(hidden.size()), *k2 = dev(key.size()), + *q2 = dev(key.size()), *g2 = dev((size_t) n_tok * H), *gd2 = dev(key.size()), + *s2 = dev((size_t) n_tok * hist.size()); + up(h2, hist); + strata::kernels::native_ple_postops_tokens(d_key, d_hidden, d_value, h2, w, {k2, q2, g2, gd2, q2, r2}, n_tok, s, + s2); + check(cudaStreamSynchronize(s), "ple tokens"); + const int br = bitwise_diff(down(r0, hidden.size()), down(r2, hidden.size()), "result"); + const int bh = bitwise_diff(down(h0, hist.size()), down(h2, hist.size()), "history"); + const int bs = bitwise_diff(down(s0, (size_t) n_tok * hist.size()), down(s2, (size_t) n_tok * hist.size()), + "histories"); + if (br || bh || bs) std::fprintf(stderr, "ple tokens: n_tok %d: results or histories differ\n", n_tok); + bad += br + bh + bs; + for (float* p : {d_key, d_hidden, d_value, h0, r0, s0, k1, q1, g1, gd1, c1, h2, r2, k2, q2, g2, gd2, s2}) + cudaFree(p); + } + for (float* p : {d_nk, d_nq, d_nc}) cudaFree(p); + cudaFree(d_taps); + std::printf("ple_postops_tokens: %s\n", + bad ? "MISMATCH" : "bitwise equal (1-23 tokens, results, history, the history after each token)"); + return bad; +} + +// ---- the decode attention (split, merged, gated) vs the prompt path's kernels (INT8 and FP16 pools; 1-2051 cells) +int test_decode_attn(std::mt19937& rng, cudaStream_t s) { + const strata::kernels::QsaShapes sh = strata::kernels::qsa_real_shapes(); + const int HD = (int) sh.head_dim, NH = (int) sh.n_head, NKV = (int) sh.n_head_kv; + const int cells = 4096, pages = cells / (int) sh.page_size; + const int64_t cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, sh); + const size_t n_codes = (size_t) cells * NKV * HD, n_scales = n_codes / strata::kernels::KV_Q8_GROUP; + std::vector table((size_t) pages); + for (int i = 0; i < pages; ++i) table[(size_t) i] = (i * 3 + 1) % pages; // a permuted page table + int32_t* d_table = dev(table.size()); + up(d_table, table); + auto half_bits = [&](float lo, float hi) { // a random positive FP16 in [lo, hi) + const float f = std::uniform_real_distribution(lo, hi)(rng); + uint32_t b; + std::memcpy(&b, &f, 4); + const int e = (int) ((b >> 23) & 0xFF) - 127 + 15; + return (uint16_t) ((e << 10) | ((b >> 13) & 0x3FF)); + }; + std::vector kq(n_codes), vq(n_codes); + for (auto* v : {&kq, &vq}) + for (auto& c : *v) c = (int8_t) ((int) (rng() % 255) - 127); + std::vector ks(n_scales), vs(n_scales), kh(n_codes), vh(n_codes); + for (auto& x : ks) x = half_bits(0.002f, 0.05f); + for (auto& x : vs) x = half_bits(0.002f, 0.05f); + for (auto& x : kh) x = (uint16_t) (half_bits(0.01f, 1.0f) | ((rng() & 1u) << 15)); + for (auto& x : vh) x = (uint16_t) (half_bits(0.01f, 1.0f) | ((rng() & 1u) << 15)); + int8_t *d_kq = dev(n_codes), *d_vq = dev(n_codes); + uint16_t *d_ks = dev(n_scales), *d_vs = dev(n_scales), *d_kh = dev(n_codes), + *d_vh = dev(n_codes); + up(d_kq, kq); up(d_vq, vq); up(d_ks, ks); up(d_vs, vs); up(d_kh, kh); up(d_vh, vh); + const std::vector widths = {1, 5, 64, 65, 700, 2051}; + const int n_q = (int) widths.size(); + std::vector ids((size_t) n_q * cap, 0), steps((size_t) n_q * strata::kernels::kStepCount, 0); + for (int i = 0; i < n_q; ++i) { + std::vector all(cells); + for (int c = 0; c < cells; ++c) all[(size_t) c] = c; + std::shuffle(all.begin(), all.end(), rng); + std::sort(all.begin(), all.begin() + widths[(size_t) i]); + std::copy(all.begin(), all.begin() + widths[(size_t) i], ids.begin() + (size_t) i * cap); + steps[(size_t) i * strata::kernels::kStepCount + strata::kernels::kStepWidth] = widths[(size_t) i]; + } + std::vector q((size_t) n_q * NH * HD), qf((size_t) n_q * NH * 2 * HD); // qf: queries, then gate logits + for (auto& x : q) x = 3.0f * std::normal_distribution(0.0f, 1.0f)(rng); + for (auto& x : qf) x = 2.0f * std::normal_distribution(0.0f, 1.0f)(rng); + int32_t *d_ids = dev(ids.size()), *d_steps = dev(steps.size()); + float *d_q = dev(q.size()), *d_a = dev(q.size()), *d_b = dev(q.size()); + float *d_qf = dev(qf.size()), *d_g = dev(q.size()), *d_c = dev(q.size()); + const uint64_t scr = strata::kernels::qsa_decode_attn_scratch_floats(cap, sh); + float* d_scratch = dev((size_t) n_q * scr); + up(d_ids, ids); up(d_steps, steps); up(d_q, q); up(d_qf, qf); + check(cudaDeviceSynchronize(), "attn inputs"); + const int64_t SC = strata::kernels::kStepCount; + int bad = 0, gate_bad = 0; + double worst = 0.0; + for (int int8 = 0; int8 < 2; ++int8) { + strata::kernels::QsaAttnPools pools; + pools.page_table = d_table; + if (int8) { pools.k_q = d_kq; pools.v_q = d_vq; pools.k_scale = d_ks; pools.v_scale = d_vs; } + else { pools.k_pool = d_kh; pools.v_pool = d_vh; } + strata::kernels::qsa_prefill_attn(d_q, pools, d_ids, d_steps, cap, sh, d_b, n_q, s); + for (int alone = 0; alone < 2; ++alone) { // the queries in one call, and each alone (the most splits) + if (alone) + for (int i = 0; i < n_q; ++i) + strata::kernels::qsa_decode_attn_batch(d_q + (size_t) i * NH * HD, pools, d_ids + (size_t) i * cap, + d_steps + i * SC, cap, sh, d_scratch, + d_a + (size_t) i * NH * HD, 1, s); + else strata::kernels::qsa_decode_attn_batch(d_q, pools, d_ids, d_steps, cap, sh, d_scratch, d_a, n_q, s); + check(cudaStreamSynchronize(s), "decode attn"); + const std::vector a = down(d_a, q.size()), b = down(d_b, q.size()); + for (int r = 0; r < n_q * NH; ++r) { // relative to the largest value of the head's output row + double big = 0.0, diff = 0.0; + for (int d = 0; d < HD; ++d) { + big = std::max(big, (double) std::fabs(a[(size_t) r * HD + d])); + diff = std::max(diff, (double) std::fabs(a[(size_t) r * HD + d] - b[(size_t) r * HD + d])); + } + const double rel = diff / std::max(big, 1e-30); + worst = std::max(worst, rel); + if (!(rel <= 2e-5)) { + if (bad < 5) + std::fprintf(stderr, "decode attn: %s pools, %s, query %d head %d: relative difference %.3g\n", + int8 ? "INT8" : "FP16", alone ? "alone" : "batched", r / NH, r % NH, rel); + ++bad; + } + } + } + // the gate in the merge against native_qsa_gate_apply on the ungated output + strata::kernels::qsa_decode_attn_batch(d_q, pools, d_ids, d_steps, cap, sh, d_scratch, d_a, n_q, s); + strata::kernels::native_qsa_gate_apply(d_a, d_qf, d_c, n_q * NH, HD, s); + strata::kernels::qsa_decode_attn_batch(d_q, pools, d_ids, d_steps, cap, sh, d_scratch, d_g, n_q, s, d_qf); + check(cudaStreamSynchronize(s), "gated attn"); + const std::vector g = down(d_g, q.size()), c = down(d_c, q.size()); + if (std::memcmp(g.data(), c.data(), g.size() * 4) != 0) { + std::fprintf(stderr, "decode attn: %s pools, the fused gate differs from native_qsa_gate_apply\n", + int8 ? "INT8" : "FP16"); + ++gate_bad; + } + } + for (void* p : {(void*) d_table, (void*) d_kq, (void*) d_vq, (void*) d_ks, (void*) d_vs, (void*) d_kh, (void*) d_vh, + (void*) d_ids, (void*) d_steps, (void*) d_q, (void*) d_a, (void*) d_b, (void*) d_qf, (void*) d_g, + (void*) d_c, (void*) d_scratch}) + cudaFree(p); + std::printf("decode_attn: %s, the gate %s (INT8 and FP16 pools, 1-2051 cells, 6 queries at once and each alone; " + "largest relative difference %.2g)\n", + bad ? "OUTSIDE TOLERANCE" : "within 2e-5 of the prompt path's kernels", + gate_bad ? "DIFFERS" : "bitwise native_qsa_gate_apply's", worst); + return bad + gate_bad; +} + +// ---- the prompt path's GDN conv and recurrence over a chunk vs the verify window's kernels, 8 tokens at a time +int test_prefill_gdn(std::mt19937& rng, cudaStream_t s) { + const int C = 10240, HK = 16, HV = 48, S = 128, ZV = HV * S; + const float eps = 1e-6f; + std::normal_distribution nd(0.0f, 1.0f); + auto uni = [&](float a, float b) { return std::uniform_real_distribution(a, b)(rng); }; + std::vector w((size_t) C * 4), hist((size_t) C * 3), state((size_t) S * HV * S), gamma((size_t) S); + for (auto& x : w) x = 0.5f * nd(rng); + for (auto& x : gamma) x = uni(0.5f, 1.5f); + float *d_w = dev(w.size()), *d_gamma = dev(gamma.size()); + up(d_w, w); + up(d_gamma, gamma); + int bad = 0; + for (int T : {1, 2, 3, 5, 37}) { + std::vector qkv((size_t) T * C), gate((size_t) T * HV), beta((size_t) T * HV), z((size_t) T * ZV); + for (auto& x : qkv) x = nd(rng); + for (auto& x : hist) x = nd(rng); + for (auto& x : state) x = 0.1f * nd(rng); + for (auto& x : gate) x = uni(-2.0f, 0.0f); + for (auto& x : beta) x = uni(0.0f, 1.0f); + for (auto& x : z) x = nd(rng); + float *d_qkv = dev(qkv.size()), *d_gate = dev(gate.size()), *d_beta = dev(beta.size()), + *d_z = dev(z.size()); + up(d_qkv, qkv); up(d_gate, gate); up(d_beta, beta); up(d_z, z); + float* d_hist[2]; + float* d_state[2]; + float* d_h[2]; + float* d_y[2]; + for (int p = 0; p < 2; ++p) { + d_hist[p] = dev(hist.size()); + d_state[p] = dev(state.size()); + d_h[p] = dev(qkv.size()); + d_y[p] = dev(z.size()); + up(d_hist[p], hist); + up(d_state[p], state); + } + uint16_t* d_y16 = dev(z.size()); + // the verify window's kernels, a window of up to 8 tokens after another, each committed + std::vector counts; + for (int t0 = 0; t0 < T; t0 += 8) counts.push_back(std::min(8, T - t0)); + int32_t* d_counts = dev(counts.size()); + up(d_counts, counts); + for (size_t i = 0; i < counts.size(); ++i) { + const size_t t0 = i * 8; + const int n = counts[i]; + strata::kernels::gdn_conv_l2_multi(d_hist[0], d_qkv + t0 * C, d_w, d_h[0] + t0 * C, C, 2 * HK, eps, n, s); + strata::kernels::gdn_conv_commit(d_hist[0], d_qkv + t0 * C, C, d_counts + i, s); + strata::kernels::gdn_step_norm_multi(d_state[0], d_h[0] + t0 * C, C, d_gate + t0 * HV, d_beta + t0 * HV, + d_z + t0 * ZV, d_gamma, eps, d_y[0] + t0 * ZV, HK, HV, n, d_counts + i, s); + } + // the prompt path, all T at once + strata::prefill::gdn_conv(d_hist[1], d_qkv, d_w, d_h[1], T, eps, s); + strata::prefill::gdn_scan(d_state[1], d_h[1], d_gate, d_beta, d_y[1], T, s); + strata::prefill::gdn_out_norm(d_y[1], d_z, d_gamma, eps, d_y16, T, s); + check(cudaStreamSynchronize(s), "prefill gdn"); + const int b = bitwise_diff(down(d_h[0], qkv.size()), down(d_h[1], qkv.size()), "conv output") + + bitwise_diff(down(d_y[0], z.size()), down(d_y[1], z.size()), "y") + + bitwise_diff(down(d_state[0], state.size()), down(d_state[1], state.size()), "state") + + bitwise_diff(down(d_hist[0], hist.size()), down(d_hist[1], hist.size()), "conv history"); + if (b) std::fprintf(stderr, "prefill gdn: T %d differs\n", T); + bad += b; + for (int p = 0; p < 2; ++p) { + cudaFree(d_hist[p]); + cudaFree(d_state[p]); + cudaFree(d_h[p]); + cudaFree(d_y[p]); + } + for (float* p : {d_qkv, d_gate, d_beta, d_z}) cudaFree(p); + cudaFree(d_y16); + cudaFree(d_counts); + } + cudaFree(d_w); + cudaFree(d_gamma); + std::printf("prefill_gdn: %s\n", bad ? "MISMATCH" : "bitwise equal to the verify window's kernels (1-37 tokens)"); + return bad; +} + +// argmax_rows against sample_tokens' greedy pick and row_top_prob_split against row_top_prob, bitwise: rows of both +// vocabularies with ties at the maximum (the lowest index wins), NaN, -inf, a row with no value above -inf, and +// launches repeated on one scratch (each must leave its counters at zero) +int test_argmax(std::mt19937& rng, cudaStream_t s) { + using namespace strata::kernels; + int bad = 0; + std::normal_distribution nd(0.0f, 4.0f); + uint8_t* sa = dev(argmax_rows_scratch_bytes(8)); // one scratch for every row count, as in the engine + uint8_t* st = dev(row_top_prob_scratch_bytes(8)); + check(cudaMemset(sa, 0, argmax_rows_scratch_bytes(8)), "memset"); + check(cudaMemset(st, 0, row_top_prob_scratch_bytes(8)), "memset"); + for (const int n : {248320, 40525, 4097, 33}) { + for (const int rows : {4, 8, 1, 3, 8, 2}) { + std::vector h((size_t) rows * n); + for (float& v : h) v = nd(rng); + for (int r = 0; r < rows; ++r) { + float* row = h.data() + (size_t) r * n; + const int a = (int) (rng() % n), b = (int) (rng() % n); + row[a] = row[b] = 40.0f; // a tie at the maximum + row[(int) (rng() % n)] = std::nanf(""); // never picked + row[(int) (rng() % n)] = -INFINITY; + if (r == 2) // no value above -inf: 0 + for (int i = 0; i < n; ++i) row[i] = i % 7 ? -INFINITY : std::nanf(""); + } + float* d_l = dev(h.size()); + up(d_l, h); + int* d_ref = dev(rows); + int32_t* d_out = dev(rows); + float* d_p1 = dev(rows); + float* d_p2 = dev(rows); + SamplerParams sp; + sp.greedy = true; + sp.temperature = 0.0f; + sample_tokens(d_l, rows, n, nullptr, 0, sp, d_ref, s); + row_top_prob(d_l, rows, n, d_ref, d_p1, s); + for (int rep = 0; rep < 3; ++rep) { + argmax_rows(d_l, rows, n, sa, d_out, s); + row_top_prob_split(d_l, rows, n, d_out, d_p2, st, s); + check(cudaStreamSynchronize(s), "argmax"); + const std::vector ref = down(d_ref, rows); + const std::vector out = down(d_out, rows); + const std::vector p1 = down(d_p1, rows), p2 = down(d_p2, rows); + for (int r = 0; r < rows; ++r) + if (ref[r] != out[r] || std::memcmp(&p1[r], &p2[r], 4) != 0) { + if (bad < 5) + std::printf(" argmax n %d rows %d row %d: %d / %d, p %.9g / %.9g\n", n, rows, r, ref[r], + out[r], p1[r], p2[r]); + ++bad; + } + } + cudaFree(d_l); cudaFree(d_ref); cudaFree(d_out); cudaFree(d_p1); cudaFree(d_p2); + } + } + cudaFree(sa); + cudaFree(st); + std::printf("argmax_rows, row_top_prob_split: %s\n", + bad ? "MISMATCH" : "bitwise equal to the one-block kernels (ties, NaN, -inf, 1-8 rows on one scratch)"); + return bad; +} + +} // namespace + +int main(int argc, char** argv) { + for (int i = 1; i < argc; ++i) { + if (std::string(argv[i]) != "--selftest") { + std::fprintf(stderr, "usage: verify_parity [--selftest]\n"); + return 2; + } + } + std::mt19937 rng(20260925); + cudaStream_t s = nullptr; + check(cudaStreamCreate(&s), "stream"); + int bad = 0; + bad += test_gather_combine(rng, s); + bad += test_hit_plan(rng, s); + bad += test_mmvf_multi(rng, s); + bad += test_router_multi(rng, s); + bad += test_verify_router(rng, s); + bad += test_router_rows(rng, s); + bad += test_gr_read(rng, s); + bad += test_rope_tokens(rng, s); + bad += test_norm_rope(rng, s); + bad += test_kv_append(rng, s); + bad += test_indexer_append(rng, s); + bad += test_ple_tokens(rng, s); + bad += test_decode_attn(rng, s); + bad += test_prefill_gdn(rng, s); + bad += test_argmax(rng, s); + cudaStreamDestroy(s); + std::printf("verify_parity: %s\n", bad ? "FAIL" : "PASS"); + return bad ? 1 : 0; +} diff --git a/src/ngram/ple_reader.cpp b/src/ngram/ple_reader.cpp index ff6d5139cd6..3c989ca1a1e 100644 --- a/src/ngram/ple_reader.cpp +++ b/src/ngram/ple_reader.cpp @@ -5,7 +5,6 @@ #include #include #include -#include #include #include #include @@ -34,14 +33,16 @@ constexpr uint32_t EMPTY = 0xFFFFFFFFu; struct RowCache { uint64_t sets = 0; std::vector keys; // sets * WAYS - std::vector data; // sets * WAYS * ROW_BYTES + std::vector data; // sets * WAYS * row_bytes std::vector next; // per-set replacement pointer uint64_t used = 0; + uint32_t row_bytes = ROW_BYTES; - void init(uint64_t rows) { + void init(uint64_t rows, uint32_t bytes) { sets = rows / WAYS; + row_bytes = bytes; keys.assign(sets * WAYS, EMPTY); - data.assign(sets * WAYS * ROW_BYTES, 0); + data.assign(sets * WAYS * row_bytes, 0); next.assign(sets, 0); used = 0; } @@ -53,7 +54,7 @@ struct RowCache { if (sets == 0) return nullptr; const uint64_t s = mix(row) % sets; for (uint32_t w = 0; w < WAYS; ++w) - if (keys[s * WAYS + w] == row) return &data[(s * WAYS + w) * ROW_BYTES]; + if (keys[s * WAYS + w] == row) return &data[(s * WAYS + w) * row_bytes]; return nullptr; } void insert(uint32_t row, const uint8_t* bytes) { @@ -63,7 +64,7 @@ struct RowCache { next[s] = (uint8_t) ((w + 1) % WAYS); if (keys[s * WAYS + w] == EMPTY) ++used; keys[s * WAYS + w] = row; - std::memcpy(&data[(s * WAYS + w) * ROW_BYTES], bytes, ROW_BYTES); + std::memcpy(&data[(s * WAYS + w) * row_bytes], bytes, row_bytes); } }; @@ -95,6 +96,7 @@ struct PleReader::Impl { DirectFile file; uint64_t table_offset = 0; uint64_t n_rows = 0; + uint32_t row_bytes = ROW_BYTES; uint32_t max_inflight = 0; uint8_t* slab = nullptr; // max_inflight slots of 2 pages std::vector free_slots; @@ -119,14 +121,6 @@ struct PleReader::Impl { uint8_t* slot_buf(uint32_t s) { return slab + (size_t) s * 2 * PAGE; } bool busy() const { return free_slots.size() < max_inflight || !delayed.empty(); } - void cancel_queued() { - for (const Job& job : queue) { - auto it = tickets.find(job.ticket); - if (it != tickets.end() && it->second.pending > 0) --it->second.pending; - } - queue.clear(); - } - void record_latency(double us) { stats.read_us_sum += us; if (stats.read_us.size() < LATENCY_RING) stats.read_us.push_back((float) us); @@ -141,14 +135,7 @@ struct PleReader::Impl { queue.pop_front(); Job& j = inflight[s]; j.issued_us = now_us(); - if (!file.submit(j.offset, slot_buf(s), j.length, s, error)) { - j.uses.clear(); - auto it = tickets.find(j.ticket); - if (it != tickets.end() && it->second.pending > 0) --it->second.pending; - free_slots.push_back(s); - cancel_queued(); - return false; - } + if (!file.submit(j.offset, slot_buf(s), j.length, s, error)) return false; stats.submit_us += now_us() - j.issued_us; ++stats.reads; } @@ -157,65 +144,41 @@ struct PleReader::Impl { bool finish(const Completion& c) { const uint32_t s = (uint32_t) c.tag; - if (s >= inflight.size()) { - error = "PleReader: invalid table-read completion"; - cancel_queued(); - return false; - } Job& j = inflight[s]; if (!c.ok) { error = "PleReader: a table read failed"; - cancel_queued(); - j.uses.clear(); - auto it = tickets.find(j.ticket); - if (it != tickets.end() && it->second.pending > 0) --it->second.pending; - free_slots.push_back(s); return false; } record_latency(now_us() - j.issued_us); stats.bytes += c.bytes; const uint8_t* buf = slot_buf(s); for (const Use& u : j.uses) { - if (u.in_page + ROW_BYTES > c.bytes) { + if (u.in_page + row_bytes > c.bytes) { error = "PleReader: short read inside the table"; - cancel_queued(); - j.uses.clear(); - auto it = tickets.find(j.ticket); - if (it != tickets.end() && it->second.pending > 0) --it->second.pending; - free_slots.push_back(s); return false; } - } - for (const Use& u : j.uses) { - std::memcpy(u.dst, buf + u.in_page, ROW_BYTES); + std::memcpy(u.dst, buf + u.in_page, row_bytes); cache.insert(u.row, buf + u.in_page); } auto it = tickets.find(j.ticket); if (it != tickets.end() && it->second.pending > 0) --it->second.pending; j.uses.clear(); free_slots.push_back(s); - return error.empty() ? pump() : true; + return pump(); } /// Completions (or wake packets) just returned by `file.wait`, applying fault injection. bool process(const Completion* got, int n) { - bool ok = true; for (int i = 0; i < n; ++i) { if (got[i].tag == DirectFile::WAKE_TAG) continue; - if (got[i].tag >= inflight.size()) { - error = "PleReader: invalid table-read completion"; - cancel_queued(); - ok = false; - continue; - } if (delay_us > 0 && now_us() - inflight[(uint32_t) got[i].tag].issued_us < delay_us) { delayed.push_back(got[i]); ++stats.late_injected; continue; } - if (!finish(got[i])) ok = false; // drain the rest of this batch so no completed slot is stranded + if (!finish(got[i])) return false; } - return ok; + return true; } bool release_delayed() { @@ -274,41 +237,31 @@ PleReader::~PleReader() { } bool PleReader::open(const std::string& path, uint64_t table_offset, uint64_t n_rows, uint32_t max_inflight, - uint64_t cache_rows, std::string& err, bool io_thread) { + uint64_t cache_rows, std::string& err, bool io_thread, uint32_t row_bytes) { close(); if (max_inflight == 0 || max_inflight > 1024) { err = "PleReader: max_inflight must be 1..1024"; return false; } + if (row_bytes == 0 || row_bytes > PAGE) { err = "PleReader: a row must fit a page"; return false; } if (!impl_->file.open(path, err)) return false; - if (table_offset + n_rows * (uint64_t) ROW_BYTES > impl_->file.size()) { + if (table_offset + n_rows * (uint64_t) row_bytes > impl_->file.size()) { err = "PleReader: the table extends past the end of " + path; close(); return false; } impl_->table_offset = table_offset; impl_->n_rows = n_rows; + impl_->row_bytes = row_bytes; impl_->max_inflight = max_inflight; impl_->slab = (uint8_t*) DirectFile::alloc_aligned((size_t) max_inflight * 2 * PAGE); if (impl_->slab == nullptr) { err = "PleReader: cannot allocate read buffers"; close(); return false; } impl_->inflight.assign(max_inflight, Job{}); impl_->free_slots.clear(); for (uint32_t s = max_inflight; s-- > 0;) impl_->free_slots.push_back(s); - impl_->cache.init(cache_rows); + impl_->cache.init(cache_rows, row_bytes); impl_->error.clear(); reset_stats(); impl_->stop = false; impl_->threaded = io_thread; - if (io_thread) { - try { - impl_->worker = std::thread([this] { impl_->worker_loop(); }); - } catch (const std::exception& e) { - err = std::string("PleReader: cannot create I/O worker: ") + e.what(); - close(); - return false; - } catch (...) { - err = "PleReader: cannot create I/O worker"; - close(); - return false; - } - } + if (io_thread) impl_->worker = std::thread([this] { impl_->worker_loop(); }); return true; } @@ -343,7 +296,7 @@ void PleReader::close() { m.delayed.clear(); m.inflight.clear(); m.free_slots.clear(); - m.cache.init(0); + m.cache.init(0, ROW_BYTES); m.threaded = false; m.stop = false; } @@ -360,20 +313,20 @@ PleReader::Ticket PleReader::issue(const uint32_t* rows, size_t n, uint8_t* out_ std::unordered_map by_page; // aligned offset -> index in `jobs` std::vector jobs; for (size_t i = 0; i < n; ++i) { - uint8_t* dst = out_raw + i * ROW_BYTES; + uint8_t* dst = out_raw + i * m.row_bytes; ++m.stats.requests; if (rows[i] >= m.n_rows) { - std::memset(dst, 0, ROW_BYTES); + std::memset(dst, 0, m.row_bytes); continue; } if (const uint8_t* hit = m.cache.find(rows[i])) { - std::memcpy(dst, hit, ROW_BYTES); + std::memcpy(dst, hit, m.row_bytes); ++m.stats.cache_hits; continue; } - const uint64_t at = m.table_offset + (uint64_t) rows[i] * ROW_BYTES; + const uint64_t at = m.table_offset + (uint64_t) rows[i] * m.row_bytes; const uint64_t first = at / PAGE * PAGE; - const uint32_t length = (uint32_t) ((at + ROW_BYTES - 1) / PAGE * PAGE - first + PAGE); + const uint32_t length = (uint32_t) ((at + m.row_bytes - 1) / PAGE * PAGE - first + PAGE); auto f = by_page.find(first); if (f != by_page.end()) { Job& j = jobs[f->second]; @@ -393,11 +346,10 @@ PleReader::Ticket PleReader::issue(const uint32_t* rows, size_t n, uint8_t* out_ // Sorted by offset: prefill chunks then read the SSD in near-sequential order. std::sort(jobs.begin(), jobs.end(), [](const Job& a, const Job& b) { return a.offset < b.offset; }); ts.pending = (uint32_t) jobs.size(); - const bool has_jobs = ts.pending > 0; for (Job& j : jobs) m.queue.push_back(std::move(j)); if (m.threaded) { lk.unlock(); - if (has_jobs) { + if (ts.pending > 0) { m.cv_work.notify_one(); m.file.wake(); // in case the worker is blocked in the port } @@ -414,26 +366,21 @@ bool PleReader::collect(Ticket t, std::string& err) { std::unique_lock lk(m.mu); auto it = m.tickets.find(t.id); if (it == m.tickets.end()) { err = "PleReader: unknown ticket"; return false; } - m.cv_done.wait(lk, [&] { - const auto current = m.tickets.find(t.id); - return current == m.tickets.end() || current->second.pending == 0; - }); - it = m.tickets.find(t.id); - if (it == m.tickets.end()) { err = "PleReader: unknown ticket"; return false; } + m.cv_done.wait(lk, [&] { return !m.error.empty() || m.tickets[t.id].pending == 0; }); if (!m.error.empty()) { err = m.error; return false; } m.stats.wait_us += now_us() - start; - m.tickets.erase(it); + m.tickets.erase(t.id); return true; } auto it = m.tickets.find(t.id); if (it == m.tickets.end()) { err = "PleReader: unknown ticket"; return false; } while (it->second.pending > 0) { - const bool drained = m.drain(-1); - if (!drained && m.error.empty()) m.error = "PleReader: read failed"; + if (!m.error.empty() || !m.drain(-1)) { + err = m.error.empty() ? "PleReader: read failed" : m.error; + return false; + } it = m.tickets.find(t.id); - if (it == m.tickets.end()) { err = "PleReader: unknown ticket"; return false; } } - if (!m.error.empty()) { err = m.error; return false; } m.stats.wait_us += now_us() - start; m.tickets.erase(it); return true; diff --git a/src/prefill/experts.cpp b/src/prefill/experts.cpp new file mode 100644 index 00000000000..392e035029d --- /dev/null +++ b/src/prefill/experts.cpp @@ -0,0 +1,596 @@ +// src/prefill/experts.cpp - see include/strata/prefill/experts.hpp. +#include "strata/prefill/experts.hpp" + +#include "strata/kernels/cpu/expert_layout.hpp" +#include "strata/kernels/iq_kernels.hpp" +#include "strata/prefill/gemm.hpp" +#include "strata/prefill/kernels.hpp" +#include "strata/prefill/moe_mmq.hpp" + +#include +#include +#include +#include + +namespace strata::prefill { +namespace { + +constexpr int64_t N = 2560, FF2 = 1280, FF = 640, K = 10; +// Batches of consecutive ids, at most G experts and ROWS rows (an expert with more rows is a batch of its own, in +// pieces of ROWS). Per-expert launches made the host, not the GPU, the limit (~10 calls for each of ~500 experts a +// layer). +constexpr int G = 16; +constexpr int64_t ROWS = 4096; +constexpr int SLOTS = 2; // batches in staging: the copies of one overlap the work on the other +constexpr size_t GEMM_WS = 32u << 20; // cuBLAS workspace +constexpr int NPTR = 7; // per-expert device pointers: blob, xs, gate/up, gu, h, down, d +constexpr int PRE_GROUP = 16; // prefetched blobs per copy event +constexpr size_t MMQ_TAIL = 4096; // zeroed after a batch's gathered experts: MMQ reads past the last one + +struct OnDevice { // a remote runner's GPU current for one call, the main one again afterwards + int main; + bool set; + OnDevice(int dev, int main_dev, bool remote) : main(main_dev), set(remote) { if (set) cudaSetDevice(dev); } + ~OnDevice() { if (set) cudaSetDevice(main); } +}; + +// Bump allocation from a region; without one it only counts (bytes_needed). +struct Alloc { + uint8_t* base = nullptr; + uint64_t cap = 0, used = 0; + template T* take(size_t n, bool& ok) { + const uint64_t bytes = ((uint64_t) n * sizeof(T) + 255) & ~255ull; + T* p = base != nullptr ? (T*) (base + used) : nullptr; + used += bytes; + if (base != nullptr && used > cap) ok = false; + return p; + } +}; + +size_t blob_cap() { return ((size_t) strata::kernels::cpu::expert_layout().max_blob + 255) & ~(size_t) 255; } + +// the pinned upload block: NPTR x n_expert pointers, then the adds (at most 3 per row, 2 per expert) with each +// batch's row bounds (2 per expert), rows, weights +size_t adds_cap(int64_t chunk, int64_t n_expert) { return (size_t) (chunk * K * 3 + n_expert * 4); } + +// a native layer whose expert types llama.cpp's MMQ covers takes it (STRATA_PREFILL_MMQ=0: FP16 for every layer) +bool mmq_layer(const strata::kernels::cpu::NativeFmt& f) { + static const bool on = [] { + const char* e = std::getenv("STRATA_PREFILL_MMQ"); + return e == nullptr || std::atoi(e) != 0; + }(); + return on && strata::kernels::cpu::expert_layout().native && mmq::supported(f.gu_type) && mmq::supported(f.d_type); +} +bool mmq_any() { + const auto& lay = strata::kernels::cpu::expert_layout(); + for (const auto& f : lay.fmt) if (mmq_layer(f)) return true; + return false; +} +size_t up_bytes(int64_t chunk, int64_t n_expert) { + return (size_t) NPTR * (size_t) n_expert * sizeof(void*) + (adds_cap(chunk, n_expert) + (size_t) chunk * K * 2) * 4; +} + +} // namespace + +struct ExpertRunner::Impl { + int dev = -1, main = 0; + bool remote = false, streams = false; + core::ExpertSource* src = nullptr; + const core::ExpertCache* cache = nullptr; + const int32_t* res = nullptr; + int64_t n_expert = 0, max_chunk = 0, T = 0; // T: the chunk the buffers are bound for + cudaStream_t s = nullptr, copy = nullptr; + cudaEvent_t copied[SLOTS] = {}, used[SLOTS] = {}; + std::vector piece_ev; // remote: a piece of the sums is in host_sum + cudaEvent_t t0 = nullptr, t1 = nullptr; // the previous layer's GPU time (remote), read at the next one + bool timed = false, live[SLOTS] = {}; + Gemm gemm; + std::unique_ptr mmq_ctx; // the MMQ layers' launches (null: none) + void* owned = nullptr; // the buffers' own allocation (no region) + uint16_t *input = nullptr, *xs = nullptr, *hh = nullptr, *dq_gu = nullptr, *dq_d = nullptr; + float *sum = nullptr, *w = nullptr, *gu = nullptr, *d = nullptr; + int32_t* rows_src = nullptr; + int32_t* adds = nullptr; // per batch: its tokens, their row lists' starts, the row lists (moe_gather_add), bounds + uint8_t* xq = nullptr; // MMQ: a batch's rows as q8_1 (gate/up's input, then down's) + int32_t* ident = nullptr; // MMQ: the identity row map + std::vector mark, first, count; + uint8_t* stage[SLOTS] = {}; + void** ptrs = nullptr; // NPTR arrays of n_expert device pointers + // a layer's uploads, staged in pinned memory (a pageable copy blocks the host behind queued expert copies): + // [pointers | adds | rows | weights], for chunks of up to max_chunk tokens + uint8_t* h_up = nullptr; + cudaEvent_t up_done = nullptr; // the previous layer's uploads have been read + bool up_live = false; + uint16_t *h_input = nullptr, *h_sum = nullptr; + // a layer's blobs copied ahead: at pre_off[e] in `pre` (-1: not there), covered by event pre_ev[pre_grp[e]] + uint8_t* pre = nullptr; + uint64_t pre_cap = 0; + int64_t pre_layer = -1; + std::vector pre_off; + std::vector pre_grp; + std::vector pre_ev; + cudaEvent_t pre_used = nullptr; // the last read of the area + bool pre_live = false; + + // the same sequence counted (bytes_needed) or carved + static void carve(Alloc& a, int64_t T, int64_t n_expert, bool remote, bool streams, Impl* m, void** ws, bool& ok) { + const size_t t = (size_t) T; + uint16_t* input = remote ? a.take(t * N, ok) : nullptr; + float* sum = remote ? a.take(t * N, ok) : nullptr; + int32_t* rows_src = a.take(t * K, ok); + float* w = a.take(t * K, ok); + int32_t* adds = a.take(adds_cap(T, n_expert), ok); + uint16_t* xs = a.take((size_t) ROWS * N, ok); + float* gu = a.take((size_t) ROWS * FF2, ok); + uint16_t* hh = a.take((size_t) ROWS * FF, ok); + float* d = a.take((size_t) ROWS * N, ok); + uint16_t* dq_gu = a.take((size_t) G * FF2 * N, ok); + uint16_t* dq_d = a.take((size_t) G * N * FF, ok); + uint8_t* stage[SLOTS] = {}; + if (streams) + for (auto& p : stage) p = a.take((size_t) G * blob_cap(), ok); + void** ptrs = a.take((size_t) NPTR * (size_t) n_expert, ok); + *ws = a.take(GEMM_WS, ok); + const bool q = mmq_any(); + uint8_t* xq = q ? a.take(mmq::q8_bytes(ROWS, N), ok) : nullptr; + int32_t* ident = q ? a.take((size_t) ROWS, ok) : nullptr; + if (m == nullptr) return; + m->xq = xq; m->ident = ident; + m->input = input; m->sum = sum; m->rows_src = rows_src; m->w = w; m->adds = adds; + m->xs = xs; m->gu = gu; m->hh = hh; m->d = d; m->dq_gu = dq_gu; m->dq_d = dq_d; + for (int i = 0; i < SLOTS; ++i) m->stage[i] = stage[i]; + m->ptrs = ptrs; + } +}; + +ExpertRunner::ExpertRunner() : impl_(new Impl) {} + +ExpertRunner::~ExpertRunner() { + Impl& m = *impl_; + if (m.dev < 0) return; + OnDevice on(m.dev, m.main, m.remote); + if (m.s) cudaStreamSynchronize(m.s); + if (m.copy) cudaStreamSynchronize(m.copy); + for (int i = 0; i < SLOTS; ++i) { + if (m.copied[i]) cudaEventDestroy(m.copied[i]); + if (m.used[i]) cudaEventDestroy(m.used[i]); + } + for (cudaEvent_t e : {m.t0, m.t1, m.pre_used, m.up_done}) if (e) cudaEventDestroy(e); + for (cudaEvent_t e : m.pre_ev) if (e) cudaEventDestroy(e); + for (cudaEvent_t e : m.piece_ev) if (e) cudaEventDestroy(e); + if (m.h_input) cudaFreeHost(m.h_input); + if (m.h_sum) cudaFreeHost(m.h_sum); + if (m.h_up) cudaFreeHost(m.h_up); + if (m.owned) cudaFree(m.owned); + if (m.copy) cudaStreamDestroy(m.copy); + if (m.remote && m.s) cudaStreamDestroy(m.s); +} + +uint64_t ExpertRunner::bytes_needed(int64_t chunk, int64_t n_expert, bool remote, bool streams) { + Alloc a; + bool ok = true; + void* ws = nullptr; + Impl::carve(a, chunk, n_expert, remote, streams, nullptr, &ws, ok); + return a.used; +} + +bool ExpertRunner::init(int device, int main_device, void* stream, core::ExpertSource* src, const core::ExpertCache* cache, + const int32_t* res, int64_t n_expert, int64_t max_chunk, bool streams, std::string& err) { + Impl& m = *impl_; + const auto& lay = strata::kernels::cpu::expert_layout(); + if (stream == nullptr && !lay.native) { + err = "prompt experts on a second GPU: need a native pack"; + return false; + } + m.dev = device; + m.main = main_device; + m.remote = stream == nullptr; + m.streams = streams; + m.src = src; + m.cache = cache; + m.res = res; + m.n_expert = n_expert; + m.max_chunk = max_chunk; + m.mark.assign((size_t) max_chunk, -1); + m.first.resize((size_t) max_chunk); + m.count.resize((size_t) max_chunk); + OnDevice on(device, main_device, m.remote); + bool ok = cudaEventCreateWithFlags(&m.up_done, cudaEventDisableTiming) == cudaSuccess && + cudaEventCreate(&m.t0) == cudaSuccess && cudaEventCreate(&m.t1) == cudaSuccess && + cudaHostAlloc((void**) &m.h_up, up_bytes(max_chunk, n_expert), cudaHostAllocDefault) == cudaSuccess; + if (m.remote) ok = ok && cudaStreamCreateWithFlags(&m.s, cudaStreamNonBlocking) == cudaSuccess; + else m.s = (cudaStream_t) stream; + if (streams) { + ok = ok && cudaStreamCreateWithFlags(&m.copy, cudaStreamNonBlocking) == cudaSuccess; + for (int i = 0; ok && i < SLOTS; ++i) + ok = cudaEventCreateWithFlags(&m.copied[i], cudaEventDisableTiming) == cudaSuccess && + cudaEventCreateWithFlags(&m.used[i], cudaEventDisableTiming) == cudaSuccess; + m.pre_ev.assign((size_t) ((n_expert + PRE_GROUP - 1) / PRE_GROUP), nullptr); + for (auto& e : m.pre_ev) ok = ok && cudaEventCreateWithFlags(&e, cudaEventDisableTiming) == cudaSuccess; + ok = ok && cudaEventCreateWithFlags(&m.pre_used, cudaEventDisableTiming) == cudaSuccess; + m.pre_off.assign((size_t) n_expert, -1); + m.pre_grp.assign((size_t) n_expert, 0); + } + if (!ok || !m.gemm.init(m.s, err, true)) { + if (err.empty()) err = "prompt experts: streams or events"; + return false; + } + if (mmq_any()) m.mmq_ctx = std::make_unique(); + // the input and the sums: pinned and portable, both GPUs copy them + if (m.remote && (cudaHostAlloc((void**) &m.h_input, (size_t) max_chunk * N * 2, cudaHostAllocPortable) != cudaSuccess || + cudaHostAlloc((void**) &m.h_sum, (size_t) max_chunk * N * 2, cudaHostAllocPortable) != cudaSuccess)) { + err = "prompt experts: pinned input and sum buffers"; + return false; + } + return true; +} + +bool ExpertRunner::bind(void* region, uint64_t bytes, int64_t chunk, std::string& err, void* area, uint64_t area_bytes) { + Impl& m = *impl_; + OnDevice on(m.dev, m.main, m.remote); + m.pre_layer = -1; + if (region == nullptr) { + if (m.owned != nullptr) return true; + chunk = m.max_chunk; + bytes = bytes_needed(chunk, m.n_expert, m.remote, m.streams); + if (cudaMalloc(&m.owned, bytes) != cudaSuccess) { + m.owned = nullptr; + err = "prompt experts: " + std::to_string(bytes >> 20) + " MiB of device buffers"; + return false; + } + region = m.owned; + } + if (chunk > m.max_chunk) { + err = "prompt experts: a chunk longer than the one they were set up for"; + return false; + } + Alloc a; + a.base = (uint8_t*) region; + a.cap = bytes; + bool ok = true; + void* ws = nullptr; + Impl::carve(a, chunk, m.n_expert, m.remote, m.streams, &m, &ws, ok); + if (!ok) { + err = "prompt experts: their buffers for a chunk of " + std::to_string(chunk) + " tokens do not fit"; + return false; + } + m.gemm.set_buffers(nullptr, 0, ws, GEMM_WS); + m.T = chunk; + if (area == nullptr && region != m.owned && bytes > a.used) { + area = (uint8_t*) region + a.used; + area_bytes = bytes - a.used; + } + m.pre = m.streams ? (uint8_t*) area : nullptr; + m.pre_cap = m.pre != nullptr ? area_bytes : 0; + return true; +} + +uint64_t ExpertRunner::bytes_for(int64_t chunk, const int32_t* skip) const { + const Impl& m = *impl_; + return bytes_needed(chunk, m.n_expert, m.remote, m.streams) + (chunk >= kPrefetchMin ? prefetch_need(skip) : 0); +} + +uint64_t ExpertRunner::prefetch_need(const int32_t* skip) const { + const Impl& m = *impl_; + if (!m.streams) return 0; + const auto& lay = strata::kernels::cpu::expert_layout(); + uint64_t most = 0; + for (int64_t l = 0; l < lay.n_layers; ++l) { + const size_t stride = ((size_t) lay.blob_bytes(l) + 255) & ~(size_t) 255; + uint64_t sum = 0; + for (int64_t e = 0; e < m.n_expert; ++e) { + const size_t i = (size_t) (l * m.n_expert + e); + if ((m.res == nullptr || m.res[i] < 0) && (skip == nullptr || skip[i] < 0)) sum += stride; + } + most = std::max(most, sum); + } + return most + most / 10; // the slots lent for the prompt make a few more experts stream +} + +bool ExpertRunner::prefetch(int64_t layer, const std::vector& experts, std::string& err) { + Impl& m = *impl_; + if (m.pre_cap == 0) return true; + OnDevice on(m.dev, m.main, m.remote); + const auto& lay = strata::kernels::cpu::expert_layout(); + const size_t bb = (size_t) lay.blob_bytes(layer); + const size_t stride = (bb + 255) & ~(size_t) 255; + std::fill(m.pre_off.begin(), m.pre_off.end(), -1); + m.pre_layer = layer; + if (m.pre_live) cudaStreamWaitEvent(m.copy, m.pre_used, 0); // the previous layer has read the area + uint64_t at = 0; + int32_t grp = 0, in_grp = 0; + const size_t n = experts.size(); + for (size_t j = 0; j < n;) { + const int32_t e = experts[j]; + const bool held = m.res && m.cache && m.res[(size_t) layer * (size_t) m.n_expert + (size_t) e] >= 0; + if (held) { ++j; continue; } + if (at + stride > m.pre_cap) break; + // a run: the next experts' blobs follow this one in the arena, and fit + const uint8_t* b0 = m.src->blob(layer, e); + if (b0 == nullptr) { err = "prompt experts: the expert source has no blob"; return false; } + size_t k = j + 1; + if (m.src->pinned(layer, e) && stride == bb) + while (k < n && experts[(size_t) k] == e + (int32_t) (k - j) && at + (k - j + 1) * stride <= m.pre_cap && + !(m.res && m.cache && m.res[(size_t) layer * (size_t) m.n_expert + (size_t) experts[k]] >= 0) && + m.src->pinned(layer, experts[k]) && m.src->blob(layer, experts[k]) == b0 + (k - j) * bb && + in_grp + (int32_t) (k - j) < PRE_GROUP) ++k; + cudaMemcpyAsync(m.pre + at, b0, (k - j) * bb, cudaMemcpyHostToDevice, m.copy); // pageable blobs: staged + for (size_t i = j; i < k; ++i) { + m.pre_off[(size_t) experts[i]] = (int64_t) (at + (i - j) * stride); + m.pre_grp[(size_t) experts[i]] = grp; + } + at += (k - j) * stride; + experts_prefetched += (int64_t) (k - j); + experts_streamed += (int64_t) (k - j); + in_grp += (int32_t) (k - j); + j = k; + if (in_grp >= PRE_GROUP) { + cudaEventRecord(m.pre_ev[(size_t) grp++], m.copy); + in_grp = 0; + } + } + if (in_grp > 0) cudaEventRecord(m.pre_ev[(size_t) grp], m.copy); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { err = std::string("prompt experts: prefetch: ") + cudaGetErrorString(e); return false; } + return true; +} + +uint16_t* ExpertRunner::host_input() const { return impl_->h_input; } +uint16_t* ExpertRunner::host_sum() const { return impl_->h_sum; } +cudaEvent_t ExpertRunner::piece_done(int64_t k) const { return impl_->piece_ev[(size_t) k]; } + +void ExpertRunner::stage_input(int64_t t0, int64_t rows, cudaEvent_t ready) { + Impl& m = *impl_; + OnDevice on(m.dev, m.main, m.remote); + cudaStreamWaitEvent(m.s, ready, 0); // behind the previous layer's sums, which pass through `input` + cudaMemcpyAsync(m.input + t0 * N, m.h_input + t0 * N, (size_t) rows * N * 2, cudaMemcpyHostToDevice, m.s); +} + +bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector& experts, const std::vector& off, + const std::vector& src, const std::vector& w, int64_t piece, + const uint16_t* input, float* sum, std::string& err) { + Impl& m = *impl_; + const size_t n_exp = experts.size(), n_rows = src.size(); + if (T > m.T || off.size() != n_exp + 1 || w.size() != n_rows || (size_t) off.back() != n_rows || + n_rows > (size_t) T * K || n_exp > (size_t) m.n_expert || (!m.remote && (input == nullptr || sum == nullptr))) { + err = "prompt experts: a layer's share is out of range"; + return false; + } + OnDevice on(m.dev, m.main, m.remote); + const auto h0 = std::chrono::steady_clock::now(); + if (m.timed && cudaEventQuery(m.t1) == cudaSuccess) { // the previous layer (the main GPU waited for it) + float ms = 0; + cudaEventElapsedTime(&ms, m.t0, m.t1); + ms_gpu += ms; + } + m.timed = false; + const auto& lay = strata::kernels::cpu::expert_layout(); + const auto& f = lay.fmt.empty() ? strata::kernels::cpu::NativeFmt{} : lay.fmt[(size_t) layer]; + const size_t bb = (size_t) lay.blob_bytes(layer); + const size_t stride = (bb + 255) & ~(size_t) 255; // staging: a run of arena blobs lands as one copy when equal + const bool use_mmq = m.mmq_ctx != nullptr && mmq_layer(f); + const size_t gub = use_mmq ? mmq::matrix_bytes(f.gu_type, FF2, N) : 0; + const size_t db = use_mmq ? mmq::matrix_bytes(f.d_type, N, FF) : 0; + auto held = [&](int32_t e) -> int32_t { + return m.res && m.cache ? m.res[(size_t) layer * (size_t) m.n_expert + (size_t) e] : -1; + }; + auto pre = [&](int32_t e) -> int64_t { return m.pre_layer == layer ? m.pre_off[(size_t) e] : -1; }; + // ---- the batches, and every expert's device pointers (one upload) + std::vector> batches; + for (size_t j0 = 0; j0 < n_exp;) { + size_t j1 = j0 + 1; + int64_t rows = off[j0 + 1] - off[j0]; + if (rows <= ROWS) + while (j1 < n_exp && j1 - j0 < (size_t) G && rows + (off[j1 + 1] - off[j1]) <= ROWS) rows += off[j1 + 1] - off[j1++]; + batches.emplace_back(j0, j1); + j0 = j1; + } + if (m.up_live) cudaEventSynchronize(m.up_done); // the host block is free again + const size_t E = (size_t) m.n_expert; // array stride + void** P = (void**) m.h_up; + int32_t* h_adds = (int32_t*) (P + NPTR * E); + int32_t* h_rows = h_adds + adds_cap(m.max_chunk, m.n_expert); + float* h_w = (float*) (h_rows + (size_t) m.max_chunk * K); + std::vector rows_of(n_exp); + for (size_t b = 0; b < batches.size(); ++b) { + const auto [j0, j1] = batches[b]; + size_t streamed = 0; + for (size_t j = j0; j < j1; ++j) { + const int64_t r = off[j] - off[j0]; + const int32_t slot = held(experts[j]); + if (slot < 0 && !m.streams) { + err = "prompt experts: an expert its cache does not hold, and no staging"; + return false; + } + P[0 * E + j] = slot >= 0 ? (void*) m.cache->device_slot(slot) + : pre(experts[j]) >= 0 ? (void*) (m.pre + pre(experts[j])) + : (void*) (m.stage[b % SLOTS] + streamed++ * stride); + P[1 * E + j] = m.xs + r * N; + P[2 * E + j] = m.dq_gu + (j - j0) * (size_t) (FF2 * N); + P[3 * E + j] = m.gu + r * FF2; + P[4 * E + j] = m.hh + r * FF; + P[5 * E + j] = m.dq_d + (j - j0) * (size_t) (N * FF); + P[6 * E + j] = m.d + r * N; + rows_of[j] = off[j + 1] - off[j]; + } + } + // per batch, its tokens and each one's rows in order (a token may be routed to several of the batch's experts): + // [tokens | starts (n + 1) | rows] at adds_at[b] + std::vector adds_at(batches.size()), bounds_at(batches.size()); + std::vector ntok(batches.size()); + size_t n_adds = 0; + for (size_t b = 0; b < batches.size(); ++b) { + const int32_t r0 = off[batches[b].first], r1 = off[batches[b].second]; + std::vector toks; + for (int32_t r = r0; r < r1; ++r) { + const int32_t t = src[(size_t) r]; + if (m.mark[(size_t) t] != (int32_t) b) { + m.mark[(size_t) t] = (int32_t) b; + m.first[(size_t) t] = (int32_t) toks.size(); + m.count[(size_t) toks.size()] = 0; + toks.push_back(t); + } + ++m.count[(size_t) m.first[(size_t) t]]; + } + const size_t nt = toks.size(), at = n_adds; + adds_at[b] = at; + ntok[b] = (int32_t) nt; + n_adds = at + nt + (nt + 1) + (size_t) (r1 - r0); + int32_t* h_tok = h_adds + at; + int32_t* h_start = h_tok + nt; + int32_t* h_list = h_start + nt + 1; + h_start[0] = 0; + for (size_t i = 0; i < nt; ++i) { + h_tok[i] = toks[i]; + h_start[i + 1] = h_start[i] + m.count[i]; + m.count[i] = h_start[i]; // the fill cursor + } + for (int32_t r = r0; r < r1; ++r) h_list[m.count[(size_t) m.first[(size_t) src[(size_t) r]]]++] = r; + for (int32_t t : toks) m.mark[(size_t) t] = -1; + bounds_at[b] = n_adds; // MMQ: the batch's experts' rows, from the batch's first + for (size_t j = batches[b].first; j <= batches[b].second; ++j) h_adds[n_adds++] = off[j] - r0; + } + if (m.remote) cudaEventRecord(m.t0, m.s); + const uint16_t* in = m.remote ? m.input : input; + float* out = m.remote ? m.sum : sum; + std::copy(src.begin(), src.end(), h_rows); + std::copy(w.begin(), w.end(), h_w); + if (n_adds > 0) cudaMemcpyAsync(m.adds, h_adds, n_adds * 4, cudaMemcpyHostToDevice, m.s); + cudaMemcpyAsync(m.ptrs, P, NPTR * E * sizeof(void*), cudaMemcpyHostToDevice, m.s); + if (n_rows > 0) { + cudaMemcpyAsync(m.rows_src, h_rows, n_rows * 4, cudaMemcpyHostToDevice, m.s); + cudaMemcpyAsync(m.w, h_w, n_rows * 4, cudaMemcpyHostToDevice, m.s); + } + if (use_mmq) mmq::iota(m.ident, ROWS, m.s); // each layer: the local runner's buffers share the dense steps' + cudaEventRecord(m.up_done, m.s); + m.up_live = true; + if (m.remote) cudaMemsetAsync(m.sum, 0, (size_t) T * N * 4, m.s); + // ---- staging: the streamed experts of batch b into slot b % SLOTS, a copy per run of adjacent arena blobs + auto stage_batch = [&](size_t b) { + const auto [j0, j1] = batches[b]; + const int sl = (int) (b % SLOTS); + bool any = false; + for (size_t j = j0; j < j1;) { + const int32_t e = experts[j]; + if (held(e) >= 0 || pre(e) >= 0) { ++j; continue; } + if (!any && m.live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0); // its previous batch is dequantized + any = true; + const uint8_t* b0 = m.src->blob(layer, e); + uint8_t* dst = (uint8_t*) P[0 * E + j]; + size_t k = j + 1; // the run: the next experts' blobs follow this one in the arena and in staging + if (m.src->pinned(layer, e) && stride == bb) + while (k < j1 && held(experts[k]) < 0 && pre(experts[k]) < 0 && m.src->pinned(layer, experts[k]) && + m.src->blob(layer, experts[k]) == b0 + (k - j) * bb && + (uint8_t*) P[0 * E + k] == dst + (k - j) * bb) ++k; + cudaMemcpyAsync(dst, b0, (k - j) * bb, cudaMemcpyHostToDevice, m.copy); // pageable blobs: staged + experts_streamed += (int64_t) (k - j); + j = k; + } + if (any) { + cudaEventRecord(m.copied[sl], m.copy); + m.live[sl] = true; + } + return any; + }; + std::vector has_stream(batches.size(), 0); + if (!batches.empty()) has_stream[0] = stage_batch(0); + auto** dp = (const uint16_t**) m.ptrs; // device arrays by index + for (size_t b = 0; b < batches.size(); ++b) { + if (b + 1 < batches.size()) has_stream[b + 1] = stage_batch(b + 1); // the next batch copies meanwhile + const auto [j0, j1] = batches[b]; + const int n = (int) (j1 - j0); + const int sl = (int) (b % SLOTS); + if (has_stream[b]) cudaStreamWaitEvent(m.s, m.copied[sl], 0); + int32_t last_grp = -1; // the batch's prefetched blobs have landed with their last group + for (size_t j = j0; j < j1; ++j) + if (held(experts[j]) < 0 && pre(experts[j]) >= 0) last_grp = std::max(last_grp, m.pre_grp[(size_t) experts[j]]); + if (last_grp >= 0) cudaStreamWaitEvent(m.s, m.pre_ev[(size_t) last_grp], 0); + const int64_t r0 = off[j0], rows = off[j1] - r0; + const bool q = use_mmq && rows <= ROWS; // an expert with more rows is multiplied in pieces, in FP16 + if (q) { + // the batch's GGUF blocks side by side in the dequantization buffers' place (MMQ reads the experts at + // one stride), a zeroed tail after them + for (size_t j = j0; j < j1; ++j) { + const uint8_t* bl = (const uint8_t*) P[0 * E + j]; + mmq::gather_native(bl, bl + f.up_off, gub / 2, bl + f.down_off, db, (uint8_t*) m.dq_gu + (j - j0) * gub, + (uint8_t*) m.dq_d + (j - j0) * db, m.s); + } + cudaMemsetAsync((uint8_t*) m.dq_gu + (size_t) n * gub, 0, MMQ_TAIL, m.s); + cudaMemsetAsync((uint8_t*) m.dq_d + (size_t) n * db, 0, MMQ_TAIL, m.s); + } else if (lay.native) + strata::kernels::iq_dequant_experts_f16(f.gu_type, f.d_type, (const uint8_t* const*) (m.ptrs + 0 * E + j0), n, + f.up_off, f.down_off, f.n_ff, f.n_embd, m.dq_gu, m.dq_d, m.s); + else + for (size_t j = j0; j < j1; ++j) + blob_dequant_f16((const uint8_t*) P[0 * E + j], (uint16_t*) P[2 * E + j], (uint16_t*) P[5 * E + j], m.s); + if (has_stream[b]) cudaEventRecord(m.used[sl], m.s); + if (last_grp >= 0) { + cudaEventRecord(m.pre_used, m.s); + m.pre_live = true; + } + for (size_t j = j0; j < j1; ++j) if (held(experts[j]) >= 0) ++experts_resident; + if (rows > ROWS) { // one expert: its rows in pieces + for (int64_t p = 0; p < rows; p += ROWS) { + const int64_t nr = std::min(ROWS, rows - p); + gather_rows16(in, m.rows_src + r0 + p, m.xs, nr, N, m.s); + m.gemm.f16(m.xs, m.dq_gu, m.gu, nr, FF2, N); + swiglu_interleaved(m.gu, m.hh, nr, m.s); + m.gemm.f16(m.hh, m.dq_d, m.d, nr, N, FF); + moe_scatter_add(out, m.d, m.w + r0 + p, m.rows_src + r0 + p, nr, m.s); + } + continue; + } + if (q) { + // the rows in FP32 in down's output place, as q8_1 for the gate/up product; h in the FP16 rows' place + gather_rows16_f32(in, m.rows_src + r0, m.d, rows, N, m.s); + mmq::quantize(m.d, nullptr, m.xq, f.gu_type, N, N, rows, m.s); + int64_t maxr = 0; + for (size_t j = j0; j < j1; ++j) maxr = std::max(maxr, off[j + 1] - off[j]); + mmq::Product p; + p.w = m.dq_gu; p.type = f.gu_type; p.w_rows = FF2; p.w_cols = N; p.expert_bytes = gub; p.n = n; + p.xq = m.xq; p.bounds = m.adds + bounds_at[b]; p.ids = m.ident; p.total_rows = rows; p.max_rows = maxr; + p.dst = m.gu; p.ld_dst = FF2; + m.mmq_ctx->run(p, m.s); + float* h = (float*) m.xs; + mmq::swiglu(m.gu, h, rows, FF, false, m.s); + mmq::quantize(h, nullptr, m.xq, f.d_type, FF, FF, rows, m.s); + p.w = m.dq_d; p.type = f.d_type; p.w_rows = N; p.w_cols = FF; p.expert_bytes = db; + p.dst = m.d; p.ld_dst = N; + m.mmq_ctx->run(p, m.s); + const int32_t* a = m.adds + adds_at[b]; + moe_gather_add(out, m.d, r0, m.w, a, a + ntok[b], a + 2 * ntok[b] + 1, ntok[b], m.s); + continue; + } + gather_rows16(in, m.rows_src + r0, m.xs, rows, N, m.s); + m.gemm.f16_grouped(dp + 1 * E + j0, dp + 2 * E + j0, (float* const*) (m.ptrs + 3 * E + j0), + (const uint16_t* const*) (P + 1 * E + j0), (const uint16_t* const*) (P + 2 * E + j0), + (float* const*) (P + 3 * E + j0), rows_of.data() + j0, n, FF2, N); + swiglu_interleaved(m.gu, m.hh, rows, m.s); + m.gemm.f16_grouped(dp + 4 * E + j0, dp + 5 * E + j0, (float* const*) (m.ptrs + 6 * E + j0), + (const uint16_t* const*) (P + 4 * E + j0), (const uint16_t* const*) (P + 5 * E + j0), + (float* const*) (P + 6 * E + j0), rows_of.data() + j0, n, N, FF); + const int32_t* a = m.adds + adds_at[b]; + moe_gather_add(out, m.d, r0, m.w, a, a + ntok[b], a + 2 * ntok[b] + 1, ntok[b], m.s); + } + if (m.remote) { // the sums as FP16 in the input's place (half the bytes over the first GPU's x4 link), in pieces + const int64_t pc = piece > 0 ? piece : T; + for (int64_t t0 = 0, k = 0; t0 < T; t0 += pc, ++k) { + const int64_t n = std::min(pc, T - t0); + if ((size_t) k == m.piece_ev.size()) { + m.piece_ev.push_back(nullptr); + cudaEventCreateWithFlags(&m.piece_ev.back(), cudaEventDisableTiming); + } + sums_to_f16(m.sum + t0 * N, m.input + t0 * N, n * N, m.s); + cudaMemcpyAsync(m.h_sum + t0 * N, m.input + t0 * N, (size_t) n * N * 2, cudaMemcpyDeviceToHost, m.s); + cudaEventRecord(m.piece_ev[(size_t) k], m.s); + } + cudaEventRecord(m.t1, m.s); + m.timed = true; + } + ms_host += std::chrono::duration(std::chrono::steady_clock::now() - h0).count(); + const cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { err = std::string("prompt experts: ") + cudaGetErrorString(e); return false; } + return true; +} + +} // namespace strata::prefill diff --git a/src/prefill/gemm.cu b/src/prefill/gemm.cu index 4742d423b3e..1363ab93f05 100644 --- a/src/prefill/gemm.cu +++ b/src/prefill/gemm.cu @@ -5,32 +5,10 @@ #include #include -#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE) -// The HIP compatibility shim maps CUDA shuffle spellings to Strata helpers. -// hipBLASLt's public headers declare native HIP shuffle functions, so keep -// those declarations from being macro-expanded in this translation unit. -#undef __shfl_xor_sync -#undef __shfl_down_sync -#undef __shfl_up_sync -#undef __shfl_sync -#undef __ballot_sync -#endif - -#include +#include #include #include -#include -#include - -#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE) -#include "hipblaslt_tuning.hpp" -#include -#include -#include -#include -#include -#include -#endif +#include namespace strata::prefill { namespace { @@ -42,299 +20,57 @@ void ck(cublasStatus_t s, const char* what) { } } -#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE) -struct HipLtCallKey { - strata::prefill::hipblaslt::InputType type; - int t; - int n; - int k; - int ldy; - uint32_t beta_bits; - - bool operator<(const HipLtCallKey& other) const { - return std::tie(type, n, k, ldy, t, beta_bits) < - std::tie(other.type, other.n, other.k, other.ldy, other.t, other.beta_bits); - } -}; - -struct HipLtCachedAlgo { - bool supported = false; - hipblasLtMatmulAlgo_t algo{}; - size_t workspace_bytes = 0; -}; - -struct HipLtState { - hipblasLtHandle_t handle = nullptr; - void* workspace = nullptr; - size_t workspace_bytes = 0; - strata::prefill::hipblaslt::TuningTable table; - std::map cache; - uint64_t lt_launches = 0; - uint64_t fallbacks = 0; - std::set> fallback_shapes; - - ~HipLtState() { - if (std::getenv("STRATA_HIPBLASLT_VERBOSE")) { - std::fprintf(stderr, "prefill gemm: hipBLASLt summary launches=%llu fallbacks=%llu unique_fallback_shapes=%zu\n", - (unsigned long long) lt_launches, (unsigned long long) fallbacks, fallback_shapes.size()); - for (const auto& shape : fallback_shapes) { - const auto type = std::get<0>(shape); - std::fprintf(stderr, "prefill gemm: fallback shape dtype=%s T=%d N=%d K=%d ldy=%d\n", - type == strata::prefill::hipblaslt::InputType::bf16 ? "bf16" : "f16", - std::get<1>(shape), std::get<2>(shape), std::get<3>(shape), std::get<4>(shape)); - } - } - if (handle) hipblasLtDestroy(handle); - } -}; - -struct HipLtDescriptors { - hipblasLtMatmulDesc_t op = nullptr; - hipblasLtMatrixLayout_t a = nullptr; - hipblasLtMatrixLayout_t b = nullptr; - hipblasLtMatrixLayout_t c = nullptr; - - ~HipLtDescriptors() { - if (op) hipblasLtMatmulDescDestroy(op); - if (a) hipblasLtMatrixLayoutDestroy(a); - if (b) hipblasLtMatrixLayoutDestroy(b); - if (c) hipblasLtMatrixLayoutDestroy(c); - } - - bool init(hipDataType type, int t, int n, int k, int ldy) { - const hipblasOperation_t trans_a = HIPBLAS_OP_T; - const hipblasOperation_t trans_b = HIPBLAS_OP_N; - if (hipblasLtMatmulDescCreate(&op, HIPBLAS_COMPUTE_32F, HIP_R_32F) != HIPBLAS_STATUS_SUCCESS || - hipblasLtMatmulDescSetAttribute(op, HIPBLASLT_MATMUL_DESC_TRANSA, &trans_a, sizeof(trans_a)) != - HIPBLAS_STATUS_SUCCESS || - hipblasLtMatmulDescSetAttribute(op, HIPBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b)) != - HIPBLAS_STATUS_SUCCESS || - hipblasLtMatrixLayoutCreate(&a, type, k, n, k) != HIPBLAS_STATUS_SUCCESS || - hipblasLtMatrixLayoutCreate(&b, type, k, t, k) != HIPBLAS_STATUS_SUCCESS || - hipblasLtMatrixLayoutCreate(&c, HIP_R_32F, n, t, ldy) != HIPBLAS_STATUS_SUCCESS) { - return false; - } - return true; - } -}; - -std::unique_ptr create_hipblaslt_state(void* workspace, size_t workspace_bytes) { - const char* path = std::getenv("STRATA_HIPBLASLT_TUNING"); - if (!path || !*path) return nullptr; - - auto state = std::make_unique(); - state->workspace = workspace; - state->workspace_bytes = workspace_bytes; - if (hipblasLtCreate(&state->handle) != HIPBLAS_STATUS_SUCCESS) { - std::fprintf(stderr, "prefill gemm: hipBLASLt handle creation failed; using hipBLASEx\n"); - return nullptr; - } - - int version = 0; - if (hipblasLtGetVersion(state->handle, &version) != HIPBLAS_STATUS_SUCCESS) { - std::fprintf(stderr, "prefill gemm: hipBLASLt version query failed; using hipBLASEx\n"); - return nullptr; - } - int device = 0; - hipDeviceProp_t properties{}; - if (hipGetDevice(&device) != hipSuccess || hipGetDeviceProperties(&properties, device) != hipSuccess) { - std::fprintf(stderr, "prefill gemm: HIP device query failed; using hipBLASEx\n"); - return nullptr; - } - std::string arch(properties.gcnArchName); - const auto suffix = arch.find(':'); - if (suffix != std::string::npos) arch.resize(suffix); - - std::string error; - if (!state->table.load(path, arch, version, error)) { - std::fprintf(stderr, "prefill gemm: %s; using hipBLASEx\n", error.c_str()); - return nullptr; - } - std::fprintf(stderr, "prefill gemm: hipBLASLt tuning enabled (%zu rows, %s, version %d)\n", - state->table.rows().size(), arch.c_str(), version); - return state; -} - -HipLtCachedAlgo resolve_hipblaslt_algo(HipLtState& state, strata::prefill::hipblaslt::InputType type, int t, - int n, int k, int ldy, float beta) { - uint32_t beta_bits = 0; - static_assert(sizeof(beta_bits) == sizeof(beta)); - std::memcpy(&beta_bits, &beta, sizeof(beta)); - const HipLtCallKey key{type, t, n, k, ldy, beta_bits}; - const auto cached = state.cache.find(key); - if (cached != state.cache.end()) return cached->second; - - HipLtCachedAlgo resolved; - const bool verbose = std::getenv("STRATA_HIPBLASLT_VERBOSE") != nullptr; - const auto* row = state.table.closest(type, n, k, ldy, t); - if (!row) { - if (verbose) { - std::fprintf(stderr, "prefill gemm: Lt fallback; no calibration for dtype=%s T=%d N=%d K=%d ldy=%d\n", - type == strata::prefill::hipblaslt::InputType::bf16 ? "bf16" : "f16", t, n, k, ldy); - } - return state.cache.emplace(key, resolved).first->second; - } - - HipLtDescriptors desc; - const hipDataType input_type = type == strata::prefill::hipblaslt::InputType::bf16 ? HIP_R_16BF : HIP_R_16F; - if (!desc.init(input_type, t, n, k, ldy)) { - return state.cache.emplace(key, resolved).first->second; - } - - std::vector solution_ids{row->solution_id}; - std::vector candidates; - if (hipblaslt_ext::getAlgosFromIndex(state.handle, solution_ids, candidates) != HIPBLAS_STATUS_SUCCESS || - candidates.empty() || candidates.front().state != HIPBLAS_STATUS_SUCCESS || - hipblaslt_ext::getIndexFromAlgo(candidates.front().algo) != row->solution_id) { - if (verbose) { - std::fprintf(stderr, "prefill gemm: Lt fallback; solution %d unavailable for T=%d N=%d K=%d ldy=%d\n", - row->solution_id, t, n, k, ldy); - } - return state.cache.emplace(key, resolved).first->second; - } - - const float alpha = 1.0f; - size_t required_workspace = 0; - auto algo = candidates.front().algo; - if (hipblaslt_ext::matmulIsAlgoSupported(state.handle, desc.op, &alpha, desc.a, desc.b, &beta, desc.c, desc.c, - algo, required_workspace) != HIPBLAS_STATUS_SUCCESS) { - if (verbose) { - std::fprintf(stderr, "prefill gemm: Lt fallback; solution %d rejects actual T=%d N=%d K=%d ldy=%d beta=%.9g\n", - row->solution_id, t, n, k, ldy, beta); - } - return state.cache.emplace(key, resolved).first->second; - } - - resolved.supported = true; - resolved.algo = algo; - resolved.workspace_bytes = required_workspace; - if (verbose) { - std::fprintf(stderr, - "prefill gemm: Lt solution=%d dtype=%s T=%d N=%d K=%d ldy=%d beta=%.9g workspace=%zu\n", - row->solution_id, type == strata::prefill::hipblaslt::InputType::bf16 ? "bf16" : "f16", t, n, - k, ldy, beta, required_workspace); - } - return state.cache.emplace(key, resolved).first->second; -} - -bool try_hipblaslt(void* opaque_state, strata::prefill::hipblaslt::InputType type, const uint16_t* x, - const uint16_t* w, float* y, int64_t t, int64_t n, int64_t k, int64_t ldy, float beta, - void* stream) { - auto* state = static_cast(opaque_state); - if (!state || t <= 0 || n <= 0 || k <= 0 || t > INT_MAX || n > INT_MAX || k > INT_MAX || ldy > INT_MAX || - ldy < n) { - return false; - } - const auto resolved = resolve_hipblaslt_algo(*state, type, (int) t, (int) n, (int) k, (int) ldy, beta); - if (!resolved.supported) { - ++state->fallbacks; - state->fallback_shapes.emplace(type, (int) t, (int) n, (int) k, (int) ldy); - return false; - } - if (resolved.workspace_bytes > state->workspace_bytes) { - ++state->fallbacks; - state->fallback_shapes.emplace(type, (int) t, (int) n, (int) k, (int) ldy); - if (std::getenv("STRATA_HIPBLASLT_VERBOSE")) { - std::fprintf(stderr, "prefill gemm: Lt fallback; solution needs %zu workspace bytes, have %zu\n", - resolved.workspace_bytes, state->workspace_bytes); - } - return false; - } - - HipLtDescriptors desc; - const hipDataType input_type = type == strata::prefill::hipblaslt::InputType::bf16 ? HIP_R_16BF : HIP_R_16F; - if (!desc.init(input_type, (int) t, (int) n, (int) k, (int) ldy)) return false; - const float alpha = 1.0f; - const hipblasStatus_t status = hipblasLtMatmul(state->handle, desc.op, &alpha, w, desc.a, x, desc.b, &beta, y, - desc.c, y, desc.c, &resolved.algo, state->workspace, - state->workspace_bytes, (hipStream_t) stream); - if (status == HIPBLAS_STATUS_SUCCESS) { - ++state->lt_launches; - return true; - } - - std::fprintf(stderr, "prefill gemm: hipBLASLt launch failed with status %d\n", (int) status); - if (beta != 0.0f) { - std::fprintf(stderr, "prefill gemm: refusing a fallback after hipBLASLt failed with nonzero beta\n"); - std::exit(1); - } - auto* mutable_state = static_cast(opaque_state); - uint32_t beta_bits = 0; - std::memcpy(&beta_bits, &beta, sizeof(beta_bits)); - auto cached = mutable_state->cache.find(HipLtCallKey{type, (int) t, (int) n, (int) k, (int) ldy, beta_bits}); - if (cached != mutable_state->cache.end()) cached->second.supported = false; - ++mutable_state->fallbacks; - mutable_state->fallback_shapes.emplace(type, (int) t, (int) n, (int) k, (int) ldy); - return false; -} -#endif - } // namespace Gemm::~Gemm() { -#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE) - delete static_cast(hipblaslt_state_); -#endif if (handle_) cublasDestroy((cublasHandle_t) handle_); - if (!external_) { - if (scratch_) cudaFree(scratch_); - if (workspace_) cudaFree(workspace_); + for (int i = 0; i < AUX; ++i) { + if (aux_handle_[i]) cublasDestroy((cublasHandle_t) aux_handle_[i]); + if (join_[i]) cudaEventDestroy((cudaEvent_t) join_[i]); + if (aux_stream_[i]) { + cudaStreamSynchronize((cudaStream_t) aux_stream_[i]); + cudaStreamDestroy((cudaStream_t) aux_stream_[i]); + } } + if (fork_) cudaEventDestroy((cudaEvent_t) fork_); } -bool Gemm::init_external(void* stream, uint16_t* scratch, int64_t scratch_elems, void* workspace, size_t ws_bytes, - std::string& err) { +bool Gemm::init(void* stream, std::string& err, bool side_streams) { cublasHandle_t h = nullptr; if (cublasCreate(&h) != CUBLAS_STATUS_SUCCESS) { err = "prefill gemm: cublasCreate failed"; return false; } handle_ = h; stream_ = stream; - external_ = true; cublasSetStream(h, (cudaStream_t) stream); - workspace_ = workspace; - cublasSetWorkspace(h, workspace_, ws_bytes); cublasSetMathMode(h, CUBLAS_DEFAULT_MATH); - scratch_ = scratch; - scratch_elems_ = scratch_elems; -#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE) - hipblaslt_state_ = create_hipblaslt_state(workspace_, ws_bytes).release(); -#endif + if (!side_streams) return true; + if (cudaEventCreateWithFlags((cudaEvent_t*) &fork_, cudaEventDisableTiming) != cudaSuccess) { + err = "prefill gemm: events"; + return false; + } + for (int i = 0; i < AUX; ++i) { + cublasHandle_t a = nullptr; + if (cudaStreamCreateWithFlags((cudaStream_t*) &aux_stream_[i], cudaStreamNonBlocking) != cudaSuccess || + cudaEventCreateWithFlags((cudaEvent_t*) &join_[i], cudaEventDisableTiming) != cudaSuccess || + cublasCreate(&a) != CUBLAS_STATUS_SUCCESS) { + err = "prefill gemm: side streams"; + return false; + } + aux_handle_[i] = a; + cublasSetStream(a, (cudaStream_t) aux_stream_[i]); + cublasSetMathMode(a, CUBLAS_DEFAULT_MATH); + } return true; } -void Gemm::rebind(uint16_t* scratch, int64_t scratch_elems, void* workspace, size_t ws_bytes) { +void Gemm::set_buffers(uint16_t* scratch, int64_t scratch_elems, void* workspace, size_t ws_bytes) { + // a fixed workspace, so the handles never allocate on the way: a slice for each stream + const int ns = aux_handle_[0] ? AUX + 1 : 1; + const size_t slice = ws_bytes / (size_t) ns / 256 * 256; + cublasSetWorkspace((cublasHandle_t) handle_, workspace, slice); + for (int i = 0; i + 1 < ns; ++i) + cublasSetWorkspace((cublasHandle_t) aux_handle_[i], (uint8_t*) workspace + (size_t) (i + 1) * slice, slice); scratch_ = scratch; scratch_elems_ = scratch_elems; - workspace_ = workspace; - cublasSetWorkspace((cublasHandle_t) handle_, workspace_, ws_bytes); -#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE) - if (hipblaslt_state_) { - auto* state = static_cast(hipblaslt_state_); - state->workspace = workspace_; - state->workspace_bytes = ws_bytes; - } -#endif -} - -bool Gemm::init(void* stream, int64_t scratch_elems, std::string& err) { - cublasHandle_t h = nullptr; - if (cublasCreate(&h) != CUBLAS_STATUS_SUCCESS) { err = "prefill gemm: cublasCreate failed"; return false; } - handle_ = h; - stream_ = stream; - cublasSetStream(h, (cudaStream_t) stream); - // A fixed workspace so the handle never allocates on the way (and graphs could capture it later). - const size_t ws = 32u << 20; - if (cudaMalloc(&workspace_, ws) != cudaSuccess) { err = "prefill gemm: workspace"; return false; } - cublasSetWorkspace(h, workspace_, ws); - cublasSetMathMode(h, CUBLAS_DEFAULT_MATH); -#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE) - hipblaslt_state_ = create_hipblaslt_state(workspace_, ws).release(); -#endif - if (scratch_elems > 0 && cudaMalloc((void**) &scratch_, (size_t) scratch_elems * 2) != cudaSuccess) { - err = "prefill gemm: dequant scratch of " + std::to_string(scratch_elems * 2 >> 20) + " MiB"; - return false; - } - scratch_elems_ = scratch_elems; - return true; } void Gemm::bf16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy, @@ -342,12 +78,6 @@ void Gemm::bf16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64 if (T <= 0 || N <= 0) return; if (ldy <= 0) ldy = N; const float alpha = 1.0f; -#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE) - if (try_hipblaslt(hipblaslt_state_, strata::prefill::hipblaslt::InputType::bf16, X, W, Y, T, N, K, ldy, - beta, stream_)) { - return; - } -#endif // Column-major view: Y^T[N, T] = W[N, K] (stored K x N col-major, transposed) . X^T[K, T]. ck(cublasGemmEx((cublasHandle_t) handle_, CUBLAS_OP_T, CUBLAS_OP_N, (int) N, (int) T, (int) K, &alpha, W, CUDA_R_16BF, (int) K, X, CUDA_R_16BF, (int) K, &beta, Y, CUDA_R_32F, (int) ldy, @@ -360,18 +90,50 @@ void Gemm::f16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_ if (T <= 0 || N <= 0) return; if (ldy <= 0) ldy = N; const float alpha = 1.0f; -#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE) - if (try_hipblaslt(hipblaslt_state_, strata::prefill::hipblaslt::InputType::f16, X, W, Y, T, N, K, ldy, - beta, stream_)) { - return; - } -#endif ck(cublasGemmEx((cublasHandle_t) handle_, CUBLAS_OP_T, CUBLAS_OP_N, (int) N, (int) T, (int) K, &alpha, W, CUDA_R_16F, (int) K, X, CUDA_R_16F, (int) K, &beta, Y, CUDA_R_32F, (int) ldy, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT), "cublasGemmEx f16"); } +void Gemm::f16_grouped(const uint16_t* const* X, const uint16_t* const* W, float* const* Y, const uint16_t* const* X_host, + const uint16_t* const* W_host, float* const* Y_host, const int* rows, int n, int64_t N, int64_t K) { + if (n <= 0) return; + if (grouped_) { + // column-major, as f16(): Y_i^T[N, rows] = W_i (K x N col-major, transposed) . X_i^T[K, rows]; a group per problem + const std::vector ta((size_t) n, CUBLAS_OP_T), tb((size_t) n, CUBLAS_OP_N); + const std::vector m((size_t) n, (int) N), k((size_t) n, (int) K), ld_wx((size_t) n, (int) K), + ld_y((size_t) n, (int) N), size((size_t) n, 1); + const std::vector alpha((size_t) n, 1.0f), beta((size_t) n, 0.0f); + const cublasStatus_t st = cublasGemmGroupedBatchedEx( + (cublasHandle_t) handle_, ta.data(), tb.data(), m.data(), rows, k.data(), alpha.data(), (const void* const*) W, + CUDA_R_16F, ld_wx.data(), (const void* const*) X, CUDA_R_16F, ld_wx.data(), beta.data(), (void* const*) Y, + CUDA_R_32F, ld_y.data(), n, size.data(), CUBLAS_COMPUTE_32F); + if (st == CUBLAS_STATUS_SUCCESS) return; + if (st != CUBLAS_STATUS_NOT_SUPPORTED && st != CUBLAS_STATUS_INVALID_VALUE) ck(st, "cublasGemmGroupedBatchedEx"); + grouped_ = false; + std::fprintf(stderr, "prefill gemm: cuBLAS does not group FP16 -> FP32 GEMMs here; one call per expert\n"); + } + // round robin over the stream and its side streams, which join it again + int64_t sum_rows = 0; + for (int i = 0; i < n; ++i) sum_rows += rows[i]; + const int ns = aux_handle_[0] && sum_rows >= 128 * (int64_t) n ? std::min(n, AUX + 1) : 1; + if (ns > 1) { + cudaEventRecord((cudaEvent_t) fork_, (cudaStream_t) stream_); + for (int s = 1; s < ns; ++s) cudaStreamWaitEvent((cudaStream_t) aux_stream_[s - 1], (cudaEvent_t) fork_, 0); + } + void* const main = handle_; + for (int i = 0; i < n; ++i) { + handle_ = i % ns == 0 ? main : aux_handle_[i % ns - 1]; + f16(X_host[i], W_host[i], Y_host[i], rows[i], N, K); + } + handle_ = main; + for (int s = 1; s < ns; ++s) { + cudaEventRecord((cudaEvent_t) join_[s - 1], (cudaStream_t) aux_stream_[s - 1]); + cudaStreamWaitEvent((cudaStream_t) stream_, (cudaEvent_t) join_[s - 1], 0); + } +} + void Gemm::native(const uint16_t* X, int ggml_type, const void* W_blocks, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy, float beta) { if (N * K > scratch_elems_) { diff --git a/src/prefill/kernels.cu b/src/prefill/kernels.cu index fb7418a6372..5908d7ad7e2 100644 --- a/src/prefill/kernels.cu +++ b/src/prefill/kernels.cu @@ -1,11 +1,12 @@ // src/prefill/kernels.cu - see include/strata/prefill/kernels.hpp. #include "strata/prefill/kernels.hpp" #include "strata/kernels/mrope.hpp" -#include "strata/kernels/router_top10.hpp" +#include "strata/kernels/verify_kernels.hpp" #include #include +#include #include #include #include @@ -70,8 +71,8 @@ __global__ void gr_norm_kernel(const float* __restrict__ R, const float* __restr xn16[row * N + d] = bf(v); } } -// F-1: the row scale only (and the BF16 image); gr_mix_r_kernel recomputes r * rs * w itself, in the same order, -// so the FP32 copy of the normalized rows (T x 10240 floats) is neither written nor read +// gr_norm_kernel's row scale and BF16 image; gr_mix_r_kernel recomputes r * rs * w in the same order, so the FP32 +// copy of the normalized rows is neither written nor read __global__ void gr_norm_rs_kernel(const float* __restrict__ R, const float* __restrict__ w, float eps, float* __restrict__ rs_out, uint16_t* __restrict__ xn16) { __shared__ float sh[32]; @@ -98,18 +99,18 @@ __global__ void gr_mix_r_kernel(const float* __restrict__ R, const float* __rest s = fmaf(x, sigm(g[j]), s); } s /= (float) HC; - mixed[i] = s; + if (mixed) mixed[i] = s; if (mixed16) mixed16[i] = bf(s); if (mixed_h) mixed_h[i] = hf(s); } -// F-2: gr_write_kernel for one row (t, c), then gr_norm_rs_kernel's reduction over it with the next half's norm -// weights - the same thread-to-element mapping (256 threads, stride 256) and block_sum, so rs and the BF16 image are -// the same bits, and R is not read back -constexpr int GRW_PER = (N + 255) / 256; -__global__ void __launch_bounds__(256) gr_write_norm_rs_kernel(float* __restrict__ R, const float* __restrict__ bo, - const float* __restrict__ inj, int64_t inj_ld, - const float* __restrict__ w, float eps, - float* __restrict__ rs_out, uint16_t* __restrict__ xn16) { +// gr_write_kernel for one row (t, c), then gr_norm_rs_kernel's reduction over it with the next read's norm weights: +// the same thread-to-element mapping (256 threads, stride 256) and block_sum, so rs and the BF16 image are the same +// bits, and R is not read back +constexpr int GRW_PER = N / 256; +__global__ void __launch_bounds__(256) +gr_write_norm_rs_kernel(float* __restrict__ R, const float* __restrict__ bo, const float* __restrict__ inj, + int64_t inj_ld, const float* __restrict__ w, float eps, float* __restrict__ rs_out, + uint16_t* __restrict__ xn16, const uint16_t* __restrict__ partial) { __shared__ float sh[32]; const int64_t row = blockIdx.x; // t * 4 + c const int64_t t = row / HC; @@ -118,19 +119,22 @@ __global__ void __launch_bounds__(256) gr_write_norm_rs_kernel(float* __restrict const float sc = 2.0f * sigm(inj[t * inj_ld + c] / (float) HC); float v[GRW_PER]; float ss = 0.0f; - int k = 0; #pragma unroll - for (int d = threadIdx.x; d < N; d += 256, ++k) { - const float x = fmaf(bo[t * N + d], sc, r[d]); + for (int k = 0; k < GRW_PER; ++k) { + const int d = threadIdx.x + 256 * k; + const float b = partial ? bo[t * N + d] + __half2float(__ushort_as_half(partial[t * N + d])) : bo[t * N + d]; + const float x = fmaf(b, sc, r[d]); r[d] = x; v[k] = x; ss += x * x; } const float rs = rsqrtf(block_sum(ss, sh) / (float) N + eps); if (threadIdx.x == 0) rs_out[row] = rs; - k = 0; #pragma unroll - for (int d = threadIdx.x; d < N; d += 256, ++k) xn16[row * N + d] = bf(v[k] * rs * w[c * N + d]); + for (int k = 0; k < GRW_PER; ++k) { + const int d = threadIdx.x + 256 * k; + xn16[row * N + d] = bf(v[k] * rs * w[c * N + d]); + } } __global__ void gr_silu_kernel(const float* __restrict__ lo, uint16_t* __restrict__ lo16, int64_t n) { const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; @@ -150,16 +154,17 @@ __global__ void gr_mix_kernel(const float* __restrict__ xn, const float* __restr s = fmaf(xn[j], sigm(g[j]), s); } s /= (float) HC; - mixed[i] = s; + if (mixed) mixed[i] = s; if (mixed16) mixed16[i] = bf(s); if (mixed_h) mixed_h[i] = hf(s); } __global__ void gr_write_kernel(float* __restrict__ R, const float* __restrict__ bo, const float* __restrict__ inj, - int64_t inj_ld, int64_t T) { + int64_t inj_ld, int64_t T, const uint16_t* __restrict__ partial) { const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; if (i >= T * D) return; const int64_t t = i / D, c = (i % D) / N, d = i % N; - R[i] = fmaf(bo[t * N + d], 2.0f * sigm(inj[t * inj_ld + c] / (float) HC), R[i]); + const float b = partial ? bo[t * N + d] + __half2float(__ushort_as_half(partial[t * N + d])) : bo[t * N + d]; + R[i] = fmaf(b, 2.0f * sigm(inj[t * inj_ld + c] / (float) HC), R[i]); } __global__ void gr_broadcast_kernel(const float* __restrict__ e, float* __restrict__ R, int64_t T) { const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; @@ -179,210 +184,173 @@ __global__ void gdn_gates_kernel(const float* __restrict__ ab, const float* __re gate[i] = (v > 20.0f ? v : log1pf(__expf(v))) * ssm_a[h]; beta[i] = sigm(ab[t * 2 * HV + HV + h]); } -// one thread per channel, walks the chunk; then a second kernel normalises -__global__ void gdn_conv_kernel(float* __restrict__ hist, const float* __restrict__ qkv, const float* __restrict__ w, - float* __restrict__ h, int64_t T) { - const int c = blockIdx.x * blockDim.x + threadIdx.x; - if (c >= C) return; - float v0 = hist[c * 3], v1 = hist[c * 3 + 1], v2 = hist[c * 3 + 2]; - const float w0 = w[c * 4], w1 = w[c * 4 + 1], w2 = w[c * 4 + 2], w3 = w[c * 4 + 3]; - for (int64_t t = 0; t < T; ++t) { - const float x = qkv[t * C + c]; - const float s = v0 * w0 + v1 * w1 + v2 * w2 + x * w3; - h[t * C + c] = s / (1.0f + __expf(-s)); - v0 = v1; v1 = v2; v2 = x; - } - hist[c * 3] = v0; hist[c * 3 + 1] = v1; hist[c * 3 + 2] = v2; -} -// C-3: the same 4-tap causal conv, tiled over tokens: thread (c, tile) reads its tile's 3 predecessors from the -// chunk (or the history before it) instead of carrying them - the conv reads inputs, not its own outputs, so the -// tiles are independent. The same expression per element (so the same bits); the history is written afterwards. -constexpr int CONV_TILE = 64; -__global__ void gdn_conv_tiled_kernel(const float* __restrict__ hist, const float* __restrict__ qkv, - const float* __restrict__ w, float* __restrict__ h, int64_t T) { - const int c = blockIdx.x * blockDim.x + threadIdx.x; - if (c >= C) return; - const int64_t t0 = (int64_t) blockIdx.y * CONV_TILE; - if (t0 >= T) return; - const int64_t t1 = t0 + CONV_TILE < T ? t0 + CONV_TILE : T; - auto input = [&](int64_t t) -> float { return t >= 0 ? qkv[t * C + c] : hist[c * 3 + (int) (t + 3)]; }; - float v0 = input(t0 - 3), v1 = input(t0 - 2), v2 = input(t0 - 1); - const float w0 = w[c * 4], w1 = w[c * 4 + 1], w2 = w[c * 4 + 2], w3 = w[c * 4 + 3]; - for (int64_t t = t0; t < t1; ++t) { - const float x = qkv[t * C + c]; - const float s = v0 * w0 + v1 * w1 + v2 * w2 + x * w3; - h[t * C + c] = s / (1.0f + __expf(-s)); - v0 = v1; v1 = v2; v2 = x; - } -} -// the history after the chunk: its last three inputs (the older history where the chunk is shorter than 3) +// The conv history after a chunk: its last three inputs, older ones moved up when the chunk is shorter. __global__ void gdn_conv_hist_kernel(float* __restrict__ hist, const float* __restrict__ qkv, int64_t T) { const int c = blockIdx.x * blockDim.x + threadIdx.x; if (c >= C) return; - float v[3]; - for (int k = 0; k < 3; ++k) { - const int64_t t = T - 3 + k; - v[k] = t >= 0 ? qkv[t * C + c] : hist[c * 3 + (int) (t + 3)]; + const float old[3] = {hist[c * 3], hist[c * 3 + 1], hist[c * 3 + 2]}; +#pragma unroll + for (int r = 0; r < 3; ++r) { + const int64_t j = T - 3 + r; + hist[c * 3 + r] = j >= 0 ? qkv[j * C + c] : old[r + T]; } - hist[c * 3] = v[0]; hist[c * 3 + 1] = v[1]; hist[c * 3 + 2] = v[2]; -} -__global__ void gdn_l2_kernel(float* __restrict__ h, float eps) { - // block (t, head) over the 32 q/k heads, 128 threads - const int64_t t = blockIdx.y; - const int head = blockIdx.x; - float* x = h + t * C + head * S; - const float v = x[threadIdx.x]; - float sq = warp_sum(v * v); - __shared__ float part[4]; - if ((threadIdx.x & 31) == 0) part[threadIdx.x >> 5] = sq; - __syncthreads(); - const float ss = part[0] + part[1] + part[2] + part[3]; - x[threadIdx.x] = v * rsqrtf(ss + eps); -} -constexpr int RG = 4, RPG = S / RG; -__global__ void __launch_bounds__(S * RG) gdn_rec_kernel(float* __restrict__ state, const float* __restrict__ h, - const float* __restrict__ gate, - const float* __restrict__ beta, const float* __restrict__ z, - const float* __restrict__ gamma, float eps, - float* __restrict__ y, uint16_t* __restrict__ y16, int64_t T) { - __shared__ float sk[S], sq[S], red[RG][S], wsum[16]; - const int head = blockIdx.x, col = threadIdx.x, rg = threadIdx.y, tid = rg * S + col; +} +// cp.async of one 4-byte word into shared memory (sm_80+), and its group bookkeeping. +__device__ __forceinline__ void cp_async4(float* dst, const float* src) { + const unsigned d = (unsigned) __cvta_generic_to_shared(dst); + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" ::"r"(d), "l"(src)); +} +__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::); } +template __device__ __forceinline__ void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +// The recurrence over the chunk, the arithmetic of gdn_step_norm_multi: a thread per (CPT columns, group of 32 +// state rows), a head's columns split over S / CB blocks. A column's four row groups sit in adjacent lanes, so +// both sums over the row groups are shuffles added in the decode kernels' order (group 0 + 1 + 2 + 3). Each +// token's q, k, v, gate and beta are copied (cp.async) into a ring of DEPTH token slots DEPTH - 1 tokens ahead, so +// the reads from memory overlap the steps before them: one barrier per token. Each thread reads its row group's +// 32 q and 32 k values into registers with 16-byte loads once for its CPT columns; a row group is padded to 36 +// words, so the four groups a warp reads at once sit in different banks. Writes o (before the norm) to `out`; +// gdn_norm_kernel finishes each token. +constexpr int RG = 4, RPG = S / RG, RP = RPG + 4, DEPTH = 8, SLOT = 2 * RG * RP; +template +__global__ void __launch_bounds__(CB / CPT * RG) gdn_scan_kernel(float* __restrict__ state, + const float* __restrict__ h, + const float* __restrict__ gate, + const float* __restrict__ beta, + float* __restrict__ out, int T) { + constexpr int NT = CB / CPT * RG, QL = (2 * S + NT - 1) / NT; // threads; q/k values each thread copies + static_assert(CB % (8 * CPT) == 0 && NT >= CB + 2, "8 column groups a warp; a thread each for v, gate, beta"); + __shared__ __align__(16) float sqk[DEPTH * SLOT]; // per slot: q, then k, [row group][row] + __shared__ float sv[DEPTH * CB], sgb[DEPTH * 2]; + const int head = blockIdx.x, tid = threadIdx.x, lane = tid & 31; + const int cl = ((tid >> 5) * 8 + (lane >> 2)) * CPT, rg = lane & 3, g0 = lane & ~3; // first local column const int qh = head % HK; - float s[RPG]; - float* base = state + ((size_t) (rg * RPG) * HV + head) * S + col; + // what this thread copies for every token: q/k value i = tid + j * NT < 2S (row i & 127 of q or k), then v of + // local column tid (tid < CB), the gate (tid == CB) or beta (tid == CB + 1) + int qk_off[QL], qk_at[QL]; +#pragma unroll + for (int j = 0; j < QL; ++j) { + const int i = tid + j * NT; + qk_off[j] = i < 2 * S ? (i < S ? 0 : HK * S) + qh * S + (i & (S - 1)) : -1; + qk_at[j] = (i >> 7) * RG * RP + ((i & (S - 1)) / RPG) * RP + (i & (RPG - 1)); + } + const int v_off = 2 * HK * S + head * S + blockIdx.y * CB + tid; + const float* gb = (tid == CB ? gate : beta) + head; + float s[CPT][RPG]; + float* base = state + ((size_t) (rg * RPG) * HV + head) * S + blockIdx.y * CB + cl; const size_t rs = (size_t) HV * S; #pragma unroll - for (int r = 0; r < RPG; ++r) s[r] = base[r * rs]; - const float g_col = gamma[col]; - for (int64_t t = 0; t < T; ++t) { - const float* ht = h + t * C; - __syncthreads(); - if (tid < S) { sq[tid] = ht[qh * S + tid]; sk[tid] = ht[HK * S + qh * S + tid]; } - __syncthreads(); - const float g = __expf(gate[t * HV + head]); - float kv = 0.0f; + for (int c = 0; c < CPT; ++c) #pragma unroll - for (int r = 0; r < RPG; ++r) kv = fmaf(s[r], sk[rg * RPG + r], kv); - red[rg][col] = kv; - __syncthreads(); - const float kv_col = red[0][col] + red[1][col] + red[2][col] + red[3][col]; - const float delta = (ht[2 * HK * S + head * S + col] - g * kv_col) * beta[t * HV + head]; - float o = 0.0f; + for (int r = 0; r < RPG; ++r) s[c][r] = base[r * rs + c]; + int next = 0; // the next token to copy + auto issue = [&]() { + if (next < T) { + const int slot = next & (DEPTH - 1); + const float* ht = h + (size_t) next * C; #pragma unroll - for (int r = 0; r < RPG; ++r) { - s[r] = fmaf(g, s[r], sk[rg * RPG + r] * delta); - o = fmaf(s[r], sq[rg * RPG + r], o); + for (int j = 0; j < QL; ++j) + if (qk_off[j] >= 0) cp_async4(&sqk[slot * SLOT + qk_at[j]], ht + qk_off[j]); + if (tid < CB) cp_async4(&sv[slot * CB + tid], ht + v_off); + else if (tid < CB + 2) cp_async4(&sgb[slot * 2 + tid - CB], gb + (size_t) next * HV); } - __syncthreads(); - red[rg][col] = o; - __syncthreads(); - float oc = 0.0f, sp = 0.0f; - if (rg == 0) { - oc = (red[0][col] + red[1][col] + red[2][col] + red[3][col]) * rsqrtf((float) S); - sp = oc * oc; - } - sp = warp_sum(sp); - if ((tid & 31) == 0) wsum[tid >> 5] = sp; - __syncthreads(); - if (rg == 0) { - const float ss = wsum[0] + wsum[1] + wsum[2] + wsum[3]; - const float v = oc * rsqrtf(ss / (float) S + eps) * g_col * sigm(z[t * HV * S + head * S + col]); - y[t * HV * S + head * S + col] = v; - y16[t * HV * S + head * S + col] = hf(v); + cp_async_commit(); + ++next; + }; + for (int t = 0; t < DEPTH - 1; ++t) issue(); + for (int t = 0; t < T; ++t) { + cp_async_wait(); + __syncthreads(); // token t's slot is filled; every thread is done with token t - 1's slot + issue(); + const int slot = t & (DEPTH - 1); + float qr[RPG], kr[RPG]; + const float4* q4 = reinterpret_cast(&sqk[slot * SLOT + rg * RP]); + const float4* k4 = reinterpret_cast(&sqk[slot * SLOT + RG * RP + rg * RP]); +#pragma unroll + for (int j = 0; j < RPG / 4; ++j) { + const float4 a = k4[j], b = q4[j]; + kr[4 * j] = a.x; kr[4 * j + 1] = a.y; kr[4 * j + 2] = a.z; kr[4 * j + 3] = a.w; + qr[4 * j] = b.x; qr[4 * j + 1] = b.y; qr[4 * j + 2] = b.z; qr[4 * j + 3] = b.w; } - } + const float g = __expf(sgb[slot * 2]), bt = sgb[slot * 2 + 1]; + float kv[CPT], o[CPT]; #pragma unroll - for (int r = 0; r < RPG; ++r) base[r * rs] = s[r]; -} - -// D-2: the recurrence with the value columns split over 4 blocks per head (4x the blocks of the kernel above, a -// quarter of its threads per __syncthreads), and the output norm - the only step that couples the head's columns - -// in its own kernel. Per column the same arithmetic in the same order (the 4 row-group partial sums added as -// red[0] + red[1] + red[2] + red[3]; the norm's warp sums over the same 32-column warps): the same bits. -constexpr int CB = 32, NCB = S / CB; -__global__ void __launch_bounds__(CB * RG) gdn_rec_cols_kernel(float* __restrict__ state, const float* __restrict__ h, - const float* __restrict__ gate, - const float* __restrict__ beta, - float* __restrict__ oc_out, int64_t T) { - __shared__ float sk[S], sq[S], red[RG][CB]; - const int head = blockIdx.x / NCB, cb = blockIdx.x % NCB; - const int c = threadIdx.x, rg = threadIdx.y, tid = rg * CB + c, col = cb * CB + c; - const int qh = head % HK; - float s[RPG]; - float* base = state + ((size_t) (rg * RPG) * HV + head) * S + col; - const size_t rs = (size_t) HV * S; + for (int c = 0; c < CPT; ++c) kv[c] = o[c] = 0.0f; #pragma unroll - for (int r = 0; r < RPG; ++r) s[r] = base[r * rs]; - for (int64_t t = 0; t < T; ++t) { - const float* ht = h + t * C; - __syncthreads(); - if (tid < S) { sq[tid] = ht[qh * S + tid]; sk[tid] = ht[HK * S + qh * S + tid]; } - __syncthreads(); - const float g = __expf(gate[t * HV + head]); - float kv = 0.0f; + for (int r = 0; r < RPG; ++r) #pragma unroll - for (int r = 0; r < RPG; ++r) kv = fmaf(s[r], sk[rg * RPG + r], kv); - red[rg][c] = kv; - __syncthreads(); - const float kv_col = red[0][c] + red[1][c] + red[2][c] + red[3][c]; - const float delta = (ht[2 * HK * S + head * S + col] - g * kv_col) * beta[t * HV + head]; - float o = 0.0f; + for (int c = 0; c < CPT; ++c) kv[c] = fmaf(s[c][r], kr[r], kv[c]); + float delta[CPT]; #pragma unroll - for (int r = 0; r < RPG; ++r) { - s[r] = fmaf(g, s[r], sk[rg * RPG + r] * delta); - o = fmaf(s[r], sq[rg * RPG + r], o); + for (int c = 0; c < CPT; ++c) { + const float k0 = __shfl_sync(0xffffffffu, kv[c], g0), k1 = __shfl_sync(0xffffffffu, kv[c], g0 + 1), + k2 = __shfl_sync(0xffffffffu, kv[c], g0 + 2), k3 = __shfl_sync(0xffffffffu, kv[c], g0 + 3); + const float kv_col = k0 + k1 + k2 + k3; + delta[c] = (sv[slot * CB + cl + c] - g * kv_col) * bt; + } +#pragma unroll + for (int r = 0; r < RPG; ++r) +#pragma unroll + for (int c = 0; c < CPT; ++c) { + s[c][r] = fmaf(g, s[c][r], kr[r] * delta[c]); + o[c] = fmaf(s[c][r], qr[r], o[c]); + } + float* ot = out + (size_t) t * HV * S + head * S + blockIdx.y * CB + cl; +#pragma unroll + for (int c = 0; c < CPT; ++c) { + const float o0 = __shfl_sync(0xffffffffu, o[c], g0), o1 = __shfl_sync(0xffffffffu, o[c], g0 + 1), + o2 = __shfl_sync(0xffffffffu, o[c], g0 + 2), o3 = __shfl_sync(0xffffffffu, o[c], g0 + 3); + if (rg == 0) ot[c] = (o0 + o1 + o2 + o3) * rsqrtf((float) S); } - __syncthreads(); - red[rg][c] = o; - __syncthreads(); - if (rg == 0) oc_out[t * HV * S + head * S + col] = (red[0][c] + red[1][c] + red[2][c] + red[3][c]) * rsqrtf((float) S); } + cp_async_wait<0>(); #pragma unroll - for (int r = 0; r < RPG; ++r) base[r * rs] = s[r]; -} -__global__ void __launch_bounds__(S) gdn_out_norm_kernel(const float* __restrict__ z, const float* __restrict__ gamma, - float eps, float* __restrict__ y, uint16_t* __restrict__ y16) { - __shared__ float wsum[4]; - const int64_t t = blockIdx.x; - const int head = blockIdx.y, col = threadIdx.x; - const size_t at = (size_t) t * HV * S + (size_t) head * S + col; - const float oc = y[at]; - float sp = warp_sum(oc * oc); - if ((col & 31) == 0) wsum[col >> 5] = sp; + for (int c = 0; c < CPT; ++c) +#pragma unroll + for (int r = 0; r < RPG; ++r) base[r * rs + c] = s[c][r]; +} +// y = rmsnorm(o) * gamma * sigmoid(z) per (token, head), in place over o; also the FP16 image +__global__ void __launch_bounds__(S) gdn_norm_kernel(float* __restrict__ y, const float* __restrict__ z, + const float* __restrict__ gamma, float eps, + uint16_t* __restrict__ y16) { + __shared__ float wsum[S / 32]; + const size_t i = ((size_t) blockIdx.y * HV + blockIdx.x) * S + threadIdx.x; + const float oc = y[i]; + const float sp = warp_sum(oc * oc); + if ((threadIdx.x & 31) == 0) wsum[threadIdx.x >> 5] = sp; __syncthreads(); const float ss = wsum[0] + wsum[1] + wsum[2] + wsum[3]; - const float v = oc * rsqrtf(ss / (float) S + eps) * gamma[col] * sigm(z[t * HV * S + head * S + col]); - y[at] = v; - y16[at] = hf(v); + const float v = oc * rsqrtf(ss / (float) S + eps) * gamma[threadIdx.x] * sigm(z[i]); + y[i] = v; + y16[i] = hf(v); } // ---------------------------------------------------------------- MoE -template __global__ void route_kernel(const float* __restrict__ logits, int32_t* __restrict__ ids, float* __restrict__ wout, int64_t T) { const int64_t t = (int64_t) blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5); if (t >= T) return; const int lane = threadIdx.x & 31; - const float* lg = logits + t * (REG * 32); - float v[REG]; + const float* lg = logits + t * 512; + float v[16]; #pragma unroll - for (int i = 0; i < REG; ++i) v[i] = lg[lane + i * 32]; + for (int i = 0; i < 16; ++i) v[i] = lg[lane + i * 32]; float mx = -INFINITY; #pragma unroll - for (int i = 0; i < REG; ++i) mx = fmaxf(mx, v[i]); + for (int i = 0; i < 16; ++i) mx = fmaxf(mx, v[i]); mx = warp_max(mx); float sum = 0.0f; #pragma unroll - for (int i = 0; i < REG; ++i) { v[i] = expf(v[i] - mx); sum += v[i]; } + for (int i = 0; i < 16; ++i) { v[i] = expf(v[i] - mx); sum += v[i]; } const float rcp = 1.0f / warp_sum(sum); #pragma unroll - for (int i = 0; i < REG; ++i) { v[i] *= rcp; if (isnan(v[i])) v[i] = -FLT_MAX; } + for (int i = 0; i < 16; ++i) { v[i] *= rcp; if (isnan(v[i])) v[i] = -FLT_MAX; } float selected = 0.0f, selected_sum = 0.0f; for (int rank = 0; rank < 10; ++rank) { float best = v[0]; int ex = lane; #pragma unroll - for (int i = 1; i < REG; ++i) if (v[i] > best) { best = v[i]; ex = lane + i * 32; } + for (int i = 1; i < 16; ++i) if (v[i] > best) { best = v[i]; ex = lane + i * 32; } #pragma unroll for (int m = 16; m; m >>= 1) { const float ob = __shfl_xor_sync(0xffffffffu, best, m); @@ -444,16 +412,31 @@ __global__ void gather_rows16_kernel(const uint16_t* __restrict__ x, const int32 const int64_t r = i / per, j = i % per; reinterpret_cast(dst)[r * per + j] = reinterpret_cast(x)[(int64_t) src[r] * per + j]; } -__global__ void moe_combine_kernel(const float* __restrict__ Dm, const int32_t* __restrict__ slot, - const float* __restrict__ w, const float* __restrict__ shared, - const float* __restrict__ sg, float* __restrict__ bo, int64_t T) { +__global__ void gather_rows16_f32_kernel(const uint16_t* __restrict__ x, const int32_t* __restrict__ src, + float* __restrict__ dst, int64_t n, int64_t width) { + const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; // 8 values + const int64_t per = width / 8; + if (i >= n * per) return; + const int64_t r = i / per, j = i % per; + const uint4 v = reinterpret_cast(x)[(int64_t) src[r] * per + j]; + const __half2* h = reinterpret_cast(&v); + float4* o = reinterpret_cast(dst + r * width + j * 8); + const float2 a = __half22float2(h[0]), b = __half22float2(h[1]), c = __half22float2(h[2]), d = __half22float2(h[3]); + o[0] = make_float4(a.x, a.y, b.x, b.y); + o[1] = make_float4(c.x, c.y, d.x, d.y); +} +__global__ void moe_shared_gate_kernel(float* __restrict__ x, const float* __restrict__ sg, int64_t T) { const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; if (i >= T * N) return; - const int64_t t = i / N, d = i % N; - float s = 0.0f; -#pragma unroll - for (int k = 0; k < 10; ++k) s = fmaf(w[t * 10 + k], Dm[(int64_t) slot[t * 10 + k] * N + d], s); - bo[i] = s + shared[i] * sigm(sg[t]); + x[i] = x[i] * sigm(sg[i / N]); +} +__global__ void moe_scatter_add_kernel(float* __restrict__ sum, const float* __restrict__ rows, + const float* __restrict__ w, const int32_t* __restrict__ src, int64_t n) { + const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n * N) return; + const int64_t r = i / N, d = i % N; + float* o = sum + (int64_t) src[r] * N + d; + *o = fmaf(w[r], rows[i], *o); } // ---------------------------------------------------------------- QSA helpers @@ -492,12 +475,10 @@ __global__ void gate_attn_kernel(const float* __restrict__ a, const float* __res o16[i] = hf(a[i] * (1.0f / (1.0f + expf(-qf[t * 24 * 512 + h * 512 + 256 + d])))); } -// one block per (token, kv head, 64-value group); 64 threads. KV streaming: the pool page only if the block is -// resident (table >= 0), and the host copy and the prompt path's staging pool (both identity layout) when given. +// one block per (token, kv head, 64-value group); 64 threads __global__ void kv_append_kernel(const float* __restrict__ K, const float* __restrict__ V, int64_t pos0, const int32_t* __restrict__ table, int64_t page_size, uint16_t* k_pool, - uint16_t* v_pool, int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, - strata::kernels::KvHostPools host, strata::kernels::KvHostPools stage) { + uint16_t* v_pool, int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale) { const int64_t t = blockIdx.x; const int kvh = blockIdx.y, g = blockIdx.z >> 1; const bool is_v = (blockIdx.z & 1) != 0; @@ -506,12 +487,8 @@ __global__ void kv_append_kernel(const float* __restrict__ K, const float* __res const int64_t pos = pos0 + t; const int64_t page = table[pos / page_size]; const int64_t row = (page * 2 + kvh) * page_size + pos % page_size; - const int64_t row_id = ((pos / page_size) * 2 + kvh) * page_size + pos % page_size; if (k_pool != nullptr) { - const uint16_t h = hf(x); - if (page >= 0) (is_v ? v_pool : k_pool)[row * 256 + d] = h; - if (host.k_pool != nullptr) (is_v ? host.v_pool : host.k_pool)[row_id * 256 + d] = h; - if (stage.k_pool != nullptr) (is_v ? stage.v_pool : stage.k_pool)[row_id * 256 + d] = h; + (is_v ? v_pool : k_pool)[row * 256 + d] = hf(x); return; } float a = fabsf(x); @@ -524,27 +501,13 @@ __global__ void kv_append_kernel(const float* __restrict__ K, const float* __res const float sf = __half2float(__ushort_as_half(sb)); int q = 0; if (sf > 0.0f) { q = __float2int_rn(x / sf); q = q < -127 ? -127 : (q > 127 ? 127 : q); } - if (page >= 0) { - (is_v ? v_q : k_q)[row * 256 + d] = (int8_t) q; - if (threadIdx.x == 0) (is_v ? v_scale : k_scale)[row * 4 + g] = sb; - } - if (host.k_q != nullptr) { - (is_v ? host.v_q : host.k_q)[row_id * 256 + d] = (int8_t) q; - if (threadIdx.x == 0) (is_v ? host.v_scale : host.k_scale)[row_id * 4 + g] = sb; - } - if (stage.k_q != nullptr) { - (is_v ? stage.v_q : stage.k_q)[row_id * 256 + d] = (int8_t) q; - if (threadIdx.x == 0) (is_v ? stage.v_scale : stage.k_scale)[row_id * 4 + g] = sb; - } + (is_v ? v_q : k_q)[row * 256 + d] = (int8_t) q; + if (threadIdx.x == 0) (is_v ? v_scale : k_scale)[row * 4 + g] = sb; } __global__ void to_f16_kernel(const float* __restrict__ x, uint16_t* __restrict__ y, int64_t n) { for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x * blockDim.x) y[i] = hf(x[i]); } -__global__ void round_f16_kernel(const float* __restrict__ x, float* __restrict__ y, int64_t n) { - for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x * blockDim.x) - y[i] = __half2float(__float2half_rn(x[i])); -} __global__ void to_bf16_kernel(const float* __restrict__ x, uint16_t* __restrict__ y, int64_t n) { for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x * blockDim.x) y[i] = bf(x[i]); @@ -554,11 +517,10 @@ __global__ void to_bf16_kernel(const float* __restrict__ x, uint16_t* __restrict void kv_append(const float* K, const float* V, int64_t T, int64_t pos0, const int32_t* page_table, int64_t page_size, uint16_t* k_pool, uint16_t* v_pool, int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, - void* stream, const strata::kernels::KvHostPools* host, const strata::kernels::KvHostPools* stage) { + void* stream) { if (T <= 0) return; - kv_append_kernel<<>>( - K, V, pos0, page_table, page_size, k_pool, v_pool, k_q, v_q, k_scale, v_scale, - host ? *host : strata::kernels::KvHostPools{}, stage ? *stage : strata::kernels::KvHostPools{}); + kv_append_kernel<<>>(K, V, pos0, page_table, page_size, k_pool, + v_pool, k_q, v_q, k_scale, v_scale); check("kv_append"); } void to_f16(const float* x, uint16_t* y, int64_t n, void* stream) { @@ -566,11 +528,6 @@ void to_f16(const float* x, uint16_t* y, int64_t n, void* stream) { to_f16_kernel<<<(unsigned) ((n + 255) / 256 < 4096 ? (n + 255) / 256 : 4096), 256, 0, (cudaStream_t) stream>>>(x, y, n); check("to_f16"); } -void round_f16(const float* x, float* y, int64_t n, void* stream) { - if (n <= 0) return; - round_f16_kernel<<<(unsigned) ((n + 255) / 256 < 4096 ? (n + 255) / 256 : 4096), 256, 0, (cudaStream_t) stream>>>(x, y, n); - check("round_f16"); -} void to_bf16(const float* x, uint16_t* y, int64_t n, void* stream) { if (n <= 0) return; to_bf16_kernel<<<(unsigned) ((n + 255) / 256 < 4096 ? (n + 255) / 256 : 4096), 256, 0, (cudaStream_t) stream>>>(x, y, n); @@ -591,10 +548,10 @@ void gr_mix_r(const float* R, const float* rs, const float* w_norm, const float* mixed_h); check("gr_mix_r"); } -void gr_write_norm_rs(float* R, const float* bo, const float* inj, int64_t inj_ld, const float* w_norm_next, float eps, - float* rs, uint16_t* xn16, int64_t T, void* stream) { - gr_write_norm_rs_kernel<<<(unsigned) (T * HC), 256, 0, (cudaStream_t) stream>>>(R, bo, inj, inj_ld, w_norm_next, eps, - rs, xn16); +void gr_write_norm_rs(float* R, const float* bo, const float* inj, int64_t inj_ld, const float* w_norm, float eps, + float* rs, uint16_t* xn16, int64_t T, void* stream, const uint16_t* partial) { + gr_write_norm_rs_kernel<<<(unsigned) (T * HC), 256, 0, (cudaStream_t) stream>>>(R, bo, inj, inj_ld, w_norm, eps, + rs, xn16, partial); check("gr_write_norm_rs"); } void gr_silu(const float* lo, uint16_t* lo16, int64_t T, void* stream) { @@ -606,8 +563,9 @@ void gr_mix(const float* xn, const float* gated, float* mixed, uint16_t* mixed16 gr_mix_kernel<<>>(xn, gated, mixed, mixed16, T, mixed_h); check("gr_mix"); } -void gr_write(float* R, const float* bo, const float* inj, int64_t inj_ld, int64_t T, void* stream) { - gr_write_kernel<<>>(R, bo, inj, inj_ld, T); +void gr_write(float* R, const float* bo, const float* inj, int64_t inj_ld, int64_t T, void* stream, + const uint16_t* partial) { + gr_write_kernel<<>>(R, bo, inj, inj_ld, T, partial); check("gr_write"); } void gr_broadcast(const float* e, float* R, int64_t T, void* stream) { @@ -619,35 +577,28 @@ void gdn_gates(const float* ab, const float* dt, const float* ssm_a, float* gate check("gdn_gates"); } void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream) { - static const bool serial = std::getenv("STRATA_GDN_CONV_SERIAL") != nullptr; // the old walk (A/B) - if (serial || T <= CONV_TILE) { - gdn_conv_kernel<<>>(history, qkv, conv_w, h, T); - } else { - gdn_conv_tiled_kernel<<>>(history, qkv, conv_w, h, T); - gdn_conv_hist_kernel<<>>(history, qkv, T); - } - gdn_l2_kernel<<>>(h, eps); + for (int64_t t0 = 0; t0 < T; t0 += 65535) + strata::kernels::gdn_conv_l2_multi(history, qkv, conv_w, h, C, 2 * HK, eps, (int) std::min(65535, T - t0), + stream, (int) t0); + gdn_conv_hist_kernel<<>>(history, qkv, T); check("gdn_conv"); } -void gdn_recurrence(float* state, const float* h, const float* gate, const float* beta, const float* z, - const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, void* stream) { - static const bool serial = std::getenv("STRATA_GDN_REC_HEADS") != nullptr; // the one-block-per-head kernel (A/B) - if (serial || T <= 0) { - gdn_rec_kernel<<>>(state, h, gate, beta, z, gamma, eps, y, y16, T); - } else { - gdn_rec_cols_kernel<<>>(state, h, gate, beta, y, T); - gdn_out_norm_kernel<<>>(z, gamma, eps, y, y16); +void gdn_scan(float* state, const float* h, const float* gate, const float* beta, float* y, int64_t T, void* stream) { + // 32 columns a block, 2 a thread (fastest of 16-128 columns, 1-2 a thread: 0.67 us a token on the 3090) + constexpr int CB = 32, CPT = 2; + gdn_scan_kernel<<>>(state, h, gate, beta, y, (int) T); + check("gdn_scan"); +} +void gdn_out_norm(float* y, const float* z, const float* gamma, float eps, uint16_t* y16, int64_t T, void* stream) { + for (int64_t t0 = 0; t0 < T; t0 += 65535) { + const int64_t n = std::min(65535, T - t0); + gdn_norm_kernel<<>>(y + t0 * HV * S, z + t0 * HV * S, gamma, + eps, y16 + t0 * HV * S); } - check("gdn_recurrence"); -} -void route(const float* logits, int32_t* ids, float* weights, int64_t T, int64_t n_expert, void* stream) { - if (n_expert == 512) - route_kernel<16><<<(unsigned) ((T + 7) / 8), 256, 0, (cudaStream_t) stream>>>(logits, ids, weights, T); - else if (n_expert == 256) - route_kernel<8><<<(unsigned) ((T + 7) / 8), 256, 0, (cudaStream_t) stream>>>(logits, ids, weights, T); - else - strata::kernels::router_top10(logits, (int) T, (int) n_expert, 10, ids, weights, stream); + check("gdn_out_norm"); +} +void route(const float* logits, int32_t* ids, float* weights, int64_t T, void* stream) { + route_kernel<<<(unsigned) ((T + 7) / 8), 256, 0, (cudaStream_t) stream>>>(logits, ids, weights, T); check("route"); } void blob_dequant(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, void* stream) { @@ -667,27 +618,53 @@ void swiglu_pair(const float* g, const float* u, uint16_t* h16, int64_t n, void* swiglu_pair_kernel<<>>(g, u, h16, n); check("swiglu_pair"); } -namespace { -__global__ void copy_i32_kernel(int32_t* __restrict__ dst, const int32_t* __restrict__ src, int64_t n) { - for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x * blockDim.x) - dst[i] = src[i]; -} -} // namespace -void copy_i32(int32_t* dst, const int32_t* src, int64_t n, void* stream) { - if (n <= 0) return; - const int64_t b = (n + 255) / 256; - copy_i32_kernel<<<(unsigned) (b < 256 ? b : 256), 256, 0, (cudaStream_t) stream>>>(dst, src, n); - check("copy_i32"); -} void gather_rows16(const uint16_t* x16, const int32_t* src, uint16_t* dst16, int64_t n, int64_t width, void* stream) { if (n <= 0) return; gather_rows16_kernel<<>>(x16, src, dst16, n, width); check("gather_rows16"); } -void moe_combine(const float* Dm, const int32_t* slot, const float* w, const float* shared, const float* sg, float* bo, - int64_t T, void* stream) { - moe_combine_kernel<<>>(Dm, slot, w, shared, sg, bo, T); - check("moe_combine"); +void gather_rows16_f32(const uint16_t* x16, const int32_t* src, float* dst, int64_t n, int64_t width, void* stream) { + if (n <= 0) return; + gather_rows16_f32_kernel<<>>(x16, src, dst, n, width); + check("gather_rows16_f32"); +} +__global__ void sums_to_f16_kernel(const float* __restrict__ x, uint16_t* __restrict__ y, int64_t n) { + const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) y[i] = hf(fminf(fmaxf(x[i], -65504.0f), 65504.0f)); +} +void sums_to_f16(const float* x, uint16_t* y, int64_t n, void* stream) { + sums_to_f16_kernel<<>>(x, y, n); + check("sums_to_f16"); +} +void moe_shared_gate(float* x, const float* sg, int64_t T, void* stream) { + moe_shared_gate_kernel<<>>(x, sg, T); + check("moe_shared_gate"); +} +__global__ void moe_gather_add_kernel(float* __restrict__ sum, const float* __restrict__ rows, int64_t r0, + const float* __restrict__ w, const int32_t* __restrict__ tok, + const int32_t* __restrict__ start, const int32_t* __restrict__ list) { + const int64_t d = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (d >= N) return; + const int b = blockIdx.y; + float* o = sum + (int64_t) tok[b] * N + d; + float s = *o; + for (int i = start[b]; i < start[b + 1]; ++i) { + const int r = list[i]; + s = fmaf(w[r], rows[(r - r0) * N + d], s); + } + *o = s; +} +void moe_gather_add(float* sum, const float* rows, int64_t r0, const float* w, const int32_t* tok, const int32_t* start, + const int32_t* list, int64_t n_tok, void* stream) { + if (n_tok <= 0) return; + moe_gather_add_kernel<<>>( + sum, rows, r0, w, tok, start, list); + check("moe_gather_add"); +} +void moe_scatter_add(float* sum, const float* rows, const float* w, const int32_t* src, int64_t n, void* stream) { + if (n <= 0) return; + moe_scatter_add_kernel<<>>(sum, rows, w, src, n); + check("moe_scatter_add"); } void rms_rows(float* x, const float* w, int64_t rows, int64_t cols, int64_t ld, float eps, void* stream) { if (rows <= 0) return; diff --git a/src/prefill/moe_mmq.cu b/src/prefill/moe_mmq.cu index ac2daee4f56..d43f5b5889d 100644 --- a/src/prefill/moe_mmq.cu +++ b/src/prefill/moe_mmq.cu @@ -91,9 +91,11 @@ bool built() { return true; } bool supported(int t) { switch ((ggml_type) t) { - case GGML_TYPE_Q2_0: case GGML_TYPE_IQ2_XXS: case GGML_TYPE_IQ2_XS: case GGML_TYPE_IQ2_S: - case GGML_TYPE_IQ3_XXS: case GGML_TYPE_IQ3_S: case GGML_TYPE_IQ4_NL: case GGML_TYPE_IQ4_XS: - case GGML_TYPE_Q8_0: // the draft layer's dense matrices (E-9) + case GGML_TYPE_Q2_0: case GGML_TYPE_IQ1_S: case GGML_TYPE_IQ2_XXS: case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: case GGML_TYPE_IQ3_XXS: case GGML_TYPE_IQ3_S: case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_Q4_K: case GGML_TYPE_Q5_K: case GGML_TYPE_Q6_K: case GGML_TYPE_Q5_1: // unsloth's files + case GGML_TYPE_Q8_0: return true; default: return false; @@ -143,6 +145,11 @@ void Context::run(const Product& p, void* stream) { case GGML_TYPE_IQ4_NL: mul_mat_q_case(ctx, a, s); break; case GGML_TYPE_IQ4_XS: mul_mat_q_case(ctx, a, s); break; case GGML_TYPE_Q8_0: mul_mat_q_case(ctx, a, s); break; + case GGML_TYPE_IQ1_S: mul_mat_q_case(ctx, a, s); break; + case GGML_TYPE_Q4_K: mul_mat_q_case(ctx, a, s); break; + case GGML_TYPE_Q5_K: mul_mat_q_case(ctx, a, s); break; + case GGML_TYPE_Q6_K: mul_mat_q_case(ctx, a, s); break; + case GGML_TYPE_Q5_1: mul_mat_q_case(ctx, a, s); break; default: std::fprintf(stderr, "prefill mmq: type %d is not covered\n", (int) t); std::exit(1); diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index 2dcdb54362b..aa712066478 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -1,64 +1,27 @@ // src/prefill/prefill.cpp - see include/strata/prefill/prefill.hpp. #include "strata/prefill/prefill.hpp" -#include "strata/core/mtp.hpp" -#include "strata/core/progress.hpp" -#include "strata/core/on_device.hpp" #include "strata/core/layout.hpp" -#include "strata/kernels/cpu/expert.hpp" +#include "strata/kernels/dequant_bf16.hpp" #include "strata/kernels/native_qsa_indexer.hpp" -#include "strata/kernels/ngram.hpp" -#include "strata/kernels/ple.hpp" #include "strata/kernels/native_ple_postops.hpp" -#include "strata/kernels/iq_kernels.hpp" -#include "strata/kernels/cpu/expert_layout.hpp" +#include "strata/kernels/ngram.hpp" #include "strata/kernels/qsa.hpp" -#include "strata/kernels/kv_stream.hpp" -#include "strata/kernels/cvec.hpp" -#include "strata/kernels/kv_q4.hpp" -#include "strata/core/layer.hpp" -#include "strata/core/native_head.hpp" -#include "strata/kernels/verify_kernels.hpp" #include "strata/kernels/qsa_decode_attn.hpp" -#include "strata/kernels/qsa_prompt_attn.hpp" #include "strata/kernels/qsa_select.hpp" +#include "strata/kernels/verify_kernels.hpp" #include "strata/prefill/gemm.hpp" -#include "strata/prefill/moe_mmq.hpp" #include "strata/prefill/kernels.hpp" #include #include #include -#include #include -#include -#include -#include -#include -#include -#include +#include #include #include -#ifndef STRATA_PREFILL_MMQ -// A build without the llama.cpp sources (no STRATA_NATIVE_EXPERTS): no MMQ, the FP16 expert path everywhere. -namespace strata::prefill::mmq { -bool built() { return false; } -bool supported(int) { return false; } -size_t matrix_bytes(int, int64_t, int64_t) { return 0; } -size_t q8_bytes(int64_t, int64_t) { return 0; } -void quantize(const float*, const int32_t*, void*, int, int64_t, int64_t, int64_t, void*) {} -Context::Context() {} -Context::~Context() {} -void Context::run(const Product&, void*) {} -void gather_native(const void*, const void*, size_t, const void*, size_t, void*, void*, void*) {} -void gather_strata_q2(const uint8_t*, void*, void*, void*) {} -void swiglu(const float*, float*, int64_t, int64_t, bool, void*) {} -void iota(int32_t*, int64_t, void*) {} -} // namespace strata::prefill::mmq -#endif - namespace strata::prefill { namespace { @@ -66,795 +29,259 @@ using Clock = std::chrono::steady_clock; constexpr float EPS = 1e-6f; constexpr int64_t N = 2560, HC = 4, D = N * HC, LR = 320, K = 10, NE = 512; constexpr int64_t C = 10240, ZV = 6144, HV = 48; -// plan v0.3 P6: staging holds the largest blob of the pack (a native pack's blobs differ per layer) -inline int64_t MAXBLOB() { return (int64_t) strata::kernels::cpu::expert_layout().max_blob; } -constexpr int STAGE = 8; // host->device expert staging ring (chunks below STREAM_ALL_MIN) -// Step 3: from this chunk size on, every non-resident expert of every layer streams in a fixed order through a -// RING_MAX-slot ring (nearly all 512 are routed at such a chunk), so the copy engine keeps working through the -// attention halves instead of waiting for each layer's routing. -constexpr int RING_MAX = 512; // the arrays; the ring itself is ring_slots() -constexpr int64_t STREAM_ALL_MIN = 2048; -double g_pinned_share = 1.0; -// The streamed ring: 384 slots when (nearly) every streamed expert is DMA'd from pinned RAM - measured on Q2_0, -// 8192-token chunks: 96 slots 1153 tok/s, 384 1294 (the next layer's experts arrive during its attention half) - -// and 96 when a large share goes through host copies (IQ3_S on 64 GB, a third unpinned: 96 slots 1216, 256 1070 - -// the host copies are the limit and the bigger ring only takes cache slots). STRATA_PREFILL_RING overrides. -inline int ring_slots(size_t T) { - const char* v = std::getenv("STRATA_PREFILL_RING"); - const int r = v ? std::atoi(v) : (g_pinned_share >= 0.9 ? 384 : 96); - const int big = r < 16 ? 16 : r > RING_MAX ? RING_MAX : r; - return (int64_t) T >= STREAM_ALL_MIN ? big : STAGE; -} -constexpr int DQ = 2; // dequantized-expert ring (FP16 gate/up + down) - -// F-1: STRATA_GR_UNFUSED=1 keeps the FP32 copy of the normalized rows (gr_norm + gr_mix), the A/B arm -inline bool gr_unfused() { - static const bool v = [] { const char* e = std::getenv("STRATA_GR_UNFUSED"); return e && e[0] == '1'; }(); - return v; -} +constexpr int64_t SUB = 2048; // tokens of a sub-chunk: everything but the routed experts +constexpr int64_t SEL_BATCH = 256; // queries per block-score launch +// A layer's quantized projections, dequantized to FP16 once for all its sub-chunks: the PLE key, or the mixer's +// (GDN: qkv, gate, out; QSA: q, k, v, out) and the shared expert's +constexpr int64_t W16 = std::max(D * N, std::max(C + ZV + ZV, 12288 + 512 + 512 + ZV) * N + 3 * 640 * N); +constexpr size_t GEMM_WS = 32u << 20; // cuBLAS workspace +constexpr int GGML_Q8_0 = 8; double ms_since(Clock::time_point t) { return std::chrono::duration(Clock::now() - t).count(); } -// Either cudaMalloc (owned, freed with the object) or a bump allocation from a borrowed region; with no base and -// no region it only counts, which is how `bytes_needed` sizes the region. +// Bump allocation from a region; without one it only counts (bytes_needed). struct Alloc { uint8_t* base = nullptr; uint64_t cap = 0, used = 0; - bool count_only = false; - std::vector* owned = nullptr; - template T* take(size_t n, bool& ok) { + bool ok = true; + template T* take(size_t n) { const uint64_t bytes = ((uint64_t) n * sizeof(T) + 256 + 255) & ~255ull; - if (count_only) { used += bytes; return nullptr; } - if (base != nullptr) { - if (used + bytes > cap) { ok = false; return nullptr; } - T* p = (T*) (base + used); - used += bytes; - return p; - } - void* p = nullptr; - if (cudaMalloc(&p, bytes) != cudaSuccess) { ok = false; return nullptr; } - owned->push_back(p); + T* p = base != nullptr ? (T*) (base + used) : nullptr; used += bytes; - return (T*) p; + if (base != nullptr && used > cap) ok = false; + return p; } }; -} // namespace +// The device buffers of one binding. +struct Bufs { + uint16_t* w16 = nullptr; + void* gemm_ws = nullptr; + // the chunk's: residual stream, the MoE's input (FP16) and output, the FFN half's injection, routing, steps, PLE + // rows, the second GPU's sums (FP16) + float *R = nullptr, *bo_moe = nullptr, *inj_f = nullptr, *w = nullptr, *ple_emb = nullptr; + uint16_t *mixed_h = nullptr, *sums = nullptr; + int32_t *ids = nullptr, *steps = nullptr; + // a sub-chunk's: the mixer's BF16 input, its output, the mixer half's injection, a read's row scales + uint16_t* mixed_bf = nullptr; + float *bo = nullptr, *inj_a = nullptr, *rs = nullptr; + // the union: what one step needs at a time + float* emb = nullptr; + float *xn = nullptr, *lo = nullptr, *gated = nullptr; + uint16_t *xn16 = nullptr, *lo16 = nullptr; + float *ple_kp = nullptr, *ple_key = nullptr, *ple_q = nullptr, *ple_gated = nullptr; + float *qkv = nullptr, *z = nullptr, *ab = nullptr, *gate = nullptr, *beta = nullptr, *hbuf = nullptr, *y = nullptr; + uint16_t* y_h = nullptr; + float *Kc = nullptr, *Vc = nullptr, *Qf = nullptr, *q = nullptr, *idx_raw = nullptr, *q_idx = nullptr, + *attn = nullptr, *sel_scores = nullptr; + uint16_t* attn_h = nullptr; + int32_t* sel_ids = nullptr; + float *logits = nullptr, *sgate = nullptr, *sup = nullptr, *sg = nullptr; + uint16_t* sh_h = nullptr; + uint8_t* local = nullptr; // this GPU's expert runner + uint64_t local_bytes = 0; +}; -// Step 4 of the prompt-speed plan: the experts the arena could not pin (a third of the streamed ones on IQ3_S) are -// copied into pinned buffers by these threads, ahead of the launches. Copied in line by the launching thread they -// left the GPU without queued work while each ~2 MB memcpy ran (~15 s of a 32K prompt on IQ3_S). Job j - a layer's -// j-th unpinned expert, in launch order - lands in host buffer j % kRing, which is free again once the DMA of job -// j - kRing (recorded by the launching thread, `issued`) is done. -struct Stager { - // D-5: the pinned ring's depth (STRATA_STAGER_RING, default 16) - how far the host copies can run ahead of the - // DMAs of the unpinned experts' blobs - int kRing = 16; - struct Job { const uint8_t* src; size_t bytes; }; - std::vector buf; - std::vector pinned; - std::vector> pageable; // the fallback when no more RAM can be pinned - std::vector dma_done; - std::vector jobs; - std::unique_ptr[]> ready; - size_t ready_cap = 0; - // gen << 32 | n << 16 | next index: a claim is a CAS on the generation it woke for (a thread late from the - // previous layer can never take a job of this one - the expert pool's issue #29 lesson) - std::atomic head{0}; - std::atomic issued{0}, active{0}; - uint32_t gen = 0; - bool quit = false; - std::mutex mu; - std::condition_variable cv; - std::vector threads; - int device = 0; +// Chunks of T tokens: the chunk-long buffers, the sub-chunk ones, then the union of the steps a layer takes one after +// the other on the stream. (Staging copies run ahead on a copy stream only while the routed experts run, after the +// host has waited for the stream, and the stream waits for them before the next step.) +Bufs layout(Alloc& a, int64_t T, int64_t cap, int64_t max_blocks, bool offload) { + const size_t t = (size_t) T, p = (size_t) std::min(T, SUB); + Bufs b; + b.w16 = a.take((size_t) W16); + b.gemm_ws = a.take(GEMM_WS); + b.R = a.take(t * D); + b.mixed_h = a.take(t * N); + b.bo_moe = a.take(t * N); + b.inj_f = a.take(t * HC); + b.ids = a.take(t * K); + b.w = a.take(t * K); + b.steps = a.take(t * strata::kernels::kStepCount); + b.ple_emb = a.take(t * N); + b.sums = offload ? a.take(t * N) : nullptr; // read by the next layer's sub-chunks + b.mixed_bf = a.take(p * N); + b.bo = a.take(p * N); + b.inj_a = a.take(p * HC); + b.rs = a.take(p * HC); + const uint64_t u0 = a.used; + uint64_t u1 = u0; + auto next = [&] { u1 = std::max(u1, a.used); a.used = u0; }; + b.emb = a.take(p * N); + next(); // the hyper-connection read + b.xn = a.take(p * D); b.xn16 = a.take(p * D); b.lo = a.take(p * LR); + b.lo16 = a.take(p * LR); b.gated = a.take(p * D); + next(); // the PLE block + b.ple_kp = a.take(p * D); b.ple_key = a.take(p * D); + b.ple_q = a.take(p * D); b.ple_gated = a.take(p * D); + next(); // GDN + b.qkv = a.take(p * C); b.z = a.take(p * ZV); b.ab = a.take(p * 2 * HV); + b.gate = a.take(p * HV); b.beta = a.take(p * HV); b.hbuf = a.take(p * C); + b.y = a.take(p * ZV); b.y_h = a.take(p * ZV); + next(); // QSA + b.Kc = a.take(p * 512); b.Vc = a.take(p * 512); b.Qf = a.take(p * 12288); + b.q = a.take(p * ZV); b.idx_raw = a.take(p * 128); b.q_idx = a.take(p * 512); + b.attn = a.take(p * ZV); b.attn_h = a.take(p * ZV); b.sel_ids = a.take(p * (size_t) cap); + b.sel_scores = a.take((size_t) SEL_BATCH * (size_t) max_blocks); + next(); // the router and the shared expert + b.logits = a.take(p * NE); b.sgate = a.take(p * 640); b.sup = a.take(p * 640); + b.sh_h = a.take(p * 640); b.sg = a.take(p); + next(); // the routed experts + b.local_bytes = ExpertRunner::bytes_needed(T, NE, false, !offload); + b.local = a.take(b.local_bytes); + next(); + a.used = u1; + return b; +} - bool init(size_t blob_bytes, int nthreads) { - if (const char* v = std::getenv("STRATA_STAGER_RING")) kRing = std::clamp(std::atoi(v), 2, 256); - buf.assign((size_t) kRing, nullptr); - pinned.assign((size_t) kRing, 0); - dma_done.assign((size_t) kRing, nullptr); - pageable.resize(kRing); - for (int i = 0; i < kRing; ++i) { - pinned[i] = cudaHostAlloc((void**) &buf[i], blob_bytes, cudaHostAllocDefault) == cudaSuccess; - if (!pinned[i]) { - cudaGetLastError(); - pageable[(size_t) i].resize(blob_bytes); - buf[i] = pageable[(size_t) i].data(); - } - if (cudaEventCreateWithFlags(&dma_done[i], cudaEventDisableTiming) != cudaSuccess) return false; - } - cudaGetDevice(&device); - for (int t = 0; t < nthreads; ++t) threads.emplace_back([this] { work(); }); - return true; - } - ~Stager() { - finish(); - { std::lock_guard lk(mu); quit = true; } - cv.notify_all(); - for (auto& t : threads) t.join(); - for (int i = 0; i < kRing; ++i) { - if (dma_done[i]) cudaEventDestroy(dma_done[i]); - if (buf[i] && pinned[i]) cudaFreeHost(buf[i]); - } - } - void work() { - cudaSetDevice(device); - uint32_t seen = 0; - for (;;) { - { - std::unique_lock lk(mu); - cv.wait(lk, [&] { return quit || gen != seen; }); - if (quit) return; - seen = gen; - } - for (;;) { - active.fetch_add(1, std::memory_order_acq_rel); - const int j = claim(seen); - if (j < 0) { active.fetch_sub(1, std::memory_order_acq_rel); break; } - const int b = j % kRing; - if (j >= kRing) { - while (issued.load(std::memory_order_acquire) <= j - kRing) std::this_thread::yield(); - cudaEventSynchronize(dma_done[b]); - } - std::memcpy(buf[b], jobs[(size_t) j].src, jobs[(size_t) j].bytes); - ready[(size_t) j].store(1, std::memory_order_release); - active.fetch_sub(1, std::memory_order_acq_rel); - } - } - } - int claim(uint32_t g) { - uint64_t cur = head.load(std::memory_order_acquire); - for (;;) { - if ((uint32_t) (cur >> 32) != g) return -1; - const int n = (int) ((cur >> 16) & 0xffff), j = (int) (cur & 0xffff); - if (j >= n) return -1; - if (head.compare_exchange_weak(cur, cur + 1, std::memory_order_acq_rel, std::memory_order_acquire)) return j; - } - } - /// A layer's jobs; the previous layer's are finished (finish()). - void start(std::vector&& js) { - if (js.empty()) return; - std::lock_guard lk(mu); - jobs = std::move(js); - if (ready_cap < jobs.size()) { - ready_cap = jobs.size() * 2; - ready.reset(new std::atomic[ready_cap]); - } - for (size_t i = 0; i < jobs.size(); ++i) ready[i].store(0, std::memory_order_relaxed); - issued.store(0); - ++gen; - head.store((uint64_t) gen << 32 | (uint64_t) jobs.size() << 16, std::memory_order_release); - cv.notify_all(); - } - /// Job j's bytes, in a pinned buffer (waits for the copy). - const uint8_t* wait(int j) { - while (!ready[(size_t) j].load(std::memory_order_acquire)) std::this_thread::yield(); - return buf[j % kRing]; - } - /// The launching thread queued job j's DMA on `copy`: its buffer is free once that is done. - void issued_one(int j, cudaStream_t copy) { - cudaEventRecord(dma_done[j % kRing], copy); - issued.store(j + 1, std::memory_order_release); - } - /// No job is running after this (the end of a layer, or an early return in the middle of one). - void finish() { - head.store((uint64_t) gen << 32, std::memory_order_release); // n = 0: nothing more to claim - issued.store(1 << 30, std::memory_order_release); - while (active.load(std::memory_order_acquire) != 0) std::this_thread::yield(); - } -}; +strata::kernels::QsaShapes shapes_of(const core::ModelGeometry& g) { + strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); + s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; + s.idx_dim = g.idx_key_dim; + return s; +} + +} // namespace struct Prefill::Impl { const core::WeightTable* wt = nullptr; const core::ModelGeometry* g = nullptr; core::SessionState* ss = nullptr; - core::ExpertSource* src = nullptr; const core::ExpertCache* cache = nullptr; const int32_t* host_res = nullptr; - int64_t T = 0, T_max = 0; - bool borrowed = false; - cudaStream_t cs = nullptr, copy = nullptr; + ExpertRunner* offload = nullptr; + ExpertRunner local; + int64_t max_chunk = 0, T = 0; // T: the chunk the buffers are bound for + strata::kernels::QsaShapes s; + int64_t cap = 0, max_blocks = 0; + cudaStream_t cs = nullptr; + // the second GPU's share, a sub-chunk at a time: the input goes out on `xfer`, the sums come back on `xsum` + cudaStream_t xfer = nullptr, xsum = nullptr; + cudaEvent_t ev_mixed = nullptr; + std::vector ev_in, ev_piece; // a sub-chunk's input is in host memory, its sums on this GPU Gemm gemm; - std::vector owned; - // chunk buffers - float *emb = nullptr, *R = nullptr, *xn = nullptr, *lo = nullptr, *gated = nullptr, *inj = nullptr; - float* grs = nullptr; // F-1: the hyper-connection read's row scales (T x 4) - uint16_t *xn16 = nullptr, *lo16 = nullptr; - float* mixed = nullptr; - uint16_t *mixed_bf = nullptr, *mixed_h = nullptr; - float* bo = nullptr; - // GDN - float *qkv = nullptr, *z = nullptr, *ab = nullptr, *gate = nullptr, *beta = nullptr, *hbuf = nullptr, *y = nullptr; - uint16_t* y_h = nullptr; - // QSA - float *Kc = nullptr, *Vc = nullptr, *Qf = nullptr, *q = nullptr, *idx_raw = nullptr, *q_idx = nullptr, *attn = nullptr; - uint16_t* attn_h = nullptr; - int32_t* steps_dev = nullptr; - std::vector steps_host; - int32_t* sel_ids = nullptr; - float* sel_scores = nullptr; // [sel_batch, max_blocks] - int64_t sel_batch = 256, max_blocks = 0; - float* attn_scratch = nullptr; - int64_t attn_batch = 32, cap = 0; - // MoE - float *logits = nullptr, *w = nullptr, *GU = nullptr, *Dm = nullptr, *sgate = nullptr, *sup = nullptr, - *shared = nullptr, *sg = nullptr; - int32_t *ids = nullptr, *slot_dev = nullptr, *src_dev = nullptr; - uint16_t *Xs = nullptr, *Hh = nullptr, *sh_h = nullptr; - // step 2b (MMQ): the activations quantized per layer, H in FP32 and its group's quantized rows, the identity - // row map, the group bounds, the group buffers of gathered experts - void *Xq = nullptr, *Hq = nullptr; - float* H = nullptr; - int32_t *ids_identity = nullptr, *bounds_dev = nullptr; - uint8_t *grp_gu = nullptr, *grp_d = nullptr; - std::vector bounds_host; - std::unique_ptr mmq_ctx; - std::vector ids_host, slot_host, src_host, cnt, off; - // The grouping tables in mapped pinned memory, [ids | slot | src] of T_max * K each, then the MMQ bounds: kernels - // read and write them in place. A cudaMemcpyAsync of them queues behind the expert blobs the copy stream already - // holds (up to `ring` of them, ~70 us each), and the GPU idles meanwhile - measured 4.2 s of a 128K prompt's - // 70 s at the default ring, 0.7 s with a 16-slot one. STRATA_GROUP_COPY=1: the copies (the A/B arm). - int32_t* grp_host = nullptr; - int32_t* grp_dev = nullptr; // its device alias - size_t grp_n = 0, grp_tk = 0; // int32s allocated; T_max * K (the offset of slot, and of src past it) - uint16_t* dq_gu[DQ] = {}; - uint16_t* dq_d[DQ] = {}; - uint8_t* stage_dev[RING_MAX] = {}; - int ring = STAGE; // the slots of this layout's ring (ring_slots) - std::unique_ptr stager; // the unpinned experts' host copies (step 4) - cudaEvent_t copied[RING_MAX] = {}, used[RING_MAX] = {}; - bool stage_live[RING_MAX] = {}; - // PLE - float* ple_emb = nullptr; - std::vector ple_pageable[2]; // the fallback when no more RAM can be pinned - float* ple_emb_host[2] = {}; // pinned, double-buffered: the next chunk's rows are read while this - cudaEvent_t ple_copied[2] = {}; // one runs; the event marks that buffer's upload done + void* owned = nullptr; // the buffers' own allocation (no region) + Bufs b; + std::vector> w16_have; // this layer's dequantized weights + int64_t w16_used = 0; + int32_t* ids_host = nullptr; // pinned [max_chunk, K] + float* w_host = nullptr; + std::vector steps_host, cnt, off, entry_of; + std::vector ex1, off1, src1, ex2, off2, src2; // this GPU's experts, the second GPU's + std::vector w1, w2; + // a chunk's PLE rows, read on a thread while the chunk before runs: two pinned buffers, their row indices + float* ple_host[2] = {}; std::vector ple_rows[2]; - float* ple_norm = nullptr; - uint8_t* region = nullptr; // the attention/MoE scratch region (idle while the PLE block runs) - uint64_t region_bytes = 0; - PrefillStats* stats = nullptr; - // KV streaming: one layer's whole K/V, staged from the host copy per layer and chunk (identity layout) - strata::kernels::KvHostPools stage; - int32_t* ident_table = nullptr; - // layer split: the device, and the hand-off to the next stage (two pinned chunk buffers, used in turn) - int device = -1; - float* hand[2] = {}; - // C-4: the chunk's token ids on the device, for one batched embedding gather - int32_t* tok_dev = nullptr; - std::vector tok_host; }; -namespace { -// the staging pool of a streamed session: every page of one layer (same sequence in init and bytes_needed) -// STRATA_KV_STAGE_OWN (A/B only): the staging pool gets its own allocation instead of borrowed expert slots, so a -// streamed run lends the prompt path exactly the slots a resident one does (a lent expert runs on the CPU, which -// rounds differently: without this an A/B compares two expert placements as well as two KV placements) -bool stage_own() { static const bool v = std::getenv("STRATA_KV_STAGE_OWN") != nullptr; return v; } -void take_stage(Alloc& o_borrowed, const core::SessionState& ss, const strata::kernels::QsaShapes& s, - strata::kernels::KvHostPools& st, bool& ok) { - const core::QsaState& q0 = ss.qsa_states[0]; - if (q0.kv_mode != 1) return; - if (stage_own() && o_borrowed.count_only) return; - Alloc own; - own.owned = o_borrowed.owned; - Alloc& o = stage_own() ? own : o_borrowed; - const size_t rows = (size_t) q0.n_pages * s.n_head_kv * s.page_size; - if (q0.kv_q4) { - st.k_q4 = o.take(rows * strata::kernels::kv_q4_bytes_per_head((int) s.head_dim), ok); - st.v_q4 = o.take(rows * strata::kernels::kv_q4_bytes_per_head((int) s.head_dim), ok); - } else if (q0.kv_int8) { - st.k_q = o.take(rows * s.head_dim, ok); - st.v_q = o.take(rows * s.head_dim, ok); - st.k_scale = o.take(rows * (s.head_dim / 64), ok); - st.v_scale = o.take(rows * (s.head_dim / 64), ok); - } else { - st.k_pool = o.take(rows * s.head_dim, ok); - st.v_pool = o.take(rows * s.head_dim, ok); - } -} -strata::kernels::QsaAttnPools pools_of(const strata::kernels::KvHostPools& h, const int32_t* table) { - strata::kernels::QsaAttnPools p; - p.k_pool = h.k_pool; p.v_pool = h.v_pool; p.k_q = h.k_q; p.v_q = h.v_q; p.k_scale = h.k_scale; p.v_scale = h.v_scale; - p.k_q4 = h.k_q4; p.v_q4 = h.v_q4; - p.page_table = table; - return p; -} -} // namespace - Prefill::Prefill() : impl_(new Impl) {} Prefill::~Prefill() { if (!impl_) return; - if (impl_->cs) cudaStreamSynchronize(impl_->cs); - if (impl_->copy) cudaStreamSynchronize(impl_->copy); - for (int i = 0; i < RING_MAX; ++i) { - if (impl_->copied[i]) cudaEventDestroy(impl_->copied[i]); - if (impl_->used[i]) cudaEventDestroy(impl_->used[i]); - } - for (int b = 0; b < 2; ++b) { - if (impl_->hand[b]) cudaFreeHost(impl_->hand[b]); - if (impl_->ple_copied[b]) cudaEventDestroy(impl_->ple_copied[b]); - if (impl_->ple_emb_host[b] && impl_->ple_pageable[b].empty()) cudaFreeHost(impl_->ple_emb_host[b]); - } - if (impl_->copy) cudaStreamDestroy(impl_->copy); - if (impl_->grp_host) cudaFreeHost(impl_->grp_host); - for (void* p : impl_->owned) cudaFree(p); -} - -namespace { -constexpr int64_t GEMM_SCRATCH = 32ll << 20; // FP16 elements for the largest dequantized dense weight -constexpr size_t GEMM_WS = 32u << 20; // cuBLAS workspace - -// THE ATTENTION HALF AND THE MoE HALF SHARE THEIR BUFFERS. A layer runs its attention (GDN or QSA), writes it back -// into the residual, and only then its MoE, so the three sets of scratch are never live at once: one region the size -// of the largest holds them all. That is ~260 KB of the ~680 KB a prompt token cost - which is what lets a chunk -// grow (every expert is streamed once per chunk, so a bigger chunk streams fewer bytes per token). The sizes are -// counted with the same `take` sequence `init` uses; a mismatch makes `init` fail with "do not fit", never overlap. -uint64_t gdn_set_bytes(size_t T) { - Alloc a; a.count_only = true; bool ok = true; - a.take(T * C, ok); a.take(T * ZV, ok); a.take(T * 2 * HV, ok); a.take(T * HV, ok); - a.take(T * HV, ok); a.take(T * C, ok); a.take(T * ZV, ok); a.take(T * ZV, ok); - return a.used; -} -uint64_t qsa_set_bytes(size_t T, int64_t cap, int64_t max_blocks, int64_t sel_batch, int64_t attn_batch, - const strata::kernels::QsaShapes& s) { - Alloc a; a.count_only = true; bool ok = true; - a.take(T * 512, ok); a.take(T * 512, ok); a.take(T * 12288, ok); a.take(T * ZV, ok); - a.take(T * 128, ok); a.take(T * 512, ok); a.take(T * ZV, ok); a.take(T * ZV, ok); - a.take(T * (size_t) cap, ok); - a.take((size_t) sel_batch * (size_t) max_blocks, ok); - a.take((size_t) attn_batch * strata::kernels::qsa_decode_attn_scratch_floats(cap, s), ok); - return a.used; -} -// Step 2b: which layers' experts go through MMQ (both weight types covered; the Strata Q2_0 pack always - its blob -// is converted to GGUF Q2_0 blocks on the gather), whether any layer keeps the FP16 path (IQ1_M), and the largest -// gate/up and down matrices a group buffer slot holds. STRATA_PREFILL_MMQ=0: the FP16 path everywhere (the A/B). -constexpr int MMQ_GROUP = 16; // experts per MMQ launch (the gather is per expert, as blobs arrive) -// MMQ reads up to one 256-value tile past a matrix's last row when the row length is not a multiple of it (the down -// product: 640 values). Those bytes meet zero activations, which is harmless only if they decode to finite numbers - -// llama.cpp zero-pads after every tensor, and so does a group buffer: this many zeroed bytes follow its last expert. -constexpr size_t MMQ_TAIL = 4096; -struct MmqPlan { - bool any = false, fallback = true; - std::vector layer; // per layer: MMQ - size_t gu_max = 0, d_max = 0; -}; -const MmqPlan& mmq_plan() { - static const MmqPlan plan = [] { - MmqPlan p; - const auto& lay = strata::kernels::cpu::expert_layout(); - const char* env = std::getenv("STRATA_PREFILL_MMQ"); - const bool on = mmq::built() && (env == nullptr || std::atoi(env) != 0); - const int64_t layers = lay.native ? (int64_t) lay.fmt.size() : lay.n_layers; - p.layer.assign((size_t) std::max(layers, 0), 0); - p.fallback = !on || layers <= 0; - for (int64_t l = 0; on && l < layers; ++l) { - const int gt = lay.native ? lay.fmt[(size_t) l].gu_type : 42, dt = lay.native ? lay.fmt[(size_t) l].d_type : 42; - if (!mmq::supported(gt) || !mmq::supported(dt)) { p.fallback = true; continue; } - p.layer[(size_t) l] = 1; - p.any = true; - p.gu_max = std::max(p.gu_max, mmq::matrix_bytes(gt, 1280, N)); - p.d_max = std::max(p.d_max, mmq::matrix_bytes(dt, N, 640)); - } - return p; - }(); - return plan; -} -uint64_t moe_set_bytes(size_t T, int64_t n_expert) { - const MmqPlan& mp = mmq_plan(); - Alloc a; a.count_only = true; bool ok = true; - a.take(T * n_expert, ok); a.take(T * K, ok); a.take(T * K, ok); a.take(T * K, ok); - a.take(T * K, ok); - if (mp.fallback) a.take(T * K * N, ok); - a.take(T * K * 1280, ok); - if (mp.fallback) a.take(T * K * 640, ok); - a.take(T * K * N, ok); a.take(T * 640, ok); - a.take(T * 640, ok); a.take(T * 640, ok); a.take(T * N, ok); a.take(T, ok); - if (mp.any) { - a.take(mmq::q8_bytes((int64_t) (T * K), N), ok); - a.take(T * K * 640, ok); - a.take(mmq::q8_bytes((int64_t) (T * K), 640), ok); - } - return a.used; -} + Impl& m = *impl_; + if (m.cs) cudaStreamSynchronize(m.cs); + if (m.xfer) cudaStreamSynchronize(m.xfer); + if (m.xsum) cudaStreamSynchronize(m.xsum); + if (m.ev_mixed) cudaEventDestroy(m.ev_mixed); + for (cudaEvent_t e : m.ev_in) if (e) cudaEventDestroy(e); + for (cudaEvent_t e : m.ev_piece) if (e) cudaEventDestroy(e); + if (m.xfer) cudaStreamDestroy(m.xfer); + if (m.xsum) cudaStreamDestroy(m.xsum); + for (float* h : m.ple_host) if (h) cudaFreeHost(h); + if (m.ids_host) cudaFreeHost(m.ids_host); + if (m.w_host) cudaFreeHost(m.w_host); + if (m.owned) cudaFree(m.owned); } bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, core::SessionState& ss, - core::ExpertSource* src, const core::ExpertCache* cache, const int32_t* host_res, int64_t chunk, - void* stream, std::string& err, void* borrow, uint64_t borrow_bytes) { + core::ExpertSource* src, const core::ExpertCache* cache, const int32_t* host_res, int64_t max_chunk, + void* stream, ExpertRunner* offload, std::string& err) { Impl& m = *impl_; - m.wt = &wt; m.g = &g; m.ss = &ss; m.src = src; m.cache = cache; m.host_res = host_res; - m.T = chunk; m.cs = (cudaStream_t) stream; m.stats = &stats_; - if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_expert < 1 || ss.k != K) { + m.wt = &wt; m.g = &g; m.ss = &ss; m.cache = cache; m.host_res = host_res; m.offload = offload; + m.max_chunk = max_chunk; m.cs = (cudaStream_t) stream; + if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_expert != NE || ss.k != K) { err = "prefill: geometry differs from the artifact's"; return false; } - cudaGetDevice(&m.device); - if (stage_le_ < 0) stage_le_ = g.n_layers; - if (stage_lb_ < 0 || stage_lb_ >= stage_le_ || stage_le_ > g.n_layers || (stage_le_ < g.n_layers) != (next_ != nullptr)) { - err = "prefill: the stage's layer range is wrong"; + m.s = shapes_of(g); + m.cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, m.s); + m.max_blocks = ss.qsa_states[0].max_cells / m.s.idx_block + 2; + m.ev_in.assign((size_t) ((max_chunk + SUB - 1) / SUB), nullptr); + m.ev_piece.assign(m.ev_in.size(), nullptr); + bool pieces = true; + for (auto* v : {&m.ev_in, &m.ev_piece}) + for (cudaEvent_t& e : *v) + pieces = pieces && cudaEventCreateWithFlags(&e, cudaEventDisableTiming) == cudaSuccess; + if (!pieces || cudaStreamCreateWithFlags(&m.xfer, cudaStreamNonBlocking) != cudaSuccess || + cudaStreamCreateWithFlags(&m.xsum, cudaStreamNonBlocking) != cudaSuccess || + cudaEventCreateWithFlags(&m.ev_mixed, cudaEventDisableTiming) != cudaSuccess || + cudaHostAlloc((void**) &m.ids_host, (size_t) max_chunk * K * 4, cudaHostAllocDefault) != cudaSuccess || + cudaHostAlloc((void**) &m.w_host, (size_t) max_chunk * K * 4, cudaHostAllocDefault) != cudaSuccess) { + err = "prefill: streams, events and pinned routing buffers"; return false; } - for (int b = 0; next_ != nullptr && b < 2; ++b) - if (!m.hand[b] && cudaHostAlloc((void**) &m.hand[b], (size_t) chunk * D * 4, cudaHostAllocPortable) != cudaSuccess) { - err = "prefill: the layer split's hand-off buffers"; - return false; - } - if (m.tok_dev == nullptr) { - if (cudaMalloc((void**) &m.tok_dev, (size_t) chunk * sizeof(int32_t)) != cudaSuccess) { - err = "prefill: the token id buffer"; - return false; - } - m.owned.push_back(m.tok_dev); - m.tok_host.resize((size_t) chunk); - } - if (cudaStreamCreateWithFlags(&m.copy, cudaStreamNonBlocking) != cudaSuccess) { err = "prefill: copy stream"; return false; } - const size_t T = (size_t) chunk; - m.T_max = chunk; - m.borrowed = borrow != nullptr; - bool ok = true; - // one-time: events, the stager, the host buffers (for the largest chunk), the identity page table - for (int i = 0; i < RING_MAX; ++i) { - if (cudaEventCreateWithFlags(&m.copied[i], cudaEventDisableTiming) != cudaSuccess) ok = false; - if (cudaEventCreateWithFlags(&m.used[i], cudaEventDisableTiming) != cudaSuccess) ok = false; - } - if (!m.stager) { - m.stager = std::make_unique(); - const int hw = (int) std::thread::hardware_concurrency(); - const char* stv = std::getenv("STRATA_STAGER_THREADS"); // D-5: the host copy threads of unpinned blobs - if (!m.stager->init((size_t) MAXBLOB(), stv ? std::clamp(std::atoi(stv), 1, 32) : std::max(2, std::min(4, hw / 4)))) - ok = false; - } + const size_t T = (size_t) max_chunk; m.steps_host.resize(T * strata::kernels::kStepCount); - m.ids_host.resize(T * K); m.slot_host.resize(T * K); m.src_host.resize(T * K); m.cnt.resize(m.g->n_expert); m.off.resize(m.g->n_expert + 1); - { - const size_t need = 3 * T * K + (size_t) (2 * (m.g->n_expert + m.g->n_expert / MMQ_GROUP + 2)); - const char* gc = std::getenv("STRATA_GROUP_COPY"); - if (m.grp_n < need && !(gc && gc[0] == '1')) { - if (m.grp_host) cudaFreeHost(m.grp_host); - m.grp_host = m.grp_dev = nullptr; - m.grp_n = m.grp_tk = 0; - void *h = nullptr, *d = nullptr; - if (cudaHostAlloc(&h, need * 4, cudaHostAllocMapped) == cudaSuccess && - cudaHostGetDevicePointer(&d, h, 0) == cudaSuccess) { - m.grp_host = (int32_t*) h; - m.grp_dev = (int32_t*) d; - m.grp_n = need; - m.grp_tk = T * K; - } else { // the copies, as before - if (h) cudaFreeHost(h); - cudaGetLastError(); - } - } - } - for (int b = 0; b < 2; ++b) { - if (!m.ple_emb_host[b] && - cudaHostAlloc((void**) &m.ple_emb_host[b], (size_t) T * N * 4, cudaHostAllocDefault) != cudaSuccess) { - cudaGetLastError(); - m.ple_pageable[b].resize(T * N); // pageable: the upload is staged before it returns - m.ple_emb_host[b] = m.ple_pageable[b].data(); + m.cnt.resize(NE); + m.off.resize(NE + 1); + m.entry_of.resize(T * K); + for (int i = 0; i < 2; ++i) { + m.ple_rows[i].resize(T * strata::kernels::PLE_N_HEADS); + if (cudaHostAlloc((void**) &m.ple_host[i], T * N * 4, cudaHostAllocDefault) != cudaSuccess) { + err = "prefill: pinned PLE row buffers"; + return false; } - if (!m.ple_copied[b] && cudaEventCreateWithFlags(&m.ple_copied[b], cudaEventDisableTiming) != cudaSuccess) - ok = false; - m.ple_rows[b].resize(T * strata::kernels::PLE_N_HEADS); - } - if (ss.qsa_states[0].kv_mode == 1) { // KV streaming: the staging pool's identity page table - const int64_t pages = ss.qsa_states[0].n_pages; - std::vector ident((size_t) pages); - for (int64_t i = 0; i < pages; ++i) ident[(size_t) i] = (int32_t) i; - if (cudaMalloc((void**) &m.ident_table, ident.size() * 4) != cudaSuccess || - cudaMemcpy(m.ident_table, ident.data(), ident.size() * 4, cudaMemcpyHostToDevice) != cudaSuccess) - ok = false; - else - m.owned.push_back(m.ident_table); - } - if (!ok) { err = "prefill: host buffers or events for a chunk of " + std::to_string(chunk) + " tokens"; return false; } - Alloc o; - o.base = (uint8_t*) borrow; - o.cap = borrow_bytes; - o.owned = &m.owned; - { - uint16_t* gs = o.take((size_t) GEMM_SCRATCH, ok); - void* ws = o.take(GEMM_WS, ok); - if (!ok) { err = "prefill: GEMM scratch does not fit"; return false; } - if (!m.gemm.init_external(stream, gs, GEMM_SCRATCH, ws, GEMM_WS, err)) return false; } - if (!carve(T, &o)) { - err = "prefill: device buffers for a chunk of " + std::to_string(chunk) + " tokens do not fit"; - return false; - } - return true; + int dev = 0; + cudaGetDevice(&dev); + return m.gemm.init(stream, err) && + m.local.init(dev, dev, stream, src, cache, host_res, NE, max_chunk, offload == nullptr, err); } -// Every device buffer of a chunk of T tokens, from the Alloc `alloc` (after the GEMM scratch and workspace): `init` -// once, and `relayout` for a request's own chunk. The order is `bytes_needed`'s. -bool Prefill::carve(size_t T, void* alloc) { - Impl& m = *impl_; - Alloc& o = *static_cast(alloc); - const core::ModelGeometry& g = *m.g; - core::SessionState& ss = *m.ss; - bool ok = true; - m.emb = o.take(T * N, ok); m.R = o.take(T * D, ok); - m.xn = gr_unfused() ? o.take(T * D, ok) : nullptr; // F-1: not needed (gr_mix_r reads R) - m.grs = o.take(T * HC, ok); - m.xn16 = o.take(T * D, ok); m.lo = o.take(T * LR, ok); m.lo16 = o.take(T * LR, ok); - m.gated = o.take(T * D, ok); m.inj = o.take(T * HC, ok); - m.mixed = o.take(T * N, ok); m.mixed_bf = o.take(T * N, ok); - m.mixed_h = o.take(T * N, ok); m.bo = o.take(T * N, ok); - m.steps_dev = o.take(T * strata::kernels::kStepCount, ok); - strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); - s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; - s.idx_dim = g.idx_key_dim; - m.cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s); - m.max_blocks = ss.qsa_states[0].max_cells / s.idx_block + 2; - { - // one region for the attention half's and the MoE half's scratch (see gdn_set_bytes) - const uint64_t region = std::max({gdn_set_bytes(T), qsa_set_bytes(T, m.cap, m.max_blocks, m.sel_batch, - m.attn_batch, s), moe_set_bytes(T, m.g->n_expert)}); - uint8_t* base = o.take((size_t) region, ok); - m.region = base; - m.region_bytes = region; - Alloc a; - a.base = base; a.cap = region; a.owned = &m.owned; - m.qkv = a.take(T * C, ok); m.z = a.take(T * ZV, ok); m.ab = a.take(T * 2 * HV, ok); - m.gate = a.take(T * HV, ok); m.beta = a.take(T * HV, ok); m.hbuf = a.take(T * C, ok); - m.y = a.take(T * ZV, ok); m.y_h = a.take(T * ZV, ok); - Alloc b; - b.base = base; b.cap = region; b.owned = &m.owned; - m.Kc = b.take(T * 512, ok); m.Vc = b.take(T * 512, ok); m.Qf = b.take(T * 12288, ok); - m.q = b.take(T * ZV, ok); m.idx_raw = b.take(T * 128, ok); m.q_idx = b.take(T * 512, ok); - m.attn = b.take(T * ZV, ok); m.attn_h = b.take(T * ZV, ok); - m.sel_ids = b.take(T * (size_t) m.cap, ok); - m.sel_scores = b.take((size_t) m.sel_batch * (size_t) m.max_blocks, ok); - m.attn_scratch = b.take((size_t) m.attn_batch * strata::kernels::qsa_decode_attn_scratch_floats(m.cap, s), ok); - Alloc c; - c.base = base; c.cap = region; c.owned = &m.owned; - m.logits = c.take(T * m.g->n_expert, ok); m.w = c.take(T * K, ok); m.ids = c.take(T * K, ok); - m.slot_dev = c.take(T * K, ok); m.src_dev = c.take(T * K, ok); - const MmqPlan& mp = mmq_plan(); - m.Xs = mp.fallback ? c.take(T * K * N, ok) : nullptr; - m.GU = c.take(T * K * 1280, ok); - m.Hh = mp.fallback ? c.take(T * K * 640, ok) : nullptr; - m.Dm = c.take(T * K * N, ok); - m.sgate = c.take(T * 640, ok); m.sup = c.take(T * 640, ok); m.sh_h = c.take(T * 640, ok); - m.shared = c.take(T * N, ok); m.sg = c.take(T, ok); - if (mp.any) { - m.Xq = c.take(mmq::q8_bytes((int64_t) (T * K), N), ok); - m.H = c.take(T * K * 640, ok); - m.Hq = c.take(mmq::q8_bytes((int64_t) (T * K), 640), ok); - } - if (base == nullptr) ok = false; - } - for (int i = 0; i < DQ; ++i) { m.dq_gu[i] = o.take(1280 * 2560, ok); m.dq_d[i] = o.take(2560 * 640, ok); } - if (mmq_plan().any) { - const MmqPlan& mp = mmq_plan(); - m.ids_identity = o.take(T * K, ok); - m.bounds_dev = o.take((size_t) (2 * (m.g->n_expert + m.g->n_expert / MMQ_GROUP + 2)), ok); - m.grp_gu = o.take(MMQ_GROUP * mp.gu_max + MMQ_TAIL, ok); - m.grp_d = o.take(MMQ_GROUP * mp.d_max + MMQ_TAIL, ok); - // (written at every run's start, not here: when serving, these are live expert-cache slots until a request - // lends them - a write now would corrupt a resident expert) - if (!m.mmq_ctx) m.mmq_ctx = std::make_unique(); - } - m.ring = ring_slots(T); - for (int i = 0; i < m.ring; ++i) { - m.stage_dev[i] = o.take((size_t) MAXBLOB(), ok); - m.stage_live[i] = false; // a new buffer: nothing of an earlier layout to wait for - } - m.ple_emb = o.take(T * N, ok); - m.ple_norm = o.take((size_t) strata::kernels::NG_HC_DIM, ok); - take_stage(o, ss, s, m.stage, ok); - m.T = (int64_t) T; - return ok; +uint64_t Prefill::bytes_needed(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk, bool offload) { + const strata::kernels::QsaShapes s = shapes_of(g); + Alloc a; + layout(a, chunk, strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s), + ss.qsa_states[0].max_cells / s.idx_block + 2, offload); + return a.used; } -bool Prefill::relayout(int64_t chunk, void* borrow, uint64_t borrow_bytes, std::string& err) { - Impl& m = *impl_; - if (!m.borrowed || borrow == nullptr || chunk <= 0 || chunk > m.T_max) { - err = "prefill: relayout needs borrowed buffers and a chunk of at most " + std::to_string(m.T_max); - return false; - } - if (cudaStreamSynchronize(m.cs) != cudaSuccess || cudaStreamSynchronize(m.copy) != cudaSuccess) { - err = "prefill: relayout: the stream failed"; - return false; - } - bool ok = true; - Alloc o; - o.base = (uint8_t*) borrow; - o.cap = borrow_bytes; - o.owned = &m.owned; - uint16_t* gs = o.take((size_t) GEMM_SCRATCH, ok); - void* ws = o.take(GEMM_WS, ok); - if (ok) m.gemm.rebind(gs, GEMM_SCRATCH, ws, GEMM_WS); - if (!ok || !carve((size_t) chunk, &o)) { - err = "prefill: device buffers for a chunk of " + std::to_string(chunk) + " tokens do not fit"; - return false; - } - return true; +uint64_t Prefill::bytes_for(int64_t chunk) const { + const Impl& m = *impl_; + return bytes_needed(*m.g, *m.ss, chunk, m.offload != nullptr) + + (m.offload == nullptr && chunk >= ExpertRunner::kPrefetchMin ? m.local.prefetch_need(nullptr) : 0); } -int64_t Prefill::chunk() const { return impl_->T; } - -bool Prefill::draft_kv(core::MtpDrafter& mtp, const float* R_rows, const int32_t* next_tokens, int64_t n, int64_t cell0, - std::string& err) { +bool Prefill::bind(void* region, uint64_t bytes, int64_t chunk, std::string& err) { Impl& m = *impl_; - static const bool off = [] { const char* v = std::getenv("STRATA_MTP_BATCH"); return v != nullptr && v[0] == '0'; }(); - core::QsaState& st = mtp.kv_state_rw(); - if (off || n <= 0 || m.g == nullptr || m.region == nullptr || st.kv_mode != 0 || st.kv_hybrid || - mtp.device() != m.device) - return false; - const auto t0 = Clock::now(); - const core::ModelGeometry& g = *m.g; - const int64_t Nn = g.n_embd, HCN = g.hc * g.n_embd, KV = g.n_head_kv * g.head_dim; - constexpr int kQ8_0 = 8; // GGML_TYPE_Q8_0 - const float* w_ne = mtp.tensor_f32("pre_fc_norm_embedding.weight"); - const void* w_fe = mtp.tensor_q8("fc_embedding.weight"); - const float* w_nh = mtp.tensor_f32("pre_fc_norm_hidden.weight"); - const void* w_fh = mtp.tensor_q8("fc_hidden.weight"); - const float* w_hn = mtp.tensor_f32("attn_hyper_connection.hc_norm.weight"); - const uint16_t* w_dn = mtp.tensor_bf16("attn_hyper_connection.input_mix_weight_down.weight"); - const uint16_t* w_up = mtp.tensor_bf16("attn_hyper_connection.input_mix_weight_up.weight"); - const void* w_k = mtp.tensor_q8("self_attn.k_proj.weight"); - const void* w_v = mtp.tensor_q8("self_attn.v_proj.weight"); - const float* w_kn = mtp.tensor_f32("self_attn.k_norm.weight"); - if (!w_ne || !w_fe || !w_nh || !w_fh || !w_hn || !w_dn || !w_up || !w_k || !w_v || !w_kn) return false; - const core::NativeEmbed* nemb = core::native_embed(); - const core::WeightRef* wemb = nemb ? nullptr : m.wt->find("token_embd.weight"); - if (!nemb && (wemb == nullptr || wemb->codebook_iq4nl || wemb->ne0 != g.n_embd || wemb->group_elems <= 0 || - (wemb->code_bits != 2 && wemb->code_bits != 4 && wemb->code_bits != 8))) - return false; - // the cells the drafter's window can still reach - const int64_t r0 = std::max(0, mtp.first_needed() - cell0); - if (r0 >= n) return true; - // per row: emb/e2 (N), en16 (N half), hn/h2/Rm/gated (HCN), hn16/xn16 (HCN half), lo (LR) + lo16, grs, mixed (N) + - // mixed_h, K and V (KV each), the token id - // E-9: the drafter's Q8_0 matrices through Q8_1 x Q8_0 MMQ - its own pass's integer dot products (mmvq), so - // its K/V stay close to what the drafter computes itself; STRATA_MTP_BATCH_F16=1: FP16 GEMMs (the A/B) - static const bool f16_only = [] { const char* v = std::getenv("STRATA_MTP_BATCH_F16"); return v && v[0] == '1'; }(); - const bool q8 = !f16_only && mmq::built() && mmq::supported(kQ8_0); - const uint64_t per_row = 4 * (2 * Nn + 4 * HCN + LR + HC + Nn + 2 * KV + 1) + 2 * (Nn + 2 * HCN + LR + Nn) + 64 + - (q8 ? (uint64_t) mmq::q8_bytes(g.hc, Nn) + 4 * g.hc : 0); - const int64_t B = std::min(n - r0, (int64_t) (m.region_bytes / per_row) & ~(int64_t) 63); - if (B < 64) return false; - uint8_t* q = m.region; - auto carve = [&](size_t bytes) { void* p = q; q += (bytes + 255) & ~(size_t) 255; return p; }; - float* emb = (float*) carve((size_t) B * Nn * 4); - float* e2 = (float*) carve((size_t) B * Nn * 4); - uint16_t* en16 = (uint16_t*) carve((size_t) B * Nn * 2); - float* hn = (float*) carve((size_t) B * HCN * 4); - float* h2 = (float*) carve((size_t) B * HCN * 4); - float* Rm = (float*) carve((size_t) B * HCN * 4); - float* gated = (float*) carve((size_t) B * HCN * 4); - uint16_t* hn16 = (uint16_t*) carve((size_t) B * HCN * 2); - uint16_t* xn16 = (uint16_t*) carve((size_t) B * HCN * 2); - float* lo = (float*) carve((size_t) B * LR * 4); - uint16_t* lo16 = (uint16_t*) carve((size_t) B * LR * 2); - float* grs = (float*) carve((size_t) B * HC * 4); - float* mixed = (float*) carve((size_t) B * Nn * 4); - uint16_t* mixed_h = (uint16_t*) carve((size_t) B * Nn * 2); - float* Kc = (float*) carve((size_t) B * KV * 4); - float* Vc = (float*) carve((size_t) B * KV * 4); - int32_t* tok = (int32_t*) carve((size_t) B * 4); - void* xq = q8 ? carve(mmq::q8_bytes(B * g.hc, Nn)) : nullptr; - int32_t* ident = q8 ? (int32_t*) carve((size_t) B * g.hc * 4) : nullptr; - int32_t* bnd = q8 ? (int32_t*) carve(16) : nullptr; - if ((uint64_t) (q - m.region) > m.region_bytes) return false; - if (!mtp.idle(err)) return false; // the drafter's own stream (its graph uploads) before this writes its K/V - strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); - s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; - s.idx_dim = g.idx_key_dim; - std::vector tk((size_t) B); - if (q8) { - if (!m.mmq_ctx) m.mmq_ctx = std::make_unique(); - mmq::iota(ident, B * g.hc, m.cs); - } - // y[rows, n_out] = x[rows, k] . w^T for a Q8_0 matrix: MMQ from the FP32 rows (bounds slot 0: rows, 1: rows*hc) - auto proj = [&](const float* x, const uint16_t* x16, const void* w, float* y, int64_t rows, int64_t n_out, - int64_t k, int slot) { - if (!q8) { m.gemm.native(x16, kQ8_0, w, y, rows, n_out, k); return; } - mmq::quantize(x, nullptr, xq, kQ8_0, k, k, rows, m.cs); - mmq::Product p; - p.w = w; p.type = kQ8_0; p.w_rows = n_out; p.w_cols = k; p.expert_bytes = mmq::matrix_bytes(kQ8_0, n_out, k); - p.n = 1; p.xq = xq; p.bounds = bnd + 2 * slot; p.ids = ident; p.total_rows = rows; p.max_rows = rows; - p.dst = y; p.ld_dst = n_out; - m.mmq_ctx->run(p, m.cs); - }; - for (int64_t b0 = r0; b0 < n; b0 += B) { - const int64_t nb = std::min(B, n - b0), c0 = cell0 + b0; - for (int64_t i = 0; i < nb; ++i) tk[(size_t) i] = next_tokens[b0 + i]; - const int32_t bh[4] = {0, (int32_t) nb, 0, (int32_t) (nb * g.hc)}; - if (q8 && cudaMemcpyAsync(bnd, bh, sizeof bh, cudaMemcpyHostToDevice, m.cs) != cudaSuccess) { - err = "prefill: the draft bounds' upload failed"; - return false; - } - if (cudaMemcpyAsync(tok, tk.data(), (size_t) nb * 4, cudaMemcpyHostToDevice, m.cs) != cudaSuccess) { - err = "prefill: the draft tokens' upload failed"; + if (region == nullptr) { + if (m.owned != nullptr) return true; // allocated once, for the longest chunk + chunk = m.max_chunk; + bytes = bytes_needed(*m.g, *m.ss, chunk, m.offload != nullptr); + if (cudaMalloc(&m.owned, bytes) != cudaSuccess) { + m.owned = nullptr; + err = "prefill: " + std::to_string(bytes >> 20) + " MiB of device buffers"; return false; } - // the input branches: the next token's embedding, and this cell's final residual rows - if (nemb) { - nemb->gather_dev(tok, nb, emb, m.cs); - } else { - const auto* codes = (const uint8_t*) wemb->data; - const auto* scales = (const float*) (codes + wemb->codes_bytes); - const auto* offsets = wemb->has_offset ? (const float*) (codes + wemb->codes_bytes + wemb->scales_bytes) - : nullptr; - strata::kernels::embedding_gather_dev(codes, scales, offsets, tok, (int) nb, wemb->ne0, wemb->code_bits, - wemb->code_bias, wemb->group_elems, - (uint64_t) (wemb->ne0 / (8 / wemb->code_bits)), - (uint64_t) (wemb->ne0 / wemb->group_elems), emb, m.cs); - } - rms_rows(emb, w_ne, nb, Nn, Nn, EPS, m.cs); - if (!q8) to_f16(emb, en16, nb * Nn, m.cs); - proj(emb, en16, w_fe, e2, nb, Nn, Nn, 0); - cudaMemcpyAsync(hn, R_rows + (size_t) b0 * HCN, (size_t) nb * HCN * 4, cudaMemcpyDeviceToDevice, m.cs); - rms_rows(hn, w_nh, nb, HCN, HCN, EPS, m.cs); - if (!q8) to_f16(hn, hn16, nb * HCN, m.cs); - proj(hn, hn16, w_fh, h2, nb * g.hc, Nn, Nn, 1); // every stream through fc_hidden - strata::kernels::add_streams_broadcast(h2, e2, Rm, Nn, (int) g.hc, (int) nb, m.cs); - // the attention hyper-connection's read (its mixed input only: this pass writes nothing back) - gr_norm_rs(Rm, w_hn, EPS, grs, xn16, nb, m.cs); - m.gemm.bf16(xn16, w_dn, lo, nb, LR, HCN); - gr_silu(lo, lo16, nb, m.cs); - m.gemm.bf16(lo16, w_up, gated, nb, HCN, LR); - gr_mix_r(Rm, grs, w_hn, gated, mixed, nullptr, nb, m.cs, mixed_h); - // K and V into the drafter's cache, as the prompt path's QSA layers append theirs - proj(mixed, mixed_h, w_k, Kc, nb, KV, Nn, 0); - proj(mixed, mixed_h, w_v, Vc, nb, KV, Nn, 0); - rms_rows(Kc, w_kn, nb * g.n_head_kv, g.head_dim, g.head_dim, EPS, m.cs); - rope(Kc, nb, g.n_head_kv, g.head_dim, KV, c0, (float) strata::kernels::qsa_freq_base(), m.cs); - if (st.kv_q4) { - strata::kernels::fwht256_inplace_cuda(Kc, nb * g.n_head_kv, m.cs); - strata::kernels::fwht256_inplace_cuda(Vc, nb * g.n_head_kv, m.cs); - strata::kernels::kv_append_q4(st.k_q4, st.v_q4, st.page_table, c0, nb, Kc, Vc, s, m.cs, &st.host); - } else { - kv_append(Kc, Vc, nb, c0, st.page_table, s.page_size, st.kv_int8 ? nullptr : st.k_pool, - st.kv_int8 ? nullptr : st.v_pool, st.k_q, st.v_q, st.k_scale, st.v_scale, m.cs, &st.host); - } + region = m.owned; } - if (cudaStreamSynchronize(m.cs) != cudaSuccess) { - err = std::string("prefill: the draft layer's K/V: ") + cudaGetErrorString(cudaGetLastError()); + if (chunk <= 0 || chunk > m.max_chunk) { + err = "prefill: a chunk of " + std::to_string(chunk) + " tokens (set up for " + std::to_string(m.max_chunk) + ")"; return false; } - mtp.ms_prefill += ms_since(t0); - return true; -} -void Prefill::set_pinned_share(double share) { g_pinned_share = share; } -double Prefill::pinned_share() { return g_pinned_share; } - -uint64_t Prefill::bytes_needed(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk) { - // the same allocation sequence as `init`, counted - const size_t T = (size_t) chunk; - bool ok = true; - Alloc o; - o.count_only = true; - o.take((size_t) GEMM_SCRATCH, ok); - o.take(GEMM_WS, ok); - auto f = [&](size_t n) { o.take(n, ok); }; - f(T * N); f(T * D); f(T * D); o.take(T * D, ok); f(T * LR); o.take(T * LR, ok); - f(T * D); f(T * HC); f(T * N); o.take(T * N, ok); o.take(T * N, ok); f(T * N); - o.take(T * strata::kernels::kStepCount, ok); - strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); - s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; - s.idx_dim = g.idx_key_dim; - const int64_t cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s); - const int64_t max_blocks = ss.qsa_states[0].max_cells / s.idx_block + 2; - o.take((size_t) std::max({gdn_set_bytes(T), qsa_set_bytes(T, cap, max_blocks, 256, 32, s), - moe_set_bytes(T, g.n_expert)}), ok); - for (int i = 0; i < DQ; ++i) { o.take(1280 * 2560, ok); o.take(2560 * 640, ok); } - if (mmq_plan().any) { - const MmqPlan& mp = mmq_plan(); - o.take(T * K, ok); - o.take((size_t) (2 * (g.n_expert + g.n_expert / MMQ_GROUP + 2)), ok); - o.take(MMQ_GROUP * mp.gu_max + MMQ_TAIL, ok); - o.take(MMQ_GROUP * mp.d_max + MMQ_TAIL, ok); + Alloc a; + a.base = (uint8_t*) region; + a.cap = bytes; + m.b = layout(a, chunk, m.cap, m.max_blocks, m.offload != nullptr); + if (!a.ok) { + err = "prefill: device buffers for a chunk of " + std::to_string(chunk) + " tokens do not fit"; + return false; } - for (int i = 0; i < ring_slots(T); ++i) o.take((size_t) MAXBLOB(), ok); - f(T * N); - f((size_t) strata::kernels::NG_HC_DIM); - strata::kernels::KvHostPools stage; - take_stage(o, ss, s, stage, ok); - return o.used + (8u << 20); // alignment slack + m.gemm.set_buffers(nullptr, 0, m.b.gemm_ws, GEMM_WS); + // the rest of a lent region takes this GPU's prefetched experts (outside the union: copies run beside the steps) + const bool area = region != m.owned && bytes > a.used; + if (!m.local.bind(m.b.local, m.b.local_bytes, chunk, err, area ? (uint8_t*) region + a.used : nullptr, + area ? bytes - a.used : 0)) + return false; + m.T = chunk; + return true; } namespace { @@ -864,12 +291,6 @@ const core::WeightRef* need(const core::LayerView& v, const char* suffix, std::s if (!r) err = v.name(suffix) + " is missing"; return r; } -bool native_proj(Gemm& gm, const core::WeightRef* w, const uint16_t* X, float* Y, int64_t T, const std::string& name, - std::string& err, int64_t ldy = 0) { - if (!w->native_data) { err = "prefill: " + name + " has no native GGUF blocks (run with --native)"; return false; } - gm.native(X, w->native_type, w->native_data, Y, T, w->ne1, w->ne0, ldy); - return true; -} bool bf16_proj(Gemm& gm, const core::WeightRef* w, const uint16_t* X, float* Y, int64_t T, const std::string& name, std::string& err, int64_t ldy = 0) { if (w->kind != core::WeightKind::Bf16InF32 || !w->data) { err = "prefill: " + name + " is not a resident BF16 tensor"; return false; } @@ -879,1010 +300,475 @@ bool bf16_proj(Gemm& gm, const core::WeightRef* w, const uint16_t* X, float* Y, } // namespace -namespace { -// STRATA_PREFILL_TIMING=1: the prompt path's GPU time by phase. Events are recorded on the compute stream in order; -// the time between two consecutive marks is charged to the phase of the first, so a gap where the GPU waits (for the -// host's expert grouping, or for an expert's copy) lands on the phase that was waiting. Events are reused: the marks -// are folded at every MoE layer's host sync, after which all of them have completed. -enum PfPhase { kPfStart, kPfHc, kPfGdn, kPfQsa, kPfQsaIdx, kPfQsaSel, kPfQsaAttn, kPfRouter, kPfHostGroup, kPfGather, - kPfWaitCopy, kPfDequant, kPfGemmGU, kPfGemmD, kPfCombine, kPfPle, kPfKvStage, kPfGdnConv, kPfGdnRec, kPfGdnOut, - kPfCount }; -const char* const kPfNames[kPfCount] = {"embed+steps", "hc read", "gdn", "qsa proj", "qsa indexer", "qsa select", - "qsa attn", "router+shared", "host grouping", "gather", "wait copy", "dequant", - "gemm gate/up", "gemm down", "combine", "ple", "kv stage", "gdn conv+gates", - "gdn recurrence", "gdn out proj"}; -struct PfTimer { - bool on = std::getenv("STRATA_PREFILL_TIMING") != nullptr; - std::vector ev; - std::vector ph; - size_t used = 0; - double ms[kPfCount] = {}; - void mark(int phase, cudaStream_t s) { - if (!on) return; - if (used == ev.size()) { - cudaEvent_t e = nullptr; - cudaEventCreate(&e); - ev.push_back(e); - ph.push_back(0); - } - ph[used] = phase; - cudaEventRecord(ev[used], s); - ++used; - } - // every recorded mark has completed (the stream was synchronized): charge the gaps, keep the last mark - void fold() { - if (!on || used < 2) return; - for (size_t i = 0; i + 1 < used; ++i) { - float t = 0.0f; - if (cudaEventElapsedTime(&t, ev[i], ev[i + 1]) == cudaSuccess) ms[ph[i]] += t; - } - std::swap(ev[0], ev[used - 1]); - std::swap(ph[0], ph[used - 1]); - used = 1; - } - ~PfTimer() { - for (cudaEvent_t e : ev) cudaEventDestroy(e); - } -}; -} // namespace - bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err) { Impl& m = *impl_; - const core::OnDevice on_device(m.device); + const Bufs& b = m.b; const core::ModelGeometry& g = *m.g; core::SessionState& ss = *m.ss; + const strata::kernels::QsaShapes& s = m.s; + constexpr int64_t SC = strata::kernels::kStepCount; + if (m.T <= 0) { err = "prefill: no buffers bound"; return false; } const auto t_start = Clock::now(); - const int64_t LB = stage_lb_, LE = stage_le_; - // the next stage reads chunk c on a thread while this one reads chunk c + 1 (declared first: an early return - // waits for it before anything it reads goes away) - std::string next_err; - std::future next_run; - int hand_buf = 0; - double host_sync_ms = 0, host_chunk_ms = 0, host_setup_ms = 0; // STRATA_PREFILL_TIMING: the host's share - strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); - s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; - s.idx_dim = g.idx_key_dim; const uint64_t gdn_floats = (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size + (uint64_t) g.ssm_conv_channels * (g.ssm_d_conv - 1); int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]}; - PfTimer pt; - const cudaStream_t cs = (cudaStream_t) m.cs; - // the MMQ row table lives in the borrowed cache slots, which the refill after a prompt overwrites with experts: - // write it again for every prompt (a layout is reused as long as the chunk and the slots are the same) - if (m.ids_identity != nullptr) mmq::iota(m.ids_identity, m.T * K, m.cs); - // The PLE rows of a chunk are read from the model file on the host (an SSD read per missed row): the chunk - // after this one is read on a thread while the GPU runs this one, into the other of two buffers. The rows - // depend only on the tokens (the two before a position name its n-grams), so this is the same data. - const bool ple_on = ss.ple.ready() && LB <= 1 && 1 < LE; - // the PLE block batched over the chunk: the pinned postops and a BF16 or GGUF-native key (else token by token); - // STRATA_PLE_BATCH=0 keeps the per-token block (the A/B) - static const bool ple_batch_env = [] { - const char* v = std::getenv("STRATA_PLE_BATCH"); - return v == nullptr || std::atoi(v) != 0; - }(); - const bool ple_batch = ple_on && ple_batch_env && strata::kernels::ple_native_postops_enabled() && - (ss.ple.w.key_bf16 != nullptr || ss.ple.w.key_native_data != nullptr) && - m.region_bytes / ((uint64_t) (3 * strata::kernels::NG_HC_DIM + N + 4) * 4 + (uint64_t) N * 2 + 4096) >= 64; - const int32_t prev0[2] = {prev[0], prev[1]}; - auto ple_gather = [&m, &ss, tokens, n, prev0](int64_t c0, int buf, std::string& e) -> bool { - const int64_t T = std::min(m.T, n - c0); - auto at = [&](int64_t i) { return i < 2 ? prev0[i] : (int32_t) tokens[i - 2]; }; // prev0, then the tokens - int32_t pv[2] = {at(c0), at(c0 + 1)}; - for (int64_t t = 0; t < T; ++t) { - const int32_t tok = (int32_t) tokens[c0 + t]; - strata::kernels::ngram_rows(&tok, pv, 1, ss.ple.consts, - m.ple_rows[buf].data() + t * strata::kernels::PLE_N_HEADS); - pv[0] = pv[1]; - pv[1] = tok; + // Y[T, rows] = X . W^T for a quantized W [rows, cols]: dequantized at its first use in the layer + auto w16 = [&](int type, const void* blocks, int64_t rows, int64_t cols) -> const uint16_t* { + for (const auto& [k, p] : m.w16_have) + if (k == blocks) return p; + if (m.w16_used + rows * cols > W16) return nullptr; + uint16_t* p = b.w16 + m.w16_used; + strata::kernels::dequant_f16(type, blocks, 0, rows, cols, p, m.cs); + m.w16_used += (rows * cols + 127) / 128 * 128; + m.w16_have.emplace_back(blocks, p); + return p; + }; + auto native_proj = [&](const core::WeightRef* w, const uint16_t* X, float* Y, int64_t T, const std::string& name, + std::string& e) { + const uint16_t* W = w->native_data ? w16(w->native_type, w->native_data, w->ne1, w->ne0) : nullptr; + if (W == nullptr) { + e = "prefill: " + name + (w->native_data ? " does not fit the dequantization buffer" + : " has no native GGUF blocks (run with --native)"); + return false; } - return ss.ple.table->gather_batch(m.ple_rows[buf].data(), (size_t) T, m.ple_emb_host[buf], e); + m.gemm.f16(X, W, Y, T, w->ne1, w->ne0); + return true; + }; + auto new_layer = [&] { + m.w16_have.clear(); + m.w16_used = 0; + }; + // --prefill-profile: an event at the start of each section; the time to the next event is charged to it + std::vector> ev; + auto mark = [&](int section) { + if (!profile) return; + cudaEvent_t e; + cudaEventCreate(&e); + cudaEventRecord(e, m.cs); + ev.push_back({section, e}); }; - std::string ple_next_err; - std::future ple_next; // declared after everything it reads: an early return waits for it - int ple_buf = 0; - for (int64_t c0 = 0; c0 < n; c0 += m.T) { - if (should_stop && should_stop()) { err = "cancelled"; return false; } - if (std::getenv("STRATA_TRACE")) { std::fprintf(stderr, "strata trace: prompt chunk %lld of %lld\n", (long long) c0, (long long) n); std::fflush(stderr); } - const int64_t T = std::min(m.T, n - c0), p0 = pos0 + c0; - ++stats_.chunks; - pt.mark(kPfStart, cs); - const auto tsetup = Clock::now(); - // ---- embeddings, broadcast to the four streams - or, in a later stage of a layer split, the rows the - // previous stage handed on - if (hand_in_ != nullptr) { - if (cudaMemcpyAsync(m.R, hand_in_ + (size_t) c0 * D, (size_t) T * D * 4, cudaMemcpyHostToDevice, m.cs) != - cudaSuccess) { - err = "prefill: the layer split's hand-off upload failed"; - return false; - } - } - // C-4: the whole chunk's rows in one gather (the same per-element arithmetic as the per-token path, so the - // same bits); a chunk with picture rows, or a token outside the table, takes the per-token path - bool batched = hand_in_ == nullptr && m.tok_dev != nullptr; - const core::NativeEmbed* nemb = core::native_embed(); - const core::WeightRef* wemb = nemb ? nullptr : m.wt->find("token_embd.weight"); - if (batched && nemb == nullptr && - (wemb == nullptr || wemb->codebook_iq4nl || wemb->ne0 != g.n_embd || wemb->group_elems <= 0 || - (wemb->code_bits != 2 && wemb->code_bits != 4 && wemb->code_bits != 8))) - batched = false; - for (int64_t t = 0; batched && t < T; ++t) { - const int64_t tok = tokens[c0 + t]; - if ((embd_rows && embd_rows[p0 + t]) || tok < 0 || (wemb && tok >= wemb->ne1)) batched = false; - else m.tok_host[(size_t) t] = (int32_t) tok; - } - if (batched) { - if (cudaMemcpyAsync(m.tok_dev, m.tok_host.data(), (size_t) T * sizeof(int32_t), cudaMemcpyHostToDevice, - m.cs) != cudaSuccess) { - err = "prefill: the token id upload failed"; - return false; - } - if (nemb) { - nemb->gather_dev(m.tok_dev, T, m.emb, m.cs); - } else { - const auto* codes = (const uint8_t*) wemb->data; - const auto* scales = (const float*) (codes + wemb->codes_bytes); - const auto* offsets = wemb->has_offset ? (const float*) (codes + wemb->codes_bytes + wemb->scales_bytes) - : nullptr; - strata::kernels::embedding_gather_dev(codes, scales, offsets, m.tok_dev, (int) T, wemb->ne0, - wemb->code_bits, wemb->code_bias, wemb->group_elems, - (uint64_t) (wemb->ne0 / (8 / wemb->code_bits)), - (uint64_t) (wemb->ne0 / wemb->group_elems), m.emb, m.cs); - } - } - for (int64_t t = 0; hand_in_ == nullptr && !batched && t < T; ++t) { - const float* row = embd_rows ? embd_rows[p0 + t] : nullptr; - if (row) { - if (cudaMemcpyAsync(m.emb + t * N, row, (size_t) N * 4, cudaMemcpyHostToDevice, m.cs) != cudaSuccess) { - err = "prefill: the image embedding upload failed"; - return false; - } - } else if (!core::embed_row(*m.wt, g, tokens[c0 + t], m.emb + t * N, m.cs, err)) { - return false; - } + // The MTP draft layer's K/V over a chunk's cells: its input norms and projections (e = fc_embedding(rms(the + // embedding of the token at t + 1)), h = fc_hidden(rms(R_t)) per stream, h + e on every stream), the attention + // hyper-connection read, then K (normed and rotated) and V into its cells, from the window's first cell on. The + // chunk's buffers are free once its last layer is done; its residual rows are normed in place. + const core::MtpPromptKv dk = draft != nullptr ? draft->prompt_kv() : core::MtpPromptKv{}; + if (draft != nullptr && (!dk.norm_emb || !dk.norm_hidden || !dk.fc_emb || !dk.fc_hidden || !dk.hc_norm || + !dk.hc_down || !dk.hc_up || !dk.k_proj || !dk.v_proj || !dk.k_norm)) { + err = "prefill: a draft layer weight is missing"; + return false; + } + auto draft_pass = [&](const int64_t* toks, int64_t T, int64_t p0) -> bool { + const int64_t from = std::clamp(dk.first_cell - p0, 0, T); + if (from == T) return true; + mark(kPsDraft); + new_layer(); + const uint16_t *We = w16(GGML_Q8_0, dk.fc_emb, N, N), *Wh = w16(GGML_Q8_0, dk.fc_hidden, N, N), + *Wk = w16(GGML_Q8_0, dk.k_proj, 512, N), *Wv = w16(GGML_Q8_0, dk.v_proj, 512, N); + if (!We || !Wh || !Wk || !Wv) { + err = "prefill: the draft layer's projections do not fit the dequantization buffer"; + return false; } - if (hand_in_ == nullptr) gr_broadcast(m.emb, m.R, T, m.cs); - // ---- the PLE rows of the whole chunk, one batched SSD request (read ahead on a thread, see ple_gather) - if (ple_on) { - const auto tp = Clock::now(); - if (!ple_next.valid()) { - if (!ple_gather(c0, ple_buf, err)) return false; - } else if (!ple_next.get()) { - err = ple_next_err; - return false; - } - cudaMemcpyAsync(m.ple_emb, m.ple_emb_host[ple_buf], (size_t) T * N * 4, cudaMemcpyHostToDevice, m.cs); - cudaEventRecord(m.ple_copied[ple_buf], m.cs); - if (c0 + m.T < n) { - cudaEventSynchronize(m.ple_copied[ple_buf ^ 1]); // the other buffer's upload (a chunk ago) is done - ple_next = std::async(std::launch::async, [&ple_gather, &ple_next_err, c1 = c0 + m.T, b = ple_buf ^ 1] { - return ple_gather(c1, b, ple_next_err); - }); - } - ple_buf ^= 1; - stats_.ms_ple += ms_since(tp); + for (int64_t t = from; t < T; ++t) m.ids_host[t - from] = (int32_t) toks[t + 1]; + cudaMemcpyAsync(b.ids, m.ids_host, (size_t) (T - from) * 4, cudaMemcpyHostToDevice, m.cs); + const core::QsaState& st = *dk.kv; + for (int64_t t0 = from; t0 < T; t0 += SUB) { + const int64_t P = std::min(SUB, T - t0), pp = p0 + t0; + float* R = b.R + t0 * D; + if (!core::embed_rows(*m.wt, g, b.ids + (t0 - from), P, b.emb, m.cs, err)) return false; + rms_rows(b.emb, dk.norm_emb, P, N, N, EPS, m.cs); + to_f16(b.emb, b.mixed_h, P * N, m.cs); + m.gemm.f16(b.mixed_h, We, b.bo, P, N, N); + rms_rows(R, dk.norm_hidden, P, D, D, EPS, m.cs); + to_f16(R, b.xn16, P * D, m.cs); + m.gemm.f16(b.xn16, Wh, b.gated, P * HC, N, N); + strata::kernels::add_streams_broadcast(b.gated, b.bo, b.gated, N, (int) HC, (int) P, m.cs); + gr_norm(b.gated, dk.hc_norm, EPS, b.xn, b.xn16, P, m.cs); + m.gemm.bf16(b.xn16, dk.hc_down, b.lo, P, LR, D); + gr_silu(b.lo, b.lo16, P, m.cs); + m.gemm.bf16(b.lo16, dk.hc_up, b.gated, P, D, LR); + gr_mix(b.xn, b.gated, nullptr, nullptr, P, m.cs, b.mixed_h); + m.gemm.f16(b.mixed_h, Wk, b.Kc, P, 512, N); + m.gemm.f16(b.mixed_h, Wv, b.Vc, P, 512, N); + rms_rows(b.Kc, dk.k_norm, P * 2, 256, 256, EPS, m.cs); + rope(b.Kc, P, 2, 256, 512, pp, (float) strata::kernels::qsa_freq_base(), m.cs); + kv_append(b.Kc, b.Vc, P, pp, st.page_table, s.page_size, st.kv_int8 ? nullptr : st.k_pool, + st.kv_int8 ? nullptr : st.v_pool, st.k_q, st.v_q, st.k_scale, st.v_scale, m.cs); } - for (int64_t t = 0; t < T; ++t) { prev[0] = prev[1]; prev[1] = (int32_t) tokens[c0 + t]; } - // ---- the QSA step records of every position in the chunk - for (int64_t t = 0; t < T; ++t) strata::kernels::qsa_step_fill(m.steps_host.data() + t * strata::kernels::kStepCount, p0 + t, s); - cudaMemcpyAsync(m.steps_dev, m.steps_host.data(), (size_t) T * strata::kernels::kStepCount * 4, - cudaMemcpyHostToDevice, m.cs); + return true; + }; - int64_t qsa_index = 0, gdn_index = 0; - for (int64_t l = 0; l < LB; ++l) (core::is_qsa_layer(g, l) ? qsa_index : gdn_index) += 1; - // step 3: this chunk's stream - every non-resident expert of every layer, layer by layer in id order (entry - // k lands in ring slot k % ring); a copy is issued once the entry `ring` before it is consumed (its slot's - // `used` event recorded), so the copy stream never waits on an event that is not queued yet - const strata::kernels::cpu::ExpertLayout& lay0 = strata::kernels::cpu::expert_layout(); - const bool stream_all = m.ring > STAGE && T >= STREAM_ALL_MIN && m.src != nullptr; - struct StreamEntry { int32_t l, e; const uint8_t* blob; int job; }; - std::vector seq; - std::vector seq_start; - size_t issued = 0, consumed = 0; - if (stream_all) { - seq_start.assign((size_t) g.n_layers + 1, 0); - std::vector js; - for (int64_t l = LB; l < LE; ++l) { - seq_start[(size_t) l] = seq.size(); - for (int32_t e = 0; e < m.g->n_expert; ++e) { - if (m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0) continue; - const uint8_t* b = m.src->blob(l, e); - if (!b) { err = "prefill: expert source has no blob"; return false; } - int job = -1; - if (!m.src->pinned(l, e)) { - job = (int) js.size(); - js.push_back({b, (size_t) lay0.blob_bytes(l)}); + // ---- a chunk's PLE rows (one batched SSD request) on a thread: the first chunk's while its layer 0 runs, each + // next chunk's while the one before runs; waited for before the PLE block at layer 1 + struct PleRead { + std::thread th; + bool ok = true; + std::string err; + bool wait(std::string& e) { + if (th.joinable()) th.join(); + if (!ok) e = err; + return ok; + } + ~PleRead() { if (th.joinable()) th.join(); } + } ple_read; + int32_t pa[2] = {prev[0], prev[1]}; // the rows' tokens before the next chunk to read + int ple_cur = 0; + auto ple_start = [&](int64_t c, int buf) { + const int64_t Tc = std::min(m.T, n - c); + uint32_t* rows = m.ple_rows[buf].data(); + for (int64_t t = 0; t < Tc; ++t) { + const int32_t tok = (int32_t) tokens[c + t]; + strata::kernels::ngram_rows(&tok, pa, 1, ss.ple.consts, rows + t * strata::kernels::PLE_N_HEADS); + pa[0] = pa[1]; + pa[1] = tok; + } + ple_read.ok = true; + ple_read.th = std::thread([&m, &ss, &ple_read, rows, Tc, out = m.ple_host[buf]] { + ple_read.ok = ss.ple.table->gather_batch(rows, (size_t) Tc, out, ple_read.err); + }); + }; + if (ss.ple.ready() && n > 0) ple_start(0, 0); + + static const bool trace = std::getenv("STRATA_TRACE") != nullptr; + processed_ = 0; + int64_t c0 = 0; + for (; c0 < n; c0 += m.T) { + if (should_stop && should_stop()) break; // the chunks before stay: `prev` and the state hold them + if (trace) { std::fprintf(stderr, "strata trace: prompt chunk %lld of %lld\n", (long long) c0, (long long) n); std::fflush(stderr); } + const int64_t T = std::min(m.T, n - c0), p0 = pos0 + c0; + ++stats_.chunks; + // ---- the embeddings, broadcast to the four streams + for (int64_t t0 = 0; t0 < T; t0 += SUB) { + const int64_t P = std::min(SUB, T - t0); + for (int64_t t = t0; t < t0 + P; ++t) { + const float* row = embd_rows ? embd_rows[p0 + t] : nullptr; + if (row) { + if (cudaMemcpyAsync(b.emb + (t - t0) * N, row, (size_t) N * 4, cudaMemcpyHostToDevice, m.cs) != cudaSuccess) { + err = "prefill: the image embedding upload failed"; + return false; } - seq.push_back({(int32_t) l, e, b, job}); + } else if (!core::embed_row(*m.wt, g, tokens[c0 + t], b.emb + (t - t0) * N, m.cs, err)) { + return false; } } - for (int64_t l = LE; l <= g.n_layers; ++l) seq_start[(size_t) l] = seq.size(); - m.stager->start(std::move(js)); + gr_broadcast(b.emb, b.R + t0 * D, P, m.cs); } - struct StagerDone { - Stager* st; - ~StagerDone() { if (st) st->finish(); } - } chunk_stager_done{stream_all ? m.stager.get() : nullptr}; - auto issue_until = [&](size_t limit) { - limit = std::min(limit, seq.size()); - while (issued < limit) { - const StreamEntry& en = seq[issued]; - const int sl = (int) (issued % (size_t) m.ring); - const auto th = Clock::now(); - const size_t bytes = (size_t) lay0.blob_bytes(en.l); - if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0); - if (en.job < 0) { - cudaMemcpyAsync(m.stage_dev[sl], en.blob, bytes, cudaMemcpyHostToDevice, m.copy); - ++stats_.experts_dma; - } else { - const uint8_t* hb = m.stager->wait(en.job); - cudaMemcpyAsync(m.stage_dev[sl], hb, bytes, cudaMemcpyHostToDevice, m.copy); - m.stager->issued_one(en.job, m.copy); - } - cudaEventRecord(m.copied[sl], m.copy); - m.stage_live[sl] = true; - stats_.ms_experts_host += ms_since(th); - ++stats_.experts_streamed; - ++issued; + const bool ple_on = ss.ple.ready(); + // ---- the QSA step records of every position in the chunk + for (int64_t t = 0; t < T; ++t) strata::kernels::qsa_step_fill(m.steps_host.data() + t * SC, p0 + t, s); + cudaMemcpyAsync(b.steps, m.steps_host.data(), (size_t) T * SC * 4, cudaMemcpyHostToDevice, m.cs); + + // ---- a layer's FFN write, deferred to the sub-chunk that reads its rows next: the second GPU's sums come back + // a sub-chunk at a time while this GPU works on the sub-chunks before (with `norm`, the next read's norm in + // the same pass) + bool pending = false; + auto ffn_write = [&](int64_t t0, int64_t P, const float* norm) { + const uint16_t* part = nullptr; + if (m.offload) { + cudaStreamWaitEvent(m.cs, m.ev_piece[(size_t) (t0 / SUB)], 0); + part = b.sums + t0 * N; } + float *R = b.R + t0 * D, *bo = b.bo_moe + t0 * N, *inj = b.inj_f + t0 * HC; + if (norm) gr_write_norm_rs(R, bo, inj, HC, norm, EPS, b.rs, b.xn16, P, m.cs, part); + else gr_write(R, bo, inj, HC, P, m.cs, part); }; - // D-5: the stream is issued by its own host thread, so the thread launching the layers' kernels never waits - // behind a host copy of an unpinned blob (that wait left the GPU idle: the 'wait copy' / 'dequant' time of the - // i-quant prompts). The same copies in the same order into the same slots, and a slot is refilled only once - // the compute stream has recorded that it is done with it: the same results. STRATA_PREFILL_ISSUER=0: inline. - static const bool issuer_on = [] { - const char* v = std::getenv("STRATA_PREFILL_ISSUER"); - return v == nullptr || std::atoi(v) != 0; - }(); - std::atomic a_issued{0}, a_consumed{0}; - std::atomic a_stop{false}; - double iss_ms = 0; - int64_t iss_streamed = 0, iss_dma = 0; - std::thread issuer; - struct IssuerJoin { - std::atomic* stop; - std::thread* t; - ~IssuerJoin() { if (t->joinable()) { stop->store(true); t->join(); } } - } issuer_join{&a_stop, &issuer}; - const bool threaded_issue = stream_all && issuer_on; - if (threaded_issue) { - issuer = std::thread([&] { - const core::OnDevice od(m.device); - for (size_t idx = 0; idx < seq.size(); ++idx) { - while (idx >= a_consumed.load(std::memory_order_acquire) + (size_t) m.ring) { - if (a_stop.load(std::memory_order_acquire)) return; - std::this_thread::yield(); - } - const StreamEntry& en = seq[idx]; - const int sl = (int) (idx % (size_t) m.ring); - const auto th = Clock::now(); - const size_t bytes = (size_t) lay0.blob_bytes(en.l); - if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0); - if (en.job < 0) { - cudaMemcpyAsync(m.stage_dev[sl], en.blob, bytes, cudaMemcpyHostToDevice, m.copy); - ++iss_dma; - } else { - const uint8_t* hb = m.stager->wait(en.job); - cudaMemcpyAsync(m.stage_dev[sl], hb, bytes, cudaMemcpyHostToDevice, m.copy); - m.stager->issued_one(en.job, m.copy); - } - cudaEventRecord(m.copied[sl], m.copy); - m.stage_live[sl] = true; - iss_ms += ms_since(th); - ++iss_streamed; - a_issued.store(idx + 1, std::memory_order_release); - } - }); - } else if (stream_all) { - issue_until((size_t) m.ring); // layer 0's first experts, behind the embedding and the PLE - } - // the consumer's side: entry k's copy is on the copy stream (the thread issued it), then k is given back - auto wait_issued = [&](size_t k) { - if (!threaded_issue) return; - while (a_issued.load(std::memory_order_acquire) <= k) std::this_thread::yield(); - }; - auto give_back = [&](size_t upto) { - if (threaded_issue) a_consumed.store(upto, std::memory_order_release); - else issue_until(upto + (size_t) m.ring); - }; - host_setup_ms += ms_since(tsetup); - bool normed = false; // F-2: the previous half's write already normed R for this half (grs, xn16) - for (int64_t l = LB; l < LE; ++l) { - core::progress_beat(); // the serve watchdog: a prompt chunk of 8192 tokens is still moving + int64_t qsa_index = 0, gdn_index = 0; + for (int64_t l = 0; l < g.n_layers; ++l) { const core::LayerView v(*m.wt, l); - // ---- the PLE block at layer 1, token by token (its conv reads the previous tokens' rows) - if (l == 1 && ple_on && ple_batch) { - // the whole chunk at once, in sub-batches carved from the idle scratch region: the key and value - // projections as GEMMs (a token at a time they re-read ~52 MB of BF16 key per token on the IQ - // files), the rest with the per-token kernels' arithmetic (native_ple_postops_batch) - pt.mark(kPfPle, cs); + const bool qsa = core::is_qsa_layer(g, l); + new_layer(); + // a long chunk uses nearly every expert: the ones this GPU's cache lacks are copied ahead, here while + // the dense steps run, or on the second GPU + if (T >= ExpertRunner::kPrefetchMin) { + m.ex1.clear(); + for (int32_t e = 0; e < NE; ++e) + if (!(m.host_res && m.cache && m.host_res[(size_t) l * NE + e] >= 0)) m.ex1.push_back(e); + if (!(m.offload ? m.offload : &m.local)->prefetch(l, m.ex1, err)) return false; + } + mark(kPsPle); + // ---- the PLE block at layer 1: the key and value projections, then the decode path's per-token + // arithmetic (the conv reads the previous tokens' normalized rows) + if (l == 1 && ple_on) { + // the chunk's rows, then the next chunk's read on the thread (the other buffer) const auto tp = Clock::now(); + if (!ple_read.wait(err)) return false; + cudaMemcpyAsync(b.ple_emb, m.ple_host[ple_cur], (size_t) T * N * 4, cudaMemcpyHostToDevice, m.cs); + if (c0 + T < n) ple_start(c0 + T, 1 - ple_cur); + ple_cur = 1 - ple_cur; + stats_.ms_ple += ms_since(tp); const strata::kernels::PleWeights& pw = ss.ple.w; - constexpr int64_t HD = strata::kernels::NG_HC_DIM; - const uint64_t per_token = (uint64_t) (3 * HD + N + 4) * 4 + (uint64_t) N * 2 + 4096; - const int64_t SB = std::min(T, (int64_t) (m.region_bytes / per_token)); - for (int64_t s0 = 0; s0 < T; s0 += SB) { - const int64_t nb = std::min(SB, T - s0); - uint8_t* q = m.region; - auto carve_f = [&](size_t n) { float* p = (float*) q; q += (n * 4 + 255) & ~(size_t) 255; return p; }; - float* key = carve_f((size_t) nb * HD); - float* qn = carve_f((size_t) nb * HD); - float* gated = carve_f((size_t) nb * HD); - float* val = carve_f((size_t) nb * N); - float* gate = carve_f((size_t) nb * 4); - uint16_t* e16 = (uint16_t*) carve_f((size_t) nb * N / 2); - const float* emb = m.ple_emb + s0 * N; + for (int64_t t0 = 0; t0 < T; t0 += SUB) { + const int64_t P = std::min(SUB, T - t0); + if (pending) ffn_write(t0, P, nullptr); + const float* ple_emb = b.ple_emb + t0 * N; + to_bf16(ple_emb, b.mixed_bf, P * N, m.cs); if (pw.key_bf16 != nullptr) { - to_bf16(emb, e16, nb * N, m.cs); - m.gemm.bf16(e16, pw.key_bf16, key, nb, HD, N); + m.gemm.bf16(b.mixed_bf, pw.key_bf16, b.ple_kp, P, D, N); + } else if (const uint16_t* W = pw.key_native_data ? w16(pw.key_native_type, pw.key_native_data, D, N) + : nullptr) { + to_f16(ple_emb, b.mixed_h, P * N, m.cs); + m.gemm.f16(b.mixed_h, W, b.ple_kp, P, D, N); } else { - to_f16(emb, e16, nb * N, m.cs); - m.gemm.native(e16, pw.key_native_type, pw.key_native_data, key, nb, HD, N); - to_bf16(emb, e16, nb * N, m.cs); + err = "prefill: the PLE key has neither a BF16 nor a native GGUF form (run with --native)"; + return false; } - m.gemm.bf16(e16, pw.value_bf16, val, nb, N, N); + m.gemm.bf16(b.mixed_bf, pw.value_bf16, b.bo, P, N, N); try { - strata::kernels::native_ple_postops_batch(key, m.R + s0 * D, val, ss.ple.hist, pw, qn, gated, - gate, (int) nb, m.cs); + strata::kernels::native_ple_postops_tokens( + b.ple_kp, b.R + t0 * D, b.bo, ss.ple.hist, pw, + {b.ple_key, b.ple_q, b.inj_a, b.ple_gated, b.ple_q, b.R + t0 * D}, (int) P, m.cs); } catch (const std::exception& e) { err = std::string("prefill PLE: ") + e.what(); return false; } } - stats_.ms_ple += ms_since(tp); - } else if (l == 1 && ple_on) { - pt.mark(kPfPle, cs); - const auto tp = Clock::now(); - for (int64_t t = 0; t < T; ++t) { - strata::kernels::PleOut po; - po.normalized = m.ple_norm; - po.result = m.R + t * D; - try { - strata::kernels::ple_block(m.ple_emb + t * N, m.R + t * D, ss.ple.hist, ss.ple.w, po, - ss.ple.scratch, m.cs); - } catch (const std::exception& e) { err = std::string("prefill PLE: ") + e.what(); return false; } - strata::kernels::ple_history_advance(ss.ple.hist, m.ple_norm, m.cs); - } - stats_.ms_ple += ms_since(tp); + pending = false; + new_layer(); } - for (int half = 0; half < 2; ++half) { - // ---- the hyper-connection read of this half - const char* pre = half == 0 ? "hc_attn_" : "hc_ffn_"; - const std::string sn = std::string(pre) + "norm.weight", sd = std::string(pre) + "down.weight", - su = std::string(pre) + "up.weight", si = std::string(pre) + "inject.weight"; - const core::WeightRef *wn = need(v, sn.c_str(), err), *wd = need(v, sd.c_str(), err), - *wu = need(v, su.c_str(), err), *wi = need(v, si.c_str(), err); - if (!wn || !wd || !wu || !wi) return false; - pt.mark(kPfHc, cs); - if (gr_unfused()) gr_norm(m.R, (const float*) wn->data, EPS, m.xn, m.xn16, T, m.cs); - else if (!normed) gr_norm_rs(m.R, (const float*) wn->data, EPS, m.grs, m.xn16, T, m.cs); - normed = false; - if (!bf16_proj(m.gemm, wd, m.xn16, m.lo, T, sd, err)) return false; - gr_silu(m.lo, m.lo16, T, m.cs); - if (!bf16_proj(m.gemm, wu, m.lo16, m.gated, T, su, err)) return false; - if (!bf16_proj(m.gemm, wi, m.xn16, m.inj, T, si, err)) return false; - if (gr_unfused()) gr_mix(m.xn, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h); - else gr_mix_r(m.R, m.grs, (const float*) wn->data, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h); - - if (half == 0 && !core::is_qsa_layer(g, l)) { - // ======================= GDN ======================= - const core::WeightRef *wqkv = need(v, "attn_qkv.weight", err), *wg = need(v, "attn_gate.weight", err), - *wo = need(v, "ssm_out.weight", err), *wa = need(v, "ssm_alpha.weight", err), - *wb = need(v, "ssm_beta.weight", err), *wc = need(v, "ssm_conv1d.weight", err), - *wnm = need(v, "ssm_norm.weight", err), *wdt = need(v, "ssm_dt.bias", err), - *wsa = need(v, "ssm_a", err); - if (!wqkv || !wg || !wo || !wa || !wb || !wc || !wnm || !wdt || !wsa) return false; - pt.mark(kPfGdn, cs); - float* state = ss.gdn_state + (size_t) gdn_index * gdn_floats; - float* conv = state + (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size; - if (!native_proj(m.gemm, wqkv, m.mixed_h, m.qkv, T, v.name("attn_qkv.weight"), err)) return false; - if (!native_proj(m.gemm, wg, m.mixed_h, m.z, T, v.name("attn_gate.weight"), err)) return false; - if (!bf16_proj(m.gemm, wa, m.mixed_bf, m.ab, T, v.name("ssm_alpha.weight"), err, 2 * HV)) return false; - if (!bf16_proj(m.gemm, wb, m.mixed_bf, m.ab + HV, T, v.name("ssm_beta.weight"), err, 2 * HV)) return false; - pt.mark(kPfGdnConv, cs); // "gdn" is the projections in; the rest on their own lines - gdn_gates(m.ab, (const float*) wdt->data, (const float*) wsa->data, m.gate, m.beta, T, m.cs); - gdn_conv(conv, m.qkv, (const float*) wc->data, m.hbuf, T, EPS, m.cs); - pt.mark(kPfGdnRec, cs); - gdn_recurrence(state, m.hbuf, m.gate, m.beta, m.z, (const float*) wnm->data, EPS, m.y, m.y_h, T, m.cs); - pt.mark(kPfGdnOut, cs); - if (!native_proj(m.gemm, wo, m.y_h, m.bo, T, v.name("ssm_out.weight"), err)) return false; - ++gdn_index; - } else if (half == 0) { - // ======================= QSA ======================= - const core::QsaState& st = ss.qsa_states[qsa_index]; - const core::WeightRef *wq = need(v, "attn_q.weight", err), *wk = need(v, "attn_k.weight", err), - *wv = need(v, "attn_v.weight", err), *wo = need(v, "attn_output.weight", err), - *wik = need(v, "indexer.k_proj.weight", err), - *wiq = need(v, "indexer.q_proj.weight", err), - *wqn = need(v, "attn_q_norm.weight", err), *wkn = need(v, "attn_k_norm.weight", err), - *wiqn = need(v, "indexer.q_norm.weight", err), - *wikn = need(v, "indexer.k_norm.weight", err); - if (!wq || !wk || !wv || !wo || !wik || !wiq || !wqn || !wkn || !wiqn || !wikn) return false; - pt.mark(kPfQsa, cs); - if (!native_proj(m.gemm, wk, m.mixed_h, m.Kc, T, v.name("attn_k.weight"), err)) return false; - if (!native_proj(m.gemm, wv, m.mixed_h, m.Vc, T, v.name("attn_v.weight"), err)) return false; - if (!native_proj(m.gemm, wq, m.mixed_h, m.Qf, T, v.name("attn_q.weight"), err)) return false; - if (!bf16_proj(m.gemm, wik, m.mixed_bf, m.idx_raw, T, v.name("indexer.k_proj.weight"), err)) return false; - if (!bf16_proj(m.gemm, wiq, m.mixed_bf, m.q_idx, T, v.name("indexer.q_proj.weight"), err)) return false; - rms_rows(m.Kc, (const float*) wkn->data, T * 2, 256, 256, EPS, m.cs); - rope(m.Kc, T, 2, 256, 512, p0, (float) strata::kernels::qsa_freq_base(), m.cs); - // KV streaming: this layer's cells [0, p0) come in from the host copy to the staging pool, and the - // chunk's cells go to the host copy, the staging pool, and the VRAM slots of resident blocks - const bool staged = st.kv_mode == 1; - if (staged) { - pt.mark(kPfKvStage, cs); - strata::kernels::kv_stage_from_host(pools_of(m.stage, m.ident_table), st.host, - core::qsa_kv_format(st), - (p0 + s.page_size - 1) / s.page_size, s, m.cs); - pt.mark(kPfQsa, cs); - } - if (st.kv_hybrid) { // K8V4: K INT8 unrotated, V rotated Q4_0 (only V and the output rotate) - strata::kernels::fwht256_inplace_cuda(m.Vc, T * 2, m.cs); - kv_append(m.Kc, m.Kc, T, p0, st.page_table, s.page_size, nullptr, nullptr, - st.k_q, st.k_q, st.k_scale, st.k_scale, m.cs, nullptr, // mode 0: no host mirror - staged ? &m.stage : nullptr); - strata::kernels::kv_append_q4(st.v_q4, st.v_q4, st.page_table, p0, T, m.Vc, m.Vc, s, m.cs, - nullptr, staged ? &m.stage : nullptr); - } else if (st.kv_q4) { // Q4_0 KV (kv_q4.hpp): rotated K and V, the queries below too, the output back - strata::kernels::fwht256_inplace_cuda(m.Kc, T * 2, m.cs); - strata::kernels::fwht256_inplace_cuda(m.Vc, T * 2, m.cs); - strata::kernels::kv_append_q4(st.k_q4, st.v_q4, st.page_table, p0, T, m.Kc, m.Vc, s, m.cs, - &st.host, staged ? &m.stage : nullptr); - } else { - kv_append(m.Kc, m.Vc, T, p0, st.page_table, s.page_size, st.kv_int8 ? nullptr : st.k_pool, - st.kv_int8 ? nullptr : st.v_pool, st.k_q, st.v_q, st.k_scale, st.v_scale, m.cs, - &st.host, staged ? &m.stage : nullptr); - } - split_q(m.Qf, m.q, T, m.cs); - rms_rows(m.q, (const float*) wqn->data, T * 24, 256, 256, EPS, m.cs); - rope(m.q, T, 24, 256, 6144, p0, (float) strata::kernels::qsa_freq_base(), m.cs); - if (st.kv_q4) strata::kernels::fwht256_inplace_cuda(m.q, T * 24, m.cs); - rms_rows(m.q_idx, (const float*) wiqn->data, T * 4, 128, 128, EPS, m.cs); - rope(m.q_idx, T, 4, 128, 512, p0, (float) strata::kernels::qsa_freq_base(), m.cs); - // the indexer appends, token by token; then scores + selection for many queries at once: - // a query reads completed blocks (final once completed) and `dead` for its own tail block - const strata::kernels::QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; - pt.mark(kPfQsaIdx, cs); - // C-2: the chunk's appends in three launches instead of one per token (the same end state: - // the queries below read it only after the whole chunk is appended). STRATA_INDEXER_PER_TOKEN=1: the old - try { - static const bool per_token = std::getenv("STRATA_INDEXER_PER_TOKEN") != nullptr; - if (!per_token) { - strata::kernels::native_qsa_indexer_append_batch(m.idx_raw, T, p0, 0, (const float*) wikn->data, - EPS, ib, s, st.max_cells, - (float) strata::kernels::qsa_freq_base(), m.cs); - } - for (int64_t t = 0; per_token && t < T; ++t) { - const int32_t* step_t = m.steps_dev + t * strata::kernels::kStepCount; - strata::kernels::native_qsa_indexer_append(m.idx_raw + t * 128, step_t + strata::kernels::kStepPos, 0, - (const float*) wikn->data, EPS, ib, s, st.max_cells, - (float) strata::kernels::qsa_freq_base(), m.cs); - } - } catch (const std::exception& e) { err = std::string("prefill indexer: ") + e.what(); return false; } - pt.mark(kPfQsaSel, cs); - for (int64_t t0 = 0; t0 < T; t0 += m.sel_batch) { - const int64_t nb = std::min(m.sel_batch, T - t0); - const int32_t* steps0 = m.steps_dev + t0 * strata::kernels::kStepCount; - // C-1: the grid reaches the batch's last query's n_bid (they rise with the position) - const int64_t active = (int64_t) m.steps_host[(size_t) ((t0 + nb - 1) * strata::kernels::kStepCount + - strata::kernels::kStepNBid)] + 1; - // the scores on tensor cores (3xTF32: FP32-level, not bitwise); STRATA_SELECT_OLD=1: the warp kernel - static const bool old_sel = std::getenv("STRATA_SELECT_OLD") != nullptr; - if (old_sel || !strata::kernels::qsa_block_scores_tc(st.idx_pooled, st.idx_dead, m.q_idx + t0 * 512, - steps0, nb, m.max_blocks, s, m.sel_scores, - m.cs, active)) - strata::kernels::qsa_block_scores(st.idx_pooled, st.idx_dead, m.q_idx + t0 * 512, steps0, nb, - m.max_blocks, s, m.sel_scores, m.cs, active); - strata::kernels::qsa_block_topk(m.sel_scores, steps0, nb, m.max_blocks, m.cap, s, - m.sel_ids + t0 * m.cap, m.cs); - } - // STRATA_SEL_OVERLAP (debug, D-1's question): how much do neighbouring queries' selections share? - // Per tile of 16 queries: the union of their selected cells against the sum of their widths. - if (static const bool ovl = std::getenv("STRATA_SEL_OVERLAP") != nullptr; ovl && qsa_index == 0) { - std::vector ids((size_t) (T * m.cap)); - cudaMemcpyAsync(ids.data(), m.sel_ids, ids.size() * 4, cudaMemcpyDeviceToHost, m.cs); - cudaStreamSynchronize(m.cs); - double sum_w = 0, sum_u = 0; - for (int64_t t0 = 0; t0 + 16 <= T; t0 += 16) { - std::vector u; - for (int64_t t = t0; t < t0 + 16; ++t) { - const int64_t w = m.steps_host[(size_t) (t * strata::kernels::kStepCount + strata::kernels::kStepWidth)]; - sum_w += (double) w; - u.insert(u.end(), ids.begin() + t * m.cap, ids.begin() + t * m.cap + w); - } - std::sort(u.begin(), u.end()); - sum_u += (double) (std::unique(u.begin(), u.end()) - u.begin()); - } - std::fprintf(stderr, "strata prefill: selection overlap at %lld: 16-query tiles read %.1f%% of the " - "cells one query at a time does\n", (long long) p0, sum_w > 0 ? 100.0 * sum_u / sum_w : 0.0); - } - // STRATA_IDX_FP16_CHECK: would FP16 pooled indexer keys select the same cells? (the KV-streaming - // design's last question). Every query is selected again from the pooled keys and `dead` rounded - // to fp16 (exactly what an fp16 store reads back); the agreement with the fp32 selection is - // printed cumulatively after each chunk's last QSA layer. Debug: syncs per layer. - if (static const bool f16chk = std::getenv("STRATA_IDX_FP16_CHECK") != nullptr; f16chk) { - static float *pooled16 = nullptr, *dead16 = nullptr; - static int32_t* ids16 = nullptr; - static double shared = 0, cells = 0; - static long long queries = 0, same = 0, sel_queries = 0; - const int64_t rows = st.idx_pooled_rows; - if (pooled16 == nullptr && - (cudaMalloc((void**) &pooled16, (size_t) rows * s.idx_dim * 4) != cudaSuccess || - cudaMalloc((void**) &dead16, (size_t) s.idx_dim * 4) != cudaSuccess || - cudaMalloc((void**) &ids16, (size_t) (m.T * m.cap) * 4) != cudaSuccess)) { - err = "STRATA_IDX_FP16_CHECK: no room for its buffers"; - return false; - } - round_f16(st.idx_pooled, pooled16, rows * s.idx_dim, m.cs); - round_f16(st.idx_dead, dead16, s.idx_dim, m.cs); - for (int64_t t0 = 0; t0 < T; t0 += m.sel_batch) { - const int64_t nb = std::min(m.sel_batch, T - t0); - const int32_t* steps0 = m.steps_dev + t0 * strata::kernels::kStepCount; - strata::kernels::qsa_block_scores(pooled16, dead16, m.q_idx + t0 * 512, steps0, nb, - m.max_blocks, s, m.sel_scores, m.cs); - strata::kernels::qsa_block_topk(m.sel_scores, steps0, nb, m.max_blocks, m.cap, s, - ids16 + t0 * m.cap, m.cs); - } - std::vector a((size_t) (T * m.cap)), b((size_t) (T * m.cap)); - cudaMemcpyAsync(a.data(), m.sel_ids, a.size() * 4, cudaMemcpyDeviceToHost, m.cs); - cudaMemcpyAsync(b.data(), ids16, b.size() * 4, cudaMemcpyDeviceToHost, m.cs); - cudaStreamSynchronize(m.cs); - for (int64_t t = 0; t < T; ++t) { - const int64_t w = m.steps_host[(size_t) (t * strata::kernels::kStepCount + strata::kernels::kStepWidth)]; - const int32_t *x = a.data() + t * m.cap, *y = b.data() + t * m.cap; - int64_t i = 0, j = 0, c = 0; - while (i < w && j < w) { - if (x[i] == y[j]) { ++c; ++i; ++j; } else if (x[i] < y[j]) ++i; else ++j; - } - ++queries; - same += c == w; - if (p0 + t + 1 > m.cap) { ++sel_queries; shared += (double) c; cells += (double) w; } - } - if (qsa_index + 1 == g.n_qsa_layers()) - std::fprintf(stderr, "strata prefill: FP16 indexer keys: %lld of %lld selections identical; " - "where the selection is sparse, %.4f%% of cells shared (%lld queries)\n", - same, queries, cells > 0 ? 100.0 * shared / cells : 100.0, sel_queries); - } - // STRATA_QSA_DUMP=: append every QSA layer's selected cells for the prompt's last - // STRATA_QSA_DUMP_LAST (4096) positions - records of int32 {qsa layer, pos0, T, cap} + T*cap cells, - // for tools/qsa_locality.py (how local the sparse attention's reads are: the KV-streaming question) - if (static const char* dump = std::getenv("STRATA_QSA_DUMP"); dump != nullptr) { - static const long long last = std::getenv("STRATA_QSA_DUMP_LAST") - ? std::atoll(std::getenv("STRATA_QSA_DUMP_LAST")) : 4096; - if (p0 + T > pos0 + n - last) { - std::vector h((size_t) (T * m.cap)); - cudaMemcpyAsync(h.data(), m.sel_ids, h.size() * 4, cudaMemcpyDeviceToHost, m.cs); - cudaStreamSynchronize(m.cs); - if (std::FILE* f = std::fopen(dump, "ab")) { - const int32_t hdr[4] = {(int32_t) qsa_index, (int32_t) p0, (int32_t) T, (int32_t) m.cap}; - std::fwrite(hdr, 4, 4, f); - std::fwrite(h.data(), 4, h.size(), f); - std::fclose(f); - } - } - } - const strata::kernels::QsaAttnPools pools = staged ? pools_of(m.stage, m.ident_table) - : core::qsa_attn_pools(st); - pt.mark(kPfQsaAttn, cs); - // perf-review D-1: the whole chunk on tensor cores, one block per (query, KV head), FP32-level - // accuracy but not bitwise (qsa_prompt_attn.hpp). Q4_0 KV, or STRATA_PROMPT_ATTN_OLD=1: the - // decode kernel, 32 queries at a time (K8V4 runs the tensor kernel's mode 3: INT8 K, - // V dequantized from its q4_0 blocks to fp16 at gather) - static const bool old_attn = std::getenv("STRATA_PROMPT_ATTN_OLD") != nullptr; - if (old_attn || !strata::kernels::qsa_prompt_attn_batch(m.q, pools, m.sel_ids, m.steps_dev, m.cap, s, - m.attn, T, m.cs)) - for (int64_t t0 = 0; t0 < T; t0 += m.attn_batch) { - const int64_t nb = std::min(m.attn_batch, T - t0); - strata::kernels::qsa_decode_attn_batch(m.q + t0 * ZV, pools, m.sel_ids + t0 * m.cap, - m.steps_dev + t0 * strata::kernels::kStepCount, m.cap, - s, m.attn_scratch, m.attn + t0 * ZV, nb, m.cs); - } - if (st.kv_q4 || st.kv_hybrid) strata::kernels::fwht256_inplace_cuda(m.attn, T * 24, m.cs); - pt.mark(kPfQsa, cs); - gate_attn(m.attn, m.Qf, m.attn_h, T, m.cs); - if (!native_proj(m.gemm, wo, m.attn_h, m.bo, T, v.name("attn_output.weight"), err)) return false; - ++qsa_index; - } else { - // ======================= MoE ======================= - const core::WeightRef *wr = need(v, "ffn_gate_inp.weight", err), - *wgi = need(v, "ffn_gate_inp_shexp.weight", err), - *wsg = need(v, "ffn_gate_shexp.weight", err), - *wsu = need(v, "ffn_up_shexp.weight", err), - *wsd = need(v, "ffn_down_shexp.weight", err); - if (!wr || !wgi || !wsg || !wsu || !wsd) return false; - pt.mark(kPfRouter, cs); - if (!bf16_proj(m.gemm, wr, m.mixed_bf, m.logits, T, v.name("ffn_gate_inp.weight"), err)) return false; - route(m.logits, m.ids, m.w, T, m.g->n_expert, m.cs); - // the shared expert and its scalar gate - if (!native_proj(m.gemm, wsg, m.mixed_h, m.sgate, T, v.name("ffn_gate_shexp.weight"), err)) return false; - if (!native_proj(m.gemm, wsu, m.mixed_h, m.sup, T, v.name("ffn_up_shexp.weight"), err)) return false; - swiglu_pair(m.sgate, m.sup, m.sh_h, T, m.cs); - if (!native_proj(m.gemm, wsd, m.sh_h, m.shared, T, v.name("ffn_down_shexp.weight"), err)) return false; - if (wgi->kind != core::WeightKind::Bf16InF32) { err = "prefill: shared gate is not BF16"; return false; } - m.gemm.bf16(m.mixed_bf, (const uint16_t*) wgi->data, m.sg, T, 1, N); - // group the (token, k) pairs by expert on the host - pt.mark(kPfHostGroup, cs); - // (the sync below also orders this layer's writes of slot/src/bounds after the previous - // layer's kernels that read them) - const bool grp_mapped = m.grp_host != nullptr; - int32_t* ids_h = grp_mapped ? m.grp_host : m.ids_host.data(); - int32_t* slot_h = grp_mapped ? m.grp_host + m.grp_tk : m.slot_host.data(); - int32_t* src_h = grp_mapped ? m.grp_host + 2 * m.grp_tk : m.src_host.data(); - if (grp_mapped) copy_i32(m.grp_dev, m.ids, T * K, m.cs); - else cudaMemcpyAsync(m.ids_host.data(), m.ids, (size_t) T * K * 4, cudaMemcpyDeviceToHost, m.cs); - cudaStreamSynchronize(m.cs); - pt.fold(); - std::fill(m.cnt.begin(), m.cnt.end(), 0); - for (int64_t i = 0; i < T * K; ++i) { - const int32_t e = ids_h[(size_t) i]; - if (e < 0 || e >= m.g->n_expert) { err = "prefill: routed id out of range"; return false; } - ++m.cnt[(size_t) e]; - } - m.off[0] = 0; - for (int64_t e = 0; e < m.g->n_expert; ++e) m.off[(size_t) e + 1] = m.off[(size_t) e] + m.cnt[(size_t) e]; - std::vector fill(m.off.begin(), m.off.end() - 1); - for (int64_t i = 0; i < T * K; ++i) { - const int32_t e = ids_h[(size_t) i]; - const int32_t p = fill[(size_t) e]++; - slot_h[(size_t) i] = p; - src_h[(size_t) p] = (int32_t) (i / K); - } - if (grp_mapped) { - copy_i32(m.slot_dev, m.grp_dev + m.grp_tk, T * K, m.cs); - copy_i32(m.src_dev, m.grp_dev + 2 * m.grp_tk, T * K, m.cs); - } else { - cudaMemcpyAsync(m.slot_dev, m.slot_host.data(), (size_t) T * K * 4, cudaMemcpyHostToDevice, m.cs); - cudaMemcpyAsync(m.src_dev, m.src_host.data(), (size_t) T * K * 4, cudaMemcpyHostToDevice, m.cs); - } - // the experts, in id order: resident ones from VRAM, the others through the staging ring - std::vector order; - for (int32_t e = 0; e < m.g->n_expert; ++e) if (m.cnt[(size_t) e] > 0) order.push_back(e); - const strata::kernels::cpu::ExpertLayout& lay = strata::kernels::cpu::expert_layout(); - const bool use_mmq = mmq_plan().any && mmq_plan().layer[(size_t) l]; - const int mmq_gt = lay.native ? lay.fmt[(size_t) l].gu_type : 42; - const int mmq_dt = lay.native ? lay.fmt[(size_t) l].d_type : 42; - const size_t mmq_gub = use_mmq ? mmq::matrix_bytes(mmq_gt, 1280, N) : 0; - const size_t mmq_db = use_mmq ? mmq::matrix_bytes(mmq_dt, N, 640) : 0; - pt.mark(kPfGather, cs); - if (use_mmq) { - // step 2b: the layer's activations as q8_1 rows in expert order, straight from `mixed` - mmq::quantize(m.mixed, m.src_dev, m.Xq, mmq_gt, N, N, T * K, m.cs); - // each group's rows: absolute bounds (gate/up reads the layer's rows), relative ones (down - // reads the group's own quantized H) - const size_t n = order.size(), ng = (n + MMQ_GROUP - 1) / MMQ_GROUP; - m.bounds_host.resize(n + 1 + ng * (MMQ_GROUP + 1)); - for (size_t j = 0; j < n; ++j) m.bounds_host[j] = m.off[(size_t) order[j]]; - m.bounds_host[n] = (int32_t) (T * K); - for (size_t g = 0; g < ng; ++g) - for (size_t i = 0; i <= MMQ_GROUP; ++i) - m.bounds_host[n + 1 + g * (MMQ_GROUP + 1) + i] = - m.bounds_host[std::min(n, g * MMQ_GROUP + i)] - m.bounds_host[g * MMQ_GROUP]; - if (grp_mapped) { - int32_t* bh = m.grp_host + 3 * m.grp_tk; - std::memcpy(bh, m.bounds_host.data(), m.bounds_host.size() * 4); - copy_i32(m.bounds_dev, m.grp_dev + 3 * m.grp_tk, (int64_t) m.bounds_host.size(), m.cs); - } else { - cudaMemcpyAsync(m.bounds_dev, m.bounds_host.data(), m.bounds_host.size() * 4, - cudaMemcpyHostToDevice, m.cs); - } - } else { - gather_rows16(m.mixed_h, m.src_dev, m.Xs, T * K, N, m.cs); - } - // Stage ahead: the copy stream moves blobs host -> device while the compute stream works. - int stage_next = 0; - std::vector stage_of(order.size(), -1); - // the unpinned ones are copied to pinned buffers by the stager's threads, in this order - std::vector job_of(order.size(), -1); - if (!stream_all) { - std::vector js; - for (size_t j = 0; j < order.size(); ++j) { - const int32_t e = order[j]; - if (m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0) continue; - if (m.src->pinned(l, e)) continue; - const uint8_t* b = m.src->blob(l, e); - if (!b) { err = "prefill: expert source has no blob"; return false; } - job_of[j] = (int) js.size(); - js.push_back({b, (size_t) lay.blob_bytes(l)}); - } - m.stager->start(std::move(js)); + const core::WeightRef *wr = nullptr, *wgi = nullptr, *wsg = nullptr, *wsu = nullptr, *wsd = nullptr; + if (!(wr = need(v, "ffn_gate_inp.weight", err)) || !(wgi = need(v, "ffn_gate_inp_shexp.weight", err)) || + !(wsg = need(v, "ffn_gate_shexp.weight", err)) || !(wsu = need(v, "ffn_up_shexp.weight", err)) || + !(wsd = need(v, "ffn_down_shexp.weight", err))) + return false; + if (wgi->kind != core::WeightKind::Bf16InF32) { err = "prefill: shared gate is not BF16"; return false; } + // ================ per sub-chunk: the mixer half, then the FFN half up to the routing ================ + for (int64_t t0 = 0; t0 < T; t0 += SUB) { + const int64_t P = std::min(SUB, T - t0), pp = p0 + t0; + uint16_t* mixed_h = b.mixed_h + t0 * N; + for (int half = 0; half < 2; ++half) { + mark(half == 0 ? kPsHcRead : kPsHcFfn); + // ---- the hyper-connection read of this half (the FFN half's norm came with the mixer's write) + const char* pre = half == 0 ? "hc_attn_" : "hc_ffn_"; + const std::string sn = std::string(pre) + "norm.weight", sd = std::string(pre) + "down.weight", + su = std::string(pre) + "up.weight", si = std::string(pre) + "inject.weight"; + const core::WeightRef *wn = need(v, sn.c_str(), err), *wd = need(v, sd.c_str(), err), + *wu = need(v, su.c_str(), err), *wi = need(v, si.c_str(), err); + if (!wn || !wd || !wu || !wi) return false; + float* inj = half == 0 ? b.inj_a : b.inj_f + t0 * HC; + if (half == 0 && pending) { + mark(kPsCombine); + ffn_write(t0, P, (const float*) wn->data); + mark(kPsHcRead); + } else if (half == 0) { + gr_norm_rs(b.R + t0 * D, (const float*) wn->data, EPS, b.rs, b.xn16, P, m.cs); } - StagerDone stager_done{stream_all ? nullptr : m.stager.get()}; - auto stage_one = [&](size_t j) -> bool { - const int32_t e = order[j]; - const bool resident = m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0; - if (resident) return true; - const int sl = stage_next; - stage_next = (stage_next + 1) % STAGE; - const auto th = Clock::now(); - const uint8_t* b = m.src->blob(l, e); - if (!b) { err = "prefill: expert source has no blob"; return false; } - if (m.src->pinned(l, e)) { - // DMA straight from the page-locked arena: the copy stream only waits for the slot - if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0); - cudaMemcpyAsync(m.stage_dev[sl], b, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, m.copy); - ++stats_.experts_dma; - } else { - // copied to a pinned buffer by the stager (waits only if it is behind), then DMA - const uint8_t* hb = m.stager->wait(job_of[j]); - if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0); - cudaMemcpyAsync(m.stage_dev[sl], hb, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, m.copy); - m.stager->issued_one(job_of[j], m.copy); - } - cudaEventRecord(m.copied[sl], m.copy); - m.stage_live[sl] = true; - stage_of[j] = sl; - stats_.ms_experts_host += ms_since(th); - ++stats_.experts_streamed; - return true; - }; - // one expert's products from its blob on the device; `slot` (a ring slot, or -1 for a resident - // expert) is released once the blob is read - auto compute = [&](size_t j, const uint8_t* blob_dev, int slot) -> bool { - const int32_t e = order[j]; - pt.mark(kPfDequant, cs); - if (use_mmq) { - // gather the expert into its group slot (GGUF blocks, unchanged or converted) - const size_t q = j % MMQ_GROUP; - if (lay.native) { - const auto& f = lay.fmt[(size_t) l]; - mmq::gather_native(blob_dev, blob_dev + f.up_off, mmq_gub / 2, blob_dev + f.down_off, - mmq_db, m.grp_gu + q * mmq_gub, m.grp_d + q * mmq_db, m.cs); - } else { - mmq::gather_strata_q2(blob_dev, m.grp_gu + q * mmq_gub, m.grp_d + q * mmq_db, m.cs); - } - if (slot >= 0) cudaEventRecord(m.used[slot], m.cs); - if (q + 1 < MMQ_GROUP && j + 1 < order.size()) return true; - // the group's products: gate/up, swiglu, the group's H to q8_1, down - const size_t j0 = j - q, g = j0 / MMQ_GROUP, n = order.size(); - const int ngx = (int) (q + 1); - const int64_t r0 = m.bounds_host[j0], nr = m.bounds_host[j + 1] - r0; - int64_t maxr = 0; - for (size_t i = j0; i <= j; ++i) maxr = std::max(maxr, m.cnt[(size_t) order[i]]); - pt.mark(kPfGemmGU, cs); - // the zeroed tail after the group's last expert (see MMQ_TAIL) - cudaMemsetAsync(m.grp_gu + (size_t) ngx * mmq_gub, 0, MMQ_TAIL, m.cs); - cudaMemsetAsync(m.grp_d + (size_t) ngx * mmq_db, 0, MMQ_TAIL, m.cs); - mmq::Product gu; - gu.w = m.grp_gu; gu.type = mmq_gt; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = mmq_gub; - gu.n = ngx; gu.xq = m.Xq; gu.bounds = m.bounds_dev + j0; gu.ids = m.ids_identity; - gu.total_rows = T * K; gu.max_rows = maxr; gu.dst = m.GU; gu.ld_dst = 1280; - m.mmq_ctx->run(gu, m.cs); - mmq::swiglu(m.GU + r0 * 1280, m.H + r0 * 640, nr, 640, !lay.native, m.cs); - pt.mark(kPfGemmD, cs); - mmq::quantize(m.H + r0 * 640, nullptr, m.Hq, mmq_dt, 640, 640, nr, m.cs); - mmq::Product dn; - dn.w = m.grp_d; dn.type = mmq_dt; dn.w_rows = N; dn.w_cols = 640; dn.expert_bytes = mmq_db; - dn.n = ngx; dn.xq = m.Hq; dn.bounds = m.bounds_dev + n + 1 + g * (MMQ_GROUP + 1); - dn.ids = m.ids_identity; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = m.Dm + r0 * N; - dn.ld_dst = N; - m.mmq_ctx->run(dn, m.cs); - return true; - } - const int q = (int) (j % DQ); - if (lay.native) { - // plan v0.3 P6: a native pack's layer, dequantized by llama.cpp's own formulas - const auto& f = lay.fmt[(size_t) l]; - strata::kernels::iq_dequant_gu_f16(f.gu_type, blob_dev, blob_dev + f.up_off, f.n_ff, f.n_embd, - m.dq_gu[q], m.cs); - strata::kernels::iq_dequant_f16(f.d_type, blob_dev + f.down_off, f.n_embd * f.n_ff, m.dq_d[q], m.cs); - } else { - blob_dequant_f16(blob_dev, m.dq_gu[q], m.dq_d[q], m.cs); - } - if (slot >= 0) cudaEventRecord(m.used[slot], m.cs); - const int64_t o0 = m.off[(size_t) e], ne = m.cnt[(size_t) e]; - pt.mark(kPfGemmGU, cs); - m.gemm.f16(m.Xs + o0 * N, m.dq_gu[q], m.GU + o0 * 1280, ne, 1280, N); - swiglu_interleaved(m.GU + o0 * 1280, m.Hh + o0 * 640, ne, m.cs); - pt.mark(kPfGemmD, cs); - m.gemm.f16(m.Hh + o0 * 640, m.dq_d[q], m.Dm + o0 * N, ne, N, 640); - return true; - }; - if (!stream_all) { - size_t staged = 0; - const size_t lookahead = STAGE - 1; - for (size_t j = 0; j < order.size(); ++j) { - while (staged < order.size() && staged <= j + lookahead) { - if (!stage_one(staged)) return false; - ++staged; - } - const int32_t e = order[j]; - if (stage_of[j] < 0) { - ++stats_.experts_resident; - if (!compute(j, m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e]), -1)) return false; - } else { - pt.mark(kPfWaitCopy, cs); - cudaStreamWaitEvent(m.cs, m.copied[stage_of[j]], 0); - if (!compute(j, m.stage_dev[stage_of[j]], stage_of[j])) return false; - } - } + if (!bf16_proj(m.gemm, wd, b.xn16, b.lo, P, sd, err)) return false; + gr_silu(b.lo, b.lo16, P, m.cs); + if (!bf16_proj(m.gemm, wu, b.lo16, b.gated, P, su, err)) return false; + if (!bf16_proj(m.gemm, wi, b.xn16, inj, P, si, err)) return false; + gr_mix_r(b.R + t0 * D, b.rs, (const float*) wn->data, b.gated, nullptr, b.mixed_bf, P, m.cs, + mixed_h); + if (half == 1) break; + + if (!qsa) { + // ======================= GDN ======================= + const core::WeightRef *wqkv = need(v, "attn_qkv.weight", err), *wg = need(v, "attn_gate.weight", err), + *wo = need(v, "ssm_out.weight", err), *wa = need(v, "ssm_alpha.weight", err), + *wb = need(v, "ssm_beta.weight", err), *wc = need(v, "ssm_conv1d.weight", err), + *wnm = need(v, "ssm_norm.weight", err), *wdt = need(v, "ssm_dt.bias", err), + *wsa = need(v, "ssm_a", err); + if (!wqkv || !wg || !wo || !wa || !wb || !wc || !wnm || !wdt || !wsa) return false; + mark(kPsGdn); + float* state = ss.gdn_state + (size_t) gdn_index * gdn_floats; + float* conv = state + (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size; + if (!native_proj(wqkv, mixed_h, b.qkv, P, v.name("attn_qkv.weight"), err)) return false; + if (!native_proj(wg, mixed_h, b.z, P, v.name("attn_gate.weight"), err)) return false; + if (!bf16_proj(m.gemm, wa, b.mixed_bf, b.ab, P, v.name("ssm_alpha.weight"), err, 2 * HV)) return false; + if (!bf16_proj(m.gemm, wb, b.mixed_bf, b.ab + HV, P, v.name("ssm_beta.weight"), err, 2 * HV)) return false; + gdn_gates(b.ab, (const float*) wdt->data, (const float*) wsa->data, b.gate, b.beta, P, m.cs); + mark(kPsGdnConv); + gdn_conv(conv, b.qkv, (const float*) wc->data, b.hbuf, P, EPS, m.cs); + mark(kPsGdnScan); + gdn_scan(state, b.hbuf, b.gate, b.beta, b.y, P, m.cs); + mark(kPsGdnNorm); + gdn_out_norm(b.y, b.z, (const float*) wnm->data, EPS, b.y_h, P, m.cs); + mark(kPsGdnOut); + if (!native_proj(wo, b.y_h, b.bo, P, v.name("ssm_out.weight"), err)) return false; } else { - // the streamed walk: this layer's entries [k, kend) in id order; an entry the routing did not - // pick only gives its slot back - size_t k = seq_start[(size_t) l]; - const size_t kend = seq_start[(size_t) l + 1]; - auto release_to = [&](int32_t e_stop) { - while (k < kend && seq[k].e < e_stop) { - cudaEventRecord(m.used[k % (size_t) m.ring], m.cs); - consumed = ++k; - give_back(consumed); - } - }; - for (size_t j = 0; j < order.size(); ++j) { - const int32_t e = order[j]; - release_to(e); - if (k < kend && seq[k].e == e) { - const int sl = (int) (k % (size_t) m.ring); - pt.mark(kPfWaitCopy, cs); - wait_issued(k); - cudaStreamWaitEvent(m.cs, m.copied[sl], 0); - if (!compute(j, m.stage_dev[sl], sl)) return false; - consumed = ++k; - give_back(consumed); - } else { - ++stats_.experts_resident; - if (!compute(j, m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e]), -1)) return false; - } + // ======================= QSA ======================= + const core::QsaState& st = ss.qsa_states[qsa_index]; + const core::WeightRef *wq = need(v, "attn_q.weight", err), *wk = need(v, "attn_k.weight", err), + *wv = need(v, "attn_v.weight", err), *wo = need(v, "attn_output.weight", err), + *wik = need(v, "indexer.k_proj.weight", err), + *wiq = need(v, "indexer.q_proj.weight", err), + *wqn = need(v, "attn_q_norm.weight", err), *wkn = need(v, "attn_k_norm.weight", err), + *wiqn = need(v, "indexer.q_norm.weight", err), + *wikn = need(v, "indexer.k_norm.weight", err); + if (!wq || !wk || !wv || !wo || !wik || !wiq || !wqn || !wkn || !wiqn || !wikn) return false; + const int32_t* steps = b.steps + t0 * SC; + mark(kPsQsaProj); + if (!native_proj(wk, mixed_h, b.Kc, P, v.name("attn_k.weight"), err)) return false; + if (!native_proj(wv, mixed_h, b.Vc, P, v.name("attn_v.weight"), err)) return false; + if (!native_proj(wq, mixed_h, b.Qf, P, v.name("attn_q.weight"), err)) return false; + if (!bf16_proj(m.gemm, wik, b.mixed_bf, b.idx_raw, P, v.name("indexer.k_proj.weight"), err)) return false; + if (!bf16_proj(m.gemm, wiq, b.mixed_bf, b.q_idx, P, v.name("indexer.q_proj.weight"), err)) return false; + rms_rows(b.Kc, (const float*) wkn->data, P * 2, 256, 256, EPS, m.cs); + rope(b.Kc, P, 2, 256, 512, pp, (float) strata::kernels::qsa_freq_base(), m.cs); + kv_append(b.Kc, b.Vc, P, pp, st.page_table, s.page_size, st.kv_int8 ? nullptr : st.k_pool, + st.kv_int8 ? nullptr : st.v_pool, st.k_q, st.v_q, st.k_scale, st.v_scale, m.cs); + split_q(b.Qf, b.q, P, m.cs); + rms_rows(b.q, (const float*) wqn->data, P * 24, 256, 256, EPS, m.cs); + rope(b.q, P, 24, 256, 6144, pp, (float) strata::kernels::qsa_freq_base(), m.cs); + rms_rows(b.q_idx, (const float*) wiqn->data, P * 4, 128, 128, EPS, m.cs); + rope(b.q_idx, P, 4, 128, 512, pp, (float) strata::kernels::qsa_freq_base(), m.cs); + // the indexer appends of the sub-chunk; then scores + selection for many queries at once: + // a query reads completed blocks (final once completed) and `dead` for its own tail block + const strata::kernels::QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; + mark(kPsQsaIndexer); + try { + strata::kernels::native_qsa_indexer_append_multi( + b.idx_raw, steps + strata::kernels::kStepPos, (int) SC, (int) P, 0, + (const float*) wikn->data, EPS, ib, s, st.max_cells, (float) strata::kernels::qsa_freq_base(), m.cs); + } catch (const std::exception& e) { err = std::string("prefill indexer: ") + e.what(); return false; } + mark(kPsQsaScores); + for (int64_t q0 = 0; q0 < P; q0 += SEL_BATCH) { + const int64_t nb = std::min(SEL_BATCH, P - q0); + // the batch's last query has completed the most blocks + const int64_t grid_blocks = + m.steps_host[(size_t) (t0 + q0 + nb - 1) * SC + strata::kernels::kStepNBid] + 1; + strata::kernels::qsa_block_scores(st.idx_pooled, st.idx_dead, b.q_idx + q0 * 512, steps + q0 * SC, + nb, m.max_blocks, s, b.sel_scores, m.cs, grid_blocks); + strata::kernels::qsa_block_topk(b.sel_scores, steps + q0 * SC, nb, m.max_blocks, m.cap, s, + b.sel_ids + q0 * m.cap, m.cs); } - release_to(m.g->n_expert); + strata::kernels::QsaAttnPools pools; + pools.page_table = st.page_table; + if (st.kv_int8) { pools.k_q = st.k_q; pools.v_q = st.v_q; pools.k_scale = st.k_scale; pools.v_scale = st.v_scale; } + else { pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; } + mark(kPsQsaAttn); + strata::kernels::qsa_prefill_attn(b.q, pools, b.sel_ids, steps, m.cap, s, b.attn, P, m.cs); + mark(kPsQsaOut); + gate_attn(b.attn, b.Qf, b.attn_h, P, m.cs); + if (!native_proj(wo, b.attn_h, b.bo, P, v.name("attn_output.weight"), err)) return false; } - pt.mark(kPfCombine, cs); - moe_combine(m.Dm, m.slot_dev, m.w, m.shared, m.sg, m.bo, T, m.cs); - // debug: STRATA_DBG_NAN=1 reports the first layer of a chunk whose MoE produced non-finite values - if (static const bool dbg = std::getenv("STRATA_DBG_NAN") != nullptr; dbg) { - cudaStreamSynchronize(m.cs); - auto bad = [&](const float* d, int64_t n) { - std::vector h((size_t) n); - cudaMemcpy(h.data(), d, (size_t) n * 4, cudaMemcpyDeviceToHost); - int64_t c = 0; - for (float v : h) c += !std::isfinite(v); - return c; - }; - const int64_t bgu = bad(m.GU, T * K * 1280), bdm = bad(m.Dm, T * K * N), bbo = bad(m.bo, T * N); - const int64_t bh = m.H ? bad(m.H, T * K * 640) : -1; - static int64_t reported = -1; - if ((bgu || bdm || bbo || bh > 0) && reported != stats_.chunks) { - reported = stats_.chunks; - std::fprintf(stderr, "strata dbg: layer %lld (mmq %d, types %d/%d, %zu experts): non-finite GU %lld " - "H %lld Dm %lld bo %lld of T %lld\n", (long long) l, (int) use_mmq, mmq_gt, mmq_dt, - order.size(), (long long) bgu, (long long) bh, (long long) bdm, (long long) bbo, - (long long) T); - } - } - } - // ---- the hyper-connection write of this half; F-2: fused with the next half's norm when nothing else - // touches R in between (not the stage's last half, not before the PLE block of layer 1, not under a - // control vector) - const int64_t nl = half == 0 ? l : l + 1; - const bool fuse = !gr_unfused() && nl < LE && !(half == 1 && nl == 1 && ple_on) && - !(half == 1 && strata::kernels::cvec().covers(l)); - const core::WeightRef* wnn = nullptr; - if (fuse) { - const core::LayerView vn(*m.wt, nl); - wnn = need(vn, half == 0 ? "hc_ffn_norm.weight" : "hc_attn_norm.weight", err); - if (!wnn) return false; + // ---- the hyper-connection write of the mixer half, with the FFN half's norm + mark(kPsHcFfn); + const core::WeightRef* wfn = need(v, "hc_ffn_norm.weight", err); + if (!wfn) return false; + gr_write_norm_rs(b.R + t0 * D, b.bo, b.inj_a, HC, (const float*) wfn->data, EPS, b.rs, b.xn16, P, + m.cs); } - if (wnn) { - gr_write_norm_rs(m.R, m.bo, m.inj, HC, (const float*) wnn->data, EPS, m.grs, m.xn16, T, m.cs); - normed = true; - } else { - gr_write(m.R, m.bo, m.inj, HC, T, m.cs); + // ---- the FFN half's input is ready: its trip to the second GPU starts here + if (m.offload) { + cudaEventRecord(m.ev_mixed, m.cs); + cudaStreamWaitEvent(m.xfer, m.ev_mixed, 0); + cudaMemcpyAsync(m.offload->host_input() + t0 * N, mixed_h, (size_t) P * N * 2, + cudaMemcpyDeviceToHost, m.xfer); + cudaEventRecord(m.ev_in[(size_t) (t0 / SUB)], m.xfer); + m.offload->stage_input(t0, P, m.ev_in[(size_t) (t0 / SUB)]); } - if (half == 1 && strata::kernels::cvec().covers(l)) // --control-vector-scaled - strata::kernels::cvec_apply(m.R, l, T, D, nullptr, 0, nullptr, 0, false, m.cs); + // ---- the router, and the shared expert (its gated output starts the MoE sum) + mark(kPsRouter); + if (!bf16_proj(m.gemm, wr, b.mixed_bf, b.logits, P, v.name("ffn_gate_inp.weight"), err)) return false; + route(b.logits, b.ids + t0 * K, b.w + t0 * K, P, m.cs); + if (!native_proj(wsg, mixed_h, b.sgate, P, v.name("ffn_gate_shexp.weight"), err)) return false; + if (!native_proj(wsu, mixed_h, b.sup, P, v.name("ffn_up_shexp.weight"), err)) return false; + swiglu_pair(b.sgate, b.sup, b.sh_h, P, m.cs); + if (!native_proj(wsd, b.sh_h, b.bo_moe + t0 * N, P, v.name("ffn_down_shexp.weight"), err)) return false; + m.gemm.bf16(b.mixed_bf, (const uint16_t*) wgi->data, b.sg, P, 1, N); + moe_shared_gate(b.bo_moe + t0 * N, b.sg, P, m.cs); + cudaMemcpyAsync(m.ids_host + t0 * K, b.ids + t0 * K, (size_t) P * K * 4, cudaMemcpyDeviceToHost, m.cs); + cudaMemcpyAsync(m.w_host + t0 * K, b.w + t0 * K, (size_t) P * K * 4, cudaMemcpyDeviceToHost, m.cs); } - } - if (issuer.joinable()) { - issuer.join(); - stats_.ms_experts_host += iss_ms; - stats_.experts_streamed += iss_streamed; - stats_.experts_dma += iss_dma; - } - stats_.tokens += T; - pt.mark(kPfStart, cs); - if (next_ != nullptr) { - // the rows to the host buffer the next stage read two chunks ago (it has finished: waited below) - float* h = m.hand[hand_buf]; - if (cudaMemcpyAsync(h, m.R, (size_t) T * D * 4, cudaMemcpyDeviceToHost, m.cs) != cudaSuccess || - cudaStreamSynchronize(m.cs) != cudaSuccess) { - err = std::string("prefill: the layer split's hand-off: ") + cudaGetErrorString(cudaGetLastError()); + pending = false; + if (!qsa) ++gdn_index; + else ++qsa_index; + // ================ the routed experts of the whole chunk ================ + if (cudaStreamSynchronize(m.cs) != cudaSuccess) { + err = std::string("prefill: ") + cudaGetErrorString(cudaGetLastError()); return false; } - // this stage's state is at the chunk's end now (synced) and moves on with the next chunk below - if (on_stage_chunk && !on_stage_chunk(p0 + T, err)) return false; - if (next_run.valid() && !next_run.get()) { err = next_err; return false; } - next_->hand_in_ = h; - next_run = std::async(std::launch::async, [this, tokens, c0, T, p0, &next_err] { - return next_->run(tokens + c0, T, p0, next_err); - }); - hand_buf ^= 1; - continue; // the last stage reports the chunk (on_chunk) - } - if (const char* dump = std::getenv("STRATA_PREFILL_DUMP_R")) { // debug: the final residuals, every 64th - cudaStreamSynchronize(m.cs); // position (A/B quality of this path) - if (std::FILE* f = std::fopen(dump, c0 == 0 ? "wb" : "ab")) { - std::vector row((size_t) D); - for (int64_t t = (64 - p0 % 64) % 64; t < T; t += 64) { - cudaMemcpy(row.data(), m.R + t * D, (size_t) D * 4, cudaMemcpyDeviceToHost); - const int64_t pos = p0 + t; - std::fwrite(&pos, sizeof pos, 1, f); - std::fwrite(row.data(), 4, row.size(), f); + mark(kPsExperts); + // group the (token, k) entries by expert: this GPU takes the experts its cache holds (all of them without + // a second GPU), the second GPU the others + std::fill(m.cnt.begin(), m.cnt.end(), 0); + for (int64_t i = 0; i < T * K; ++i) { + const int32_t e = m.ids_host[i]; + if (e < 0 || e >= NE) { err = "prefill: routed id out of range"; return false; } + ++m.cnt[(size_t) e]; + } + m.off[0] = 0; + for (int64_t e = 0; e < NE; ++e) m.off[(size_t) e + 1] = m.off[(size_t) e] + m.cnt[(size_t) e]; + { + std::vector fill(m.off.begin(), m.off.end() - 1); + for (int64_t i = 0; i < T * K; ++i) m.entry_of[(size_t) fill[(size_t) m.ids_host[i]]++] = (int32_t) i; + } + for (auto* v2 : {&m.ex1, &m.src1, &m.ex2, &m.src2}) v2->clear(); + m.w1.clear(); + m.w2.clear(); + m.off1.assign(1, 0); + m.off2.assign(1, 0); + for (int32_t e = 0; e < NE; ++e) { + if (m.cnt[(size_t) e] == 0) continue; + const bool here = !m.offload || (m.host_res && m.cache && m.host_res[(size_t) l * NE + e] >= 0); + std::vector& ex = here ? m.ex1 : m.ex2; + std::vector& of = here ? m.off1 : m.off2; + std::vector& sr = here ? m.src1 : m.src2; + std::vector& wv = here ? m.w1 : m.w2; + ex.push_back(e); + for (int32_t p = m.off[(size_t) e]; p < m.off[(size_t) e + 1]; ++p) { + const int32_t i = m.entry_of[(size_t) p]; + sr.push_back(i / (int32_t) K); + wv.push_back(m.w_host[i]); } - std::fclose(f); + of.push_back((int32_t) sr.size()); } - } - if (on_chunk || on_stage_chunk) { - const auto toc = Clock::now(); - if (cudaStreamSynchronize(m.cs) != cudaSuccess) { - err = std::string("prefill: ") + cudaGetErrorString(cudaGetLastError()); + if (m.offload && !m.offload->run_layer(l, T, m.ex2, m.off2, m.src2, m.w2, SUB, nullptr, nullptr, err)) return false; + if (!m.local.run_layer(l, T, m.ex1, m.off1, m.src1, m.w1, SUB, b.mixed_h, b.bo_moe, err)) return false; + // ---- the MoE output (+ the second GPU's sums) goes into the residual stream as the next layer reads it + if (m.offload) { + for (int64_t t0 = 0; t0 < T; t0 += SUB) { + cudaStreamWaitEvent(m.xsum, m.offload->piece_done(t0 / SUB), 0); + const size_t bytes = (size_t) std::min(SUB, T - t0) * N * 2; + cudaMemcpyAsync(b.sums + t0 * N, m.offload->host_sum() + t0 * N, bytes, cudaMemcpyHostToDevice, + m.xsum); + cudaEventRecord(m.ev_piece[(size_t) (t0 / SUB)], m.xsum); + } + } + pending = true; + } + mark(kPsCombine); // the last layer's write + for (int64_t t0 = 0; t0 < T; t0 += SUB) ffn_write(t0, std::min(SUB, T - t0), nullptr); + if (draft != nullptr && !draft_pass(tokens + c0, T, p0)) return false; + if (profile) { + mark(kPsPle); + cudaEventSynchronize(ev.back().second); + for (size_t i = 0; i + 1 < ev.size(); ++i) { + float ms = 0; + cudaEventElapsedTime(&ms, ev[i].second, ev[i + 1].second); + stats_.ms_section[ev[i].first] += ms; } - const auto toc2 = Clock::now(); - if (on_stage_chunk && !on_stage_chunk(p0 + T, err)) return false; - if (on_chunk && !on_chunk(m.R, T, p0, err)) return false; - host_sync_ms += std::chrono::duration(toc2 - toc).count(); - host_chunk_ms += ms_since(toc2); + for (auto& e : ev) cudaEventDestroy(e.second); + ev.clear(); } + stats_.tokens += T; + for (int64_t t = 0; t < T; ++t) { prev[0] = prev[1]; prev[1] = (int32_t) tokens[c0 + t]; } } - if (next_run.valid() && !next_run.get()) { err = next_err; return false; } ss.ple_prev[0] = prev[0]; ss.ple_prev[1] = prev[1]; - if (std::getenv("STRATA_DBG_NAN") != nullptr) { // debug: the state the prompt leaves for the token path - cudaStreamSynchronize(m.cs); - auto bad = [&](const float* d, int64_t n) { - std::vector h((size_t) n); - cudaMemcpy(h.data(), d, (size_t) n * 4, cudaMemcpyDeviceToHost); - int64_t c = 0; - double mx = 0; - for (float v : h) { c += !std::isfinite(v); if (std::isfinite(v)) mx = std::max(mx, (double) std::fabs(v)); } - std::fprintf(stderr, " %lld non-finite (max |x| %.3g)", (long long) c, mx); - }; - const int64_t last = (n - 1) % m.T; - std::fprintf(stderr, "strata dbg: prompt end: last residual row"); - bad(m.R + last * D, D); - if (ss.ple.ready()) { std::fprintf(stderr, "; PLE history"); bad(ss.ple.hist, (int64_t) strata::kernels::NG_HIST * strata::kernels::NG_HC_DIM); } - std::fprintf(stderr, "; GDN state 0"); - bad(ss.gdn_state, 64 * 1024); - std::fprintf(stderr, "\n"); - } if (cudaStreamSynchronize(m.cs) != cudaSuccess) { err = std::string("prefill: ") + cudaGetErrorString(cudaGetLastError()); return false; } - // (PR #121) an expert copy that failed on the copy stream surfaces here, not in the next request - if (const cudaError_t cst = cudaStreamSynchronize(m.copy); cst != cudaSuccess) { - err = std::string("prefill: expert copy stream: ") + cudaGetErrorString(cst); - return false; - } + stats_.experts_streamed = m.local.experts_streamed; + stats_.experts_prefetched = m.local.experts_prefetched; + stats_.experts_resident = m.local.experts_resident; + stats_.ms_experts_host = m.local.ms_host; stats_.ms_total += ms_since(t_start); - if (pt.on) { - pt.fold(); - double total = 0.0; - for (double v : pt.ms) total += v; - std::string line; - char b[96]; - for (int i = 0; i < kPfCount; ++i) { - if (pt.ms[i] <= 0.0) continue; - std::snprintf(b, sizeof b, " %s %.0f (%.1f%%)", kPfNames[i], pt.ms[i], total > 0 ? 100.0 * pt.ms[i] / total : 0.0); - line += b; - } - std::fprintf(stderr, "strata prefill timing: %lld tokens, GPU timeline %.0f ms, wall %.0f ms, host staging %.0f ms:%s\n", - (long long) n, total, ms_since(t_start), stats_.ms_experts_host, line.c_str()); - std::fprintf(stderr, "strata prefill timing: host: chunk setup (PLE rows, the expert stream plan) %.0f ms, " - "waiting for each chunk %.0f ms, after each chunk (the draft layer, progress) %.0f ms, " - "PLE %.0f ms\n", host_setup_ms, host_sync_ms, host_chunk_ms, stats_.ms_ple); - } - if (std::getenv("STRATA_STATE_HASH_GDN") != nullptr) { // debug: the GDN states as the prompt path leaves them - cudaStreamSynchronize(m.cs); - std::vector b((size_t) gdn_floats * 4); - std::string line; - char h[8]; - for (int64_t i = 0; i < g.n_gdn_layers(); ++i) { - cudaMemcpy(b.data(), ss.gdn_state + (size_t) i * gdn_floats, b.size(), cudaMemcpyDeviceToHost); - uint64_t x = 1469598103934665603ull; - for (uint8_t c : b) x = (x ^ c) * 1099511628211ull; - std::snprintf(h, sizeof(h), "%04llx ", (unsigned long long) (x & 0xffff)); - line += h; - } - std::fprintf(stderr, "strata prefill: GDN_HASH %s\n", line.c_str()); + processed_ = std::min(c0, n); + if (c0 < n) { + err = "cancelled"; + return false; } return true; } +const char* prefill_section_name(int section) { + static const char* const names[kPsCount] = { + "PLE block", "HC read (mixer)", "GDN projections", "GDN conv", "GDN recurrence", "GDN output norm", + "GDN output projection", + "QSA projections, norms, KV", "QSA indexer appends", + "QSA block scores + top-k", "QSA attention", "QSA gate + output", "HC write + HC read (FFN)", + "router, shared expert", "grouping, experts", "combine + HC write", "draft layer K/V"}; + return section >= 0 && section < kPsCount ? names[section] : "?"; +} + } // namespace strata::prefill diff --git a/src/program/generate.cpp b/src/program/generate.cpp index 0bbc92e03e7..c03fded4742 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -14,10 +14,12 @@ // AND IT IS PHASE 2, so hit rate is `h = 0` and the number it prints is slow on purpose // (`phase-2-correct-engine.md:5-9`). What it is FOR is the honest tok/s figure and the logit dump. +#include "strata/artifact/gguf_split.hpp" +#include "strata/core/adaptive_tier.hpp" #include "strata/core/expert_cache.hpp" +#include "strata/core/second_gpu.hpp" +#include "strata/core/split_head.hpp" #include "strata/core/expert_source.hpp" -#include "strata/core/remote_experts.hpp" -#include "strata/core/on_device.hpp" #include "strata/core/layer.hpp" #include "strata/core/layout.hpp" #include "strata/core/session.hpp" @@ -28,44 +30,30 @@ #include "strata/kernels/ngram.hpp" #include "strata/kernels/s2_expert_grouped.hpp" #include "strata/kernels/sampler.hpp" -#include "strata/kernels/verify_kernels.hpp" #include "strata/kernels/shared_expert.hpp" +#include "strata/kernels/iq_kernels.hpp" #include "strata/kernels/native_moe.hpp" +#include "strata/kernels/native_mmvq.hpp" #include "strata/kernels/native_gdn.hpp" #include "strata/kernels/native_router.hpp" #include "strata/kernels/native_qsa.hpp" #include "strata/kernels/native_qsa_indexer.hpp" #include "strata/kernels/native_rope.hpp" #include "strata/kernels/mrope.hpp" -#include "strata/kernels/kv_q4.hpp" #include "strata/kernels/qsa.hpp" +#include "strata/kernels/verify_kernels.hpp" #include "strata/core/native_head.hpp" #include "strata/core/verify.hpp" #include "strata/core/mtp.hpp" +#include "strata/core/prompt_cache.hpp" +#include "strata/prefill/experts.hpp" #include "strata/prefill/prefill.hpp" +#include "strata/spec/suffix_drafter.hpp" #include "strata/core/native_dense.hpp" #include "strata/program/logits_selection.hpp" -#include "strata/program/conv_cache.hpp" -#include "strata/spec/draft_policy.hpp" -#include "strata/spec/suffix_drafter.hpp" -#include "strata/kernels/cvec.hpp" -#include "strata/core/progress.hpp" -#ifndef NOMINMAX -#define NOMINMAX // gguf_reader.hpp includes windows.h -#endif -#include "strata/artifact/gguf_reader.hpp" -#if defined(_WIN32) -#include -#include -#include -#else -#include -#include -#endif #include -#include #include #include #include @@ -73,27 +61,21 @@ #include #include #include -#include #include #include #include -#include #include #include #include -#include #include +#include +#include #include #include #include #include namespace { -// perf-review D-4: the lent slots are refilled with queued copies and one wait; STRATA_REFILL_BLOCKING=1 waits on each -bool refill_blocking() { - static const bool v = std::getenv("STRATA_REFILL_BLOCKING") != nullptr; - return v; -} using Clock = std::chrono::steady_clock; @@ -146,7 +128,6 @@ struct Options { double ple_delay_us = 0; ///< fault injection: every row read completes no earlier than this bool ple_sync_submit = false; ///< A/B arm: submit reads on the token thread, no I/O worker std::string kv = "fp16"; ///< plan v0.3 P7: KV storage, fp16 (default) or int8 (half the VRAM) - int64_t kv_resident = 0; ///< KV streaming: resident cells per QSA layer (0: all in VRAM) std::string dump_residual; /// The head input, `bb.mixed`. It exists so the head can be SPLIT: steps 1-4 (the per-stream norm, the two /// bf16 projections and the stream mean) recompute cheaply in Python, and only the 794 MB GEMV does not. @@ -182,7 +163,6 @@ struct Options { /// is reserved for a host thread that has nothing to do while the drain runs. bool no_host_worker = false; bool mmap_experts = false; ///< R2.1: opt OUT of the resident arena, back to MapViewOfFile - bool resident_cpu_experts = false; ///< mmap-backed static-cache misses copied into ordinary RAM /// R4: slots of VRAM-resident experts. **0 = off, and off is the default.** /// **THE COMMENT THAT USED TO BE HERE WAS FALSE AND ROUND 328 MEASURED IT.** It said "the cache has no /// consumer yet - `moe_hit_grouped_s2` does not exist - so switching it on costs the fill traffic and @@ -192,8 +172,6 @@ struct Options { /// true is that the ADMISSION POLICY gave every slot to the first position, which is why the earlier /// measurement found nothing - see `expert_cache_per_layer`. int expert_cache = 0; - std::array expert_cache_remote{}; ///< CUDA1..3 slots; CUDA0 keeps dense/state/MTP - std::string expert_cache_remote_placement = "stripe"; ///< stripe experts or assign complete layers to CUDA1..3 bool expert_cache_cpu_order = false; /// **R4.2g. ROUND 328 MEASURED THAT THE GLOBAL ADMISSION POLICY CANNOT WORK, AND THIS IS THE FIX.** /// The default policy hands out slots in arrival order from one counter shared by all 48 layers, so the @@ -230,13 +208,13 @@ struct Options { int vram_reserve_mib = 700; /// Plan v0.3 P5: batched prompt processing in chunks of this many tokens (0 = the token path). int64_t prefill_chunk = 0; - /// `--prefill auto`: the largest chunk (up to 8192) whose buffers the expert cache can lend. Every expert a chunk - /// routes to is streamed once per chunk, so a bigger chunk streams fewer bytes per token (the "ubatch" effect). - bool prefill_auto = false; bool no_split_rows = false; ///< plan v0.3 P4 A/B: one whole expert per pool thread - /// Plan v0.3 P5: the prompt path borrows the top expert-cache slots for its buffers and refills them after - /// the prompt (default); `--no-prefill-borrow` reserves the buffers' VRAM for the whole session instead. + /// Plan v0.3 P5: the prompt path borrows the last expert-cache slots of each GPU for its buffers and refills + /// them after the prompt (default); `--no-prefill-borrow` reserves the buffers' VRAM for the whole session. bool no_prefill_borrow = false; + /// With --second-gpu, the prompt path's experts the main GPU's cache does not hold run on the second GPU + /// (default); `--no-prompt-offload` streams them to the main GPU instead. + bool no_prompt_offload = false; /// Plan v0.3 P5 validation: batch only positions [0, P) and run the rest of the prompt through the token path /// (teacher-forced), so the logits of positions >= P - which depend on the batched state - can be scored /// against the oracle at many positions. 0 = the whole prompt but the last position. @@ -250,9 +228,21 @@ struct Options { int spec = 0; std::string spec_oracle; int spec_corrupt = 0; + /// Benchmarks: the run emits this continuation (token ids) instead of the window's argmax, and a draft is + /// accepted when it matches it, so runs with different speculation settings process the same text. + std::string spec_follow; /// Plan v0.3 P6: the MTP draft layer's runtime directory (tools/mtp_rt.py); drafts come from it. std::string mtp; int64_t mtp_window = 32768; ///< the draft layer attends to the last N cells (0 = every cell) + int main_gpu = 0; ///< the CUDA device (of the visible ones) the engine runs on + /// A second GPU (device ordinal) as another expert tier, with this much VRAM for experts (0 = all but + /// `second_gpu_reserve_mib`). + int second_gpu = -1; + double second_gpu_gib = 0.0; + int second_gpu_reserve_mib = 2048; + double second_gpu_min_mb = 4.0; ///< a layer's misses from which it takes its share (smaller: the CPU is quicker) + double second_gpu_prefetch_mb = 12.0; ///< per layer, the likeliest experts no GPU holds copied to it ahead (0 = off) + double head_split = 0.45; ///< the second GPU's share of the head's rows (0: all on the main GPU) /// Plan v0.3 P6: the share (0..1) of each layer's distinct missed experts the GPU reads over PCIe from the /// pinned arena while the CPU computes the rest (verify windows). double pcie_frac = -1.0; ///< < 0: the model's default (0.2 direct for the Q2_0 pack, 0.55 DMA for native packs) @@ -263,51 +253,38 @@ struct Options { /// Plan v0.3 P6: a draft enters the verify window only while every draft before it (and itself) has at least /// this probability under the draft layer; 0 = always --spec-1 drafts. double spec_min_p = 0.0; + /// Prompt lookup: when the text's last N+ tokens occurred before and the tokens after that occurrence are more + /// drafts than --spec gives the draft layer, the next window verifies up to kVerifyMaxT - 1 of them instead; + /// 0 = off. + int spec_lookup = 0; + /// The verify windows' capacity: --spec, or the kernels' maximum with prompt lookup. + int max_window() const { return spec_lookup > 0 ? std::max(spec, strata::kernels::kVerifyMaxT) : spec; } /// Stop when the model emits an end-of-turn token (<|endoftext|> 248044, <|im_end|> 248046, or --eos-ids). bool stop_eos = false; std::vector eos_ids = {248044, 248046}; bool spec_split = false; ///< opt-in split verify window (the overlap study: exact, ~7% slower) - /// --serve, multi-GPU layer split: "K" or "K1,K2,.." (the first layer of each later stage) or "auto" (placed - /// from each GPU's free VRAM); empty = one GPU - std::string layer_split; - /// the later stages' devices "D1,D2,.." (default: the next visible GPUs; "0" with one K: both stages on this - /// GPU, sharing everything - the bit-exact A/B of the hand-off) - std::string split_device; + bool window_profile = false; ///< GPU timestamps between the verify window's stages, printed by --stats + std::string window_hashes; ///< per verify window: a hash of its final residual rows and its argmaxes + std::string window_logits; ///< per emitted position: the 64 likeliest tokens and their log-probabilities + bool prefill_profile = false; ///< GPU time per section of the batched prompt path /// Plan v0.3 P8: stay resident and take requests on stdin (see the --serve block in main). bool serve = false; /// The vision path: keep a per-cell (t, h, w) rotary position table so --serve can take GENI requests. bool vision = false; + /// --serve prompt cache: checkpoints of the sequence state kept in host memory (0 = every request starts + /// from an empty sequence). A request continues from the longest prefix the session or a checkpoint holds. + int prompt_cache = 16; + /// Host memory for sequences the live one replaced (their K/V cells), in MiB. + int64_t cache_ram_mib = 16384; + /// Checkpoint interval inside long prompts, in whole prompt chunks (counted from the request's first new cell). + int64_t cache_every = 16384; + /// Prompt parts of up to this many tokens go through verify windows instead of the batched prompt path: a new + /// part that short, and a last prompt chunk that short. A batched call has a fixed cost (it lends and refills + /// expert-cache slots and reads every expert its tokens use): ~0.55 s with a second GPU, ~1.6 s with IQ3_XXS on + /// the 3090 alone; verify windows of 8 tokens (--spec-lookup) cost 5-8 ms per token. They break even at ~73 + /// (UD-Q4_K_XL) and ~113 (IQ3_XXS) tokens with a second GPU, ~400 without (bench/feed_test.py). + int64_t feed_max = 128; int adapt_swaps = 96; - /// --serve: how many conversation checkpoints to keep between requests (0 = every request reads its whole - /// prompt again, the v0.1.2 behaviour). One is the GDN recurrence of the 36 layers, the QSA indexer tails and - /// the PLE history (~118 MB of host RAM); the KV cache itself is positional and stays where it is. - int prompt_cache = 6; - /// --serve: also keep a checkpoint every N freshly read prompt tokens (0 = only at the last turn boundary) - int64_t prompt_cache_every = 16384; - /// --serve: a prompt read from token 0 is also checkpointed at its first turn boundary - the end of the system - /// prompt, which every chat of the same client shares - when that is at least N tokens (0 = never) - int64_t prompt_cache_root = 2048; - /// --serve: the token that opens a chat turn (<|im_start|>). The last one in a prompt is where the chat's - /// history ends and the new assistant turn begins, which is the checkpoint the next request can reuse. - int64_t turn_token = 248045; - /// --serve: a text part of the prompt of at most N tokens (a chat message, the assistant header, a short tool - /// result) goes through the verify windows, S tokens at a time, instead of the batched prompt path (0 = always - /// the batched path) - int64_t short_read = 64; - /// The suffix drafter (prompt lookup): when the text being written repeats an earlier stretch of the context (code - /// edits, quoted input, tool-call JSON) by at least this many tokens, the window may be filled with what followed - /// it there instead of the MTP's drafts, where the MTP's own first guess agrees and the draft policy expects it to - /// pay (strata/spec/draft_policy.hpp). On by default; 0 = MTP only. - int suffix_draft = 3; - /// The MTP's own window cap (0 = --spec): with --spec 6 --mtp-max-t 4 the long windows come from suffix matches. - int mtp_max_t = 0; - /// A control vector on the residual stream (strata/kernels/cvec.hpp), with llama.cpp's flags: the - /// `experimental-speed-projection` profile passes `--control-vector-scaled FILE:1.0 --control-vector-layer-range - /// 4 44 --cvec-mode project --cvec-dir per-layer`. None by default; --serve switches a loaded one per request. - std::vector> cvec_files; - int cvec_first = -1, cvec_last = -1; ///< llama.cpp's defaults: 1 .. the last layer - int cvec_mode = 1; ///< 0 = project, 1 = add (llama.cpp's default) - int cvec_single = -1; ///< --cvec-dir single:L (project mode): layer L's direction everywhere }; void usage() { @@ -317,23 +294,18 @@ void usage() { " --pack DIR the pack directory (default pack/full)\n" " --tokens LIST the prompt as comma-separated token IDS (required)\n" " --tokens-file PATH pretokenized prompt, commas or whitespace (alternative to --tokens)\n" - " --ple-gguf PATH required PLE table (original second GGUF shard)\n" + " --ple-gguf PATH the PLE table: the model shard holding per_layer_token_embd (with --native,\n" + " the --native model's shard that holds it by default)\n" " --no-ple explicit diagnostic ablation of the PLE layer\n" - " --ple-io direct|mmap n-gram table reads (plan v0.3 P2). direct (default): unbuffered SSD\n" - " reads, the table never enters RAM or the file cache; mmap: A/B arm\n" + " --ple-io direct|mmap|ram n-gram table reads (plan v0.3 P2). direct (default): unbuffered SSD\n" + " reads, the table never enters RAM or the file cache; mmap: A/B arm; ram:\n" + " read into RAM once the experts are loaded (28.8 GB), direct reads until then\n" " --ple-row-cache N bounded cache of fetched rows, 90 B each (default 1048576; 0 = off)\n" " --ple-inflight N outstanding SSD reads (default 64)\n" " --ple-delay-us U fault injection: each row read completes no earlier than U us\n" " --ple-sync-submit A/B arm: submit table reads on the token thread (default: an I/O thread)\n" " --kv fp16|int8 KV storage (plan v0.3 P7): int8 codes + fp16 scale per 64 values, half the\n" " VRAM; default fp16 until gate G-C accepts int8\n" - " --kv q4_0 4-bit K/V after a Hadamard rotation (PR #21): half of int8's memory,\n" - " slightly lower precision (see bench/results/2026-09-27-kv-q4)\n" - " --kv k8v4 hybrid: INT8 K (exact attention scores) + rotated Q4_0 V, 816 B/cell\n" - " (vs int8's 1,056); not with --kv-resident\n" - " --kv-resident N KV streaming: keep N cells of each QSA layer in VRAM (min 20480) and the\n" - " whole K/V in pinned RAM; the freed VRAM goes to expert slots. 0 (default):\n" - " all of it in VRAM. A context of N cells or fewer is not streamed\n" " --stream-token enqueue token work on the session stream (experimental)\n" " --check-logits copy and check all logits in the stream-token path\n" " --gr-fp32-activations experimental CUDA-oracle GR activation precision\n" @@ -350,11 +322,12 @@ void usage() { " --native-ple-postops experimental pinned PLE postprojection arithmetic\n" " --native-router experimental pinned CUDA 512-expert top-10 routing\n" " --cpu-oracle-q8-0 experimental pinned CPU expert quantization and dot reduction\n" - " --native SHARD1 every full-context native path at once (plan v0.3 P1): stream-token,\n" + " --native SHARD1 the model's first GGUF shard (the others of a split model are found by\n" + " name) and every full-context native path at once (plan v0.3 P1): stream-token,\n" " GR MMVF, BF16, head, dense + PLE key, MoE combine, GDN, router, QSA,\n" " indexer, RoPE, PLE postops, and the CPU q8_0 contract unless the\n" " expert cache is on. Individual --native-* flags stay for A/B.\n" - " --native-head-gguf PATH native Q5_K head from model shard 1; requires --stream-token\n" + " --native-head-gguf PATH native head from the model with this first shard; requires --stream-token\n" " --native-dense-gguf PATH native GDN/QSA/shared projections; repeat for each source model shard\n" " --expert-cache-cpu-order experimental GPU expert reduction matching CPU order\n" " --max-new N tokens to generate (default 16)\n" @@ -372,28 +345,46 @@ void usage() { " --shared-late A/B: shared expert after the CPU pool (default: overlapped with it)\n" " --keep-canonical A/B: also load canonical copies of natively served tensors (more VRAM)\n" " --vision --serve takes images too (GENI requests; embeddings from strata-vision)\n" - " --prompt-cache N --serve: keep N conversation checkpoints between requests (default 6, ~118 MB\n" - " of RAM each; 0 = read every prompt from the start)\n" - " --prompt-cache-every N --serve: also checkpoint every N fresh prompt tokens (default 16384, 0 = off)\n" - " --turn-token ID --serve: the token that opens a chat turn (default 248045, <|im_start|>)\n" - " --short-read N --serve: read at most N fresh text tokens through the decode windows instead\n" - " of the batched prompt path (default 64, 0 = off)\n" - " --suffix-draft N prompt lookup: draft from an earlier repeat of the last N+ tokens of context\n" - " when it pays (default 3; 0 = MTP only)\n" - " --mtp-max-t M cap the MTP's windows at M tokens (0 = --spec; longer ones come from suffixes)\n" - " --control-vector-scaled FILE:SCALE[,...] a control vector GGUF on the residual stream (llama.cpp's\n" - " format; --control-vector FILE = scale 1). --serve: requests switch it (cvec=0|1)\n" - " --control-vector-layer-range A B the layers it follows (inclusive; default 1 .. the last)\n" - " --cvec-mode add|project h += s v (default) or h -= s (h.v) v with v unit\n" - " --cvec-dir per-layer|single:L each layer's own direction (default) or layer L's everywhere (project)\n" + " --prompt-cache N --serve: keep N checkpoints of the sequence state in host memory and\n" + " continue each request from the longest cached prefix (default 16;\n" + " 0 = every request starts from an empty sequence)\n" + " --cache-every N checkpoint interval inside long prompts, rounded down to whole --prefill\n" + " chunks (default 16384)\n" + " --cache-ram MIB --serve: host memory for the K/V cells of sequences a request replaced,\n" + " so switching back to them continues where they were (default 16384)\n" + " --feed-max N --serve: new prompt parts and last prompt chunks of up to N tokens run\n" + " through verify windows instead of the batched prompt path (default 128)\n" + " --spec T verify windows of up to T tokens: the last one and T-1 MTP drafts (2..8)\n" + " --spec-min-p P a draft enters the window only while it and the drafts before it have at\n" + " least probability P under the draft layer (default 0: always T-1 drafts)\n" + " --spec-lookup N when the text's last N+ tokens occurred before, the next window verifies\n" + " up to 7 of the tokens that followed there instead, if they are more than\n" + " --spec allows the draft layer (windows of up to 8 tokens; 0 = off)\n" + " --mtp DIR the MTP draft layer's runtime files (tools/mtp_rt.py)\n" + " --mtp-window N the draft layer attends to the last N cells (default 32768; 0 = every cell)\n" + " --main-gpu N run on CUDA device N of the visible ones, in nvidia-smi's order unless\n" + " CUDA_DEVICE_ORDER says otherwise (default 0)\n" + " --second-gpu N CUDA device N (of the visible ones, not --main-gpu) as another\n" + " expert tier: it holds the profile's next experts, then the conversation's,\n" + " and computes them beside the CPU pool (native packs, verify windows)\n" + " --second-gpu-gib G its VRAM for experts (default: all but --second-gpu-reserve-mib, 2048)\n" + " --second-gpu-min-mb M a layer's missed experts from which it takes its share (default 4; below\n" + " it the CPU is quicker than its ~100 us round trip)\n" + " --second-gpu-prefetch-mb M per layer, copy its likeliest experts that no GPU holds (the next\n" + " layer's router on this layer's input), up to M MiB, to it while the RAM is idle\n" + " (default 12: ~0.28 ms at its ~48 GB/s; 0 = off)\n" + " --head-split F with --second-gpu: its share of the output head's rows (default 0.45;\n" + " each GPU streams its part after the last layer; 0 = all on the main GPU)\n" + " --no-prompt-offload with --second-gpu: the prompt path streams the experts the main GPU's cache\n" + " lacks to the main GPU instead of computing them on the second\n" + " --spec-follow PATH benchmarks: emit this continuation (token ids) instead of the argmax and\n" + " accept the drafts that match it, so speculation settings compare on the\n" + " same text (not --serve)\n" " --no-token-graph A/B: two graphs per layer (the host launches each) instead of one per token\n" " --no-fused-gr A/B: the six-kernel hyper-connection read and a separate write (native)\n" - " --prefill CHUNK batched prompt processing in chunks of CHUNK tokens (needs --native); auto =\n" - " the largest chunk up to 8192 whose buffers the expert cache can lend\n" + " --prefill CHUNK batched prompt processing in chunks of CHUNK tokens (needs --native; auto: 16384)\n" " --no-pool skip the CPU expert pool (the GPU-only floor)\n" " --sync-every-layer debug: synchronise after every layer\n" - " --ple-gguf PATH the n-gram/PLE shard. WITHOUT IT LAYER 1's PLE IS SILENTLY SKIPPED,\n" - " which changes every number downstream - pass it for any real run\n" " --dump-mixed PATH write the post-attention residual (n_embd, f32)\n" " --stage-timing per-stage KERNEL-COUNT shares. NOT a time profile: an uncaptured\n" " event interval includes host gaps, so run with --gpu-only-full first\n" @@ -402,6 +393,16 @@ void usage() { " --gpu-only-full MEASURE: replay pre+post for all 48 layers plus the LM head, no pool.\n" " THE TRUE PER-TOKEN GPU FLOOR. Quote this one, not --graph-only.\n" " --stats print the per-stage breakdown\n" + " --window-hashes PATH per verify window, a line with its round, first position and size, a 64-bit hash\n" + " of its final residual rows and its argmaxes: two builds that do the same\n" + " arithmetic write the same file (fixed text through --spec-follow,\n" + " --adapt-every 0; not --serve)\n" + " --window-logits PATH per token the verify windows emit, a line with its index among the generated\n" + " tokens, the token and the 64 likeliest as id:log-probability, from the\n" + " head's logits (comparisons with llama.cpp, bench/parity.py; not --serve)\n" + " --window-profile with --stats: the GPU time of each stage of the verify window (timestamps\n" + " inside the graph; each costs ~2 us)\n" + " --prefill-profile the GPU time of each section of the batched prompt path (not --serve)\n" " --gpu-stages R0.9: capture the layer as three graphs (mixer / ffn+router / post)\n" " and time them from OUTSIDE the capture. The per-stage table on the\n" " real graph that --stage-timing cannot give. Prints and exits.\n" @@ -416,11 +417,6 @@ void usage() { " GPU via `moe_hit_grouped_s2`. DEFAULT 0. Measured at 4096 slots\n" " with --expert-cache-per-layer: 54.4%% hits, CPU pool drain 19.1 -> 10.3\n" " ms/token, -2.7 ms/token end to end.\n" - " --expert-cache-device1 N pre-fill N experts on CUDA1 (experimental)\n" - " --expert-cache-device2 N pre-fill N more experts on CUDA2\n" - " --expert-cache-device3 N pre-fill N more experts on CUDA3\n" - " --expert-cache-remote-placement stripe|layer distribute expert ranks or whole\n" - " layers across CUDA1..3 (default: stripe)\n" " --expert-cache-per-layer R4.2g: give each layer its OWN slots instead of letting the first\n" " position take all of them. The default policy fills in arrival order\n" " from one shared counter, so 256 slots went to ~26 layers of position 0\n" @@ -434,26 +430,7 @@ void usage() { " deviation from `cpu_s2` is attributed.\n" " --mmap-experts R2.1: opt OUT of the resident expert arena, back to MapViewOfFile.\n" " The A/B arm: the mmap's rate depends on the OS page cache holding\n" - " 34 GB, and measured 71.97 vs 34.78 ms/token cold vs warm.\n" - " --resident-cpu-experts with mmap and a static profile, keep CPU misses resident in ordinary RAM.\n" - " Borrowed GPU-cache entries may read from mmap during prompt prefill.\n"); -} - -/// All shards of a split GGUF, from shard 1's path ("...-00001-of-00002.gguf"); just the path when it is not split. -std::vector model_shards(const std::string& first) { - const std::string tag = "-00001-of-"; - const size_t at = first.rfind(tag); - if (at == std::string::npos || first.size() < at + tag.size() + 10) return {first}; - const int total = std::atoi(first.substr(at + tag.size(), 5).c_str()); - std::vector out; - for (int i = 1; i <= total && i <= 99; ++i) { - char num[8]; - std::snprintf(num, sizeof num, "%05d", i); - std::string p = first; - p.replace(at + 1, 5, num); - if (std::ifstream(p, std::ios::binary)) out.push_back(p); - } - return out.empty() ? std::vector{first} : out; + " 34 GB, and measured 71.97 vs 34.78 ms/token cold vs warm.\n"); } bool parse_i64_list(const char* s, std::vector& out, std::string& err) { @@ -486,6 +463,48 @@ struct Drive { /// needs no locking. `d.layers` is the CURRENT layer on entry (the adapter increments it as it walks the /// blob), which is why the layer index comes from there rather than from a counter of our own. std::FILE* routing = nullptr; + /// The first GPU's adaptive tier: its copies cross that GPU's link, and while they ran they slowed the window's + /// own transfers there (the router's rows and doorbell to the host, the CPU's rows to the combine) by 20-60 us a + /// layer. So they go out in pieces while the link has nothing else to carry: once a layer's CPU rows are in, as + /// much as lands before the next layer's router writes to the host (that layer's running average of the time from + /// the previous layer's CPU rows to its ring, by window size, less the time since and `kLinkReserveUs`). A + /// window's moves are staged before it (`drive_window`). Without a second GPU the pieces cannot carry what the + /// tier moves (the CPU takes every miss; pieces during the pool too slowed its RAM reads): `burst`, a window's + /// start sends all the queued moves, and an update waits for the previous one's. + strata::core::AdaptiveTier* tier1 = nullptr; + /// Where the host and the pool's workers run: checked before each window, once a second (CorePlacement::tick). + strata::kernels::cpu::CorePlacement* placement = nullptr; + bool burst = false; + std::vector ring_us; // by window size and layer + Clock::time_point rows_at{}; // this window's last CPU rows (none yet: the epoch) + int win_t = 0; // this window's tokens + int64_t n_layers = 0; + /// The second GPU's adaptive tier, paced: its copies go between windows, where the RAM is idle (the head, the + /// commit, the drafts); through the layers they slowed the CPU pool's RAM reads, and on that GPU's one copy engine + /// they held up the prefetch copies. An update is ranked at a window's tail, and its copies go out as far as they + /// end before the next window's first pool (`gap_ms`, a running average of that gap, at the tier's measured copy + /// rate), the rest at the last tail before the next update. + strata::core::AdaptiveTier* tier2 = nullptr; + double gap_ms = 0; + Clock::time_point tail_at{}, copies_end{}; // the last window's CPU rows done; the gap's queued copies end + bool gap_open = false; // no pool since the tail + bool adapt_next = false; // an update is ranked at this window's tail + bool flush_next = false; // the next window's tail ranks one: this gap takes the rest + std::string tier_err; + /// The adaptive tiers' update (ranking, the second tier's first copies) on a thread of its own, started at the + /// window's tail and joined before the next window. + std::function adapt; + std::thread adapt_thr; + bool adapt_ok = true; + void start_adapt() { + adapt_ok = true; + adapt_thr = std::thread([this] { adapt_ok = adapt(); }); + } + bool join_adapt() { + if (adapt_thr.joinable()) adapt_thr.join(); + return adapt_ok; + } + ~Drive() { join_adapt(); } }; void drive_pool(void* user, const float* x_f, const int32_t* ids, const float* weights, int64_t n_embd, int64_t k, @@ -517,166 +536,117 @@ void drive_pool(void* user, const float* x_f, const int32_t* ids, const float* w } } -/// Plan v0.3 P6: the pool for a verify window. -void drive_pool_multi(void* user, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k, float* out, - int64_t layer) { +/// The second GPU's prefetch from a layer's predicted routing (Verifier::set_predict). +void drive_predict(void* user, int64_t layer, const int32_t* ids, const float* w, int64_t n_tok, int64_t k) { Drive* t = (Drive*) user; - t->d.layers = layer; - const Clock::time_point a = Clock::now(); - strata::core::expert_pool_dispatch_multi(t->d, x_f, ids, n_tok, k, out); - t->cpu_ms += std::chrono::duration(Clock::now() - a).count(); - ++t->calls; - // the routing trace for the serve path: the same record format drive_pool writes (layer, k, ids, weights), - // one record per token. The multi dispatch fuses the router weights into the kernel and does not surface - // them, so records carry unit weights: tools/make_profile.py ranks pairs by routed frequency, which is the - // signal that matters; a one-shot --dump-routing run records true weights if a weighted ranking is wanted. - if (t->routing != nullptr && layer >= 0 && layer < 48) { - for (int64_t tok = 0; tok < n_tok; ++tok) { - const int32_t rec[2] = {(int32_t) layer, (int32_t) k}; - std::fwrite(rec, sizeof rec, 1, t->routing); - std::fwrite(ids + tok * k, sizeof(int32_t), (size_t) k, t->routing); - static const float one[64] = {}; // k <= 64 in a verify window; zeros read as unit weights - std::fwrite(one, sizeof(float), (size_t) k, t->routing); - } - } + strata::core::expert_prefetch_multi(t->d, layer, ids, w, n_tok, k); } -/// Layer split: every verify stage shares one Drive (its counters, usage and failure flags); the GPU plan, the expert -/// cache and the PCIe share the pool uses for a layer are those of the stage that runs it. -struct SplitDrive { - static constexpr int kMax = 8; - Drive* base = nullptr; - int n = 0; ///< stages - int64_t end[kMax] = {}; ///< stage i runs the layers from end[i - 1] (0) below end[i] - strata::core::GpuPlanSink* plan[kMax] = {}; - const uint8_t* cache_base[kMax] = {}; - const uint64_t* cache_slot_off[kMax] = {}; - int pcie_num[kMax] = {}; -}; -void drive_pool_split(void* user, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k, float* out, - int64_t layer) { - SplitDrive* s = (SplitDrive*) user; - int st = 0; - while (st + 1 < s->n && layer >= s->end[st]) ++st; - Drive& d = *s->base; - d.d.plan = s->plan[st]; - d.d.cache_base = s->cache_base[st]; - d.d.cache_slot_off = s->cache_slot_off[st]; - d.d.pcie_num = s->pcie_num[st]; - drive_pool_multi(s->base, x_f, ids, n_tok, k, out, layer); +// The first tier's piece ends this long before the next ring: the router (20-30 us), the doorbell, and the copy's +// start ~20 us after the call under WDDM (50 and 30 us were slower; pieces between windows too, slower again). +constexpr double kLinkReserveUs = 80; + +/// The first tier's copy rate in bytes per ms (a first guess until measured). +double tier1_rate(const Drive* t) { + const double r = t->tier1->copy_rate(); + return r > 0 ? r : 5e6; } -/// Layer split across GPUs: a later stage on its own device, with its own copy of the dense weights, a session, an -/// expert cache for its layers, a verify window and a prompt path; the last one also holds the head (the drafter -/// lives on its device too). -struct GpuStage { - int dev = 0; - int64_t lb = 0, le = 0; - double pcie_frac = 0.0; - strata::core::WeightTable wt; - strata::core::NativeDense dense; - strata::core::NativeHead head; - strata::core::SessionState ss; - cudaStream_t stream = nullptr; - strata::core::ExpertCache cache; - std::vector> profile; ///< its layers' share of the profile, hottest first - int32_t* d_res = nullptr; ///< the residency table on its device - strata::core::Verifier ver; - strata::prefill::Prefill sp; - cudaStream_t adapt_stream = nullptr; - cudaEvent_t adapt_ev = nullptr; - bool adapt_live = false; ///< swaps of this request are in flight on it - int32_t* mrope = nullptr; ///< --vision: the image-position table on its device -}; +/// Before a window of `T` tokens: the tiers' landed moves admitted, and the first tier's next moves staged for what +/// this window's layers take (their residents leave before the window's residency snapshot). +void drive_window(Drive* t, int T) { + t->win_t = T; + t->rows_at = Clock::time_point{}; + if (t->placement != nullptr && t->d.pool != nullptr) { + const std::string moved = t->placement->tick(*t->d.pool); + if (!moved.empty()) std::fprintf(stderr, "strata generate: %s\n", moved.c_str()); + } + if (t->tier2 != nullptr) t->tier2->apply_pending(false); + if (t->tier1 == nullptr) return; + t->tier1->apply_pending(false); + if (t->burst) { // all of them, at once + t->tier1->stage(UINT64_MAX); + if (!t->tier1->pump_bytes(UINT64_MAX, t->tier_err)) { + t->d.failed = true; + t->d.fail = t->tier_err.c_str(); + } + return; + } + double us = 0; + for (int64_t l = 1; l < t->n_layers; ++l) + us += std::max(0.0, t->ring_us[(size_t) (T * t->n_layers + l)] - kLinkReserveUs); + t->tier1->stage((uint64_t) (1.25 * tier1_rate(t) * us / 1e3)); +} -// ---- issue #31: what the watchdog prints before it stops a stalled engine -struct MemSample { - unsigned long long faults = 0, rss_mib = 0, avail_mib = 0, commit_mib = 0; // faults: hard (Linux) / all (Windows) -}; -MemSample mem_sample() { - MemSample m; -#if defined(_WIN32) - PROCESS_MEMORY_COUNTERS pmc{}; - if (K32GetProcessMemoryInfo(GetCurrentProcess(), &pmc, sizeof pmc)) { - m.faults = pmc.PageFaultCount; - m.rss_mib = pmc.WorkingSetSize >> 20; - m.commit_mib = pmc.PagefileUsage >> 20; +/// A layer's CPU rows are in (Verifier::set_rows_in): the first tier's piece that lands before the next router. +void drive_rows_in(void* user, int64_t layer) { + Drive* t = (Drive*) user; + if (t->tier1 == nullptr || t->burst || t->d.failed || layer + 1 >= t->n_layers) return; + const double us = t->ring_us[(size_t) (t->win_t * t->n_layers + layer + 1)] - kLinkReserveUs - + std::chrono::duration(Clock::now() - t->rows_at).count(); + if (us > 0 && !t->tier1->pump_bytes((uint64_t) (tier1_rate(t) * us / 1e3), t->tier_err)) { + t->d.failed = true; + t->d.fail = t->tier_err.c_str(); } - MEMORYSTATUSEX ms{}; - ms.dwLength = sizeof ms; - if (GlobalMemoryStatusEx(&ms)) m.avail_mib = ms.ullAvailPhys >> 20; -#else - if (std::FILE* f = std::fopen("/proc/self/stat", "r")) { - char buf[4096]; - const size_t n = std::fread(buf, 1, sizeof buf - 1, f); - buf[n] = 0; - std::fclose(f); - const char* s = std::strrchr(buf, ')'); // fields after the command name: 3 state ... 12 majflt - for (int field = 2; s && field < 12; ++field) s = std::strchr(s + 1, ' '); - if (s) m.faults = std::strtoull(s + 1, nullptr, 10); +} + +/// The second tier's queued moves that end before the next window's first pool, or all of them. +bool pump_gap(Drive* t, std::string& err, bool all = false) { + if (t->tier2 == nullptr) return true; + const Clock::time_point start = std::max(Clock::now(), t->copies_end); + const double left = t->gap_ms - std::chrono::duration(start - t->tail_at).count(); + const double rate = t->tier2->copy_rate(); + const bool paced = !all && rate > 0 && t->gap_ms > 0; + if (paced && left <= 0) return true; + uint64_t sent = 0; + if (!t->tier2->pump(paced ? (uint64_t) (left * rate) : UINT64_MAX, sent, err)) return false; + if (rate > 0) + t->copies_end = start + std::chrono::duration_cast( + std::chrono::duration((double) sent / rate)); + return true; +} + +/// Whether the second GPU's work so far succeeded (Verifier::set_watch). +bool drive_watch(void* user, std::string& err) { + Drive* t = (Drive*) user; + return t->d.gpu2 == nullptr || t->d.gpu2->healthy(err); +} + +/// A window's last CPU rows are in (Verifier::set_tail). +void drive_tail(void* user) { + Drive* t = (Drive*) user; + t->tail_at = Clock::now(); + t->copies_end = t->tail_at; + t->gap_open = true; + if (t->d.failed) return; + if (t->adapt_next) { + t->start_adapt(); + } else if (!pump_gap(t, t->tier_err, t->flush_next)) { + t->d.failed = true; + t->d.fail = t->tier_err.c_str(); } - auto kb = [](const char* path, const char* key) -> unsigned long long { - unsigned long long v = 0; - if (std::FILE* f = std::fopen(path, "r")) { - char line[256]; - const size_t kl = std::strlen(key); - while (std::fgets(line, sizeof line, f)) - if (std::strncmp(line, key, kl) == 0) { v = std::strtoull(line + kl, nullptr, 10); break; } - std::fclose(f); - } - return v; - }; - m.rss_mib = kb("/proc/self/status", "VmRSS:") >> 10; - m.commit_mib = kb("/proc/self/status", "VmSwap:") >> 10; - m.avail_mib = kb("/proc/meminfo", "MemAvailable:") >> 10; -#endif - return m; } -void stall_report(std::FILE* f, uint64_t layers_during) { - strata::core::Progress& p = strata::core::progress(); - std::fprintf(f, "strata serve: stall report (engine %s): stage \"%s %lld\" for %lld s; %llu layers served since the " - "last finished step (0 = stopped, more = slow)\n", STRATA_VERSION, p.where.load(), - (long long) p.detail.load(), (long long) ((strata::core::progress_now_ms() - p.since_ms.load()) / 1000), - (unsigned long long) layers_during); - for (int pass = 0; pass < 2; ++pass) { - if (pass == 1) { - std::this_thread::sleep_for(std::chrono::seconds(2)); - std::fprintf(f, " 2 s later:\n"); - } - if (auto fn = strata::core::diag_pool_fn().load()) fn(f); - if (auto fn = strata::core::diag_verify_fn().load()) fn(f); - const MemSample m = mem_sample(); - std::fprintf(f, " memory: %llu MiB resident, %llu MiB %s, %llu MiB RAM available; %llu %s\n", m.rss_mib, - m.commit_mib, -#if defined(_WIN32) - "committed", m.avail_mib, m.faults, "page faults so far" -#else - "in swap", m.avail_mib, m.faults, "major page faults so far" -#endif - ); - std::fflush(f); +/// Plan v0.3 P6: the pool for a verify window. +void drive_pool_multi(void* user, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k, float* out, + int64_t layer) { + Drive* t = (Drive*) user; + t->d.layers = layer; + const Clock::time_point a = Clock::now(); + const size_t at = (size_t) (t->win_t * t->n_layers + std::min(layer, t->n_layers - 1)); + if (t->rows_at != Clock::time_point{}) { // the last CPU rows to this ring + const double us = std::chrono::duration(a - t->rows_at).count(); + t->ring_us[at] = t->ring_us[at] > 0 ? 0.8 * t->ring_us[at] + 0.2 * us : us; } -#if defined(_WIN32) - // every thread's stack, to read against this build's PDB: a few MB beside the engine's working directory - if (HMODULE dbg = LoadLibraryA("dbghelp.dll")) { - using Fn = BOOL(WINAPI*)(HANDLE, DWORD, HANDLE, int, void*, void*, void*); - if (auto write = (Fn) GetProcAddress(dbg, "MiniDumpWriteDump")) { - char path[64]; - std::snprintf(path, sizeof path, "strata-stall-%lu.dmp", (unsigned long) GetCurrentProcessId()); - HANDLE h = CreateFileA(path, GENERIC_WRITE, 0, nullptr, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr); - if (h != INVALID_HANDLE_VALUE) { - const int kThreadInfo = 0x1000; // MiniDumpWithThreadInfo; MiniDumpNormal = 0 - const BOOL ok = write(GetCurrentProcess(), GetCurrentProcessId(), h, kThreadInfo, nullptr, nullptr, nullptr); - CloseHandle(h); - char full[MAX_PATH]; - if (!GetFullPathNameA(path, MAX_PATH, full, nullptr)) std::snprintf(full, sizeof full, "%s", path); - std::fprintf(f, " %s the thread stacks to %s (attach it to the issue)\n", ok ? "wrote" : "could not write", - full); - } - } + if (t->gap_open) { // the first pool after a window: the RAM's idle time between them + t->gap_open = false; + const double gap = std::chrono::duration(a - t->tail_at).count(); + if (gap < 20.0) t->gap_ms = t->gap_ms > 0 ? 0.8 * t->gap_ms + 0.2 * gap : gap; } -#endif + strata::core::expert_pool_dispatch_multi(t->d, x_f, ids, n_tok, k, out); + t->rows_at = Clock::now(); + t->cpu_ms += std::chrono::duration(t->rows_at - a).count(); + ++t->calls; } /// STRATA_TRACE=1: the VRAM left at a step of the startup (finds what fills the card after the cache is sized) @@ -695,191 +665,108 @@ int argmax(const std::vector& v) { return best; } -// ---- --serve's conversation cache. A chat or an agent sends the whole conversation again with every request, and -// reading it again is what made a long session wait minutes for every turn. What a sequence leaves behind splits in -// two, and only one half needs copying: -// * POSITIONAL state - the KV cache of the 12 QSA layers and their pooled indexer keys, the draft layer's KV. A -// cell is written once for its position and read only by later positions (the block scores take `dead` for the -// block being filled, never its pooled row), so rewinding to a position just means writing from there again. -// * RUNNING state - the 36 GDN recurrences and conv histories, each QSA layer's indexer tail (the unfinished -// block's raw keys) and the PLE's normalized history. Each describes "everything so far" and cannot be -// rewound, so a checkpoint is a copy of exactly these: ~118 MB, the same set the verifier snapshots to roll -// back rejected drafts. -// A checkpoint is only valid while the positional cells below it still hold ITS tokens, so the serve loop keeps -// just the checkpoints that are prefixes of the tokens the session holds now. -struct ImgKey { - int64_t start = 0; ///< the image's first <|image_pad|> position - uint64_t hash = 0; ///< its embeddings and grid: the pad tokens alone are the same for every picture - bool operator==(const ImgKey& o) const { return start == o.start && hash == o.hash; } -}; - -struct ConvCheckpoint { - std::vector ids; ///< the tokens this state has consumed - std::vector imgs; ///< the images among them - std::vector gdn, ple, tails; - uint64_t used = 0; ///< last-use stamp for the retention policy (conv_cache.hpp) - std::vector stage_parts; ///< a layer split's later stages: their sessions' running state -}; - -uint64_t fnv1a(const void* data, size_t n, uint64_t h = 1469598103934665603ull) { - const uint8_t* p = (const uint8_t*) data; - for (size_t i = 0; i < n; ++i) { h ^= p[i]; h *= 1099511628211ull; } - return h; -} - -struct ConvStateSizes { - size_t gdn = 0, ple = 0, tail = 0; -}; - -ConvStateSizes conv_state_sizes(const strata::core::ModelGeometry& g) { - ConvStateSizes z; - z.gdn = (size_t) g.n_gdn_layers() * - ((size_t) g.ssm_state_size * (size_t) g.ssm_v_heads * (size_t) g.ssm_state_size + - (size_t) g.ssm_conv_channels * (size_t) (g.ssm_d_conv - 1)) * sizeof(float); - z.ple = (size_t) strata::kernels::NG_HIST * (size_t) strata::kernels::NG_HC_DIM * sizeof(float); - z.tail = (size_t) (strata::kernels::qsa_real_shapes().idx_block - 1) * (size_t) g.idx_key_dim * sizeof(float); - return z; -} - -/// Copies the running state out. The caller has synchronized the device. -bool checkpoint_save(ConvCheckpoint& c, const strata::core::SessionState& ss, const strata::core::ModelGeometry& g) { - const ConvStateSizes z = conv_state_sizes(g); - c.gdn.resize(z.gdn); - c.ple.resize(ss.ple_hist != nullptr ? z.ple : 0); - c.tails.resize(z.tail * (size_t) g.n_qsa_layers()); - if (cudaMemcpy(c.gdn.data(), ss.gdn_state, z.gdn, cudaMemcpyDeviceToHost) != cudaSuccess) return false; - if (!c.ple.empty() && cudaMemcpy(c.ple.data(), ss.ple_hist, z.ple, cudaMemcpyDeviceToHost) != cudaSuccess) - return false; - for (int64_t i = 0; i < g.n_qsa_layers(); ++i) - if (cudaMemcpy(c.tails.data() + (size_t) i * z.tail, ss.qsa_states[i].idx_tail, z.tail, cudaMemcpyDeviceToHost) != - cudaSuccess) - return false; - return true; -} - -/// Puts a checkpoint's running state back; the positional cells below it are the caller's to guarantee. -bool checkpoint_restore(const ConvCheckpoint& c, strata::core::SessionState& ss, const strata::core::ModelGeometry& g) { - const ConvStateSizes z = conv_state_sizes(g); - if (c.gdn.size() != z.gdn || c.tails.size() != z.tail * (size_t) g.n_qsa_layers()) return false; - if (cudaMemcpy(ss.gdn_state, c.gdn.data(), z.gdn, cudaMemcpyHostToDevice) != cudaSuccess) return false; - if (!c.ple.empty() && cudaMemcpy(ss.ple_hist, c.ple.data(), z.ple, cudaMemcpyHostToDevice) != cudaSuccess) - return false; - for (int64_t i = 0; i < g.n_qsa_layers(); ++i) - if (cudaMemcpy(ss.qsa_states[i].idx_tail, c.tails.data() + (size_t) i * z.tail, z.tail, cudaMemcpyHostToDevice) != - cudaSuccess) - return false; - // the PLE's token window is the last two tokens, OLDEST FIRST, -1 where there is none (as session_zero leaves it) - const size_t L = c.ids.size(); - ss.ple_prev[0] = L >= 2 ? c.ids[L - 2] : -1; - ss.ple_prev[1] = L >= 1 ? c.ids[L - 1] : -1; - return cudaDeviceSynchronize() == cudaSuccess; -} - -// --control-vector-scaled: llama.cpp's `common_control_vector_load` (every file's `direction.` times its scale, -// summed; layer 0 has none) and `llama_adapter_cvec::apply` with the projection-mode patch (project: the unit -// direction and its norm as the scale), into the tables `cvec_upload` takes. `summary` is what INFO reports. -bool load_control_vectors(const Options& o, const strata::core::ModelGeometry& g, std::string& summary, std::string& err) { - const int64_t L = g.n_layers, N = g.n_embd; - std::vector data((size_t) (L * N), 0.0f); - std::vector have((size_t) L, false); - for (const auto& [path, scale] : o.cvec_files) { - try { - strata::GgufFile f(path); - const strata::MetaValue* arch = f.get("general.architecture"); - if (arch == nullptr || arch->s != "controlvector") { - err = path + ": not a control vector GGUF (general.architecture is not 'controlvector')"; - return false; - } - const strata::MetaValue* hint = f.get("controlvector.model_hint"); - if (hint != nullptr && hint->s != "qwen4exp") - std::fprintf(stderr, "strata generate: %s was made for '%s', not qwen4exp\n", path.c_str(), hint->s.c_str()); - int found = 0; - for (const strata::TensorInfo& t : f.tensors()) { - if (t.name.rfind("direction.", 0) != 0) continue; - const long l = std::strtol(t.name.c_str() + 10, nullptr, 10); - if (l < 1 || l >= L) continue; // layer 0 has no vector; past the model is ignored, as in llama.cpp - if (t.type != 0 || t.elements() != (uint64_t) N) { - err = path + ": " + t.name + " must be " + std::to_string((long long) N) + " f32"; - return false; - } - const float* src = reinterpret_cast(f.tensor_data(t)); - for (int64_t j = 0; j < N; ++j) data[(size_t) (l * N + j)] += scale * src[j]; - have[(size_t) l] = true; - ++found; +/// Expert-cache slots lent to the batched prompt path for its buffers: the cache's last slots, as many as hold the +/// buffers, while at least 128 stay. `lend` takes their experts out of the residency table (the prompt path +/// streams them instead), `give_back` copies them in again from the arena once the prompt is done. +struct SlotLoan { + strata::core::ExpertCache* cache = nullptr; + std::vector* res = nullptr; + int32_t* d_res = nullptr; // the table's device copy, or null + int device = 0; // the cache's GPU; `main_device` is current again after `give_back` + int main_device = 0; + int32_t first = -1; // the first lent slot, -1 when the cache cannot spare the bytes + std::vector> lent; // (residency index, slot) while lent + + uint64_t offset(int64_t slot) const { + return cache->slot_offsets() ? cache->slot_offsets()[slot] + : (uint64_t) slot * (uint64_t) (cache->bytes() / cache->slots()); + } + /// The last slots that hold `need` bytes; false when the cache cannot spare them. + bool plan(uint64_t need) { + int64_t k = 0; + while (k < cache->slots() && (uint64_t) cache->bytes() - offset(cache->slots() - k) < need) ++k; + first = k + 128 <= cache->slots() ? (int32_t) (cache->slots() - k) : -1; + return first >= 0; + } + void* base() const { return first >= 0 ? (void*) cache->device_slot(first) : nullptr; } + uint64_t bytes() const { return first >= 0 ? (uint64_t) cache->bytes() - offset(first) : 0; } + int64_t slots() const { return first >= 0 ? cache->slots() - first : 0; } + void lend() { + if (first < 0) return; + std::vector& r = *res; + for (size_t i = 0; i < r.size(); ++i) + if (r[i] >= first) { + lent.emplace_back((int32_t) i, r[i]); + r[i] = strata::core::kNotResident; + } + if (d_res != nullptr) cudaMemcpy(d_res, r.data(), r.size() * sizeof(int32_t), cudaMemcpyHostToDevice); + } + /// After the borrower's last use: its GPU work is waited for first. + bool give_back(strata::core::ExpertSource& src, int64_t n_expert, std::string& err) { + if (lent.empty()) return true; + cudaSetDevice(device); + cudaDeviceSynchronize(); + bool ok = true; + const auto& lay = strata::kernels::cpu::expert_layout(); + for (const auto& [i, slot] : lent) { + const int64_t l = i / n_expert; + const uint8_t* b = src.blob(l, i % n_expert); + if (b == nullptr || !cache->fill_slot_blocking(slot, b, err, (int64_t) lay.blob_bytes(l))) { + if (b == nullptr) err = "no arena blob for a lent slot"; + ok = false; + break; } - if (found == 0) { err = path + ": no direction. tensors"; return false; } - } catch (const std::exception& e) { - err = e.what(); - return false; - } - } - const int first = o.cvec_first <= 0 ? 1 : o.cvec_first; - const int last = (o.cvec_last <= 0 || o.cvec_last >= L) ? (int) L - 1 : o.cvec_last; - const int single = o.cvec_mode == 0 ? o.cvec_single : -1; - if (single >= 0 && (single >= L || !have[(size_t) single])) { - err = "--cvec-dir single:" + std::to_string(single) + ": the vector has no direction for that layer"; - return false; - } - std::vector dir((size_t) (L * N), 0.0f), s((size_t) L, 0.0f); - int steered = 0; - for (int64_t l = first; l <= last; ++l) { - const int64_t src = single >= 0 ? single : l; - if (!have[(size_t) src]) continue; - const float* d = data.data() + (size_t) (src * N); - if (o.cvec_mode == 0) { - double nrm = 0.0; - for (int64_t j = 0; j < N; ++j) nrm += (double) d[j] * d[j]; - nrm = std::sqrt(nrm); - if (nrm <= 0.0) continue; - s[(size_t) l] = (float) nrm; - for (int64_t j = 0; j < N; ++j) dir[(size_t) (l * N + j)] = (float) (d[j] / nrm); - } else { - s[(size_t) l] = 1.0f; - std::copy(d, d + N, dir.begin() + (size_t) (l * N)); + (*res)[(size_t) i] = slot; } - ++steered; + cudaSetDevice(main_device); + if (ok && d_res != nullptr) + cudaMemcpy(d_res, res->data(), res->size() * sizeof(int32_t), cudaMemcpyHostToDevice); + lent.clear(); + return ok; } - if (steered == 0) { err = "the control vector has no direction in layers " + std::to_string(first) + ".." + std::to_string(last); return false; } - if (!strata::kernels::cvec_upload(dir, s, o.cvec_mode, first, last, N, g.hc, err)) return false; - summary = std::string(o.cvec_mode == 0 ? "project" : "add") + ":" + std::to_string(first) + "-" + std::to_string(last) + - (single >= 0 ? ":single" + std::to_string(single) : ""); - // the line llama.cpp's patched build prints, so a log shows the same thing - std::fprintf(stderr, "strata generate: control vector mode = %s, dir = %s, layers %d..%d (%d steered)\n", - o.cvec_mode == 0 ? "project" : "add", single >= 0 ? "single" : "per-layer", first, last, steered); - return true; -} +}; -// The effective host->device bandwidth of the PCIe link: copies from pinned host memory, as the expert arena's -// reads are. The native default share (0.55) was measured on x16 links (~26-28 GB/s); a x8 card in a x8 slot -// carries about half of that. Returns < 0 when the probe cannot run (then the caller keeps the default). -double probe_pcie_h2d_gbps() { - constexpr size_t kBytes = 256ull << 20; - constexpr int kIters = 4; - void* h = nullptr; - void* d = nullptr; - cudaEvent_t ev0, ev1; - if (cudaMallocHost(&h, kBytes) != cudaSuccess) return -1.0; - if (cudaMalloc(&d, kBytes) != cudaSuccess || cudaEventCreate(&ev0) != cudaSuccess || - cudaEventCreate(&ev1) != cudaSuccess) { - if (d != nullptr) cudaFree(d); - cudaFreeHost(h); - return -1.0; +/// The batched prompt path's device buffers on each GPU: expert-cache slots lent for every prompt, as many as its +/// chunk needs, else allocated once for the longest chunk. +struct PromptBuffers { + strata::prefill::Prefill* prefill = nullptr; + strata::prefill::ExpertRunner* offload = nullptr; // the second GPU's share, or null + const strata::core::ModelGeometry* g = nullptr; + const strata::core::SessionState* ss = nullptr; + SlotLoan loan, loan2; // each GPU's cache (a null cache lends nothing) + bool lend = false, lend2 = false; + + /// After the prompt path and its second GPU's runner have been initialized. + bool init(int64_t max_chunk, bool borrow, std::string& err) { + lend = borrow && loan.cache != nullptr && loan.res != nullptr && !loan.res->empty() && + loan.plan(strata::prefill::Prefill::bytes_needed(*g, *ss, max_chunk, offload != nullptr)); + lend2 = offload != nullptr && borrow && loan2.cache != nullptr && loan2.res != nullptr && !loan2.res->empty() && + loan2.plan(strata::prefill::ExpertRunner::bytes_needed(max_chunk, g->n_expert, true, true)); + return (lend || prefill->bind(nullptr, 0, max_chunk, err)) && + (offload == nullptr || lend2 || offload->bind(nullptr, 0, max_chunk, err)); + } + /// Before a prompt whose chunks hold up to `chunk` tokens. + bool take(int64_t chunk, std::string& err) { + // sized before either loan, from the residency both GPUs' prefetch areas are sized for + const uint64_t need = lend ? prefill->bytes_for(chunk) : 0; + const uint64_t need2 = lend2 ? offload->bytes_for(chunk, loan.res ? loan.res->data() : nullptr) : 0; + if (lend) { + loan.plan(need); + loan.lend(); + if (!prefill->bind(loan.base(), loan.bytes(), chunk, err)) return false; + } + if (lend2) { + loan2.plan(need2); + loan2.lend(); + if (!offload->bind(loan2.base(), loan2.bytes(), chunk, err)) return false; + } + return true; + } + /// After it: the lent slots refilled. + bool give_back(strata::core::ExpertSource& src, std::string& err) { + return loan.give_back(src, g->n_expert, err) && loan2.give_back(src, g->n_expert, err); } - std::memset(h, 0, kBytes); // fault the pages in before timing - cudaMemcpyAsync(d, h, kBytes, cudaMemcpyHostToDevice); // warmup: context up, copy engine primed - cudaEventRecord(ev0); - for (int i = 0; i < kIters; ++i) cudaMemcpyAsync(d, h, kBytes, cudaMemcpyHostToDevice); - cudaEventRecord(ev1); - const bool ok = cudaEventSynchronize(ev1) == cudaSuccess; - float ms = 0.f; - const bool timed = ok && cudaEventElapsedTime(&ms, ev0, ev1) == cudaSuccess && ms > 0.01f; - const double bw = timed ? ((double) kIters * (double) kBytes / (ms * 1e-3)) / 1e9 : -1.0; - cudaEventDestroy(ev0); - cudaEventDestroy(ev1); - cudaFree(d); - cudaFreeHost(h); - return bw; -} +}; } // namespace @@ -888,14 +775,12 @@ int main(int argc, char** argv) { // pipe or a file is block-buffered, so a program that dies loses every line it had already printed - which // turns "it crashed at step 7" into "it crashed somewhere", and the difference is a debugging session. std::setvbuf(stdout, nullptr, _IONBF, 0); - // Load every CUDA kernel when the context is created, before the expert cache takes the free VRAM. With the - // default lazy loading, a kernel first used mid-prompt (MMQ for IQ3_XXS at 64K+ on a 12 GB card) found no VRAM - // left for its code and the engine ended ("out of memory: cudaFuncSetAttribute"). Costs ~30 MB of VRAM. - if (std::getenv("CUDA_MODULE_LOADING") == nullptr) { + // device ordinals in nvidia-smi's order, not CUDA's fastest-first guess (before the first CUDA call) + if (std::getenv("CUDA_DEVICE_ORDER") == nullptr) { #if defined(_WIN32) - _putenv_s("CUDA_MODULE_LOADING", "EAGER"); + _putenv_s("CUDA_DEVICE_ORDER", "PCI_BUS_ID"); #else - setenv("CUDA_MODULE_LOADING", "EAGER", 0); + setenv("CUDA_DEVICE_ORDER", "PCI_BUS_ID", 0); #endif } Options o; @@ -955,7 +840,6 @@ int main(int argc, char** argv) { else if (a == "--ple-delay-us") o.ple_delay_us = std::atof(next("--ple-delay-us")); else if (a == "--ple-sync-submit") o.ple_sync_submit = true; else if (a == "--kv") o.kv = next("--kv"); - else if (a == "--kv-resident") o.kv_resident = std::atoll(next("--kv-resident")); else if (a == "--stream-token") o.stream_token = true; else if (a == "--check-logits") o.check_logits = true; else if (a == "--gr-fp32-activations") o.gr_fp32_activations = true; @@ -988,75 +872,46 @@ int main(int argc, char** argv) { const std::string v = next("--expert-cache"); o.expert_cache = (v == "auto") ? -1 : std::atoi(v.c_str()); } - else if (a == "--expert-cache-device1") o.expert_cache_remote[0] = std::atoi(next("--expert-cache-device1")); - else if (a == "--expert-cache-device2") o.expert_cache_remote[1] = std::atoi(next("--expert-cache-device2")); - else if (a == "--expert-cache-device3") o.expert_cache_remote[2] = std::atoi(next("--expert-cache-device3")); - else if (a == "--expert-cache-remote-placement") - o.expert_cache_remote_placement = next("--expert-cache-remote-placement"); else if (a == "--vram-reserve-mib") o.vram_reserve_mib = std::atoi(next("--vram-reserve-mib")); - else if (a == "--prefill") { + else if (a == "--prefill") { // auto: 16K-token chunks, halved until their buffers fit the lendable slots const std::string v = next("--prefill"); - o.prefill_auto = v == "auto"; - o.prefill_chunk = o.prefill_auto ? 8192 : std::atoll(v.c_str()); + o.prefill_chunk = v == "auto" ? 16384 : std::atoll(v.c_str()); } else if (a == "--no-split-rows") o.no_split_rows = true; else if (a == "--no-prefill-borrow") o.no_prefill_borrow = true; + else if (a == "--no-prompt-offload") o.no_prompt_offload = true; else if (a == "--prefill-until") o.prefill_until = std::atoll(next("--prefill-until")); else if (a == "--dump-final-r") o.dump_final_r = next("--dump-final-r"); else if (a == "--spec") o.spec = std::atoi(next("--spec")); else if (a == "--spec-oracle") o.spec_oracle = next("--spec-oracle"); + else if (a == "--spec-follow") o.spec_follow = next("--spec-follow"); else if (a == "--spec-corrupt") o.spec_corrupt = std::atoi(next("--spec-corrupt")); else if (a == "--mtp") o.mtp = next("--mtp"); else if (a == "--mtp-window") o.mtp_window = std::atoll(next("--mtp-window")); + else if (a == "--main-gpu") o.main_gpu = std::atoi(next("--main-gpu")); + else if (a == "--second-gpu") o.second_gpu = std::atoi(next("--second-gpu")); + else if (a == "--second-gpu-gib") o.second_gpu_gib = std::atof(next("--second-gpu-gib")); + else if (a == "--second-gpu-reserve-mib") o.second_gpu_reserve_mib = std::atoi(next("--second-gpu-reserve-mib")); + else if (a == "--second-gpu-min-mb") o.second_gpu_min_mb = std::atof(next("--second-gpu-min-mb")); + else if (a == "--second-gpu-prefetch-mb") o.second_gpu_prefetch_mb = std::atof(next("--second-gpu-prefetch-mb")); + else if (a == "--head-split") o.head_split = std::atof(next("--head-split")); else if (a == "--pcie-frac") o.pcie_frac = std::atof(next("--pcie-frac")); else if (a == "--adapt-every") o.adapt_every = std::atoi(next("--adapt-every")); else if (a == "--spec-min-p") o.spec_min_p = std::atof(next("--spec-min-p")); + else if (a == "--spec-lookup") o.spec_lookup = std::max(0, std::atoi(next("--spec-lookup"))); else if (a == "--stop-eos") o.stop_eos = true; else if (a == "--spec-split") o.spec_split = true; - else if (a == "--layer-split") o.layer_split = next("--layer-split"); - else if (a == "--split-device") o.split_device = next("--split-device"); + else if (a == "--window-profile") o.window_profile = true; + else if (a == "--prefill-profile") o.prefill_profile = true; + else if (a == "--window-hashes") o.window_hashes = next("--window-hashes"); + else if (a == "--window-logits") o.window_logits = next("--window-logits"); else if (a == "--pcie-mode") o.pcie_mode = next("--pcie-mode"); else if (a == "--serve") o.serve = true; else if (a == "--vision") o.vision = true; else if (a == "--prompt-cache") o.prompt_cache = std::max(0, std::atoi(next("--prompt-cache"))); - else if (a == "--prompt-cache-every") o.prompt_cache_every = std::max(0LL, std::atoll(next("--prompt-cache-every"))); - else if (a == "--prompt-cache-root") o.prompt_cache_root = std::max(0LL, std::atoll(next("--prompt-cache-root"))); - else if (a == "--turn-token") o.turn_token = std::atoll(next("--turn-token")); - else if (a == "--short-read") o.short_read = std::max(0LL, std::atoll(next("--short-read"))); - else if (a == "--suffix-draft") o.suffix_draft = std::max(0, std::atoi(next("--suffix-draft"))); - else if (a == "--mtp-max-t") o.mtp_max_t = std::max(0, std::atoi(next("--mtp-max-t"))); - else if (a == "--control-vector") o.cvec_files.push_back({next("--control-vector"), 1.0f}); - else if (a == "--control-vector-scaled") { - // FILE:SCALE, comma-separated; the LAST colon splits, so a Windows path (C:\...) keeps its drive - std::stringstream list(next("--control-vector-scaled")); - std::string item; - while (std::getline(list, item, ',')) { - const size_t colon = item.rfind(':'); - char* end = nullptr; - const float sc = colon == std::string::npos ? 0.0f : std::strtof(item.c_str() + colon + 1, &end); - if (colon == std::string::npos || colon == 0 || end == item.c_str() + colon + 1 || *end != '\0') { - std::fprintf(stderr, "--control-vector-scaled: expected FILE:SCALE, got '%s'\n", item.c_str()); - return 2; - } - o.cvec_files.push_back({item.substr(0, colon), sc}); - } - } - else if (a == "--control-vector-layer-range") { - o.cvec_first = std::atoi(next("--control-vector-layer-range")); - o.cvec_last = std::atoi(next("--control-vector-layer-range")); - } - else if (a == "--cvec-mode") { - const std::string m = next("--cvec-mode"); - if (m == "project") o.cvec_mode = 0; - else if (m == "add") o.cvec_mode = 1; - else { std::fprintf(stderr, "--cvec-mode: add or project, got '%s'\n", m.c_str()); return 2; } - } - else if (a == "--cvec-dir") { - const std::string d = next("--cvec-dir"); - if (d == "per-layer") o.cvec_single = -1; - else if (d.rfind("single:", 0) == 0) o.cvec_single = std::atoi(d.c_str() + 7); - else { std::fprintf(stderr, "--cvec-dir: per-layer or single:L, got '%s'\n", d.c_str()); return 2; } - } + else if (a == "--cache-every") o.cache_every = std::max(1, std::atoll(next("--cache-every"))); + else if (a == "--cache-ram") o.cache_ram_mib = std::max(0, std::atoll(next("--cache-ram"))); + else if (a == "--feed-max") o.feed_max = std::max(0, std::atoll(next("--feed-max"))); else if (a == "--no-spec-split") o.spec_split = false; else if (a == "--eos-ids") { std::string e; @@ -1070,7 +925,6 @@ int main(int argc, char** argv) { else if (a == "--expert-profile") o.expert_profile = next("--expert-profile"); else if (a == "--gpu-stages") o.gpu_stages = true; else if (a == "--mmap-experts") o.mmap_experts = true; - else if (a == "--resident-cpu-experts") o.resident_cpu_experts = true; else if (a == "--stats") o.stats = true; else if (a == "--shared-late") o.shared_late = true; else if (a == "--keep-canonical") o.keep_canonical = true; @@ -1088,100 +942,6 @@ int main(int argc, char** argv) { return 2; } } - // Layer split (multi-GPU): the later stages run layers [K_i, K_i+1) on their own GPUs (--split-device, default - // the next visible ones); "auto" places the K from each GPU's free VRAM once the weights are in (below). Across - // GPUs, not yet: KV streaming, images, control vectors, the helper caches (--expert-cache-remote), and lending - // cache slots to the prompt path (each stage's prompt path has its own buffers). - const bool pcie_given = o.pcie_frac >= 0.0; - std::vector split_at; - std::vector split_devs; - bool split_auto = false, split_same = false; - if (!o.layer_split.empty()) { - int n_dev = 1; - if (cudaGetDeviceCount(&n_dev) != cudaSuccess || n_dev < 1) n_dev = 1; - cudaGetLastError(); - auto ints = [](const std::string& str, auto& out) -> bool { - using V = typename std::decay_t::value_type; - size_t a = 0; - while (a < str.size()) { - size_t b = str.find(',', a); - if (b == std::string::npos) b = str.size(); - const std::string t = str.substr(a, b - a); - if (t.empty() || t.find_first_not_of("0123456789") != std::string::npos) return false; - out.push_back((V) std::atoll(t.c_str())); - a = b + 1; - } - return !out.empty(); - }; - split_auto = o.layer_split == "auto"; - bool ok = o.serve && (split_auto || ints(o.layer_split, split_at)); - if (ok && !o.split_device.empty()) ok = ints(o.split_device, split_devs); - else if (ok) - for (int d = 1; d < n_dev && (split_auto || split_devs.size() < split_at.size()); ++d) split_devs.push_back(d); - if (ok && !split_auto && split_devs.empty() && split_at.size() == 1) split_devs.push_back(0); // one GPU - split_same = ok && split_devs.size() == 1 && split_devs[0] == 0 && !split_auto; - if (ok && split_auto && split_devs.empty()) { - std::fprintf(stderr, "strata generate: --layer-split auto: one GPU visible, so no split\n"); - o.layer_split.clear(); - split_auto = false; - } else if (ok) { - ok = (split_auto || split_at.size() == split_devs.size()) && split_devs.size() < (size_t) SplitDrive::kMax; - for (size_t i = 0; ok && i < split_at.size(); ++i) ok = split_at[i] >= 2 && (i == 0 || split_at[i] > split_at[i - 1]); - for (size_t i = 0; ok && !split_same && i < split_devs.size(); ++i) { - ok = split_devs[i] > 0 && split_devs[i] < n_dev; - for (size_t j = 0; ok && j < i; ++j) ok = split_devs[i] != split_devs[j]; - } - } - if (!ok) { - std::fprintf(stderr, "strata generate: --layer-split K[,K2..]|auto needs --serve, rising K from 2, and one " - "distinct GPU per K in --split-device (1..%d; or 0 with one K: the same GPU)\n", n_dev - 1); - return 2; - } - } - const bool multi_gpu = !split_devs.empty() && !split_same; - if (o.resident_cpu_experts && (!o.mmap_experts || o.expert_profile.empty() || o.adapt_every != 0)) { - std::fprintf(stderr, "strata generate: --resident-cpu-experts requires --mmap-experts, a static --expert-profile and --adapt-every 0\n"); - return 2; - } - if (o.resident_cpu_experts && - (!o.layer_split.empty() || o.expert_cache_remote[0] > 0 || o.expert_cache_remote[1] > 0 || - o.expert_cache_remote[2] > 0)) { - std::fprintf(stderr, "strata generate: --resident-cpu-experts does not support layer splits or remote expert caches\n"); - return 2; - } - // the helper-GPU expert caches (--expert-cache-remote, docs/SECOND_GPU.md): CUDA1..3 on one GPU; with a layer - // split, the visible GPUs no stage runs on, in order - int remote_dev[3] = {1, 2, 3}; - if (multi_gpu) { - if (o.expert_profile.empty()) { - std::fprintf(stderr, "strata generate: a layer split across GPUs needs --expert-profile\n"); - return 2; - } - o.no_prefill_borrow = true; // each stage's prompt path has its own buffers - int n_vis = 1; - if (cudaGetDeviceCount(&n_vis) != cudaSuccess || n_vis < 1) n_vis = 1; - cudaGetLastError(); - int next_free = 1; - for (int r = 0; r < 3; ++r) { - if (o.expert_cache_remote[(size_t) r] <= 0) continue; - while (next_free < n_vis && - std::find(split_devs.begin(), split_devs.end(), next_free) != split_devs.end()) ++next_free; - if (next_free >= n_vis) { - std::fprintf(stderr, "strata generate: --expert-cache-remote with a layer split needs a GPU that runs no " - "stage (%d visible, %zu used by the split)\n", n_vis, split_devs.size() + 1); - return 2; - } - remote_dev[r] = next_free++; - } - std::string devs; - for (const int d : split_devs) devs += (devs.empty() ? "" : ",") + std::to_string(d); - std::fprintf(stderr, "strata generate: layer split across %zu GPUs: CUDA0, then CUDA%s (split %s)\n", - split_devs.size() + 1, devs.c_str(), o.layer_split.c_str()); - } - if (o.prefill_auto && (o.no_prefill_borrow || o.expert_profile.empty())) { - o.prefill_auto = false; // nothing to lend from: the buffers are reserved for the session, so keep them small - o.prefill_chunk = 2048; - } if (!have_tokens && o.serve) { // plan v0.3 P8: requests bring their own tokens o.tokens = {248045}; o.max_new = 1; @@ -1195,39 +955,45 @@ int main(int argc, char** argv) { return 2; } - if ((o.ple_io != "direct" && o.ple_io != "mmap") || o.ple_row_cache < 0 || o.ple_inflight < 1 || + if ((o.ple_io != "direct" && o.ple_io != "mmap" && o.ple_io != "ram") || o.ple_row_cache < 0 || o.ple_inflight < 1 || o.ple_inflight > 1024 || !(o.ple_delay_us >= 0)) { std::fprintf(stderr, "strata generate: invalid --ple-io/--ple-row-cache/--ple-inflight/--ple-delay-us\n"); return 2; } - if (o.kv == "q4") o.kv = "q4_0"; - if (o.kv != "fp16" && o.kv != "int8" && o.kv != "q4_0" && o.kv != "k8v4") { - std::fprintf(stderr, "strata generate: --kv must be fp16, int8, q4_0 or k8v4\n"); + if (o.kv != "fp16" && o.kv != "int8") { + std::fprintf(stderr, "strata generate: --kv must be fp16 or int8\n"); return 2; } - strata::core::qsa_set_kv_int8(o.kv == "int8"); - strata::core::qsa_set_kv_q4(o.kv == "q4_0"); // PR #21: 4-bit codes after a Hadamard rotation (kv_q4.hpp) - strata::core::qsa_set_kv_hybrid(o.kv == "k8v4"); // K8V4: INT8 K + rotated Q4_0 V, 816 B/cell - if (o.kv_resident < 0) { - std::fprintf(stderr, "strata generate: --kv-resident must be >= 0\n"); + // every CUDA call from here on uses this device; the adaptive tier sets it on its own thread too + int n_dev = 0; + if (cudaGetDeviceCount(&n_dev) != cudaSuccess || o.main_gpu < 0 || o.main_gpu >= n_dev || + cudaSetDevice(o.main_gpu) != cudaSuccess) { + std::fprintf(stderr, "strata generate: --main-gpu %d is not a visible CUDA device (%d visible)\n", o.main_gpu, + n_dev); return 2; } - if (o.kv == "k8v4" && o.kv_resident > 0) { - std::fprintf(stderr, "strata generate: --kv k8v4 does not support --kv-resident streaming (yet)\n"); + if (o.second_gpu >= 0 && (o.second_gpu == o.main_gpu || o.second_gpu >= n_dev)) { + std::fprintf(stderr, "strata generate: --second-gpu %d is not another visible CUDA device\n", o.second_gpu); return 2; } - strata::core::qsa_set_kv_resident(o.kv_resident); - // Prompt lookup (the suffix drafter, on by default): the MTP keeps its --spec windows and a lookup window may be - // up to 2 tokens longer; the draft policy (strata/spec/draft_policy.hpp) takes one only where it pays. Code - // edits +6-11%, ordinary text unchanged (bench/results/2026-09-27-spec). --suffix-draft 0 turns it off. - if (o.suffix_draft > 0 && o.spec >= 2 && o.mtp_max_t == 0) { - o.mtp_max_t = o.spec; - o.spec = std::min(o.spec + 2, 8); // kVerifyMaxT - } + strata::core::qsa_set_kv_int8(o.kv == "int8"); strata::core::layer_set_shared_early(!o.shared_late); + std::vector native_shards; // the --native model's GGUF shards if (!o.native_preset.empty()) { - if (o.no_ple || o.ple_gguf.empty()) { - std::fprintf(stderr, "strata generate: --native requires --ple-gguf (the PLE key is native too)\n"); + if (o.no_ple) { + std::fprintf(stderr, "strata generate: --native requires the PLE (the PLE key is native too)\n"); + return 2; + } + try { + native_shards = strata::gguf_split_paths(o.native_preset); + } catch (const std::exception& e) { + std::fprintf(stderr, "strata generate: --native: %s\n", e.what()); + return 2; + } + if (o.ple_gguf.empty()) o.ple_gguf = strata::kernels::ple_table_shard(native_shards); + if (o.ple_gguf.empty()) { + std::fprintf(stderr, "strata generate: no shard of the --native model holds per_layer_token_embd; " + "name the file with --ple-gguf\n"); return 2; } o.stream_token = true; @@ -1237,10 +1003,8 @@ int main(int argc, char** argv) { o.native_qsa = o.native_qsa_indexer = o.native_rope = o.native_ple_postops = true; if (o.native_head_gguf.empty()) o.native_head_gguf = o.native_preset; if (o.native_dense_gguf.empty()) { - // every shard of the model (-0000N-of-0000M.gguf beside --native), then the PLE shard: a split - // may put any layer in any shard (Swift's GGUFs: layers 13-47 in shard 2, the PLE table in shard 1) - o.native_dense_gguf = model_shards(o.native_preset); - if (std::find(o.native_dense_gguf.begin(), o.native_dense_gguf.end(), o.ple_gguf) == o.native_dense_gguf.end()) + o.native_dense_gguf = native_shards; + if (std::find(native_shards.begin(), native_shards.end(), o.ple_gguf) == native_shards.end()) o.native_dense_gguf.push_back(o.ple_gguf); } // Plan v0.3 (24 Sep): the CPU experts stay on the VNNI kernel. The llama.cpp-CPU-exact q8_0 contract @@ -1273,18 +1037,10 @@ int main(int argc, char** argv) { return 2; } if (!std::isfinite(o.temperature) || o.temperature < 0 || !std::isfinite(o.top_p) || - o.top_p <= 0 || o.top_p > 1 || o.top_k < 0 || o.expert_cache < -1 || - o.pool_workers < 0 || std::any_of(o.expert_cache_remote.begin(), o.expert_cache_remote.end(), - [](int slots) { return slots < 0; }) || - (o.expert_cache_remote[1] > 0 && o.expert_cache_remote[0] == 0) || - (o.expert_cache_remote[2] > 0 && o.expert_cache_remote[1] == 0)) { + o.top_p <= 0 || o.top_p > 1 || o.top_k < 0 || o.expert_cache < -1 || o.pool_workers < 0) { std::fprintf(stderr, "strata generate: invalid sampling or resource parameter\n"); return 2; } - if (o.expert_cache_remote_placement != "stripe" && o.expert_cache_remote_placement != "layer") { - std::fprintf(stderr, "strata generate: --expert-cache-remote-placement must be stripe or layer\n"); - return 2; - } if (o.native_flash_attn_short && o.max_context > 256) { std::fprintf(stderr, "strata generate: --native-flash-attn-short requires --max-context <=256\n"); @@ -1324,44 +1080,21 @@ int main(int argc, char** argv) { } } const bool native_pack = strata::kernels::cpu::expert_layout().native; - if (o.expert_cache_remote[0] > 0) { - // Keep CUDA1's proven startup order: initialise its context before - // allocating GPU0 weights or mapping the large host expert arena. - double free_gib = 0; - if (!strata::core::RemoteExperts::preflight(remote_dev[0], free_gib, err)) { - std::fprintf(stderr, "strata generate: %s\n", err.c_str()); - return 1; - } - std::fprintf(stderr, "strata generate: CUDA%d context ready, %.2f GiB free before expert arena registration\n", - remote_dev[0], free_gib); - } - // plan v0.3 P6: the PCIe share of the missed experts, measured per kind of pack (the paper, finding on PCIe). - // PR #44: a x8 link carries half of what the native default assumes - the GPU's SMs read that share over the - // link (the copy kernel, since 0.1.14), so on a slower link it must shrink or the window waits for it. The - // real H2D bandwidth is probed once; from 20 GB/s up (x16 PCIe 4/5) the measured default stays. The canonical - // pack's 0.2 was never measured against the link, so it is left alone. `--calibrate` measures it outright. - if (o.pcie_frac < 0.0) { - const double base = native_pack ? 0.55 : 0.2; - const double bw = native_pack ? probe_pcie_h2d_gbps() : -1.0; - if (!native_pack) { - o.pcie_frac = base; - } else if (bw > 0.0) { - // below ~4 GB/s (an x1 link: ~0.9 GB/s) a missed expert's 1.4 MB takes longer to cross than the CPU - // pool takes to compute it, so none of them go over the link - o.pcie_frac = bw >= 20.0 ? base : bw < 4.0 ? 0.0 : std::min(base, std::max(0.05, base * (bw / 26.0))); - std::fprintf(stderr, "strata generate: PCIe probe: %.1f GB/s host->device -> pcie_frac %.2f (default %.2f)\n", - bw, o.pcie_frac, base); - } else { - o.pcie_frac = base; - std::fprintf(stderr, "strata generate: PCIe probe failed -> pcie_frac default %.2f\n", base); + if (native_pack) + for (int64_t l = 0; l < strata::kernels::cpu::expert_layout().n_layers; ++l) { + const auto& f = strata::kernels::cpu::expert_layout().fmt[(size_t) l]; + if (!strata::kernels::native_expert_supported(f.gu_type, f.d_type, f.n_embd, f.n_ff)) { + std::fprintf(stderr, "strata generate: layer %lld's experts (GGML types %d / %d) have no GPU kernels\n", + (long long) l, f.gu_type, f.d_type); + return 1; + } } - } + // plan v0.3 P6: the PCIe share of the missed experts, measured per kind of pack (the paper, finding on PCIe) + if (o.pcie_frac < 0.0) o.pcie_frac = native_pack ? 0.55 : 0.2; // the canonical Q2_0 pack's CPU kernels are AVX-512 only; a native pack runs on AVX2 CPUs as well if (!native_pack) strata::kernels::cpu::cpu_require_expert_support(); else if (!strata::kernels::cpu::cpu_avx512_ok()) - std::fprintf(stderr, "strata generate: this CPU has no AVX-512: the expert kernels run on %s " - "(multi-token for the i-quant gate/up rows)\n", - std::getenv("STRATA_NO_IQ256") == nullptr ? "AVX-2" : "ggml-cpu vec_dot (STRATA_NO_IQ256 set)"); + std::fprintf(stderr, "strata generate: this CPU has no AVX-512: the expert kernels run on AVX2\n"); strata::core::NativeEmbed native_embed; if (native_pack) { if (o.native_preset.empty() || o.spec < 2 || o.keep_canonical || @@ -1371,7 +1104,7 @@ int main(int argc, char** argv) { return 2; } const strata::core::ModelGeometry g0; - if (!native_embed.load(o.native_preset, g0.n_embd, 248320, err)) { + if (!native_embed.load(native_shards, g0.n_embd, 248320, err)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } @@ -1463,30 +1196,8 @@ int main(int argc, char** argv) { strata::kernels::mrope_table_set(d_mrope); } strata::kernels::ple_set_native_postops(o.native_ple_postops); - strata::core::ModelGeometry g; // canonical defaults; the model file overrides the MoE shape below - int64_t K = 10; - if (!o.native_preset.empty()) { - // a pruned variant (GSQ-RCO Coder) ships fewer experts than the canonical 512x10; the model file - // is the authority on its own MoE shape - everything else in the geometry is unchanged - try { - strata::GgufFile model_gguf(o.native_preset); - if (const strata::MetaValue* v = model_gguf.get("qwen4exp.expert_count")) g.n_expert = (int64_t) v->u; - if (const strata::MetaValue* v = model_gguf.get("qwen4exp.expert_used_count")) K = (int64_t) v->u; - } catch (const std::exception& e) { - std::fprintf(stderr, "strata generate: reading the model's expert shape from %s: %s\n", - o.native_preset.c_str(), e.what()); - return 1; - } - } - // before session_init: every graph captured from here on has the vector's kernels where it applies - std::string cvec_summary = "0"; - if (!o.cvec_files.empty()) { - std::string ce; - if (!load_control_vectors(o, g, cvec_summary, ce)) { - std::fprintf(stderr, "strata generate: control vector: %s\n", ce.c_str()); - return 2; - } - } + const strata::core::ModelGeometry g; + const int64_t K = 10; if (o.max_context < (int64_t) o.tokens.size() + o.max_new) { std::fprintf(stderr, "strata generate: --max-context %lld cannot hold %zu prompt + %lld new tokens\n", (long long) o.max_context, o.tokens.size(), (long long) o.max_new); @@ -1516,10 +1227,6 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: session_init failed\n"); return 1; } - if (g.n_qsa_layers() > 0 && ss.qsa_states[0].kv_mode == 1) - std::fprintf(stderr, "strata generate: KV streaming: %lld of %lld cells per QSA layer in VRAM, the K/V in " - "%.2f GiB of pinned RAM\n", (long long) (ss.qsa_states[0].n_slots * 4), - (long long) o.max_context, (double) strata::core::qsa_kv_host_bytes() / 1073741824.0); // ---- **THE HALF-LEVEL DUMP HAS TO BE ARMED BEFORE `session_capture`, AND THE LADDER MUST NOT BE.** The // half copies are issued from inside `block_layer_pre`/`block_layer_post`, so they are only ever enqueued @@ -1564,7 +1271,9 @@ int main(int argc, char** argv) { float* ple_scratch = nullptr; if (!o.ple_gguf.empty()) { strata::kernels::PleIoOptions pio; - pio.mode = o.ple_io == "mmap" ? strata::kernels::PleIo::Mmap : strata::kernels::PleIo::Direct; + pio.mode = o.ple_io == "mmap" ? strata::kernels::PleIo::Mmap + : o.ple_io == "ram" ? strata::kernels::PleIo::Ram + : strata::kernels::PleIo::Direct; pio.max_inflight = (uint32_t) o.ple_inflight; pio.cache_rows = (uint64_t) o.ple_row_cache; pio.io_thread = !o.ple_sync_submit; @@ -1593,7 +1302,8 @@ int main(int argc, char** argv) { ss.ple.w.key_scales = (const float*) ((const uint8_t*) wk->data + wk->codes_bytes); } if (o.native_ple_key && wk->quantized()) { - if (!wk->native_data || (wk->native_type != 42 && wk->native_type != 18 && wk->native_type != 23) || !wk->native_q8_1) { + // Q2_0 in the Q2_0 file, Q8_0 in UD-Q4_K_XL + if (!wk->native_data || !strata::kernels::native_mmvq_supported(wk->native_type) || !wk->native_q8_1) { std::fprintf(stderr, "strata generate: native PLE key is absent or incompatible\n"); return 1; } @@ -1639,107 +1349,6 @@ int main(int argc, char** argv) { return 1; } - // ---- layer split across GPUs: each later stage's own copy of the dense weights, its session and (the last) the - // head, made on its device before the host arena is mapped (as the drafter below, for the same WDDM reason) - std::vector> stages; - for (size_t i = 0; multi_gpu && i < split_devs.size(); ++i) { - stages.push_back(std::make_unique()); - GpuStage& st = *stages.back(); - st.dev = split_devs[i]; - const bool last = i + 1 == split_devs.size(); - double free_gib = 0; - if (!strata::core::RemoteExperts::preflight(st.dev, free_gib, err)) { - std::fprintf(stderr, "strata generate: %s\n", err.c_str()); - return 1; - } - const strata::core::OnDevice on(st.dev); - void* arena_s = nullptr; - if (cudaMalloc(&arena_s, pool_bytes) != cudaSuccess || - !st.wt.load(o.pack, arena_s, pool_bytes, err, skip.empty() ? nullptr : &skip)) { - std::fprintf(stderr, "strata generate: layer split, CUDA%d weights: %s\n", st.dev, - err.empty() ? "the weight arena does not fit" : err.c_str()); - return 1; - } - if (!o.native_dense_gguf.empty() && !st.dense.load(o.native_dense_gguf, st.wt, err, o.native_ple_key)) { - std::fprintf(stderr, "strata generate: layer split, CUDA%d native dense projections: %s\n", st.dev, - err.c_str()); - return 1; - } - void* sbuf_s = nullptr; - if (cudaMalloc(&sbuf_s, strata::core::session_bytes(g, o.max_context, K)) != cudaSuccess || - strata::core::session_init(g, o.max_context, K, sbuf_s, st.ss) == 0 || - cudaStreamCreateWithFlags(&st.stream, cudaStreamNonBlocking) != cudaSuccess || - cudaStreamCreateWithFlags(&st.adapt_stream, cudaStreamNonBlocking) != cudaSuccess || - cudaEventCreateWithFlags(&st.adapt_ev, cudaEventDisableTiming) != cudaSuccess) { - std::fprintf(stderr, "strata generate: layer split, CUDA%d: the session state failed\n", st.dev); - return 1; - } - const strata::core::WeightRef* wo_s = st.wt.find("output.weight"); - if (wo_s == nullptr || - (last && !o.native_head_gguf.empty() && !st.head.load(o.native_head_gguf, g.n_embd, wo_s->ne1, err))) { - std::fprintf(stderr, "strata generate: layer split, CUDA%d head: %s\n", st.dev, - wo_s == nullptr ? "output.weight is missing" : err.c_str()); - return 1; - } - // a control vector (the speed projection): its tables on this device too - the stage's layers apply it here - if (!strata::kernels::cvec_replicate(err)) { - std::fprintf(stderr, "strata generate: layer split, CUDA%d: %s\n", st.dev, err.c_str()); - return 1; - } - // --vision: this device's image-position table (the identity until a picture request), read by every rope - // kernel its stage runs - set before any of its graphs is captured - if (o.vision) { - if (cudaMalloc(&st.mrope, mrope_host.size() * sizeof(int32_t)) != cudaSuccess || - cudaMemcpy(st.mrope, mrope_host.data(), mrope_host.size() * sizeof(int32_t), cudaMemcpyHostToDevice) != - cudaSuccess) { - std::fprintf(stderr, "strata generate: layer split, CUDA%d: the image position table failed\n", st.dev); - return 1; - } - strata::kernels::mrope_table_set(st.mrope); - } - // its own PCIe share of the missed experts (the same rule as CUDA0's above: its link is probed) - st.pcie_frac = o.pcie_frac; - if (!pcie_given && native_pack) { - const double bw = probe_pcie_h2d_gbps(); - if (bw > 0.0) st.pcie_frac = bw >= 20.0 ? 0.55 : bw < 4.0 ? 0.0 : std::min(0.55, std::max(0.05, 0.55 * (bw / 26.0))); - std::fprintf(stderr, "strata generate: layer split: CUDA%d PCIe probe %.1f GB/s -> pcie_frac %.2f\n", st.dev, - bw, st.pcie_frac); - } - size_t fb = 0, tb = 0; - cudaMemGetInfo(&fb, &tb); - std::fprintf(stderr, "strata generate: layer split: CUDA%d holds its weights, session%s; %.2f GiB free\n", - st.dev, last ? " and head" : "", (double) fb / 1073741824.0); - } - GpuStage* const last_st = stages.empty() ? nullptr : stages.back().get(); - - // Secure MTP's CUDA0 allocations before the large host arena is registered with both CUDA contexts. - // In particular WDDM can refuse the draft weights after mapping tens of GiB of host pages. - strata::core::MtpDrafter mtp; - if (!o.mtp.empty()) { - if (o.spec < 2) { - std::fprintf(stderr, "strata generate: --mtp is ignored without --spec T (T >= 2)\n"); - o.mtp.clear(); - } - if (!o.mtp.empty()) mtp.set_prompt_len((int64_t) o.tokens.size()); - // the draft layer is the canonical model's MTP head (512 experts) even when the target is pruned, - // so it always sees the canonical geometry; `static` because MtpDrafter keeps a reference - static const strata::core::ModelGeometry draft_geometry{}; - // with a layer split across GPUs the drafter reads the last stage's residual: it lives on that device - const strata::core::OnDevice on_mtp(last_st ? last_st->dev : -1); - if (!o.mtp.empty() && !mtp.load(o.mtp, draft_geometry, last_st ? last_st->ss : ss, o.spec, err, o.mtp_window)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } - } - // Create the additional contexts after MTP has secured CUDA0 memory, but - // before the host arena maps its expert pages into their address spaces. - for (int r = 1; r < 3; ++r) if (o.expert_cache_remote[(size_t) r] > 0) { - double free_gib = 0; - if (!strata::core::RemoteExperts::preflight(remote_dev[r], free_gib, err)) { - std::fprintf(stderr, "strata generate: %s\n", err.c_str()); - return 1; - } - std::fprintf(stderr, "strata generate: CUDA%d context ready, %.2f GiB free before expert arena registration\n", - remote_dev[r], free_gib); - } - // ---- the CPU expert pool // // R2.1: the experts are loaded into a RESIDENT ARENA by default. The mmap path is kept behind @@ -1764,16 +1373,6 @@ int main(int argc, char** argv) { strata::core::ArenaExpertSource arena_src; strata::core::ExpertSource* srcp = nullptr; if (o.mmap_experts) { - // FileExpertSource maps the pack's experts.bin: a canonical pack has it; a native (IQ) pack has it when - // built with `tools/iq_pack.py --experts-bin` (the per-layer blob sizes of its layout, PR #121). The low-RAM - // mode: the experts come from the file through the OS cache instead of a pinned copy in RAM, for a PC whose - // GPU holds most of them but whose RAM cannot hold them all. - if (native_pack && !std::filesystem::exists(std::filesystem::path(o.pack) / "experts.bin")) { - std::fprintf(stderr, "strata generate: --mmap-experts needs the pack's experts.bin; %s is a native (IQ) pack " - "built without it: python tools/iq_pack.py --gguf --out %s --experts-bin\n", - o.pack.c_str(), o.pack.c_str()); - return 2; - } if (!src.open(o.pack, g.n_layers, g.n_expert, err)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; @@ -1781,13 +1380,8 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: experts via mmap (--mmap-experts; the A/B arm of R2.1)\n"); srcp = &src; } else { - arena_src.set_gguf(o.native_preset); // plan v0.3 P6: a native pack may take its experts from shard 1 - // On the multi-GPU Windows experiment, start with at most 8 GiB of mapped host pages. - // Unregistered layers remain in the resident arena and use the CPU expert path. - // (a layer split across GPUs too: pinning all of it into two contexts leaves WDDM refusing every later - // allocation - measured on the 5080 + 3090 rig: cudaMemGetInfo and the next cudaMalloc fail) - const uint64_t pin_limit = (o.expert_cache_remote[0] > 0 || multi_gpu) ? (8ull << 30) : 0; - if (!arena_src.open(o.pack, g.n_layers, g.n_expert, /*threads=*/6, err, pin_limit)) { + arena_src.set_gguf(native_shards); // plan v0.3 P6: a native pack may take its experts from the GGUF + if (!arena_src.open(o.pack, g.n_layers, g.n_expert, /*threads=*/6, err)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } @@ -1795,194 +1389,91 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: loaded %.2f GiB at %.2f GiB/s\n", (double) strata::kernels::cpu::expert_layout().total / (1024.0 * 1024 * 1024), arena_src.load_gib_per_second()); - // A rate under ~0.2 GiB/s is not the hardware. Task Scheduler / service contexts throttle this - // read+fill about 24x (measured 0.05 vs 1.42 GiB/s for the same binary, args and cache state; the - // scheduler's defaults - Below normal priority and a least-privilege token - were the only - // difference between the runs). Say so instead of letting the user blame the disk; see - // docs/DETAILS.md, "Running it at startup (Task Scheduler)". -#ifdef _WIN32 // a Windows launch context; elsewhere a load this slow is the disk - if (arena_src.load_gib_per_second() > 0.0 && arena_src.load_gib_per_second() < 0.2) { - std::fprintf(stderr, - "strata generate: hint: ~24x below what this hardware streams from a normal " - "launch. If Strata is started by Task Scheduler or a service, register the task " - "with Priority 4 (Normal) and 'Run with highest privileges' - the scheduler's " - "defaults (Below normal + a least-privilege token) throttle the load. See " - "docs/DETAILS.md ('Running it at startup').\n"); - } -#endif srcp = &arena_src; } - strata::kernels::cpu::ExpertPool pool(o.pool_workers, /*pin=*/true, /*host_works=*/!o.no_host_worker); - if (o.no_ple_prefetch) strata::kernels::ple_prefetch_enable(false); - // ---- R4's slot storage. Allocated AFTER the weights and the session, so `cudaMemGetInfo` inside `open` - // sees the memory this process actually has left rather than the card's idle figure - and refuses with both - // numbers if the slots do not fit, instead of handing back a cache smaller than it was asked for. - mem_mark("the weights, the session and the drafter"); - strata::core::ExpertCache xcache; - std::vector> profile; - if (!o.expert_profile.empty()) { - int64_t pslots = 0; - if (!strata::core::read_expert_profile(o.expert_profile, g.n_layers, g.n_expert, profile, pslots, err)) { + // Plan v0.3 P6: the MTP draft layer, loaded before the VRAM expert tier is sized from what is left. + strata::core::MtpDrafter mtp; + if (!o.mtp.empty()) { + if (o.spec < 2) { + std::fprintf(stderr, "strata generate: --mtp is ignored without --spec T (T >= 2)\n"); + o.mtp.clear(); + } + if (!o.mtp.empty()) mtp.set_prompt_len((int64_t) o.tokens.size()); + if (!o.mtp.empty() && !mtp.load(o.mtp, g, ss, o.max_window(), err, o.mtp_window)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } - // An explicit number truncates the ranked list ("what would 2,000 slots give" without rebuilding the - // file). `--expert-cache 0` used to take the count the profile was built for; the profile now ranks - // every pair (issue #46: a card that holds more than the old 8,000 used to stop there), so it means auto. - if (o.expert_cache == 0) o.expert_cache = -1; - std::fprintf(stderr, "strata generate: profile %s: %zu ranked pairs, built for %lld slots\n", - o.expert_profile.c_str(), profile.size(), (long long) pslots); - } - // ---- layer split across GPUs: "auto" places the split points by a cost model of one decode window, measured on - // the 5080 + 3090 rig (bench/results/2026-09-29-layer-split): - // - every layer costs its GPU a time inversely proportional to SMs x clock (0.33 ms on an RTX 5080, 0.50 on a - // 3090: the per-layer round trip and kernels, not the bytes - both cards have ~950 GB/s); - // - an expert no cache holds costs ~190 ms per unit of routed mass: the CPU pool in decode and the PCIe stream - // in prompts (fitted: the sweep's best K, 26-28, is where one more layer on the faster card stops paying - // for the ~0.1% of the mass it pushes out of its cache); - // - which experts a cache holds: its layers' profiled pairs, hottest first, until its free VRAM (less the - // reserve, the prompt path's buffers and, on a later GPU, 1 GiB for its windows and the drafter) is used; - // the routed mass of rank r is taken as (r+1)^-1.2 (fits the sweep's hit rates: K=24/26/28 predicted - // 99.53/99.34/99.15%, measured 99.5/99.4/99.0%). - // Up to 3 GPUs every placement is tried; beyond, the layers are shared in proportion to speed. - // STRATA_SPLIT_MISS_MS tunes the miss cost (a slower CPU: higher). - const int64_t split_pf_mib = o.prefill_chunk > 0 ? 160 + (o.prefill_chunk * 680) / 1024 : 0; - auto stage_room = [&](int dev, bool later) -> int64_t { - const strata::core::OnDevice on(dev); - size_t fb = 0, tb = 0; - if (const cudaError_t e = cudaMemGetInfo(&fb, &tb); e != cudaSuccess) - std::fprintf(stderr, "strata generate: layer split: CUDA%d free memory: %s\n", dev < 0 ? 0 : dev, - cudaGetErrorString(e)); - const int64_t reserve = ((int64_t) o.vram_reserve_mib + split_pf_mib + (later ? 1024 : 0)) << 20; - return std::max((int64_t) fb - reserve, 0); - }; - if (multi_gpu && split_auto) { - const auto& lay = strata::kernels::cpu::expert_layout(); - const int ns = (int) stages.size() + 1; - std::vector cap((size_t) ns), used((size_t) ns); - std::vector layer_ms((size_t) ns); - for (int i = 0; i < ns; ++i) { - const int dev = i == 0 ? 0 : stages[(size_t) i - 1]->dev; - cap[(size_t) i] = stage_room(i == 0 ? -1 : dev, i > 0); - int sms = 0, khz = 0; - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); - if (cudaDeviceGetAttribute(&khz, cudaDevAttrClockRate, dev) != cudaSuccess || khz <= 0) khz = 1800000; - cudaGetLastError(); - const double speed = std::max(1.0, (double) sms * (double) khz / 1e6); // SMs x GHz - layer_ms[(size_t) i] = 0.33 * (84.0 * 2.617) / speed; - std::fprintf(stderr, "strata generate: layer split auto: CUDA%d %d SMs at %.2f GHz -> %.2f ms per layer, " - "%.2f GiB for experts\n", dev, sms, khz / 1e6, layer_ms[(size_t) i], - (double) cap[(size_t) i] / 1073741824.0); - } - const double miss_ms = std::getenv("STRATA_SPLIT_MISS_MS") ? std::atof(std::getenv("STRATA_SPLIT_MISS_MS")) : 190.0; - std::vector mass(profile.size()); - double total_mass = 0; - for (size_t r = 0; r < profile.size(); ++r) total_mass += (mass[r] = std::pow((double) r + 1.0, -1.2)); - auto cost = [&](int64_t l) -> int64_t { - return native_pack ? ((int64_t) lay.blob_bytes(l) + 255) / 256 * 256 : (int64_t) lay.max_blob; - }; - // the predicted window time (ms) of a placement, and the routed mass its caches hold - auto predict = [&](const std::vector& at, double& held_mass, int64_t& held) -> double { - std::fill(used.begin(), used.end(), 0); - held_mass = 0; - held = 0; - std::vector full((size_t) ns, false); - for (size_t r = 0; r < profile.size(); ++r) { - const int64_t l = profile[r].first; - int st = 0; - while (st + 1 < ns && l >= at[(size_t) st]) ++st; - if (full[(size_t) st]) continue; - if (used[(size_t) st] + cost(l) > cap[(size_t) st]) { full[(size_t) st] = true; continue; } // as the fill - used[(size_t) st] += cost(l); - held_mass += mass[r]; - ++held; - } - held_mass /= std::max(total_mass, 1e-9); - double ms = miss_ms * (1.0 - held_mass); - for (int i = 0; i < ns; ++i) { - const int64_t lb = i == 0 ? 0 : at[(size_t) i - 1], le = i + 1 < ns ? at[(size_t) i] : g.n_layers; - ms += (double) (le - lb) * layer_ms[(size_t) i]; - } - return ms; - }; - std::vector best, at((size_t) ns - 1); - double best_ms = 1e30, best_mass = 0; - int64_t best_held = 0; - auto consider = [&]() { - double hm = 0; - int64_t held = 0; - const double ms = predict(at, hm, held); - if (ms < best_ms) { best = at; best_ms = ms; best_mass = hm; best_held = held; } - }; - const int64_t L = g.n_layers; - if (ns == 2) { - for (int64_t k = 2; k < L; ++k) { at[0] = k; consider(); } - } else if (ns == 3) { - for (int64_t k1 = 2; k1 + 1 < L; ++k1) - for (int64_t k2 = k1 + 1; k2 < L; ++k2) { at[0] = k1; at[1] = k2; consider(); } - } else { - double total = 0; - for (const double c : layer_ms) total += 1.0 / c; - double acc = 0; - for (int i = 0; i + 1 < ns; ++i) { - acc += 1.0 / layer_ms[(size_t) i]; - at[(size_t) i] = std::clamp((int64_t) std::llround(acc / total * (double) L), - i == 0 ? 2 : at[(size_t) i - 1] + 1, L - (ns - 1 - i)); - } - consider(); - } - split_at = best; - std::string ks; - for (const int64_t k : split_at) ks += (ks.empty() ? "" : ",") + std::to_string(k); - std::fprintf(stderr, "strata generate: layer split auto: K=%s - predicted %.1f ms per decode window; the caches " - "hold %lld of %zu profiled pairs (~%.1f%% of the routed mass)\n", ks.c_str(), best_ms, - (long long) best_held, profile.size(), 100.0 * best_mass); - } - for (size_t i = 0; i < split_at.size(); ++i) - if (split_at[i] >= g.n_layers) { - std::fprintf(stderr, "strata generate: --layer-split: layer %lld is past the last (%lld)\n", - (long long) split_at[i], (long long) (g.n_layers - 1)); - return 2; - } - // the stage that runs a layer (0: CUDA0's) - auto stage_of = [&](int64_t l) -> int { - int st = 0; - while (st < (int) split_at.size() && l >= split_at[(size_t) st]) ++st; - return st; - }; - if (multi_gpu) { - std::vector> mine; - for (const auto& pr : profile) { - const int st = stage_of(pr.first); - (st == 0 ? mine : stages[(size_t) st - 1]->profile).push_back(pr); - } - profile.swap(mine); - for (size_t i = 0; i < stages.size(); ++i) { - stages[i]->lb = split_at[i]; - stages[i]->le = i + 1 < stages.size() ? split_at[i + 1] : g.n_layers; - } + mtp.set_max_drafts(o.spec - 1); } - // THE HEAD BEFORE THE CACHE. The expert cache takes what is free minus the reserve, so everything allocated - // after it comes out of the reserve. The native head (~0.5 GB with IQ3_S) was loaded after it and ate most of - // the 700 MiB: 128K IQ3_S ended with 30 MiB free, the driver paged, and a request stalled for good at its first - // verify window. Loaded first, the cache is sized around it. + // --ple-io ram: the table goes into RAM once the SSD has delivered the experts and the draft layer + if (ple_table.is_open() && !ple_table.start_ram_load(err)) + std::fprintf(stderr, "strata generate: %s; its rows stay on the SSD\n", err.c_str()); + // where the host loop and the pool's workers run, moved off the processors that take interrupts as generation + // goes (drive_window); with a second GPU, whose work raises most of them, the first core starts free + strata::kernels::cpu::CorePlacement placement(/*spare_first=*/o.second_gpu >= 0); + strata::kernels::cpu::ExpertPool pool(o.pool_workers, /*pin=*/true, /*host_works=*/!o.no_host_worker, + placement.workers()); + if (o.no_ple_prefetch) strata::kernels::ple_prefetch_enable(false); + std::fprintf(stderr, "strata generate: session is up; locating the head\n"); const strata::core::WeightRef* wo = wt.find("output.weight"); if (wo == nullptr) { std::fprintf(stderr, "strata generate: output.weight is missing\n"); return 1; } const int64_t n_vocab = wo->ne1; strata::core::NativeHead native_head; - if (!o.native_head_gguf.empty() && !multi_gpu) { // a layer split's head is on its last stage - if (!native_head.load(o.native_head_gguf, g.n_embd, n_vocab, err)) { + std::unique_ptr split_head; + if (!o.native_head_gguf.empty()) { + std::vector head_shards; + try { + head_shards = strata::gguf_split_paths(o.native_head_gguf); + } catch (const std::exception& e) { + std::fprintf(stderr, "strata generate: native head: %s\n", e.what()); + return 1; + } + if (!native_head.load(head_shards, g.n_embd, n_vocab, err)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } - std::fprintf(stderr, "strata generate: experimental native Q5_K head, %llu bytes\n", + std::fprintf(stderr, "strata generate: native head, type %d, %llu bytes\n", native_head.type(), (unsigned long long) native_head.weight_bytes()); + // the head's last rows on the second GPU (split_head.hpp): each GPU streams its part after the last layer. + // The drafter gathers its draft head's rows first; without one it needs the whole head here. + if (o.second_gpu >= 0 && o.head_split > 0.0 && o.head_split < 1.0 && o.spec >= 2) { + const int64_t split = (int64_t) ((1.0 - o.head_split) * (double) n_vocab) / 256 * 256; + auto sh = std::make_unique(); + std::string e2; + const bool ok = (o.mtp.empty() || mtp.bind_head(&native_head, e2)) && + (o.mtp.empty() || mtp.has_draft_head()) && + sh->init(o.second_gpu, o.main_gpu, head_shards, g.n_embd, n_vocab, split, o.max_window(), + e2) && + native_head.keep_rows(split, e2); + if (ok) { + std::fprintf(stderr, "strata generate: the head's rows %lld-%lld on the second GPU (%.0f MiB), the " + "first %lld here\n", (long long) split, (long long) n_vocab, + (double) sh->bytes() / 1048576.0, (long long) split); + split_head = std::move(sh); + } else { + std::fprintf(stderr, "strata generate: the head stays on the main GPU: %s\n", + e2.empty() ? "the draft layer has no draft head of its own" : e2.c_str()); + } + } } - std::vector logits((size_t) n_vocab); - float* d_logits = nullptr; - if (cudaMalloc(&d_logits, (size_t) n_vocab * 4) != cudaSuccess) { - std::fprintf(stderr, "strata generate: the logits buffer failed\n"); - return 1; + // ---- R4's slot storage. Allocated AFTER the weights, the session and the head, so `cudaMemGetInfo` inside + // `open` sees the memory this process actually has left rather than the card's idle figure - and refuses with + // both numbers if the slots do not fit, instead of handing back a cache smaller than it was asked for. + mem_mark("the weights, the session, the drafter and the head"); + strata::core::ExpertCache xcache; + std::vector> profile; + if (!o.expert_profile.empty()) { + int64_t pslots = 0; + if (!strata::core::read_expert_profile(o.expert_profile, g.n_layers, g.n_expert, profile, pslots, err)) { + std::fprintf(stderr, "strata generate: %s\n", err.c_str()); + return 1; + } + // The profile knows how many slots it was built for. `--expert-cache 0` means "take the profile's"; + // an explicit smaller number is allowed and simply truncates the ranked list, which is the right + // behaviour for asking "what would 2,000 slots give" without rebuilding the file. + if (o.expert_cache == 0) o.expert_cache = (int) pslots; + std::fprintf(stderr, "strata generate: profile %s: %zu ranked pairs, built for %lld slots\n", + o.expert_profile.c_str(), profile.size(), (long long) pslots); } const bool auto_cache = o.expert_cache < 0; if (o.expert_cache < 0) { @@ -1992,27 +1483,16 @@ int main(int argc, char** argv) { // under WDDM an over-subscribed allocation does not fail, it pages to system memory and crawls. // (with borrowing - the default with a profile - the prompt path lends cache slots instead) const bool borrow = !o.no_prefill_borrow && !o.expert_profile.empty(); - const int64_t prefill_mib = (o.prefill_chunk > 0 && !borrow) ? 160 + (o.prefill_chunk * 680) / 1024 : 0; + const int64_t prefill_mib = (o.prefill_chunk > 0 && !borrow) + ? (int64_t) (strata::prefill::Prefill::bytes_needed(g, ss, o.prefill_chunk, + o.second_gpu >= 0 && !o.no_prompt_offload) >> 20) + : 0; const int64_t reserve = ((int64_t) o.vram_reserve_mib + prefill_mib) << 20; int64_t slots = ((int64_t) free_b - reserve) / (int64_t) strata::kernels::cpu::expert_layout().max_blob; if (!profile.empty()) slots = std::min(slots, (int64_t) profile.size()); o.expert_cache = (int) std::max(slots, 0); std::fprintf(stderr, "strata generate: expert cache auto: %.2f GiB free, %d MiB reserved -> %d slots\n", (double) free_b / 1073741824.0, o.vram_reserve_mib, o.expert_cache); - } else if (multi_gpu && o.expert_cache > 0) { - // a layer split's prompt path has its own buffers (it borrows no slots): an explicit cache size leaves room - // for them and the reserve, or the first prompt fails with "device buffers ... do not fit" - size_t free_b = 0, total_b = 0; - cudaMemGetInfo(&free_b, &total_b); - const int64_t prefill_mib = o.prefill_chunk > 0 ? 160 + (o.prefill_chunk * 680) / 1024 : 0; - const int64_t reserve = ((int64_t) o.vram_reserve_mib + prefill_mib) << 20; - const int64_t fit = std::max(((int64_t) free_b - reserve) / (int64_t) strata::kernels::cpu::expert_layout().max_blob, 0); - if (o.expert_cache > fit) { - std::fprintf(stderr, "strata generate: layer split: --expert-cache %d leaves no room for the prompt path's " - "buffers (%lld MiB) on CUDA0: %lld slots\n", o.expert_cache, (long long) prefill_mib, - (long long) fit); - o.expert_cache = (int) fit; - } } // plan v0.3 P6: a native pack's blobs differ per layer, so with a profile its slots are sized per pair: the // same VRAM holds ~30% more IQ3_XXS experts than slots of the largest blob would @@ -2034,73 +1514,20 @@ int main(int argc, char** argv) { o.expert_cache = (int) sized_slots.size(); } if (o.expert_cache > 0) { - // keep the first `keep_bytes` of the cache (the profile's hottest experts first); false when nothing is left - auto shrink_to = [&](int64_t keep_bytes) -> bool { - if (keep_bytes <= 0) { o.expert_cache = 0; sized_slots.clear(); return false; } - if (!sized_slots.empty()) { - int64_t used = 0; - size_t keep = 0; - while (keep < sized_slots.size() && used + (sized_slots[keep] + 255) / 256 * 256 <= keep_bytes) - used += (sized_slots[keep++] + 255) / 256 * 256; - sized_slots.resize(keep); - o.expert_cache = (int) keep; - } else { - o.expert_cache = (int) (keep_bytes / (int64_t) strata::kernels::cpu::expert_layout().max_blob); - } - if (o.expert_cache <= 0) { o.expert_cache = 0; sized_slots.clear(); return false; } - return true; - }; - auto cache_bytes = [&]() -> int64_t { - if (sized_slots.empty()) return (int64_t) o.expert_cache * (int64_t) strata::kernels::cpu::expert_layout().max_blob; - int64_t b = 0; - for (const int64_t s : sized_slots) b += (s + 255) / 256 * 256; - return b; - }; // With `--expert-cache auto` the reserve must still be free once the slots are WRITTEN: under WDDM an // allocation is not resident until it is touched, and the free figure read before it can be ~1 GB too // high. A cache sized from it filled the card to 0 MiB, the driver then paged, and a request that needed a // page back while the verify graph spun on a host flag never finished. So the slots are zeroed and the // free figure read again; while it is short of the reserve the cache is reopened smaller. - // STRATA_TEST_CACHE_FAIL=N: the first N opens fail as an out-of-commit cudaMalloc does (tests the retry) - int fake_fails = std::getenv("STRATA_TEST_CACHE_FAIL") ? std::atoi(std::getenv("STRATA_TEST_CACHE_FAIL")) : 0; - int failed = 0; for (int attempt = 0;; ++attempt) { - bool ok = false; - if (fake_fails > 0) { - --fake_fails; - err = "ExpertCache: cudaMalloc failed: out of memory (STRATA_TEST_CACHE_FAIL)"; - } else { - ok = sized_slots.empty() - ? xcache.open(o.expert_cache, g.n_layers, g.n_expert, (int64_t) strata::kernels::cpu::expert_layout().max_blob, err) - : xcache.open_sized(sized_slots, g.n_layers, g.n_expert, err); - } + const bool ok = sized_slots.empty() + ? xcache.open(o.expert_cache, g.n_layers, g.n_expert, (int64_t) strata::kernels::cpu::expert_layout().max_blob, err) + : xcache.open_sized(sized_slots, g.n_layers, g.n_expert, err); if (!ok) { - // Issue #60: on Windows a device allocation is also charged to the system commit (RAM + page file), - // so with a small page file the cache's one big cudaMalloc fails while the VRAM is free. An auto - // cache then tries three quarters of the size, a few times, instead of stopping the engine. - char commit[96] = ""; -#if defined(_WIN32) - MEMORYSTATUSEX ms{}; - ms.dwLength = sizeof ms; - if (GlobalMemoryStatusEx(&ms)) - std::snprintf(commit, sizeof commit, " (Windows has %.1f GiB of commit left: RAM + page file)", - (double) ms.ullAvailPageFile / 1073741824.0); -#endif - if (auto_cache && failed < 8 && shrink_to(cache_bytes() / 4 * 3)) { - ++failed; - std::fprintf(stderr, "strata generate: %s%s; trying a smaller expert cache: %d slots\n", err.c_str(), - commit, o.expert_cache); - continue; - } std::fprintf(stderr, "strata generate: %s\n", err.c_str()); -#if defined(_WIN32) - std::fprintf(stderr, "strata generate: on Windows the graphics card's memory also needs room in the page " - "file: set it to \"System managed\" (System > About > Advanced system settings > " - "Performance > Advanced > Virtual memory), or lower --expert-cache\n"); -#endif return 1; } - if (!auto_cache || attempt - failed >= 6) break; + if (!auto_cache || attempt >= 6) break; cudaMemset(xcache.device_slot(0), 0, (size_t) xcache.bytes()); cudaDeviceSynchronize(); size_t free_b = 0, total_b = 0; @@ -2114,28 +1541,38 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: only %lld MiB free once the slots are written (reserve %d MiB); " "shrinking the expert cache\n", (long long) (free_b >> 20), o.vram_reserve_mib); xcache.close(); - if (!shrink_to(keep_bytes)) break; + if (keep_bytes <= 0) { o.expert_cache = 0; sized_slots.clear(); break; } + if (!sized_slots.empty()) { + int64_t used = 0; + size_t keep = 0; + while (keep < sized_slots.size() && used + (sized_slots[keep] + 255) / 256 * 256 <= keep_bytes) + used += (sized_slots[keep++] + 255) / 256 * 256; + sized_slots.resize(keep); + o.expert_cache = (int) keep; + } else { + o.expert_cache = (int) (keep_bytes / (int64_t) strata::kernels::cpu::expert_layout().max_blob); + } + if (o.expert_cache <= 0) { o.expert_cache = 0; sized_slots.clear(); break; } } - if (failed > 0 && o.expert_cache > 0) - std::fprintf(stderr, "strata generate: expert cache: %d slots (%.2f GiB) after %d smaller tries - a bigger " - "page file lets it use more of the free VRAM\n", - o.expert_cache, (double) xcache.bytes() / 1073741824.0, failed); } if (o.expert_cache > 0) { std::fprintf(stderr, "strata generate: expert cache %lld slots, %.2f GiB of VRAM; policy is\n", (long long) xcache.slots(), xcache.gib()); mem_mark("opening the expert cache"); xcache.set_per_layer_admission(o.expert_cache_per_layer); - // Round 328 warned here that the GPU hit path was wrong (tokens diverged from a cache-off run from - // token 0). That fault was fixed long since (native_expert_parity, expert_parity, the grouped kernels' - // tests), and the warning outlived it (issue #23). What remains is rounding: a GPU expert and the CPU's - // compute the same quantized expert with different float order, so a near-tie can flip. Measured teacher- - // forced on 2,557 tokens (bench/results/2026-09-27-cache-parity): 95-98% same top-1, and perplexity equal - // (on - off = -0.005 +- 0.005 nats). Neither output is more correct than the other. + // **ROUND 328: THE HIT PATH IS PROVABLY WRONG, AND THIS SAYS SO OUT LOUD RATHER THAN LETTING IT + // CORRUPT A RUN QUIETLY.** With the cache on, the generated tokens DIVERGE from the cache-off run: + // at 256 global slots (2.97% hits) the first difference is at **token 40**; at 4096 per-layer slots + // (54.4% hits) it is at **token 0**. The cache-off run is deterministic across repeated runs, so + // this is a real fault in `moe_hit_grouped_s2`'s inputs or the fill - not noise. It also explains + // what R4 recorded as "a better profile makes the token worse": more hits means more wrong rows, so + // the payoff is non-monotone BY CONSTRUCTION rather than by any memory-system effect. + // The cache stays opt-in and this warning is not a refusal, because the divergence IS the diagnostic. std::fprintf(stderr, - "strata generate: the GPU computes the experts in the cache; it rounds differently from the CPU,\n" - " so a reply can differ slightly from a run without the cache (same quality:\n" - " bench/results/2026-09-27-cache-parity).\n"); + "strata generate: *** WARNING: --expert-cache is enabled and the GPU hit path is NOT\n" + " CORRECT. The generated tokens diverge from a cache-off run (measured:\n" + " first difference at token 40 at 2.97%% hits, token 0 at 54.4%%). Any\n" + " timing from this run is real; any OUTPUT from it is not. ***\n"); if (o.expert_cache_per_layer) { int64_t lo = 0, hi = 0; xcache.layer_slot_range(0, lo, hi); @@ -2180,151 +1617,11 @@ int main(int argc, char** argv) { (long long) prefilled, (long long) want); } - if (o.resident_cpu_experts) { - if (!src.pin_cache_complement(xcache, err, /*pin=*/false)) { - std::fprintf(stderr, "strata generate: CPU expert residency: %s\n", err.c_str()); - return 1; - } - std::fprintf(stderr, "strata generate: steady-state CPU cache misses are resident in ordinary RAM; borrowed cache entries may use mmap during prompt prefill\n"); - } - - for (auto& stp : stages) { - GpuStage& st = *stp; - const auto& lay = strata::kernels::cpu::expert_layout(); - const int64_t room = stage_room(st.dev, true); - const strata::core::OnDevice on(st.dev); - std::vector sized; - int64_t used = 0; - for (const auto& pr : st.profile) { - const int64_t b = native_pack ? ((int64_t) lay.blob_bytes(pr.first) + 255) / 256 * 256 : (int64_t) lay.max_blob; - if (used + b > room) break; - used += b; - sized.push_back((int64_t) lay.blob_bytes(pr.first)); - } - if (sized.empty() || - !(native_pack ? st.cache.open_sized(sized, g.n_layers, g.n_expert, err) - : st.cache.open((int64_t) sized.size(), g.n_layers, g.n_expert, (int64_t) lay.max_blob, err))) { - std::fprintf(stderr, "strata generate: layer split, CUDA%d expert cache: %s\n", st.dev, - sized.empty() ? "no room" : err.c_str()); - return 1; - } - int64_t filled = 0; - for (const auto& pr : st.profile) { - if (filled >= st.cache.slots()) break; - const int32_t slot = st.cache.admit(pr.first, pr.second); - if (slot == strata::core::kNotResident) break; - const uint8_t* b = srcp->blob(pr.first, pr.second); - if (b == nullptr || !st.cache.fill_slot_blocking(slot, b, err, (int64_t) lay.blob_bytes(pr.first))) { - std::fprintf(stderr, "strata generate: layer split, CUDA%d profile fill failed at pair %lld: %s\n", - st.dev, (long long) filled, err.c_str()); - return 1; - } - ++filled; - } - if (filled == 0 || !st.cache.verify_slot(st.cache.slot_of(st.profile[0].first, st.profile[0].second), - srcp->blob(st.profile[0].first, st.profile[0].second), err, - (int64_t) lay.blob_bytes(st.profile[0].first))) { - std::fprintf(stderr, "strata generate: layer split, CUDA%d expert cache: %s\n", st.dev, - filled == 0 ? "nothing filled" : err.c_str()); - return 1; - } - std::fprintf(stderr, "strata generate: layer split: CUDA%d runs layers %lld-%lld, expert cache %lld slots " - "(%.2f GiB), %lld of its %zu profiled pairs; slot 0 verified\n", - st.dev, (long long) st.lb, (long long) (st.le - 1), (long long) st.cache.slots(), st.cache.gib(), - (long long) filled, st.profile.size()); - } - if (multi_gpu) - std::fprintf(stderr, "strata generate: layer split: CUDA0 runs layers 0-%lld\n", (long long) (split_at[0] - 1)); - - std::array remote_experts; - const bool multi_remote = o.expert_cache_remote[1] > 0 || o.expert_cache_remote[2] > 0; - if (o.expert_cache_remote[0] > 0) { - if (o.expert_cache <= 0 || profile.empty() || o.no_pool) { - std::fprintf(stderr, "strata generate: remote experts need --expert-profile, " - "a CUDA0 expert cache and the expert pool\n"); - return 2; - } - std::vector> ranked = profile; - if (!stages.empty()) { // a layer split: CUDA0's share of the profile, then the later stages' pairs no cache holds - for (auto& st : stages) - for (const auto& pr : st->profile) - if (st->cache.slot_of(pr.first, pr.second) < 0) ranked.push_back(pr); - } - if (multi_remote) { - // The shipped frequency profile names only 8000 of 24576 experts. Once exhausted, - // fill remaining VRAM from unranked pairs in expert-then-layer order: this spreads - // the tail across all layers instead of concentrating it on layer zero. - std::vector seen((size_t) g.n_layers * (size_t) g.n_expert, 0); - for (const auto& pair : ranked) - if (pair.first >= 0 && pair.first < g.n_layers && pair.second >= 0 && pair.second < g.n_expert) - seen[(size_t) pair.first * (size_t) g.n_expert + (size_t) pair.second] = 1; - for (int64_t e = 0; e < g.n_expert; ++e) - for (int64_t l = 0; l < g.n_layers; ++l) - if (!seen[(size_t) l * (size_t) g.n_expert + (size_t) e]) - ranked.emplace_back((int32_t) l, (int32_t) e); - std::fprintf(stderr, "strata generate: remote ranking: %zu profiled pairs, " - "%zu other pairs to fill CUDA1..3\n", profile.size(), ranked.size() - profile.size()); - } - std::array>, 3> by_device; - if (multi_remote) { - // Either stripe experts for parallel GPU work, or give each layer one - // secondary GPU to reduce switches and transfers over shared USB4. - std::vector assigned((size_t) g.n_layers * (size_t) g.n_expert, 0); - const int devices = 1 + (o.expert_cache_remote[1] > 0) + (o.expert_cache_remote[2] > 0); - int next = 0; - for (const auto& pair : ranked) { - if (pair.first < 0 || pair.first >= g.n_layers || pair.second < 0 || pair.second >= g.n_expert || - xcache.slot_of(pair.first, pair.second) >= 0) continue; - const size_t index = (size_t) pair.first * (size_t) g.n_expert + (size_t) pair.second; - if (assigned[index]) continue; - int target = -1; - if (o.expert_cache_remote_placement == "layer") { - target = pair.first % devices; - // Other layers' owners may still have room: keep scanning ranks. - if (by_device[(size_t) target].size() >= - (size_t) o.expert_cache_remote[(size_t) target]) continue; - } else { - for (int i = 0; i < devices; ++i) { - const int r = (next + i) % devices; - if (by_device[(size_t) r].size() < (size_t) o.expert_cache_remote[(size_t) r]) { - target = r; - break; - } - } - if (target < 0) break; - } - assigned[index] = 1; - by_device[(size_t) target].push_back(pair); - next = (target + 1) % devices; - } - std::fprintf(stderr, "strata generate: remote ranks %s across %d CUDA devices\n", - o.expert_cache_remote_placement == "layer" ? "grouped by layer" : "striped", devices); - } else { - by_device[0] = std::move(ranked); - } - std::vector claimed((size_t) g.n_layers * (size_t) g.n_expert, 0); - for (auto& st : stages) // a layer split: what a stage's cache holds is no helper's - for (const auto& pr : st->profile) - if (st->cache.slot_of(pr.first, pr.second) >= 0) - claimed[(size_t) pr.first * (size_t) g.n_expert + (size_t) pr.second] = 1; - for (int r = 0; r < 3; ++r) if (o.expert_cache_remote[(size_t) r] > 0) { - if (!remote_experts[(size_t) r].open(remote_dev[r], o.expert_cache_remote[(size_t) r], - g.n_layers, g.n_expert, by_device[(size_t) r], xcache, *srcp, claimed, err)) { - std::fprintf(stderr, "strata generate: %s\n", err.c_str()); - return 1; - } - std::fprintf(stderr, "strata generate: CUDA%d: %lld additional experts, %.2f GiB; " - "results return through pinned host rows\n", remote_dev[r], - (long long) remote_experts[(size_t) r].resident(), remote_experts[(size_t) r].gib()); - } - } - Drive drive; - for (int r = 0; r < 3; ++r) if (o.expert_cache_remote[(size_t) r] > 0) - drive.d.remote[drive.d.remote_count++] = &remote_experts[(size_t) r]; drive.d.hit_cpu_order = o.expert_cache_cpu_order; drive.d.split_rows = !o.no_split_rows; drive.d.pool = &pool; + drive.placement = &placement; drive.d.src = srcp; drive.d.n_expert = g.n_expert; drive.d.jobs.resize((size_t) K); @@ -2433,8 +1730,11 @@ int main(int argc, char** argv) { strata::core::HitFn hit_fn = (o.no_pool || o.expert_cache <= 0) ? nullptr : &strata::core::expert_hit_run; void* pool_user = o.no_pool ? nullptr : (void*) &drive; - std::fprintf(stderr, "strata generate: %d expert-pool workers%s%s\n", pool.workers(), - pool.host_works() ? " + the host thread" : "", + std::string on; + for (int i = 0; i < pool.workers() && i < (int) placement.workers().size(); ++i) + on += (i > 0 ? "," : "") + std::to_string(placement.workers()[(size_t) i]); + std::fprintf(stderr, "strata generate: %d pool workers on logical processors %s, the host thread on %d%s%s\n", + pool.workers(), on.c_str(), placement.host(), pool.host_works() ? " (draining too)" : "", o.no_pool ? " (UNUSED: --no-pool)" : ""); // **THE MISALIGNMENT WARNING THAT STOOD HERE IS GONE, BECAUSE THE MISALIGNMENT IS FIXED.** @@ -2493,7 +1793,12 @@ int main(int argc, char** argv) { } mem_mark("the expert cache and the graphs"); - std::fprintf(stderr, "strata generate: session is up (engine %s)\n", STRATA_VERSION); + std::vector logits((size_t) n_vocab); + float* d_logits = nullptr; + if (cudaMalloc(&d_logits, (size_t) n_vocab * 4) != cudaSuccess) { + std::fprintf(stderr, "strata generate: the logits buffer failed\n"); + return 1; + } auto run_head = [&](void* stream) -> bool { if (!native_head.loaded()) return strata::core::lm_head(wt, g, ss.block, d_logits, stream, err); @@ -2544,16 +1849,6 @@ int main(int argc, char** argv) { sp.top_k = o.top_k; sp.top_p = o.top_p; sp.temperature = o.temperature; - // what this run actually samples with (the speculative loop below gets the same parameters); serve samples - // per request instead - if (!o.serve) { - if (sp.greedy || sp.temperature <= 0.0f) - std::fprintf(stderr, "strata generate: sampling greedy\n"); - else - std::fprintf(stderr, "strata generate: sampling temperature=%g top_k=%d top_p=%g seed=%llu\n", - (double) sp.temperature, sp.top_k > 0 && sp.top_k < 64 ? sp.top_k : 64, (double) sp.top_p, - (unsigned long long) sp.seed); - } std::FILE* dump = nullptr; const int64_t dump_positions = (int64_t) o.tokens.size() - 1 + o.max_new; @@ -2872,7 +2167,7 @@ int main(int argc, char** argv) { // Initialised unconditionally, including under --no-pool: the loop validates the scratch it is handed, so // passing a default-constructed one is an error rather than a fallback. (It was, and the guard caught it - // which is the point of the guard.) One allocation at setup either way. - if (!loop_scratch.init((size_t) K * g.n_embd * 4, err)) { + if (!loop_scratch.init((size_t) K * g.n_embd * 4, err, placement.host())) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } @@ -2895,8 +2190,7 @@ int main(int argc, char** argv) { int64_t resident = 0; for (int64_t l = 0; l < g.n_layers; ++l) for (int64_t e = 0; e < g.n_expert; ++e) { - const int st = multi_gpu ? stage_of(l) : 0; - const int32_t slot = st > 0 ? stages[(size_t) st - 1]->cache.slot_of(l, e) : xcache.slot_of(l, e); + const int32_t slot = xcache.slot_of(l, e); host_res[(size_t) (l * g.n_expert + e)] = slot; if (slot != strata::core::kNotResident) ++resident; } @@ -2908,15 +2202,6 @@ int main(int argc, char** argv) { } thits.d_res = d_res; thits.n_expert = g.n_expert; - for (auto& st : stages) { // layer split across GPUs: the same table on every device - const strata::core::OnDevice on(st->dev); - if (cudaMalloc((void**) &st->d_res, host_res.size() * sizeof(int32_t)) != cudaSuccess || - cudaMemcpy(st->d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice) != - cudaSuccess) { - std::fprintf(stderr, "strata generate: layer split: CUDA%d residency table failed\n", st->dev); - return 1; - } - } thits.cache_base = drive.d.cache_base; thits.blob = drive.d.cache_blob; thits.d_slot = drive.d.d_slot; @@ -2930,6 +2215,143 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: token graph hit path: %lld resident experts, decided on the device\n", (long long) resident); } + // plan v0.3 P6: the VRAM tier follows the conversation in both decode loops (--serve and the speculative loop): + // moves every --adapt-every rounds on the adapt thread, admitted once their copies have landed + strata::core::AdaptiveTier tier; + const bool adaptive = !host_res.empty() && o.adapt_every > 0 && o.adapt_swaps > 0; + if (adaptive && !tier.init(xcache, *srcp, host_res, g.n_layers, g.n_expert, o.adapt_swaps, err, + o.main_gpu, o.main_gpu)) { + std::fprintf(stderr, "strata generate: %s\n", err.c_str()); + return 1; + } + // ---- a second GPU as another expert tier: the profile's pairs the first GPU's cache does not hold, then empty + // slots for every layer, which its own adaptive tier fills with the conversation's experts the first tier has not. + // Without free VRAM there (another process holds it) the engine runs on the main GPU alone. + strata::core::SecondGpu gpu2; + std::vector host_res2; + strata::core::AdaptiveTier tier2; + if (o.second_gpu >= 0) { + if (!native_pack || host_res.empty() || o.spec < 2) { + std::fprintf(stderr, "strata generate: --second-gpu needs a native pack, --expert-profile and --spec\n"); + return 2; + } + const int main_dev = o.main_gpu; + cudaDeviceProp prop{}; + cudaGetDeviceProperties(&prop, o.second_gpu); + // prefetch slots for the most of the smallest experts the budget holds + const auto& lay0 = strata::kernels::cpu::expert_layout(); + uint64_t min_blob = lay0.max_blob; + for (int64_t l = 0; l < g.n_layers; ++l) min_blob = std::min(min_blob, lay0.blob_bytes(l)); + const uint64_t pf_bytes = (uint64_t) (std::max(0.0, o.second_gpu_prefetch_mb) * 1048576.0); + drive.d.gpu2_prefetch_bytes = pf_bytes; + bool ok = gpu2.init(o.second_gpu, main_dev, g.n_embd, g.n_ff, o.max_window(), ss.k, err) && + gpu2.init_prefetch((int) std::min(pf_bytes / min_blob, 64), lay0.max_blob, err); + const auto& lay = strata::kernels::cpu::expert_layout(); + std::vector sizes; + std::vector> pre; + std::vector empty2; + if (ok) { + cudaSetDevice(o.second_gpu); + size_t free_b = 0, total_b = 0; + cudaMemGetInfo(&free_b, &total_b); + // the prompt path's buffers there borrow cache slots, else they are reserved here + const bool offload_own = o.prefill_chunk > 0 && !o.no_prompt_offload && o.no_prefill_borrow; + const uint64_t reserve = ((uint64_t) o.second_gpu_reserve_mib << 20) + + (offload_own ? strata::prefill::ExpertRunner::bytes_needed(o.prefill_chunk, g.n_expert, true, true) : 0); + const uint64_t budget = o.second_gpu_gib > 0 ? (uint64_t) (o.second_gpu_gib * 1073741824.0) + : (free_b > reserve ? (uint64_t) free_b - reserve : 0); + uint64_t used = 0; + for (size_t i = (size_t) prefilled; i < profile.size(); ++i) { + const auto& pr = profile[i]; + if (host_res[(size_t) (pr.first * g.n_expert + pr.second)] >= 0) continue; // the first GPU holds it + const uint64_t b = (lay.blob_bytes(pr.first) + 255) / 256 * 256; + if (used + b > budget) break; + used += b; + sizes.push_back((int64_t) lay.blob_bytes(pr.first)); + pre.push_back(pr); + } + { // the rest in empty slots (without the adaptive tier only the prompt path's loans use them) + uint64_t round = 0; + for (int64_t l = 0; l < g.n_layers; ++l) round += (lay.blob_bytes(l) + 255) / 256 * 256; + for (uint64_t q = (budget - used) / round; q > 0; --q) + for (int64_t l = 0; l < g.n_layers; ++l) { + sizes.push_back((int64_t) lay.blob_bytes(l)); + empty2.push_back((int32_t) l); + } + } + if (sizes.empty()) err = "no VRAM for experts (" + std::to_string(free_b >> 20) + " MiB free)"; + ok = !sizes.empty() && gpu2.cache().open_sized(sizes, g.n_layers, g.n_expert, err); + // as on the main GPU, the reserve must still be free once the slots are written (WDDM makes memory + // resident when it is first touched): short of it, the cache reopens without its last (empty) slots + for (int attempt = 0; ok && o.second_gpu_gib <= 0 && attempt < 6; ++attempt) { + cudaMemset(gpu2.cache().device_slot(0), 0, (size_t) gpu2.cache().bytes()); + cudaDeviceSynchronize(); + cudaMemGetInfo(&free_b, &total_b); + if ((uint64_t) free_b + (64ull << 20) >= reserve) break; + std::fprintf(stderr, "strata generate: second GPU: only %lld MiB free once its slots are written " + "(reserve %d MiB); shrinking its cache\n", (long long) (free_b >> 20), + o.second_gpu_reserve_mib); + gpu2.cache().close(); + for (uint64_t cut = 0; !sizes.empty() && cut < reserve - free_b + (64ull << 20);) { + cut += ((uint64_t) sizes.back() + 255) / 256 * 256; + sizes.pop_back(); + if (!empty2.empty()) empty2.pop_back(); + else pre.pop_back(); + } + if (sizes.empty()) err = "no VRAM for experts once the slots are written"; + ok = !sizes.empty() && gpu2.cache().open_sized(sizes, g.n_layers, g.n_expert, err); + } + for (size_t i = 0; ok && i < pre.size(); ++i) { + const int32_t slot = gpu2.cache().admit(pre[i].first, pre[i].second); + const uint8_t* b = srcp->blob(pre[i].first, pre[i].second); + ok = slot != strata::core::kNotResident && b != nullptr && + gpu2.cache().fill_slot_blocking(slot, b, err, (int64_t) lay.blob_bytes(pre[i].first)); + } + cudaSetDevice(main_dev); + } + if (ok) { + host_res2.assign((size_t) (g.n_layers * g.n_expert), strata::core::kNotResident); + for (int64_t l = 0; l < g.n_layers; ++l) + for (int64_t e = 0; e < g.n_expert; ++e) host_res2[(size_t) (l * g.n_expert + e)] = gpu2.cache().slot_of(l, e); + } + if (ok && adaptive) { + ok = tier2.init(gpu2.cache(), *srcp, host_res2, g.n_layers, g.n_expert, o.adapt_swaps, err, + o.second_gpu, main_dev); + if (ok) { + // lowest slots first: the last ones, which the prompt path borrows, stay empty longest + for (size_t i = empty2.size(); i-- > 0;) tier2.add_free(empty2[i], (int32_t) (pre.size() + i)); + tier2.set_upper(&tier); + tier.set_lower(&tier2); + tier2.set_paced(gpu2.done_event()); // between windows; a window's last share may still read its slots + drive.tier2 = &tier2; + } + } + if (ok) { + drive.d.gpu2 = &gpu2; + drive.d.host_res2 = host_res2.data(); + drive.d.gpu2_min_bytes = (uint64_t) (o.second_gpu_min_mb * 1048576.0); + std::fprintf(stderr, "strata generate: second GPU %d (%s): %zu experts from the profile and %zu empty slots, " + "%.2f GiB\n", o.second_gpu, prop.name, pre.size(), empty2.size(), gpu2.cache().gib()); + } else { + // another process may hold its VRAM (a model, a desktop app): the engine runs on the main GPU alone + std::fprintf(stderr, "strata generate: second GPU %d unused: %s\n", o.second_gpu, err.c_str()); + o.second_gpu = -1; + host_res2.clear(); + } + } + // the second tier ranks after the first on the same counts, and decays them + std::string adapt_err; + auto adapt = [&]() -> bool { + const bool two = tier2.on(); + return tier.adapt(drive.d.usage, adapt_err, !two) && + (!two || (tier2.adapt(drive.d.usage, adapt_err) && pump_gap(&drive, adapt_err))); + }; + drive.adapt = adapt; + if (tier.on()) drive.tier1 = &tier; // its copies in pieces, or at once without a second GPU (Drive::tier1) + drive.burst = !tier2.on(); + tier.set_wait(drive.burst); + drive.n_layers = g.n_layers; + drive.ring_us.assign((size_t) ((strata::kernels::kVerifyMaxT + 1) * g.n_layers), 0.0); if (!o.no_capture && !o.no_token_graph && layer_dump == nullptr && half_dump == nullptr && (hit_fn == nullptr || thits.on()) && !native_pack) { if (!strata::core::session_capture_token(wt, g, ss, d_parts, loop_scratch.y_miss, loop_scratch.parts_bytes, @@ -2940,6 +2362,48 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: token graph captured (48 layers, one launch per token)\n"); } + // ---- the batched prompt path, and the second GPU's share of its experts (not --no-prompt-offload) + auto setup_prompt_path = [&](strata::prefill::Prefill& prefill, strata::prefill::ExpertRunner& offload, + PromptBuffers& pbuf, std::string& e) -> bool { + const bool off = o.second_gpu >= 0 && !o.no_prompt_offload; + // a cache too small to lend the buffers would make them allocate on top of a card the cache filled to its + // reserve, which then pages: the chunk halves until they fit in the lendable slots (it only reads slower) + if (!o.no_prefill_borrow && d_res != nullptr && o.expert_cache > 0) { + SlotLoan probe; + probe.cache = &xcache; + int64_t chunk = o.prefill_chunk; + while (chunk > 256 && !probe.plan(strata::prefill::Prefill::bytes_needed(g, ss, chunk, off))) chunk /= 2; + if (chunk != o.prefill_chunk && probe.plan(strata::prefill::Prefill::bytes_needed(g, ss, chunk, off))) { + std::fprintf(stderr, "strata generate: prompt chunk %lld -> %lld tokens so its buffers fit in the expert " + "cache\n", (long long) o.prefill_chunk, (long long) chunk); + o.prefill_chunk = chunk; + } + } + if (off && !offload.init(o.second_gpu, o.main_gpu, nullptr, srcp, &gpu2.cache(), + host_res2.empty() ? nullptr : host_res2.data(), g.n_expert, o.prefill_chunk, true, e)) + return false; + if (!prefill.init(wt, g, ss, srcp, o.expert_cache > 0 ? &xcache : nullptr, + host_res.empty() ? nullptr : host_res.data(), o.prefill_chunk, main_cs, off ? &offload : nullptr, e)) + return false; + pbuf.prefill = &prefill; + pbuf.offload = off ? &offload : nullptr; + pbuf.g = &g; + pbuf.ss = &ss; + if (d_res != nullptr) { + pbuf.loan.cache = &xcache; + pbuf.loan.res = &host_res; + pbuf.loan.d_res = d_res; + pbuf.loan.device = pbuf.loan.main_device = o.main_gpu; + } + if (off) { + pbuf.loan2.cache = &gpu2.cache(); + pbuf.loan2.res = &host_res2; + pbuf.loan2.device = o.second_gpu; + pbuf.loan2.main_device = o.main_gpu; + } + return pbuf.init(o.prefill_chunk, !o.no_prefill_borrow, e); + }; + // ================================ WHERE THE HOST TERM GOES, PER TOKEN ================================ // // **`--gpu-only-full` MEASURES THE 48 LAYER GRAPHS AND THE LM HEAD AND NOTHING ELSE.** It never enters @@ -2964,455 +2428,74 @@ int main(int argc, char** argv) { // // and each generated token is written to stdout as `T ` as soon as its verify window is done, followed by // - // DONE - // + // DONE // - // Before that, `RESUME ` (n prompt tokens are not read again), `PP ` - // after every prompt chunk, and `REUSED ` once the prompt is read. (`ERR ` instead when a request - // cannot run; `STOP` ends the running request at its next step; `QUIT` ends the process.) A request continues - // from the live session or the longest conversation checkpoint its prompt starts with (see ConvCheckpoint), - // otherwise from an empty sequence (`session_zero`); the rest of the prompt goes through the batched prompt path - // and its last token through the first verify window - the path all three model files share. Decoding is greedy. - // The expert-cache slots (from the end of the cache) that hold the prompt path's buffers for a chunk, and the - // bytes from the first of them to the end. - // the share of expert bytes the arena could pin (sizes the prompt path's streamed ring and its lend cap) - if (srcp != nullptr && o.prefill_chunk > 0) { - uint64_t pinned = 0, total = 0; - const auto& lay = strata::kernels::cpu::expert_layout(); - for (int64_t l = 0; l < g.n_layers; ++l) - for (int64_t e = 0; e < g.n_expert; ++e) { - const uint64_t b = lay.blob_bytes(l); - total += b; - if (srcp->pinned(l, e)) pinned += b; - } - strata::prefill::Prefill::set_pinned_share(total ? (double) pinned / (double) total : 1.0); - } - auto lend_slots = [&](int64_t c) -> int64_t { - const uint64_t need = strata::prefill::Prefill::bytes_needed(g, ss, c); - const int64_t blob = (int64_t) strata::kernels::cpu::expert_layout().max_blob; - int64_t k = (int64_t) ((need + (uint64_t) blob - 1) / (uint64_t) blob); - if (xcache.slot_offsets() != nullptr) { // sized slots: take slots from the end until they hold `need` - k = 0; - while (k < xcache.slots() && - (uint64_t) (xcache.bytes() - (int64_t) xcache.slot_offsets()[xcache.slots() - k]) < need) ++k; - } - return k; - }; - auto lend_bytes = [&](int32_t first) -> uint64_t { - return xcache.slot_offsets() ? (uint64_t) (xcache.bytes() - (int64_t) xcache.slot_offsets()[first]) - : (uint64_t) (xcache.slots() - first) * - (uint64_t) strata::kernels::cpu::expert_layout().max_blob; - }; - auto request_chunk = [](int64_t tokens, int64_t max_chunk) -> int64_t { - if (tokens <= 0 || max_chunk <= 0) return 0; - const int64_t rounded = tokens > std::numeric_limits::max() - 255 - ? tokens - : ((tokens + 255) / 256) * 256; - return std::min(max_chunk, rounded); - }; - // The prompt path's chunk and the slots it borrows for its buffers: the requested chunk halved until it fits, - // or with --prefill auto the largest of kAutoChunks whose buffers take at most kAutoLendPct % of the slots (a - // lent slot's expert is streamed during the prompt and refilled after it; measured on a 12 GB card, 32K Q2_0 - // prompt: 4096 791 tok/s, 6144 878, 8192 973 with 69% of the slots lent). A request lends only what its own - // prompt needs (Prefill::relayout), so a big chunk costs short prompts nothing. 0 = none fits. - auto plan_lend = [&](int64_t& chunk) -> int64_t { - static constexpr int64_t kAutoChunks[] = {8192, 6144, 4096, 3072, 2048, 1024, 512, 256}; - // at 8192-token chunks nearly every expert streams anyway, so a lent slot costs little: 90% when the - // copies are DMA from pinned RAM (Q2_0 8192 + a 384-slot ring: 1283 tok/s), 85% when host copies are the - // limit (lending more only streams more through them). STRATA_PREFILL_LEND_PCT overrides (tuning). - const int64_t kAutoLendPct = [] { - const char* v = std::getenv("STRATA_PREFILL_LEND_PCT"); - return v ? (int64_t) std::atoi(v) - : (int64_t) (strata::prefill::Prefill::pinned_share() >= 0.9 ? 90 : 85); - }(); - auto slots_for = lend_slots; - if (o.prefill_auto) { - for (const int64_t c : kAutoChunks) { - const int64_t k = slots_for(c); - if (k + 128 <= xcache.slots() && k * 100 <= kAutoLendPct * xcache.slots()) { chunk = c; return k; } - } - return 0; - } - for (int64_t c = chunk; c >= 256; c /= 2) { - const int64_t k = slots_for(c); - if (k + 128 <= xcache.slots()) { chunk = c; return k; } - } - return 0; - }; + // (`ERR ` instead when a request cannot run; `QUIT` ends the process). Every request starts from an + // empty sequence (`session_zero`): the prompt goes through the batched prompt path and its last token through + // the first verify window - the path all three model files share. Decoding is greedy. if (o.serve) { - if (o.spec < 2 || o.mtp.empty() || o.prefill_chunk <= 0 || - (graph_hits && (thits.d_res == nullptr || host_res.empty()))) { - std::fprintf(stderr, "strata serve: needs --spec T, --mtp DIR and --prefill CHUNK (and a fillable " - "--expert-cache; the graphed hit path additionally needs --expert-profile P)\n"); + if (o.spec < 2 || o.mtp.empty() || o.prefill_chunk <= 0 || thits.d_res == nullptr || host_res.empty()) { + std::fprintf(stderr, "strata serve: needs --spec T, --mtp DIR, --prefill CHUNK, --expert-profile P and " + "--expert-cache\n"); return 2; } - strata::prefill::Prefill sp; - void* borrow = nullptr; - uint64_t borrow_bytes = 0; - int32_t lend_first = -1; // the first slot the prompt path may borrow (its largest chunk) - int32_t lend_first_now = -1; // where its buffers are laid out now - // a cache too small to lend the prompt path its buffers would make it allocate them on top - on a card the - // cache already filled to its reserve, that is the over-subscription the auto sizing avoids - so the - // prompt chunk is halved until its buffers fit in the lendable slots (a smaller chunk only reads slower) - if (!o.no_prefill_borrow && d_res != nullptr) { - int64_t chunk = o.prefill_chunk; - if (const int64_t k = plan_lend(chunk); k > 0) { - if (o.prefill_auto) - std::fprintf(stderr, "strata serve: prompt chunk auto: %lld tokens\n", (long long) chunk); - else if (chunk != o.prefill_chunk) - std::fprintf(stderr, "strata serve: prompt chunk %lld -> %lld tokens so its buffers fit in the " - "expert cache\n", (long long) o.prefill_chunk, (long long) chunk); - o.prefill_chunk = chunk; - lend_first = (int32_t) (xcache.slots() - k); - lend_first_now = lend_first; - borrow = xcache.device_slot(lend_first); - borrow_bytes = xcache.slot_offsets() - ? (uint64_t) (xcache.bytes() - (int64_t) xcache.slot_offsets()[lend_first]) - : (uint64_t) k * (uint64_t) strata::kernels::cpu::expert_layout().max_blob; - } else if (o.prefill_auto) { - o.prefill_chunk = 1024; // nothing lendable: small buffers of its own - } - } else if (o.prefill_auto && d_res == nullptr) { - o.prefill_chunk = 1024; // #85: no expert cache at all (a full 8 GB card): small buffers of its own - } - if (borrow != nullptr) - std::fprintf(stderr, "strata serve: the prompt path borrows %lld cache slots (%.2f GiB)\n", - (long long) (xcache.slots() - lend_first), (double) borrow_bytes / 1073741824.0); - else - std::fprintf(stderr, "strata serve: the prompt path allocates its own buffers (too few cache slots to borrow)\n"); - // layer split across GPUs: a prompt path per stage, each handing its chunk's rows to the next - for (size_t i = 0; i < stages.size(); ++i) { - GpuStage& st = *stages[i]; - st.sp.set_stage(st.lb, i + 1 < stages.size() ? st.le : -1, i + 1 < stages.size() ? &stages[i + 1]->sp : nullptr); - const strata::core::OnDevice on(st.dev); - if (!st.sp.init(st.wt, g, st.ss, srcp, &st.cache, host_res.data(), o.prefill_chunk, (void*) st.stream, err)) { - std::fprintf(stderr, "strata serve: layer split, CUDA%d prompt path: %s\n", st.dev, err.c_str()); - return 1; - } - } - if (multi_gpu) sp.set_stage(0, split_at[0], &stages[0]->sp); - if (!sp.init(wt, g, ss, srcp, &xcache, host_res.data(), o.prefill_chunk, main_cs, err, borrow, borrow_bytes)) { + strata::prefill::Prefill prefill; + strata::prefill::ExpertRunner offload; + PromptBuffers pbuf; + if (!setup_prompt_path(prefill, offload, pbuf, err)) { std::fprintf(stderr, "strata serve: %s\n", err.c_str()); - if (err.find("fit") != std::string::npos) // #85: say what frees VRAM - std::fprintf(stderr, "strata serve: the GPU has too little free VRAM for the prompt path: turn images " - "off (setup: --vision no), close other programs using the GPU, use a shorter " - "context, or read prompts in smaller chunks (--prefill 512)\n"); return 1; } + if (pbuf.lend || pbuf.lend2) + std::fprintf(stderr, "strata serve: the prompt path borrows up to %lld cache slots (%.2f GiB), %lld (%.2f GiB) " + "on the second GPU\n", (long long) pbuf.loan.slots(), (double) pbuf.loan.bytes() / 1073741824.0, + (long long) pbuf.loan2.slots(), (double) pbuf.loan2.bytes() / 1073741824.0); + else + std::fprintf(stderr, "strata serve: the prompt path allocates its own buffers (too few cache slots to borrow)\n"); mem_mark("the head and the prompt path"); - // the penalty-history buffer: one row per verify-window row (`penalty_rows`), each the last - // `penalty_last_n` tokens that row's pick follows, -1 padded in front. Allocated once at the cap for - // the widest window; a request without penalties gets a null buffer and takes the byte-for-byte - // neutral path (no upload, no buffer handed to the sampler). - constexpr int kPenaltyWindowCap = 4096; - constexpr size_t kHistSlots = (size_t) kPenaltyWindowCap * (size_t) strata::kernels::kVerifyMaxT; - int32_t* d_hist = nullptr; - std::vector hist_stage(kHistSlots, -1); - const int hist_dev = last_st ? last_st->dev : -1; // with the head: the last stage's device - if (const strata::core::OnDevice on_h(hist_dev); cudaMalloc(&d_hist, kHistSlots * sizeof(int32_t)) != cudaSuccess) { - std::fprintf(stderr, "strata serve: the penalty-history allocation failed\n"); - return 1; - } strata::core::Verifier ver; strata::core::VerifyHits vh; - vh.d_res = thits.d_res; + vh.res = host_res.data(); vh.cache_base = thits.cache_base; + vh.slot_off = xcache.slot_offsets(); + vh.slots = xcache.slots(); vh.blob = thits.blob; - vh.slot_off = xcache.slot_offsets(); // E-6: the device plan's pointers - vh.n_slots = xcache.slots(); - // Layer split: `ver` runs layers [0, K1) and hands its residual to the next stage's verifier, and so on; the - // last runs the head. The hand-offs are mapped pinned memory, portable: a stage on another GPU reads it. - // (--split-device 0: the second stage on this GPU, sharing its weights, session and cache - the A/B.) - strata::core::Verifier ver_same; - SplitDrive split_drive; - auto stage_ver = [&](int st) -> strata::core::Verifier& { - return st == 0 ? ver : split_same ? ver_same : stages[(size_t) st - 1]->ver; - }; - auto pcie_num_of = [](double f) { return std::max(0, std::min(256, (int) (f * 256.0 + 0.5))); }; - const int n_stages = split_devs.empty() ? 1 : (int) split_at.size() + 1; - if (n_stages > 1) { - const size_t hb = (size_t) strata::kernels::kVerifyMaxT * - (size_t) strata::core::Verifier::handoff_floats(g) * sizeof(float); - std::vector hand((size_t) n_stages - 1, nullptr); - for (float*& h : hand) { - float* hh = nullptr; - if (cudaHostAlloc((void**) &hh, hb, cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess || - cudaHostGetDevicePointer((void**) &h, hh, 0) != cudaSuccess) { - std::fprintf(stderr, "strata serve: the layer-split hand-off allocation failed\n"); - return 1; - } - std::memset(hh, 0, hb); - } - split_drive.base = &drive; - split_drive.n = n_stages; - for (int st = 0; st < n_stages; ++st) { - stage_ver(st).set_stage(st == 0 ? 0 : split_at[(size_t) st - 1], st + 1 < n_stages ? split_at[(size_t) st] : -1, - st == 0 ? nullptr : hand[(size_t) st - 1], st + 1 < n_stages ? hand[(size_t) st] : nullptr); - split_drive.end[st] = st + 1 < n_stages ? split_at[(size_t) st] : g.n_layers; - split_drive.cache_base[st] = drive.d.cache_base; - split_drive.cache_slot_off[st] = drive.d.cache_slot_off; - split_drive.pcie_num[st] = pcie_num_of(o.pcie_frac); - } - for (int st = 1; st < n_stages; ++st) { - bool ok_s = false; - if (split_same) { - ok_s = ver_same.init(wt, g, ss, vh, native_head.loaded() ? &native_head : nullptr, o.spec, err); - } else { - GpuStage& gs = *stages[(size_t) st - 1]; - const strata::core::OnDevice on(gs.dev); - strata::core::VerifyHits vs; - vs.d_res = gs.d_res; - vs.cache_base = gs.cache.device_slot(0); - vs.blob = thits.blob; - vs.slot_off = gs.cache.slot_offsets(); - vs.n_slots = gs.cache.slots(); - ok_s = gs.ver.init(gs.wt, g, gs.ss, vs, gs.head.loaded() ? &gs.head : nullptr, o.spec, err); - split_drive.cache_base[st] = gs.cache.device_slot(0); - split_drive.cache_slot_off[st] = gs.cache.slot_offsets(); - split_drive.pcie_num[st] = pcie_num_of(gs.pcie_frac); - } - if (!ok_s) { - std::fprintf(stderr, "strata serve: layer split, stage %d: %s\n", st + 1, err.c_str()); - return 1; - } - } - for (int st = 0; st + 1 < n_stages; ++st) stage_ver(st).set_next(&stage_ver(st + 1), &split_drive); - std::string plan_s = "0-" + std::to_string(split_at[0] - 1) + " (CUDA0)"; - for (int st = 1; st < n_stages; ++st) - plan_s += ", " + std::to_string(split_at[(size_t) st - 1]) + "-" + std::to_string(split_drive.end[st] - 1) + - " (CUDA" + std::to_string(split_same ? 0 : stages[(size_t) st - 1]->dev) + ")"; - std::fprintf(stderr, "strata serve: layer split: layers %s, one hand-off per window\n", plan_s.c_str()); - } - if (!ver.init(wt, g, ss, vh, native_head.loaded() ? &native_head : nullptr, o.spec, err) || - !mtp.bind(last_st ? last_st->wt : wt, last_st ? &last_st->head : &native_head, ver.final_R_all(), err)) { + ver.set_split_head(split_head.get()); + if (!ver.init(wt, g, ss, vh, native_head.loaded() ? &native_head : nullptr, o.max_window(), err) || + !mtp.bind(wt, &native_head, ver.final_R_all(), err)) { std::fprintf(stderr, "strata serve: %s\n", err.c_str()); return 1; } - for (int st = 0; st < n_stages && n_stages > 1; ++st) { - split_drive.plan[st] = stage_ver(st).plan_sink(); - if (st > 0) { - stage_ver(st).set_split(o.spec_split); - stage_ver(st).set_pcie_mode(o.pcie_mode == "dma" ? 0 : o.pcie_mode == "direct" ? 1 : 2); - } - } - // the pool the verify windows call: with a layer split, the wrapper that routes each layer to its stage - const strata::core::PoolMultiFn win_pool_fn = n_stages > 1 ? &drive_pool_split : &drive_pool_multi; - void* const win_pool_user = n_stages > 1 ? (void*) &split_drive : (void*) &drive; mem_mark("the verifier and the drafter's binding"); ver.set_split(o.spec_split); - // auto: the copy kernel for every pack. DMA (the native packs' default until 0.1.13) has the host call - // cudaMemcpyAsync + cudaLaunchHostFunc inside a verify window while the GPU spins on the flag they raise; - // issue #31's thread dumps show the host stuck in that cudaMemcpyAsync on a driver lock for good. The copy - // kernel needs no host CUDA call there, and costs ~1-3% decode on IQ3_S (45.3 -> 44.8 tok/s, 8 requests). - ver.set_pcie_mode(o.pcie_mode == "dma" ? 0 : o.pcie_mode == "direct" ? 1 : 2); - std::vector cur; - // ---- the conversation cache (see ConvCheckpoint). `live` is what the session holds right now: the tokens - // it has consumed, so a request that starts with exactly them continues without any copy. `checks` are the - // saved points; every one of them is a prefix of `live` (the loop drops the rest), so they form a chain - - // the radix cache's tree collapsed onto the one branch of history whose cells the session holds. The - // chain's root is the deepest point every request so far shared (the end of the system prompt, in - // practice); the retention policy pins it and rotates the rest LRU (conv_cache.hpp), so a NEW chat that - // shares that prefix mounts through it instead of reading it again. - std::vector live; - std::vector live_imgs, req_imgs; - bool live_ok = false; - std::vector checks; - uint64_t check_clock = 0; // the checkpoints' LRU clock; creation and every use advance it - bool cvec_cached = true; // the control vector's state the live session and the checkpoints were read with - int64_t pp_total = 0, pp_from = 0, pp_next_check = 0; - Clock::time_point pp_t0 = Clock::now(); - auto imgs_below = [&](const std::vector& all, int64_t L) { - std::vector v; - for (const ImgKey& k : all) if (k.start < L) v.push_back(k); - return v; - }; - // a checkpoint of the state after `cur[0, L)`; false only when the copy itself failed - // A layer split's mid-prompt checkpoints: when the last stage reports a chunk, the earlier ones already read - // the next, so each stage saves its own part when IT reaches a checkpoint position (the same rule as below: - // every `prompt_cache_every` tokens from where the request resumed), and the last stage puts them together. - std::mutex part_mu; - std::map> part_at; // position -> one part per stage - std::vector part_next(stages.size() + 1, INT64_MAX); - // a checkpoint of the state after `cur[0, L)`; false only when the copy itself failed. `parts`: the stages' - // states saved at L (a split's mid-prompt checkpoint); without, they are read now (everything is at L) - auto checkpoint_at = [&](int64_t L, std::vector* parts = nullptr) -> bool { - if (o.prompt_cache <= 0 || L < 1) return true; - for (ConvCheckpoint& c : checks) - if ((int64_t) c.ids.size() == L) { c.used = ++check_clock; return true; } - ConvCheckpoint c; - c.ids.assign(cur.begin(), cur.begin() + L); - c.imgs = imgs_below(req_imgs, L); - if (parts != nullptr) { - if (parts->size() != stages.size() + 1) return false; - c.gdn = std::move((*parts)[0].gdn); - c.ple = std::move((*parts)[0].ple); - c.tails = std::move((*parts)[0].tails); - for (size_t i = 1; i < parts->size(); ++i) c.stage_parts.push_back(std::move((*parts)[i])); - } else { - if (cudaDeviceSynchronize() != cudaSuccess || !checkpoint_save(c, ss, g)) return false; - for (auto& st : stages) { // a layer split's later stages: their sessions' part - const strata::core::OnDevice on(st->dev); - ConvCheckpoint part; - if (cudaDeviceSynchronize() != cudaSuccess || !checkpoint_save(part, st->ss, g)) return false; - c.stage_parts.push_back(std::move(part)); - } - } - c.used = ++check_clock; - checks.push_back(std::move(c)); - while ((int) checks.size() > o.prompt_cache) { - std::vector stamps; - stamps.reserve(checks.size()); - for (const ConvCheckpoint& k : checks) stamps.push_back(k.used); - const size_t victim = strata::program::conv_cache::eviction_victim(stamps.data(), stamps.size(), - o.prompt_cache); - checks.erase(checks.begin() + (std::ptrdiff_t) victim); - } - return true; - }; - sp.on_chunk = [&](const float* R_rows, int64_t T, int64_t p0, std::string& e) -> bool { - std::vector nxt((size_t) T); - for (int64_t t = 0; t < T; ++t) nxt[(size_t) t] = (int32_t) cur[(size_t) (p0 + t + 1)]; - // E-9: batched through the prompt path when it can (one GPU: a layer split's drafter is on the last stage) - const bool batched = !multi_gpu && sp.draft_kv(mtp, R_rows, nxt.data(), T, p0, e); - if (!e.empty() || (!batched && !mtp.prefill(R_rows, nxt.data(), T, p0, e))) return false; - // progress for the server window: PP - const int64_t done = p0 + T; - const double ms = std::chrono::duration(Clock::now() - pp_t0).count(); - std::printf("PP %lld %lld %.0f %.1f\n", (long long) done, (long long) pp_total, ms, - ms > 0.0 ? 1000.0 * (double) (done - pp_from) / ms : 0.0); - strata::core::progress_at("reading the prompt (batched), done up to token", done); - strata::core::progress_beat(); - std::fflush(stdout); - if (o.prompt_cache_every > 0 && done >= pp_next_check) { - bool saved = false; - if (multi_gpu) { // the stages' parts, saved when each of them read this chunk - std::vector parts; - { - std::lock_guard lk(part_mu); - auto it = part_at.find(done); - if (it != part_at.end()) parts = std::move(it->second); - part_at.erase(part_at.begin(), part_at.upper_bound(done)); - } - bool complete = parts.size() == stages.size() + 1; - for (const ConvCheckpoint& k : parts) complete = complete && !k.gdn.empty(); - saved = !complete || checkpoint_at(done, &parts); // an incomplete set: no checkpoint here - } else { - saved = checkpoint_at(done); - } - if (!saved) { e = "saving a conversation checkpoint failed"; return false; } - pp_next_check = done + o.prompt_cache_every; - } - return true; - }; - if (multi_gpu) { // the batched prompt is reported by its last stage (the drafter's rows are there) - stages.back()->sp.on_chunk = std::move(sp.on_chunk); - sp.on_chunk = nullptr; - for (size_t i = 0; i <= stages.size(); ++i) { - strata::prefill::Prefill& stage_sp = i == 0 ? sp : stages[i - 1]->sp; - strata::core::SessionState& stage_ss = i == 0 ? ss : stages[i - 1]->ss; - stage_sp.on_stage_chunk = [&, i](int64_t done, std::string& e) -> bool { - if (o.prompt_cache <= 0 || o.prompt_cache_every <= 0 || done < part_next[i]) return true; - part_next[i] = done + o.prompt_cache_every; - ConvCheckpoint part; // this stage's state at `done` (its stream is synchronized) - if (!checkpoint_save(part, stage_ss, g)) { e = "saving a checkpoint part failed"; return false; } - std::lock_guard lk(part_mu); - auto& v = part_at[done]; - v.resize(stages.size() + 1); - v[i] = std::move(part); - return true; - }; - } - } + ver.set_profile(o.window_profile); + ver.set_pcie_mode(o.pcie_mode == "dma" ? 0 : o.pcie_mode == "direct" ? 1 : o.pcie_mode == "kernel" ? 2 + : native_pack ? 0 : 2); // auto: DMA for the native packs, the copy kernel for Q2_0 + if (drive.d.gpu2 != nullptr && gpu2.prefetch_slots() > 0) + ver.set_predict(&drive_predict, &drive, o.adapt_every > 0 && o.adapt_swaps > 0); + if (drive.d.gpu2 != nullptr) { + ver.set_gpu2(true); + ver.set_watch(&drive_watch, &drive); + } + if (drive.tier1 != nullptr || drive.tier2 != nullptr) { // the tiers' updates start at the tail + ver.set_tail(&drive_tail, &drive); + ver.set_rows_in(&drive_rows_in, &drive); + } + prefill.draft = &mtp; drive.d.plan = ver.plan_sink(); + drive.d.host_res = ver.residency(); // the snapshot the windows decide their GPU hits by drive.d.pcie_num = std::max(0, std::min(256, (int) (o.pcie_frac * 256.0 + 0.5))); + ver.set_pcie_share(drive.d.pcie_num > 0); if (o.adapt_every > 0 && o.adapt_swaps > 0) drive.d.usage.assign((size_t) (g.n_layers * g.n_expert), 0.0f); - cudaStream_t adapt_stream = nullptr; - if (cudaStreamCreateWithFlags(&adapt_stream, cudaStreamNonBlocking) != cudaSuccess) { - std::fprintf(stderr, "strata serve: cannot create the refill stream\n"); - return 1; - } - // plan v0.3 P6: swaps in flight - (residency index, slot) admitted when adapt_ev has completed - std::vector> pending; - cudaEvent_t adapt_ev = nullptr; - cudaEventCreateWithFlags(&adapt_ev, cudaEventDisableTiming); - // a layer split's later stages keep a copy of the residency table on their devices, and swap on their own - auto res_upload = [&]() { - if (d_res != nullptr) - cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice); - for (auto& st : stages) { - const strata::core::OnDevice on(st->dev); - cudaMemcpy(st->d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice); - } - }; - auto apply_pending = [&](bool wait) { - if (pending.empty()) return; - if (wait) cudaEventSynchronize(adapt_ev); - else if (cudaEventQuery(adapt_ev) != cudaSuccess) return; - for (auto& st : stages) - if (st->adapt_live) { - if (wait) cudaEventSynchronize(st->adapt_ev); - else if (cudaEventQuery(st->adapt_ev) != cudaSuccess) return; - } - for (auto& st : stages) st->adapt_live = false; - for (const auto& [i, slot] : pending) host_res[(size_t) i] = slot; - pending.clear(); - res_upload(); - }; - // the VRAM tier follows the conversation (the same rule as the speculative loop below) - auto adapt = [&]() -> bool { - if (!pending.empty()) return true; // the previous swaps are still in flight - struct Swap { float gain; int32_t layer, in, out; }; - std::vector swaps; - std::vector> cand, vict; - for (int64_t l = 0; l < g.n_layers; ++l) { - cand.clear(); - vict.clear(); - const float* u = drive.d.usage.data() + l * g.n_expert; - const int32_t* r = host_res.data() + l * g.n_expert; - for (int32_t e = 0; e < (int32_t) g.n_expert; ++e) { - if (r[e] < 0) { if (u[e] >= 2.0f) cand.emplace_back(u[e], e); } - else vict.emplace_back(u[e], e); - } - if (cand.empty() || vict.empty()) continue; - std::sort(cand.begin(), cand.end(), [](auto& a, auto& b) { return a.first > b.first; }); - const size_t nc = std::min(cand.size(), vict.size()); - std::partial_sort(vict.begin(), vict.begin() + (ptrdiff_t) nc, vict.end(), - [](auto& a, auto& b) { return a.first < b.first; }); - for (size_t i = 0; i < nc; ++i) { - if (cand[i].first < vict[i].first + 1.5f) break; - swaps.push_back({cand[i].first - vict[i].first, (int32_t) l, cand[i].second, vict[i].second}); - } - } - std::sort(swaps.begin(), swaps.end(), [](const Swap& a, const Swap& b) { return a.gain > b.gain; }); - if ((int) swaps.size() > o.adapt_swaps) swaps.resize((size_t) o.adapt_swaps); - bool main_live = false; - for (const Swap& s : swaps) { - const size_t in = (size_t) s.layer * g.n_expert + s.in, out = (size_t) s.layer * g.n_expert + s.out; - const int32_t slot = host_res[out]; - const uint8_t* b = srcp->blob(s.layer, s.in); - const int stn = multi_gpu ? stage_of(s.layer) : 0; // the swap stays in the layer's own cache - GpuStage* gs = stn > 0 ? stages[(size_t) stn - 1].get() : nullptr; - const strata::core::OnDevice on(gs ? gs->dev : -1); - if (slot < 0 || b == nullptr || - cudaMemcpyAsync(gs ? gs->cache.device_slot(slot) : xcache.device_slot(slot), b, - (size_t) strata::kernels::cpu::expert_layout().blob_bytes(s.layer), - cudaMemcpyHostToDevice, gs ? gs->adapt_stream : adapt_stream) != cudaSuccess) - return false; - if (gs) gs->adapt_live = true; - else main_live = true; - host_res[out] = strata::core::kNotResident; // evicted now: the CPU computes it meanwhile - pending.emplace_back((int32_t) in, slot); // resident once the copy has landed - } - if (!swaps.empty()) cudaEventRecord(adapt_ev, adapt_stream); - (void) main_live; - for (auto& st : stages) - if (st->adapt_live) { - const strata::core::OnDevice on(st->dev); - cudaEventRecord(st->adapt_ev, st->adapt_stream); - } - for (float& v : drive.d.usage) v *= 0.7f; - return true; - }; + // ---- the prompt cache (strata/core/prompt_cache.hpp) + std::vector kv_states; + for (int64_t i = 0; i < g.n_qsa_layers(); ++i) kv_states.push_back(&ss.qsa_states[i]); + kv_states.push_back(&mtp.kv_state()); + strata::core::PromptCache pcache(g, ss, kv_states, main_cs, o.prompt_cache, (uint64_t) o.cache_ram_mib << 20); + if (pcache.enabled()) + std::fprintf(stderr, "strata serve: prompt cache on: up to %d checkpoints of %.1f MiB, %lld MiB for stashed " + "sequences\n", o.prompt_cache, (double) pcache.ckpt_bytes() / 1048576.0, + (long long) o.cache_ram_mib); // stdin is read on its own thread, so a STOP line reaches a request that is still running (the client went // away, or pressed Esc): the flag is checked between prompt chunks and between verify windows. std::atomic stop_req{false}; @@ -3421,35 +2504,8 @@ int main(int argc, char** argv) { std::deque in_lines; bool in_eof = false; std::thread([&] { - // read(2) on the descriptor, not std::cin: glibc's exit() flushes every stdio stream and waits for - // stdin's lock, which getline holds while it waits for input - an engine ending on an error (every - // std::exit) would hang in exit() on Linux, and the server would wait for it forever - std::string l, buf; - char chunk[4096]; - auto getline_fd = [&](std::string& out) -> bool { - for (;;) { - const size_t nlpos = buf.find('\n'); - if (nlpos != std::string::npos) { - out.assign(buf, 0, nlpos); - buf.erase(0, nlpos + 1); - return true; - } -#if defined(_WIN32) - const int n = _read(0, chunk, (unsigned) sizeof chunk); -#else - const ssize_t n = ::read(0, chunk, sizeof chunk); - if (n < 0 && errno == EINTR) continue; -#endif - if (n <= 0) { - if (buf.empty()) return false; - out.swap(buf); - buf.clear(); - return true; - } - buf.append(chunk, (size_t) n); - } - }; - while (getline_fd(l)) { + std::string l; + while (std::getline(std::cin, l)) { if (!l.empty() && l.back() == '\r') l.pop_back(); if (l == "STOP") { stop_req.store(true); continue; } std::lock_guard lk(in_mu); @@ -3468,7 +2524,21 @@ int main(int argc, char** argv) { in_lines.pop_front(); return true; }; - sp.should_stop = [&] { return stop_req.load(); }; + // PP : the prompt's progress, one line per chunk (the server's progress and + // keep-alives); `pp_at` is where the running segment starts, `pp_total` 0 outside a prompt + int64_t pp_from = 0, pp_at = 0, pp_total = 0; + Clock::time_point pp_t0 = Clock::now(); + auto pp_line = [&](int64_t done) { + if (pp_total <= 0) return; + const double ms = std::chrono::duration(Clock::now() - pp_t0).count(); + std::printf("PP %lld %lld %.0f %.1f\n", (long long) (done - pp_from), (long long) pp_total, ms, + ms > 0 ? 1000.0 * (double) (done - pp_from) / ms : 0.0); + std::fflush(stdout); + }; + prefill.should_stop = [&] { + pp_line(pp_at + prefill.processed()); + return stop_req.load(); + }; // STRATA_TRACE=1: one stderr line per step of a request (the log shows where a request stops) const bool trace = std::getenv("STRATA_TRACE") != nullptr; auto tr = [&](const char* what, long long a = -1, long long b = -1) { @@ -3479,70 +2549,49 @@ int main(int argc, char** argv) { { // what is left once everything is allocated: under WDDM a GPU filled to the brim does not fail, it pages - // and a page-in while the verify graph spins on a host flag stalls the request for good - size_t free_b = 0, total_b = 0; + size_t free_b = 0, total_b = 0, free2 = 0; cudaMemGetInfo(&free_b, &total_b); - // below ~256 MiB a later allocation (a first-used window's buffers, the desktop, another program) can make - // the driver page GPU memory, and a verify graph spinning on a host flag then never finishes - const int64_t free_mib = (int64_t) (free_b >> 20); - if (free_mib >= 256) { - std::fprintf(stderr, "strata serve: %lld MiB of VRAM free with everything loaded\n", (long long) free_mib); - } else { - std::fprintf(stderr, "strata serve: %lld MiB of VRAM free with everything loaded - LOW: requests may stall;" - " add --vram-reserve-mib %lld to the config's args (or lower --max-context)\n", - (long long) free_mib, (long long) (o.vram_reserve_mib + 512 - free_mib)); - } + if (o.second_gpu >= 0) { + cudaSetDevice(o.second_gpu); + cudaMemGetInfo(&free2, &total_b); + cudaSetDevice(o.main_gpu); + } + std::fprintf(stderr, "strata serve: %lld MiB of VRAM free with everything loaded%s%s\n", + (long long) (free_b >> 20), + o.second_gpu >= 0 ? (", " + std::to_string(free2 >> 20) + " MiB on the second GPU").c_str() : "", + free_b < ((size_t) 128 << 20) + ? " - LOW: requests may stall; lower --max-context or raise --vram-reserve-mib" : ""); + } + // the penalty-history buffer: one row per window row (`penalty_rows`), allocated once for the widest window + constexpr int kPenaltyWindowCap = 4096; + constexpr size_t kHistSlots = (size_t) kPenaltyWindowCap * (size_t) strata::kernels::kVerifyMaxT; + int32_t* d_hist = nullptr; + std::vector hist_stage(kHistSlots, -1); + if (cudaMalloc(&d_hist, kHistSlots * sizeof(int32_t)) != cudaSuccess) { + std::fprintf(stderr, "strata serve: the penalty-history allocation failed\n"); + return 1; } - // what the server's Monitor tab shows (servers before 0.1.8 skip unknown lines until READY) { + // what the server's Monitor tab shows size_t free_b = 0, total_b = 0; cudaMemGetInfo(&free_b, &total_b); - std::printf("INFO context=%lld kv=%s kv_resident=%lld expert_slots=%lld expert_cache_mib=%lld spec=%d " - "mtp_max=%d lookup=%d vram_free_mib=%lld cvec=%s arena_mib=%lld pool_workers=%d pcie_frac=%.2f " - "spec_min_p=%.2f engine=" STRATA_VERSION "\n", - (long long) o.max_context, o.kv.c_str(), - (long long) (g.n_qsa_layers() > 0 && ss.qsa_states[0].kv_mode == 1 - ? ss.qsa_states[0].n_slots * 4 : 0), - (long long) xcache.slots(), (long long) (xcache.bytes() >> 20), o.spec, o.mtp_max_t, - o.suffix_draft, (long long) (free_b >> 20), cvec_summary.c_str(), + std::printf("INFO context=%lld kv=%s kv_resident=0 expert_slots=%lld expert_cache_mib=%lld spec=%d " + "mtp_max=%d lookup=%d vram_free_mib=%lld cvec=0 arena_mib=%lld pool_workers=%d pcie_frac=%.2f " + "spec_min_p=%.2f second_gpu=%d engine=" STRATA_VERSION " fork=custom\n", + (long long) o.max_context, o.kv.c_str(), (long long) xcache.slots(), + (long long) (xcache.bytes() >> 20), o.spec, o.spec, o.spec_lookup, (long long) (free_b >> 20), (long long) (strata::kernels::cpu::expert_layout().total >> 20), pool.workers(), o.pcie_frac, - o.spec_min_p); - } - // issue #29: a request whose heartbeat (tokens, prompt chunks, verify windows) stops for this long is stuck on - // a flag nobody will raise - end the engine with where it was, so the server starts it again instead of the - // GPU spinning forever. STRATA_WATCHDOG_S=0 turns it off. Issue #31: before it does, it reports what every - // part was doing (stall_report), so one occurrence says where the wait is. - { - const char* ws = std::getenv("STRATA_WATCHDOG_S"); - const int limit = ws ? std::atoi(ws) : 60; // one step (a prompt layer, a verify window) takes seconds - if (limit > 0) - std::thread([limit] { - strata::core::Progress& p = strata::core::progress(); - uint64_t last = p.beats.load(), ticks_at = p.ticks.load(); - auto since = std::chrono::steady_clock::now(); - for (;;) { - std::this_thread::sleep_for(std::chrono::seconds(1)); - const auto now = std::chrono::steady_clock::now(); - const uint64_t b = p.beats.load(); - if (!p.busy.load() || b != last) { last = b; ticks_at = p.ticks.load(); since = now; continue; } - if (now - since < std::chrono::seconds(limit)) continue; - std::fprintf(stderr, "strata serve: no progress for %d s during a request (%s %lld) - stopping " - "the engine so the server starts it again (issue #29)\n", - limit, p.where.load(), (long long) p.detail.load()); - stall_report(stderr, p.ticks.load() - ticks_at); - std::fflush(stderr); - std::abort(); - } - }).detach(); + o.spec_min_p, o.second_gpu); } std::printf("READY %lld stop\n", (long long) o.max_context); // "stop": this engine honours STOP std::fflush(stdout); std::string line; int64_t rounds = 0; - const int S = o.spec; - const int S_mtp = o.mtp_max_t > 0 ? std::min(o.mtp_max_t, S) : S; // the MTP's windows; suffixes go up to S - if (S_mtp < S) mtp.set_max_drafts(S_mtp - 1); - strata::spec::SuffixDrafter sfx(std::max(1, o.suffix_draft), 64, (size_t) o.max_context + 4096); - strata::spec::DraftPolicy policy(S); // MTP or lookup window, learned over the whole process + const int S = o.spec, W = o.max_window(); // the draft layer's windows; with prompt lookup's + // prompt lookup: the text so far, kept across requests (a conversation's next request shares its prefix) + strata::spec::SuffixDrafter lookup(o.spec_lookup, 32, o.spec_lookup > 0 ? (size_t) o.max_context + 16 : 16); + std::vector ids32, ldrafts((size_t) W, 0); + constexpr int64_t kImStart = 248045; // <|im_start|>: a new turn begins here // The vision path (--vision): GENI carries images. The file is one // or more strata-vision records (int32 'SVE1', n, nx, ny, n_embd, then n x n_embd floats) in prompt order; // each image's rows go to its run of <|image_pad|> tokens, whose M-RoPE positions are mtmd's: t = p, @@ -3553,11 +2602,6 @@ int main(int argc, char** argv) { std::vector row_ptr; while (next_line(line)) { if (line == "QUIT") break; - // the watchdog watches a request from here until this iteration ends, whichever way it ends - struct BusyScope { - BusyScope() { strata::core::progress().busy.store(true); strata::core::progress_at("request"); } - ~BusyScope() { strata::core::progress().busy.store(false); strata::core::progress_at("idle"); } - } busy_scope; stop_req.store(false); // a STOP that arrived between requests is stale const bool geni = line.rfind("GENI ", 0) == 0; if (!geni && line.rfind("GEN ", 0) != 0) { @@ -3566,43 +2610,37 @@ int main(int argc, char** argv) { } char* endp = nullptr; const long long max_new = std::strtoll(line.c_str() + (geni ? 5 : 4), &endp, 10); - // optional sampling keys between max_new and the ids: temperature=F, top_p=F, top_k=N, min_p=F, - // penalty_last_n=N, penalty_repeat=F, penalty_freq=F, penalty_present=F, seed=N (text requests - // only). Absent keys keep today's behavior: greedy, no penalties. - float req_temperature = 0.0f, req_top_p = 1.0f; - int req_top_k = 20; // the sampler's own default; the sampled path REQUIRES top_k in 1..64 + // optional keys between max_new and the ids: temperature=F, top_p=F, top_k=N, min_p=F, penalty_last_n=N, + // penalty_repeat=F, penalty_freq=F, penalty_present=F, seed=N, spec_min_p=F (this request's draft floor); + // unknown keys (cvec, pcie_frac) are skipped. Absent keys: greedy, no penalties. + float req_temperature = 0.0f, req_top_p = 1.0f, req_min_p = 0.0f; + int req_top_k = 20; // the sampled path keeps 1..64 candidates unsigned long long req_seed = 0; - float req_min_p = 0.0f, req_penalty_repeat = 1.0f, req_penalty_freq = 0.0f, req_penalty_present = 0.0f; + float req_penalty_repeat = 1.0f, req_penalty_freq = 0.0f, req_penalty_present = 0.0f; int req_penalty_last_n = 0; - int req_cvec = 1; // cvec=0|1: a loaded control vector for this request (on when absent) - // tuning keys (setup's calibration measures settings without restarting the engine): the PCIe share of - // the missed experts and the draft-probability floor, for this request only - double req_pcie_frac = o.pcie_frac, req_spec_min_p = o.spec_min_p; - if (endp != nullptr) { // GENI takes the same keys (#75: image requests were always greedy); its - // embedding file path is the first token without an = + double req_spec_min_p = o.spec_min_p; + if (endp != nullptr) { for (;;) { while (*endp == ' ') ++endp; const char* start = endp; while (*endp != '\0' && *endp != ' ') ++endp; if (endp == start) break; - const std::string tok(start, (size_t) (endp - start)); - const size_t eq = tok.find('='); + const std::string kv(start, (size_t) (endp - start)); + const size_t eq = kv.find('='); if (eq == std::string::npos) { endp = const_cast(start); break; } - const std::string key = tok.substr(0, eq); - const float fv = std::strtof(tok.c_str() + eq + 1, nullptr); - if (key == "cvec") req_cvec = std::atoi(tok.c_str() + eq + 1); - else if (key == "temperature") req_temperature = fv; + const std::string key = kv.substr(0, eq); + const char* v = kv.c_str() + eq + 1; + const float fv = std::strtof(v, nullptr); + if (key == "temperature") req_temperature = fv; else if (key == "top_p") req_top_p = fv; - else if (key == "top_k") req_top_k = std::atoi(tok.c_str() + eq + 1); + else if (key == "top_k") req_top_k = std::atoi(v); else if (key == "min_p") req_min_p = fv; - else if (key == "penalty_last_n") req_penalty_last_n = std::atoi(tok.c_str() + eq + 1); + else if (key == "penalty_last_n") req_penalty_last_n = std::atoi(v); else if (key == "penalty_repeat") req_penalty_repeat = fv; else if (key == "penalty_freq") req_penalty_freq = fv; else if (key == "penalty_present") req_penalty_present = fv; - else if (key == "seed") req_seed = std::strtoull(tok.c_str() + eq + 1, nullptr, 10); - else if (key == "pcie_frac") req_pcie_frac = std::clamp((double) fv, 0.0, 1.0); + else if (key == "seed") req_seed = std::strtoull(v, nullptr, 10); else if (key == "spec_min_p") req_spec_min_p = std::clamp((double) fv, 0.0, 1.0); - // unknown keys are skipped: the ids start at the first token without '=' } } std::string emb_path; @@ -3618,7 +2656,6 @@ int main(int argc, char** argv) { continue; } const int64_t n = (int64_t) ids.size(); - req_imgs.clear(); if (geni && !o.vision) { std::printf("ERR this engine was started without --vision\n"); continue; } if (geni || !mrope_identity) { // positions for every cell this request can reach; the identity again for a text request @@ -3660,12 +2697,6 @@ int main(int argc, char** argv) { if (ids[(size_t) i] != kImagePad) { put(i, p, p, p); ++p; ++i; continue; } if (k >= imgs.size()) { ve = "the prompt has more images than the embeddings file"; break; } const Img& im = imgs[k++]; - { // what the conversation cache compares: a picture is its grid and its embeddings - const int64_t grid[3] = {im.n, im.nx, im.ny}; - uint64_t h = fnv1a(grid, sizeof grid); - h = fnv1a(img_rows.data() + im.off, (size_t) im.n * (size_t) g.n_embd * sizeof(float), h); - req_imgs.push_back({i, h}); - } for (int64_t j = 0; j < im.n && ve.empty(); ++j) if (i + j >= n || ids[(size_t) (i + j)] != kImagePad) ve = "image " + std::to_string(k) + " has " + std::to_string((long long) im.n) + @@ -3685,23 +2716,13 @@ int main(int argc, char** argv) { tr("positions built", (long long) img_rows.size()); cudaDeviceSynchronize(); tr("device idle"); - // CUDA0's table and, with a layer split, every later stage's (each device reads its own) - auto upload_mrope = [&]() -> bool { - bool ok = cudaMemcpy(d_mrope, mrope_host.data(), mrope_host.size() * sizeof(int32_t), - cudaMemcpyHostToDevice) == cudaSuccess; - for (auto& st : stages) { - const strata::core::OnDevice on(st->dev); - cudaDeviceSynchronize(); - ok = ok && cudaMemcpy(st->mrope, mrope_host.data(), mrope_host.size() * sizeof(int32_t), - cudaMemcpyHostToDevice) == cudaSuccess; - } - return ok; - }; - if (ve.empty() && !upload_mrope()) ve = "the image position upload failed"; + if (ve.empty() && cudaMemcpy(d_mrope, mrope_host.data(), mrope_host.size() * sizeof(int32_t), + cudaMemcpyHostToDevice) != cudaSuccess) + ve = "the image position upload failed"; if (!ve.empty()) { // leave the table as the identity so the next text request is untouched for (int64_t c = 0; c < cells; ++c) put(c, c, c, c); - upload_mrope(); + cudaMemcpy(d_mrope, mrope_host.data(), mrope_host.size() * sizeof(int32_t), cudaMemcpyHostToDevice); mrope_identity = true; std::printf("ERR %s\n", ve.c_str()); std::fflush(stdout); @@ -3709,7 +2730,7 @@ int main(int argc, char** argv) { } mrope_identity = !geni; } - sp.embd_rows = geni ? row_ptr.data() : nullptr; + prefill.embd_rows = geni ? row_ptr.data() : nullptr; if (n + max_new + 8 > o.max_context) { std::printf("ERR prompt (%lld tokens) + max_new (%lld) exceeds the context (%lld)\n", (long long) n, (long long) max_new, (long long) o.max_context); @@ -3718,218 +2739,6 @@ int main(int argc, char** argv) { bool bad = false; for (int64_t t : ids) bad = bad || t < 0 || t >= n_vocab; if (bad) { std::printf("ERR a token id is outside the vocabulary\n"); continue; } - std::array remote_before{}; - std::array launches_before{}; - std::array compact_before{}, full_before{}; - for (int r = 0; r < 3; ++r) if (o.expert_cache_remote[(size_t) r] > 0) - { - remote_before[(size_t) r] = remote_experts[(size_t) r].computed(); - launches_before[(size_t) r] = remote_experts[(size_t) r].launched_layers(); - compact_before[(size_t) r] = remote_experts[(size_t) r].returned_bytes(); - full_before[(size_t) r] = remote_experts[(size_t) r].full_row_bytes(); - } - cur = ids; - const Clock::time_point r0 = Clock::now(); - // ---- where this request starts reading: the live session, or a checkpoint, whose tokens AND pictures are - // exactly the start of this prompt - at most n - 1 of them, the last token is always the first window - auto starts_with = [&](const std::vector& pre, const std::vector& pre_imgs) -> bool { - const int64_t L = (int64_t) pre.size(); - if (L < 1 || L > n - 1) return false; - for (int64_t i = 0; i < L; ++i) - if ((int32_t) ids[(size_t) i] != pre[(size_t) i]) return false; - return imgs_below(req_imgs, L) == pre_imgs; - }; - // the control vector for this request. The live session and the checkpoints were read one way, so a - // switch reads the prompt again from the start - if (strata::kernels::cvec().loaded()) { - const bool want = req_cvec != 0; - if (want != cvec_cached) { - live_ok = false; - checks.clear(); - cvec_cached = want; - } - strata::kernels::cvec_set_enabled(want); - } - int64_t resume = 0; - bool from_live = false; - if (o.prompt_cache > 0) { - if (live_ok && starts_with(live, live_imgs)) { resume = (int64_t) live.size(); from_live = true; } - for (const ConvCheckpoint& c : checks) - if ((int64_t) c.ids.size() > resume && starts_with(c.ids, c.imgs)) { - resume = (int64_t) c.ids.size(); - from_live = false; - } - } - // this request rewrites every cell from `resume` on, so a checkpoint past it (or not on this prompt's - // path) no longer has its cells; the ones kept are prefixes of both the old tokens and the new - checks.erase(std::remove_if(checks.begin(), checks.end(), [&](const ConvCheckpoint& c) { - return (int64_t) c.ids.size() > resume || !starts_with(c.ids, c.imgs); - }), checks.end()); - live_ok = false; // until this request has finished, the session is in between - int64_t reread_to = -1; // STRATA_CKPT_REREAD only: read [0, reread_to) again instead of restoring - if (resume == 0) { - strata::core::session_zero(ss, g, nullptr, main_cs); - cudaStreamSynchronize(main_stream); - for (auto& st : stages) { - const strata::core::OnDevice on(st->dev); - strata::core::session_zero(st->ss, g, nullptr, (void*) st->stream); - cudaStreamSynchronize(st->stream); - } - checks.clear(); - } else if (!from_live) { - ConvCheckpoint* c = nullptr; - for (ConvCheckpoint& k : checks) if ((int64_t) k.ids.size() == resume) c = &k; - if (c != nullptr) c->used = ++check_clock; // mounting through it is the use LRU counts - static const bool reread = std::getenv("STRATA_CKPT_REREAD") != nullptr; - if (reread && c != nullptr) { - // THE CHECK OF THE CHECKPOINT: instead of restoring it, read its tokens again from position 0 in - // one run (below, with the prompt path's slots lent like any read) - the same chunks the request - // that saved it read them in, when that request started at 0. With the VRAM expert set fixed - // (--adapt-swaps 0) the answer must match the restored one token for token; anything the - // checkpoint missed shows up as a difference. - strata::core::session_zero(ss, g, nullptr, main_cs); - cudaStreamSynchronize(main_stream); - for (auto& st : stages) { - const strata::core::OnDevice on(st->dev); - strata::core::session_zero(st->ss, g, nullptr, (void*) st->stream); - cudaStreamSynchronize(st->stream); - } - reread_to = resume; - std::fprintf(stderr, "strata serve: STRATA_CKPT_REREAD: reading %lld tokens again instead of " - "restoring\n", (long long) resume); - } else if (c == nullptr || !checkpoint_restore(*c, ss, g) || c->stage_parts.size() != stages.size() || - [&] { - for (size_t i = 0; i < stages.size(); ++i) { - const strata::core::OnDevice on(stages[i]->dev); - if (!checkpoint_restore(c->stage_parts[i], stages[i]->ss, g)) return true; - } - return false; - }()) { - std::printf("ERR restoring a conversation checkpoint failed\n"); - return 1; - } - } - // KV streaming: the drafter's ring may hold cells past `resume` from a longer turn; the main layers' - // host copies and slots are always current (every writer writes both), so they need nothing - if (resume > 0 && reread_to <= 0) mtp.kv_restore(resume); - tr("request", n, geni ? 1 : 0); - mtp.set_prompt_len(n); - const int64_t read_from = reread_to > 0 ? 0 : resume; - pp_total = n; - pp_from = read_from; - pp_t0 = r0; - pp_next_check = reread_to > 0 ? INT64_MAX : resume + o.prompt_cache_every; - { - std::lock_guard lk(part_mu); - part_at.clear(); - std::fill(part_next.begin(), part_next.end(), pp_next_check); - } - std::printf("RESUME %lld\n", (long long) resume); // before reading: this many prompt tokens are reused - strata::core::progress_at("reading the prompt, from token", read_from); - std::fflush(stdout); - // A SHORT PART OF THE PROMPT - the new message of a chat that continues from a checkpoint, the assistant - // header - goes through the verify windows, S tokens at a time, as decode reads them. The batched path - // costs ~300 ms per run however few tokens it has (it streams every expert the chunk routes to that is - // not in VRAM over PCIe), and it borrows slots it must refill after (~180 ms); a window costs ~16 ms a - // token, with the misses on the CPU. Each part below is decided on its own, so a long first message is - // read batched and its header still goes through the windows. Picture rows need the batched path. - // STRATA_CKPT_REREAD compares a restored checkpoint with a batched re-read, so it keeps every read batched. - static const bool no_short = std::getenv("STRATA_CKPT_REREAD") != nullptr; - auto windows_ok = [&](int64_t a, int64_t b) -> bool { - if (no_short || b - a > o.short_read) return false; - if (sp.embd_rows != nullptr) - for (int64_t i = a; i < b; ++i) - if (sp.embd_rows[i] != nullptr) return false; - return true; - }; - // tokens [a, b) through the windows: commit all of them, then give the draft layer their residuals - auto read_windows = [&](int64_t a, int64_t b, std::string& e) -> bool { - // every token is committed and the picks are discarded: no head sampling (see set_head_sampling) - struct NoHeadSampling { - strata::core::Verifier& v; - explicit NoHeadSampling(strata::core::Verifier& x) : v(x) { v.set_head_sampling(false); } - ~NoHeadSampling() { v.set_head_sampling(true); } - } no_head_sampling(ver); - std::vector win((size_t) S), outw((size_t) S), nxt((size_t) S); - for (int64_t q = a; q < b;) { - if (stop_req.load()) { e = "cancelled"; return false; } - const int T = (int) std::min(S, b - q); - for (int t = 0; t < T; ++t) { - win[(size_t) t] = (int32_t) cur[(size_t) (q + t)]; - nxt[(size_t) t] = (int32_t) cur[(size_t) (q + t + 1)]; - } - drive.d.layers = 0; - drive.d.experts = 0; - drive.d.failed = false; - if (!ver.run(T, win.data(), q, win_pool_fn, win_pool_user, outw.data(), e) || drive.d.failed) { - if (drive.d.failed && drive.d.fail) e = drive.d.fail; - return false; - } - if (!ver.commit(T, e) || !mtp.prefill(ver.final_R_all(), nxt.data(), T, q, e)) return false; - q += T; - } - const double ms = std::chrono::duration(Clock::now() - pp_t0).count(); - std::printf("PP %lld %lld %.0f %.1f\n", (long long) b, (long long) pp_total, ms, - ms > 0.0 ? 1000.0 * (double) (b - pp_from) / ms : 0.0); - strata::core::progress_beat(); - std::fflush(stdout); - return true; - }; - // the batched path's slots are lent just before its first run and given back (refilled) before a window - // reads - so the windows always see the whole expert cache - or once the prompt is read - std::vector> lent_now; - int64_t lent_chunk = 0; // the chunk the lent slots hold the prompt path's buffers for - auto refill = [&](std::string& e) -> bool { - if (lent_now.empty()) return true; - tr("refill start", (long long) lent_now.size()); - for (const auto& [i, slot] : lent_now) { // D-4: queued, one wait (STRATA_REFILL_BLOCKING=1: each) - const uint8_t* b = srcp->blob(i / g.n_expert, i % g.n_expert); - const int64_t nb = (int64_t) strata::kernels::cpu::expert_layout().blob_bytes(i / g.n_expert); - if (b == nullptr || !(refill_blocking() ? xcache.fill_slot_blocking(slot, b, e, nb) - : xcache.fill_slot_queued(slot, b, e, nb))) - return false; - host_res[(size_t) i] = slot; - } - if (!xcache.sync_queued(e)) return false; - cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice); - lent_now.clear(); - lent_chunk = 0; - return true; - }; - // lend the slots `tokens` batched prompt tokens need: the prompt path's buffers for min(chunk, tokens - // rounded up to 256), laid out in the last of the slots it may borrow - auto lend = [&](int64_t tokens, std::string& e) -> bool { - if (lend_first < 0) return true; // its own buffers: nothing to lend - const int64_t want = request_chunk(tokens, o.prefill_chunk); - if (want <= 0) { - e = "prefill: cannot lend buffers for an empty request segment"; - return false; - } - if (!lent_now.empty()) { - if (want <= lent_chunk) return true; - if (!refill(e)) return false; - } - const int64_t slots = lend_slots(want); - if (slots <= 0 || slots > xcache.slots() - lend_first) { - e = "prefill: request-sized buffers exceed the configured lend region"; - return false; - } - const int32_t first = std::max(lend_first, (int32_t) (xcache.slots() - slots)); - if (want != sp.chunk() || first != lend_first_now) { - if (!sp.relayout(want, xcache.device_slot(first), lend_bytes(first), e)) return false; - lend_first_now = first; - } - for (size_t i = 0; i < host_res.size(); ++i) - if (host_res[i] >= first) { - lent_now.emplace_back((int32_t) i, host_res[i]); - host_res[i] = strata::core::kNotResident; - } - cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice); - lent_chunk = want; - return true; - }; - apply_pending(true); - // per-request sampling for the verify window's head (greedy when temperature is absent) strata::kernels::SamplerParams req_sp; req_sp.greedy = req_temperature <= 0.0f; req_sp.temperature = req_temperature; @@ -3942,200 +2751,215 @@ int main(int argc, char** argv) { req_sp.penalty_repeat = req_penalty_repeat; req_sp.penalty_freq = req_penalty_freq; req_sp.penalty_present = req_penalty_present; - req_sp.counter = 0; ver.set_sampling(req_sp); - drive.d.pcie_num = std::max(0, std::min(256, (int) (req_pcie_frac * 256.0 + 0.5))); - // a layer split: CUDA0's share as asked; a later GPU keeps its own (its link) unless the request sets one - for (int st = 0; st < split_drive.n; ++st) - split_drive.pcie_num[st] = (st == 0 || split_same || req_pcie_frac != o.pcie_frac) - ? drive.d.pcie_num : pcie_num_of(stages[(size_t) st - 1]->pcie_frac); const int hist_n = std::min(req_sp.penalty_last_n, kPenaltyWindowCap); ver.set_history(hist_n > 0 ? d_hist : nullptr, hist_n); - bool cancelled = false; - tr("prompt start", n - 1); - // The prompt is read in two parts when it has a turn boundary past `resume`: up to the last <|im_start|> - // (the conversation so far), a checkpoint there, then the new turn's header. The next request of the same - // chat renders the same history - but not always the same header or the thinking of this reply - so that - // checkpoint is the one it reuses. - int64_t turn_at = -1; - if (o.prompt_cache > 0 && o.turn_token >= 0) - for (int64_t i = n - 1; i > resume; --i) - if (ids[(size_t) i] == o.turn_token) { turn_at = i; break; } - // A prompt read from token 0 also stops at its FIRST turn boundary: the end of the system prompt (with - // the tools), which every new chat of the same client shares. That checkpoint becomes the chain's root, - // which the retention policy pins (conv_cache.hpp), so the next new chat reads only what comes after it. - // (PR #65, code-martin.) Only for a system prompt of --prompt-cache-root tokens or more: a small one - // is cheaper to read again than the extra part costs (~0.3 s). - int64_t root_at = -1; - if (o.prompt_cache > 0 && o.turn_token >= 0 && o.prompt_cache_root > 0 && read_from == 0) - for (int64_t i = 1; i < turn_at; ++i) - if (ids[(size_t) i] == o.turn_token) { - if (i >= o.prompt_cache_root) root_at = i; + const Clock::time_point r0 = Clock::now(); + auto report = [&](const std::string& e) { + std::fprintf(stderr, "strata serve: %s\n", e.c_str()); + std::printf("ERR %s\n", e.c_str()); + }; + // ---- the prompt cache: continue from the longest prefix it holds (or start from an empty sequence). + // Image requests are not cached: their cells are not identified by the token ids alone. + const bool track = pcache.enabled() && !geni; + const int64_t reuse = pcache.begin(ids, !geni); + tr("request", n, reuse); + mtp.set_prompt_len(n); + std::vector drafts((size_t) W, 0), window((size_t) W), outv((size_t) W); + std::vector dprob((size_t) W, 0.0f); + // ---- the rest of the prompt, [reuse, n - 1). A turn usually shares everything before its prompt's + // last <|im_start|> with the next one (chat templates re-render the answer, e.g. without its + // reasoning), so that position gets a checkpoint, and so do the prompt's end (a regenerated answer) + // and the end of every `step` cells of a long prompt. Short parts go through verify windows. + int64_t q = 0; + if (track) + for (int64_t i = n - 2; i > reuse; --i) + if (ids[(size_t) i] == kImStart) { q = i; break; } + int64_t batched_end = reuse; // [reuse, batched_end) through the batched prompt path + if (geni) batched_end = n - 1; // image rows reach the batched path only + else if (n - 1 - reuse > o.feed_max) batched_end = (q > reuse && n - 1 - q <= o.feed_max) ? q : n - 1; + // Every batched chunk streams each expert its tokens use, however few tokens it holds, so a short last + // chunk goes through verify windows too. Chunks and checkpoint segments start at `reuse`. + const int64_t step = std::max(1, o.cache_every / o.prefill_chunk) * o.prefill_chunk; + if (!geni && batched_end > reuse) { + const int64_t last = reuse + (batched_end - 1 - reuse) / o.prefill_chunk * o.prefill_chunk; + if (last > reuse && batched_end - last <= o.feed_max) batched_end = last; + } + if (batched_end > reuse) { + tier.apply_pending(true); // no copy may still land in a lent slot + tier2.apply_pending(true); + if (!pbuf.take(std::min(o.prefill_chunk, batched_end - reuse), err)) { report(err); return 1; } + tr("prompt start", reuse, batched_end); + pp_from = reuse; + pp_total = n - 1 - reuse; + pp_t0 = r0; + for (int64_t a0 = reuse; a0 < batched_end;) { + const int64_t b0 = track ? std::min(batched_end, a0 + step) : batched_end; + pp_at = a0; + if (!prefill.run(ids.data() + a0, b0 - a0, a0, err)) { + if (!stop_req.load()) { report(err); return 1; } + // STOP: the session holds the chunks read before it + if (track && prefill.processed() > 0) pcache.set(a0, ids.data() + a0, prefill.processed()); break; } - int64_t at = read_from; - for (const int64_t to : {reread_to, root_at, turn_at, n - 1}) { - if (to <= at) continue; - const bool win = windows_ok(at, to); - if (win && !refill(err)) { - std::printf("ERR refilling a lent slot failed: %s\n", err.c_str()); - return 1; + if (track) { + pcache.set(a0, ids.data() + a0, b0 - a0); + pcache.checkpoint(b0); + } + a0 = b0; } - if (!win && !lend(to - at, err)) { - std::printf("ERR lending the prompt path its slots failed: %s\n", err.c_str()); + if (!pbuf.give_back(*srcp, err)) { + report("refilling a lent slot failed: " + err); return 1; } - const auto tsp = Clock::now(); - const bool sp_ok = win ? read_windows(at, to, err) : sp.run(ids.data() + at, to - at, at, err); - if (trace) { - std::fprintf(stderr, "strata trace: read %lld tokens (%s) in %.1f ms\n", (long long) (to - at), - win ? "windows" : "batched", - std::chrono::duration(Clock::now() - tsp).count()); - std::fflush(stderr); - } - if (!sp_ok) { - if (!stop_req.load()) { - std::fprintf(stderr, "strata serve: %s\n", err.c_str()); - std::printf("ERR %s\n", err.c_str()); - return 1; + tr("prompt done (slots refilled)"); + } + // prompt tokens through verify windows as wide as the verifier's: each window is committed whole, and the + // draft layer catches up on the same cells (its drafts are not used, so its chain stops at the first) + auto feed = [&](int64_t from, int64_t to) -> bool { + mtp.set_max_drafts(1); + for (int64_t j = from; j < to;) { + const int T = (int) std::min(W, to - j); + for (int i = 0; i < T; ++i) window[(size_t) i] = (int32_t) ids[(size_t) (j + i)]; + drive.d.layers = 0; + drive.d.experts = 0; + drive.d.failed = false; + drive.adapt_next = drive.flush_next = false; + drive_window(&drive, T); + if (!ver.run(T, window.data(), j, &drive_pool_multi, &drive, outv.data(), err) || drive.d.failed) { + if (drive.d.failed && drive.d.fail) err = drive.d.fail; + return false; } - cancelled = true; // stopped while reading the prompt: refill the lent slots below, then DONE cancel - break; + if (!ver.commit(T, err, false)) return false; // beside the draft + for (int i = 0; i < T; ++i) outv[(size_t) i] = (int32_t) ids[(size_t) (j + i + 1)]; + if (!mtp.draft(T, outv.data(), j, T - 1, drafts.data(), err, dprob.data(), (float) o.spec_min_p) || + !ver.wait_commit(err)) + return false; + if (track) pcache.set(j, ids.data() + j, T); + j += T; } - at = to; - if ((to == turn_at || to == root_at) && !checkpoint_at(to)) { - std::printf("ERR saving a conversation checkpoint failed\n"); - return 1; + return true; + }; + if (!stop_req.load()) { + struct NoHeadSampling { // every token is committed and the picks are discarded + strata::core::Verifier& v; + explicit NoHeadSampling(strata::core::Verifier& x) : v(x) { v.set_head_sampling(false); } + ~NoHeadSampling() { v.set_head_sampling(true); } + } no_head_sampling(ver); + if (pp_total == 0) { + pp_from = reuse; + pp_total = n - 1 - reuse; + pp_t0 = r0; } - } - if (!refill(err)) { - std::printf("ERR refilling a lent slot failed: %s\n", err.c_str()); - return 1; - } - tr("prompt done (slots refilled)"); + const int64_t split = q > batched_end ? q : batched_end; + if (!feed(batched_end, split)) { report(err); return 1; } + if (split > batched_end) pp_line(split); + if (track && split == q) pcache.checkpoint(q); + if (!feed(split, n - 1)) { report(err); return 1; } + if (n - 1 > split) pp_line(n - 1); + if (track) pcache.checkpoint(n - 1); + } + pp_total = 0; const double prompt_ms = std::chrono::duration(Clock::now() - r0).count(); - std::printf("REUSED %lld\n", (long long) resume); // the prompt is read; the first window comes next - std::fflush(stdout); // the verify windows: the first holds the last prompt token alone int64_t p = n - 1; int32_t x = (int32_t) ids[(size_t) (n - 1)]; - std::vector drafts((size_t) S, 0), window((size_t) S), outv((size_t) S); - std::vector dprob((size_t) S, 0.0f); - std::vector sbuf((size_t) S, 0); - if (o.suffix_draft > 0) { - sfx.reset(); - for (int64_t t : ids) sfx.append((int32_t) t); - } bool first_window = true; - int64_t produced_n = 0, sfx_windows = 0, sfx_drafts = 0, sfx_ok = 0; - int64_t draft_offered = 0, draft_accepted = 0; - // what the session holds once this request is done: the prompt read so far, then every committed token - std::vector consumed; - consumed.reserve((size_t) (n + max_new + S)); - for (int64_t i = 0; i < n - 1; ++i) consumed.push_back((int32_t) ids[(size_t) i]); + int64_t produced_n = 0; const char* finish = "length"; const Clock::time_point d0 = Clock::now(); - // STRATA_DECODE_TIMING=1: where a request's decode time goes (one line per request) - static const bool dec_timing = std::getenv("STRATA_DECODE_TIMING") != nullptr; - struct DecSnap { - double wait, pool, host, plan, actq, jobs, run; - int64_t misses, entries, hits, pcie; - }; - auto dec_snap = [&]() { - return DecSnap{ver.ms_wait, ver.ms_pool, ver.ms_host, drive.d.ms_plan, drive.d.ms_actq, drive.d.ms_jobs, - drive.d.ms_run, drive.d.multi_misses, drive.d.multi_entries, drive.d.cache_hits, - drive.d.pcie_experts}; + int64_t draft_offered = 0, draft_accepted = 0; + // the expert-cache hits of the decode: rows either GPU computes from its VRAM cache, of all routed rows + const int64_t hits0 = drive.d.cache_hits + drive.d.gpu2_entries, look0 = hits0 + drive.d.cache_refused; + std::vector consumed; // what the state has read, for the penalties + if (hist_n > 0) + for (int64_t i = 0; i + 1 < n; ++i) consumed.push_back((int32_t) ids[(size_t) i]); + int lk = 0; // the next window's drafts from the lookup (0: the draft layer's) + if (o.spec_lookup > 0) { + ids32.resize((size_t) n); + for (int64_t i = 0; i < n; ++i) ids32[(size_t) i] = (int32_t) ids[(size_t) i]; + lookup.assign(ids32.data(), ids32.size()); + } + mtp.on_draft = [&](int j, int32_t tok, float prob) { // the drafts the next window will verify + if (lk == 0 && prob >= (float) req_spec_min_p) ver.ple_ahead(j + 1, tok); }; - const DecSnap ds0 = dec_snap(); - double dt_run = 0, dt_commit = 0, dt_draft = 0; - int64_t dec_windows = 0, dec_T = 0; - const int64_t decode_hits0 = drive.d.cache_hits; - const int64_t decode_look0 = drive.d.cache_hits + drive.d.cache_admitted + drive.d.cache_refused; - if (cancelled) finish = "cancel"; - while (!cancelled && produced_n < max_new) { - int T = S_mtp; + while (produced_n < max_new) { + if (stop_req.load()) { finish = "cancel"; break; } // the client went away + int T = S; if (req_spec_min_p > 0.0) { T = 1; - while (T < S_mtp && dprob[(size_t) T - 1] >= (float) req_spec_min_p) ++T; + while (T < S && dprob[(size_t) T - 1] >= (float) req_spec_min_p) ++T; } + if (lk > 0) T = 1 + lk; if (first_window) T = 1; - // a repeat of earlier context (prompt lookup) where the MTP's own first guess agrees: the policy takes it - // when its expected tokens per ms, from the measured acceptance and window costs, beat the MTP window's - bool from_sfx = false; - int sfx_match = 0; - if (o.suffix_draft > 0 && !first_window) { - const int k = sfx.propose(S - 1, sbuf.data()); - sfx_match = sfx.last_match(); - if (k > 0 && sbuf[0] == drafts[0]) { - const strata::spec::DraftPolicy::Pick pk = policy.choose(T, k, sfx_match); - if (pk.lookup) { T = pk.t; from_sfx = true; } + const std::vector& dr = lk > 0 ? ldrafts : drafts; + // Never verify past what this request emits: the committed cells then hold exactly the tokens the + // client saw, so its next turn continues from them. So the window stops at the output limit and + // at a drafted end-of-turn token. + T = (int) std::min(T, max_new - produced_n); + for (int i = 1; i < T; ++i) + if (std::find(o.eos_ids.begin(), o.eos_ids.end(), (int64_t) dr[(size_t) i - 1]) != + o.eos_ids.end()) { + T = i + 1; + break; } - } - const bool timed_round = !first_window; - const Clock::time_point round0 = Clock::now(); if (p + T > o.max_context) break; window[0] = x; - for (int i = 1; i < T; ++i) window[(size_t) i] = from_sfx ? sbuf[(size_t) i - 1] : drafts[(size_t) i - 1]; - drive.d.layers = 0; - drive.d.experts = 0; - drive.d.failed = false; - apply_pending(false); + for (int i = 1; i < T; ++i) window[(size_t) i] = dr[(size_t) i - 1]; if (hist_n > 0) { - // the tails the penalties count over, ONE PER ROW: the tokens the state has consumed, the - // fed-back head `x` (it joins `consumed` only after this window commits), then the drafts - // before that row - what plain decode would have counted there. (Until 0.1.19 only row 0 - // was staged, and the drafted rows read unwritten slots.) - strata::kernels::penalty_rows(consumed.data(), (int64_t) consumed.size(), window.data(), T, - hist_n, hist_stage.data()); - const strata::core::OnDevice on_h(hist_dev); + // row t counts what plain decode would have: the state's tokens, then window[0..t] + strata::kernels::penalty_rows(consumed.data(), (int64_t) consumed.size(), window.data(), T, hist_n, + hist_stage.data()); cudaMemcpy(d_hist, hist_stage.data(), (size_t) T * (size_t) hist_n * sizeof(int32_t), cudaMemcpyHostToDevice); } + drive.d.layers = 0; + drive.d.experts = 0; + drive.d.failed = false; + drive_window(&drive, T); tr("window", p, T); - const Clock::time_point tw0 = Clock::now(); - if (!ver.run(T, window.data(), p, win_pool_fn, win_pool_user, outv.data(), err) || drive.d.failed) { + drive.adapt_next = !drive.d.usage.empty() && ((rounds + 1) % o.adapt_every) == 0; + drive.flush_next = !drive.d.usage.empty() && ((rounds + 2) % o.adapt_every) == 0; + if (!ver.run(T, window.data(), p, &drive_pool_multi, &drive, outv.data(), err) || drive.d.failed) { + drive.join_adapt(); std::printf("ERR %s\n", drive.d.failed && drive.d.fail ? drive.d.fail : err.c_str()); return 1; } int a = 0; while (a < T - 1 && window[(size_t) a + 1] == outv[(size_t) a]) ++a; - if (from_sfx) { ++sfx_windows; sfx_drafts += T - 1; sfx_ok += a; } - const Clock::time_point tw1 = Clock::now(); - std::thread adapt_thr; // the adaptive tier beside the commit and the draft (as in generate) - bool adapt_ok = true; - if (!drive.d.usage.empty() && ((rounds + 1) % o.adapt_every) == 0) - adapt_thr = std::thread([&] { adapt_ok = adapt(); }); - if (!ver.commit(a + 1, err)) { - if (adapt_thr.joinable()) adapt_thr.join(); + if (!ver.commit(a + 1, err, false)) { // beside the draft + drive.join_adapt(); std::printf("ERR %s\n", err.c_str()); return 1; } - // the window's first a + 1 tokens are in the session now (the last output is not: it is next x) - for (int i = 0; i <= a; ++i) consumed.push_back(window[(size_t) i]); + if (track) pcache.set(p, window.data(), a + 1); // the committed cells p .. p + a + if (hist_n > 0) consumed.insert(consumed.end(), window.begin(), window.begin() + a + 1); draft_offered += T - 1; draft_accepted += a; + ver.ple_ahead(0, outv[(size_t) a]); first_window = false; bool eos = false; for (int i = 0; i <= a && produced_n < max_new && !eos; ++i) { std::printf("T %d\n", (int) outv[(size_t) i]); - strata::core::progress_beat(); + if (o.spec_lookup > 0) lookup.append(outv[(size_t) i]); ++produced_n; - if (o.suffix_draft > 0) sfx.append(outv[(size_t) i]); eos = std::find(o.eos_ids.begin(), o.eos_ids.end(), (int64_t) outv[(size_t) i]) != o.eos_ids.end(); } std::fflush(stdout); ++rounds; - const Clock::time_point tw2 = Clock::now(); - const bool drafted = eos || produced_n >= max_new || - mtp.draft(T, outv.data(), p, a, drafts.data(), err, dprob.data(), (float) req_spec_min_p); - { - const Clock::time_point tw3 = Clock::now(); - auto msd = [](Clock::time_point a0, Clock::time_point b0) { return std::chrono::duration(b0 - a0).count(); }; - dt_run += msd(tw0, tw1); dt_commit += msd(tw1, tw2); dt_draft += msd(tw2, tw3); - ++dec_windows; dec_T += T; + // the next window's drafts: the lookup's when they are more than the draft layer's chain allows (the + // chain stops at its first draft then) + lk = 0; + if (o.spec_lookup > 0 && !eos && produced_n < max_new) { + lk = lookup.propose(W - 1, ldrafts.data()); + if (lk < S) lk = 0; + for (int j = 0; j < lk; ++j) ver.ple_ahead(j + 1, ldrafts[(size_t) j]); } - if (adapt_thr.joinable()) adapt_thr.join(); - if (!adapt_ok) { + mtp.set_max_drafts(lk > 0 ? 1 : S - 1); + const bool drafted = (eos || produced_n >= max_new || + mtp.draft(T, outv.data(), p, a, drafts.data(), err, dprob.data(), (float) req_spec_min_p)) && + ver.wait_commit(err); + if (!drive.join_adapt()) { std::printf("ERR an adaptive refill failed\n"); return 1; } @@ -4143,176 +2967,25 @@ int main(int argc, char** argv) { std::printf("ERR %s\n", err.c_str()); return 1; } - if (timed_round && !eos) - policy.observe(from_sfx, T, a, sfx_match, - std::chrono::duration(Clock::now() - round0).count()); if (eos) { finish = "stop"; break; } - if (stop_req.load()) { finish = "cancel"; break; } x = outv[(size_t) a]; p += a + 1; } + mtp.on_draft = nullptr; const double decode_ms = std::chrono::duration(Clock::now() - d0).count(); - if (dec_timing && dec_windows > 0) { - const DecSnap d1 = dec_snap(); - const double w = (double) dec_windows, L = (double) g.n_layers; - std::fprintf(stderr, "strata decode timing: %lld windows, avg T %.2f, %.2f tokens/window, %.2f ms/window = " - "verify %.2f (GPU-reach wait %.2f + per-layer host %.2f [plan %.2f actq %.2f jobs %.2f " - "CPU %.2f] + stage %.2f) + commit/emit %.2f + draft %.2f; per layer-window: CPU experts " - "%.2f (%.2f entries), VRAM hits %.2f, PCIe %.2f\n", - (long long) dec_windows, dec_T / w, produced_n / w, decode_ms / w, dt_run / w, - (d1.wait - ds0.wait) / w, (d1.pool - ds0.pool) / w, (d1.plan - ds0.plan) / w, - (d1.actq - ds0.actq) / w, (d1.jobs - ds0.jobs) / w, (d1.run - ds0.run) / w, - (d1.host - ds0.host) / w, dt_commit / w, dt_draft / w, (d1.misses - ds0.misses) / (w * L), - (d1.entries - ds0.entries) / (w * L), (d1.hits - ds0.hits) / (w * L), (d1.pcie - ds0.pcie) / (w * L)); - const std::string pr = ver.profile_report(); - if (!pr.empty()) std::fprintf(stderr, "strata decode GPU stages (ms/window):%s\n", pr.c_str()); - } - if (!cancelled) { - // a prompt stopped halfway leaves the session somewhere between two chunks: nothing to continue from - // (the checkpoints taken while reading it are still good) - live.swap(consumed); - live_imgs = imgs_below(req_imgs, (int64_t) live.size()); - live_ok = o.prompt_cache > 0; - } - static const bool state_hash = std::getenv("STRATA_STATE_HASH") != nullptr; - if (state_hash && live_ok) { - // DEBUG: a fingerprint of every part of the session over the positions it holds ([0, L)), and - // separately of what lies past them in the last KV page (stale cells, fine unless something reads them) - cudaDeviceSynchronize(); - const int64_t L = (int64_t) live.size(); - const strata::kernels::QsaShapes qs = [&] { - strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); - s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_dim = g.idx_key_dim; - return s; - }(); - auto hash_dev = [&](const void* p, size_t bytes, uint64_t h) { - std::vector b(bytes); - if (bytes) cudaMemcpy(b.data(), p, bytes, cudaMemcpyDefault); // VRAM or a streamed host copy - return fnv1a(b.data(), b.size(), h); - }; - // the cells [c0, c1) of one int8 K or V pool ([page][kv_head][page_size][head_dim]), bytes per value `w` - auto hash_cells = [&](const void* pool, int64_t per_cell, int64_t c0, int64_t c1, uint64_t h) { - const int64_t ps = qs.page_size; - for (int64_t pg = c0 / ps; pg * ps < c1; ++pg) - for (int64_t hd = 0; hd < qs.n_head_kv; ++hd) { - const int64_t a = std::max(c0, pg * ps) - pg * ps, e = std::min(c1, (pg + 1) * ps) - pg * ps; - const size_t off = (size_t) (((pg * qs.n_head_kv + hd) * ps + a) * per_cell); - h = hash_dev((const uint8_t*) pool + off, (size_t) ((e - a) * per_cell), h); - } - return h; - }; - const ConvStateSizes z = conv_state_sizes(g); - uint64_t h_gdn = hash_dev(ss.gdn_state, z.gdn, 1469598103934665603ull); - if (std::getenv("STRATA_STATE_HASH_GDN") != nullptr) { // per GDN layer: which one differs first - const size_t per = z.gdn / (size_t) g.n_gdn_layers(); - std::string s; - char b[8]; - for (int64_t i = 0; i < g.n_gdn_layers(); ++i) { - std::snprintf(b, sizeof(b), "%04llx ", (unsigned long long) (hash_dev((const uint8_t*) ss.gdn_state + i * per, per, 1469598103934665603ull) & 0xffff)); - s += b; - } - std::fprintf(stderr, "strata serve: STATE_HASH_GDN %s\n", s.c_str()); - } - uint64_t h_ple = hash_dev(ss.ple_hist, z.ple, 1469598103934665603ull); - uint64_t h_tail = 1469598103934665603ull, h_pool = h_tail, h_kv = h_tail, h_stale = h_tail; - const int64_t kvb = qs.head_dim, scb = (qs.head_dim / 64) * 2; - // a state's K/V arrays and their bytes per (cell, head) row: the host copy when it has one - auto kv_arrays = [&](const strata::core::QsaState& st) { - const bool h = st.kv_mode != 0; - std::vector> a; - if (st.kv_q4) { - const int64_t q4b = (int64_t) strata::kernels::kv_q4_bytes_per_head((int) qs.head_dim); - a = {{h ? st.host.k_q4 : st.k_q4, q4b}, {h ? st.host.v_q4 : st.v_q4, q4b}}; - } else if (st.kv_hybrid) { - const int64_t q4b = (int64_t) strata::kernels::kv_q4_bytes_per_head((int) qs.head_dim); - a = {{h ? st.host.k_q : st.k_q, kvb}, {h ? st.host.v_q4 : st.v_q4, q4b}, - {h ? st.host.k_scale : st.k_scale, scb}}; - } else { - a = {{h ? st.host.k_q : st.k_q, kvb}, {h ? st.host.v_q : st.v_q, kvb}, - {h ? st.host.k_scale : st.k_scale, scb}, {h ? st.host.v_scale : st.v_scale, scb}}; - } - return a; - }; - const int64_t end_cell = std::min(((L + qs.page_size - 1) / qs.page_size) * qs.page_size, - ss.qsa_states[0].max_cells); - for (int64_t i = 0; i < g.n_qsa_layers(); ++i) { - const strata::core::QsaState& st = ss.qsa_states[i]; - h_tail = hash_dev(st.idx_tail, z.tail, h_tail); - h_pool = hash_dev(st.idx_pooled, (size_t) (L / qs.idx_block) * qs.idx_dim * 4, h_pool); - // KV streaming: the host copy is the identity layout and holds every cell - for (const auto& [pool, w] : kv_arrays(st)) { - h_kv = hash_cells(pool, w, 0, L, h_kv); - h_stale = hash_cells(pool, w, L, end_cell, h_stale); - } - } - const strata::core::QsaState& ms = mtp.kv_state(); - uint64_t h_mtp = 1469598103934665603ull; - const int64_t mL = std::min(L, ms.max_cells); - for (const auto& [pool, w] : kv_arrays(ms)) - if (pool != nullptr) h_mtp = hash_cells(pool, w, 0, mL, h_mtp); - std::fprintf(stderr, "strata serve: STATE_HASH L=%lld gdn=%016llx ple=%016llx tail=%016llx pooled=%016llx " - "kv=%016llx mtp=%016llx stale=%016llx ple_prev=%d,%d\n", (long long) L, - (unsigned long long) h_gdn, (unsigned long long) h_ple, (unsigned long long) h_tail, - (unsigned long long) h_pool, (unsigned long long) h_kv, (unsigned long long) h_mtp, - (unsigned long long) h_stale, ss.ple_prev[0], ss.ple_prev[1]); - } - const int64_t req_hits = drive.d.cache_hits - decode_hits0; - const int64_t req_look = (drive.d.cache_hits + drive.d.cache_admitted + drive.d.cache_refused) - decode_look0; - // DONE [hits] [lookups] - std::printf("DONE %lld %lld %.1f %.1f %s %lld %lld %lld %lld %lld\n", (long long) produced_n, (long long) n, prompt_ms, - decode_ms, finish, (long long) draft_accepted, (long long) draft_offered, (long long) resume, - (long long) req_hits, (long long) req_look); + const int64_t hits1 = drive.d.cache_hits + drive.d.gpu2_entries, look1 = hits1 + drive.d.cache_refused; + // DONE + // + std::printf("DONE %lld %lld %.1f %.1f %s %lld %lld %lld %lld %lld\n", (long long) produced_n, (long long) n, + prompt_ms, decode_ms, finish, (long long) draft_accepted, (long long) draft_offered, + (long long) reuse, (long long) (hits1 - hits0), (long long) (look1 - look0)); std::fflush(stdout); - if (drive.routing != nullptr) std::fflush(drive.routing); // the routing trace survives a crash and is watchable mid-session - const int64_t fresh = n - resume; - std::fprintf(stderr, "strata serve: prompt %lld tokens = %lld reused + %lld read in %.0f ms (%.1f tok/s), " - "%lld generated in %.0f ms (%.1f tok/s), drafts accepted %lld of %lld, %zu checkpoints%s\n", - (long long) n, (long long) resume, (long long) fresh, prompt_ms, + const int64_t fresh = n - reuse; + std::fprintf(stderr, "strata serve: %lld prompt tokens (%lld cached) in %.0f ms (%.1f tok/s), %lld generated " + "in %.0f ms (%.1f tok/s)\n", (long long) n, (long long) reuse, prompt_ms, prompt_ms > 0 ? 1000.0 * fresh / prompt_ms : 0.0, (long long) produced_n, decode_ms, - decode_ms > 0 ? 1000.0 * produced_n / decode_ms : 0.0, (long long) draft_accepted, - (long long) draft_offered, checks.size(), cancelled ? " (cancelled)" : ""); - // the VRAM share of the experts the pool looked up while decoding; experts it sent over PCIe for the GPU - // to read (--pcie-frac) are in neither count - if (req_look > 0) { - std::fprintf(stderr, "strata serve: decode expert cache hit rate: %.1f%% (%lld hits / %lld lookups)\n", - 100.0 * (double) req_hits / (double) req_look, - (long long) req_hits, (long long) req_look); - } - // STRATA_SPLIT_TIMING: where each verify stage's host time went, cumulative per window since the start - // (waiting for its GPU to ring a layer, the CPU pool and plan per layer, staging the window) - if (static const bool st_timing = std::getenv("STRATA_SPLIT_TIMING") != nullptr; st_timing) - for (int st = 0; st < n_stages; ++st) { - const strata::core::Verifier& v = stage_ver(st); - const double w = v.windows > 0 ? (double) v.windows : 1.0; - std::fprintf(stderr, "strata serve: stage %d: %lld windows; per window: wait for the GPU %.3f ms, " - "pool + plan %.3f ms, host staging %.3f ms, commit %.3f ms\n", st, - (long long) v.windows, v.ms_wait / w, v.ms_pool / w, v.ms_host / w, v.ms_commit / w); - } - if (g.n_qsa_layers() > 0 && ss.qsa_states[0].kv_mode == 1) { - // KV streaming, cumulative over the process: blocks the selections named vs blocks read from RAM - uint64_t miss = 0, look = 0; - bool over = false; - for (int64_t i = 0; i < g.n_qsa_layers(); ++i) { - const strata::kernels::KvStreamCounters c = strata::kernels::kv_stream_counters(ss.qsa_states[i].map); - miss += c.misses; look += c.lookups; over = over || c.overflow; - } - std::fprintf(stderr, "strata serve: KV streaming: %.2f%% of %llu block reads hit VRAM, %.1f MiB read " - "from RAM%s\n", look ? 100.0 * (double) (look - miss) / (double) look : 100.0, - (unsigned long long) look, (double) miss * 4224.0 / 1048576.0, - over ? " - OVERFLOW (too few resident cells)" : ""); - } - if (sfx_windows > 0) - std::fprintf(stderr, "strata serve: suffix drafts: %lld windows, %lld of %lld drafts accepted\n", - (long long) sfx_windows, (long long) sfx_ok, (long long) sfx_drafts); - for (int r = 0; r < 3; ++r) if (o.expert_cache_remote[(size_t) r] > 0) - std::fprintf(stderr, "strata serve: CUDA%d: %lld expert entries, %lld active layer launches, %.1f MiB returned " - "(%.1f MiB with full rows) in this request; host %.0f ms staging+launching, %.0f ms " - "waiting for it (since start)\n", r + 1, - (long long) (remote_experts[(size_t) r].computed() - remote_before[(size_t) r]), - (long long) (remote_experts[(size_t) r].launched_layers() - launches_before[(size_t) r]), - (double) (remote_experts[(size_t) r].returned_bytes() - compact_before[(size_t) r]) / 1048576.0, - (double) (remote_experts[(size_t) r].full_row_bytes() - full_before[(size_t) r]) / 1048576.0, - remote_experts[(size_t) r].ms_begin(), remote_experts[(size_t) r].ms_wait()); + decode_ms > 0 ? 1000.0 * produced_n / decode_ms : 0.0); + if (pcache.enabled()) std::fprintf(stderr, "strata serve: prompt cache: %s\n", pcache.summary().c_str()); } return 0; } @@ -4322,93 +2995,46 @@ int main(int argc, char** argv) { int64_t pos_start = 0; int64_t spec_pos = 0; // plan v0.3 P6: where the speculative loop starts (0 = not used) strata::prefill::Prefill prefill; + strata::prefill::ExpertRunner offload; + PromptBuffers pbuf; double prefill_batched_ms = 0; std::FILE* final_r = o.dump_final_r.empty() ? nullptr : std::fopen(o.dump_final_r.c_str(), "wb"); std::vector final_r_host(final_r ? (size_t) (g.hc * g.n_embd) : 0); - std::vector> lent; // (residency index, slot) lent to the prompt path - const int64_t n_batched = (o.prefill_until > 0 && o.prefill_until < n_prompt - 1) ? o.prefill_until : n_prompt - 1; if (o.prefill_chunk > 0 && n_prompt > 1) { - void* borrow = nullptr; - uint64_t borrow_bytes = 0; - if (!o.no_prefill_borrow && !host_res.empty() && d_res != nullptr) { - int64_t chunk = o.prefill_chunk; - int64_t k = plan_lend(chunk); // auto: the largest chunk that fits; fixed: halved to fit - const int64_t request_sized = request_chunk(n_batched, chunk); - if (k > 0 && request_sized < chunk) { // no bigger than this prompt segment needs - chunk = request_sized; - k = lend_slots(chunk); - if (k + 128 > xcache.slots()) k = 0; - if (!o.prefill_auto) o.prefill_chunk = chunk; - } - if (o.prefill_auto) { - o.prefill_chunk = k > 0 ? chunk : request_chunk(n_batched, 1024); - std::fprintf(stderr, "strata generate: prompt chunk auto: %lld tokens\n", (long long) o.prefill_chunk); - } else if (chunk != o.prefill_chunk) { - k = 0; // a fixed chunk that does not fit: its own buffers, as before - } - const int64_t blob = (int64_t) strata::kernels::cpu::expert_layout().max_blob; - if (k > 0) { // the lent slots are refilled after the prompt - const int32_t first = (int32_t) (xcache.slots() - k); - for (size_t i = 0; i < host_res.size(); ++i) - if (host_res[i] >= first) { - lent.emplace_back((int32_t) i, host_res[i]); - host_res[i] = strata::core::kNotResident; - } - cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice); - borrow = xcache.device_slot(first); - borrow_bytes = xcache.slot_offsets() ? (uint64_t) (xcache.bytes() - (int64_t) xcache.slot_offsets()[first]) - : (uint64_t) k * (uint64_t) blob; - std::fprintf(stderr, "strata generate: prompt path borrows %lld cache slots (%.2f GiB)\n", (long long) k, - (double) borrow_bytes / 1073741824.0); - } - } - if (borrow == nullptr) - std::fprintf(stderr, "strata generate: prompt path allocates its own buffers (no cache slots to borrow)\n"); - if (!prefill.init(wt, g, ss, srcp, o.expert_cache > 0 ? &xcache : nullptr, - host_res.empty() ? nullptr : host_res.data(), o.prefill_chunk, main_cs, err, borrow, - borrow_bytes)) { + const int64_t n_batched = (o.prefill_until > 0 && o.prefill_until < n_prompt - 1) ? o.prefill_until : n_prompt - 1; + if (!setup_prompt_path(prefill, offload, pbuf, err) || !pbuf.take(std::min(o.prefill_chunk, n_batched), err)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } + if (pbuf.lend || pbuf.lend2) + std::fprintf(stderr, "strata generate: prompt path borrows %lld cache slots (%.2f GiB), %lld (%.2f GiB) on the " + "second GPU\n", (long long) pbuf.loan.slots(), (double) pbuf.loan.bytes() / 1073741824.0, + (long long) pbuf.loan2.slots(), (double) pbuf.loan2.bytes() / 1073741824.0); + else + std::fprintf(stderr, "strata generate: prompt path allocates its own buffers (no cache slots to borrow)\n"); if (!o.mtp.empty()) { if (!mtp.bind(wt, &native_head, nullptr, err)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } - prefill.on_chunk = [&](const float* R_rows, int64_t T, int64_t p0, std::string& e) -> bool { - // cell i pairs R_i with the token at i + 1 (every such token is in the prompt) - std::vector nxt((size_t) T); - for (int64_t t = 0; t < T; ++t) nxt[(size_t) t] = (int32_t) o.tokens[(size_t) (p0 + t + 1)]; - if (prefill.draft_kv(mtp, R_rows, nxt.data(), T, p0, e)) return true; // E-9 - return e.empty() && mtp.prefill(R_rows, nxt.data(), T, p0, e); - }; + prefill.draft = &mtp; // cell i pairs R_i with the token at i + 1 (every such token is in the prompt) } + prefill.profile = o.prefill_profile; const Clock::time_point tp0 = Clock::now(); if (!prefill.run(o.tokens.data(), n_batched, 0, err)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } - // refill the lent slots from the arena and give them back to the decode tier - if (!lent.empty()) { - const Clock::time_point tr = Clock::now(); - for (const auto& [i, slot] : lent) { // D-4: queued, one wait (STRATA_REFILL_BLOCKING=1: each) - const uint8_t* b = srcp->blob(i / g.n_expert, i % g.n_expert); - const int64_t nb = (int64_t) strata::kernels::cpu::expert_layout().blob_bytes(i / g.n_expert); - if (b == nullptr || !(refill_blocking() ? xcache.fill_slot_blocking(slot, b, err, nb) - : xcache.fill_slot_queued(slot, b, err, nb))) { - std::fprintf(stderr, "strata generate: refilling a lent slot failed: %s\n", err.c_str()); - return 1; - } - host_res[(size_t) i] = slot; - } - if (!xcache.sync_queued(err)) { - std::fprintf(stderr, "strata generate: refilling the lent slots failed: %s\n", err.c_str()); - return 1; - } - cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice); - std::fprintf(stderr, "strata generate: %zu lent slots refilled in %.1f ms\n", lent.size(), - std::chrono::duration(Clock::now() - tr).count()); + // refill the lent slots from the arena and give them back to the decode tiers + const size_t n_lent[2] = {pbuf.loan.lent.size(), pbuf.loan2.lent.size()}; + const Clock::time_point tr = Clock::now(); + if (!pbuf.give_back(*srcp, err)) { + std::fprintf(stderr, "strata generate: refilling a lent slot failed: %s\n", err.c_str()); + return 1; } + if (n_lent[0] + n_lent[1] > 0) + std::fprintf(stderr, "strata generate: lent slots refilled (%zu, %zu on the second GPU) in %.1f ms\n", + n_lent[0], n_lent[1], std::chrono::duration(Clock::now() - tr).count()); prefill_batched_ms = std::chrono::duration(Clock::now() - tp0).count(); prefill_ms += prefill_batched_ms; pos_start = n_batched; @@ -4418,10 +3044,22 @@ int main(int argc, char** argv) { ss.ple_prev[1] = pos_start >= 1 ? (int32_t) o.tokens[(size_t) (pos_start - 1)] : -1; const strata::prefill::PrefillStats& ps = prefill.stats(); std::fprintf(stderr, "strata generate: prefill %lld tokens in %lld chunks, %.1f ms (%.1f tok/s); experts " - "streamed %lld (%lld by DMA, host %.1f ms), resident %lld; PLE %.1f ms\n", + "streamed %lld (%lld ahead), host %.1f ms, resident %lld; PLE %.1f ms\n", (long long) ps.tokens, (long long) ps.chunks, ps.ms_total, ps.ms_total > 0 ? 1000.0 * (double) ps.tokens / ps.ms_total : 0.0, (long long) ps.experts_streamed, - (long long) ps.experts_dma, ps.ms_experts_host, (long long) ps.experts_resident, ps.ms_ple); + (long long) ps.experts_prefetched, ps.ms_experts_host, (long long) ps.experts_resident, ps.ms_ple); + if (pbuf.offload != nullptr) + std::fprintf(stderr, "strata generate: prefill on the second GPU: experts streamed %lld (%lld ahead), " + "resident %lld; %.1f ms queuing, %.1f ms GPU\n", (long long) offload.experts_streamed, + (long long) offload.experts_prefetched, (long long) offload.experts_resident, offload.ms_host, + offload.ms_gpu); + if (o.prefill_profile) { + std::fprintf(stderr, "strata generate: prefill GPU ms by section:"); + for (int i = 0; i < strata::prefill::kPsCount; ++i) + std::fprintf(stderr, "%s %s %.0f", i ? ";" : "", strata::prefill::prefill_section_name(i), + ps.ms_section[i]); + std::fprintf(stderr, "\n"); + } } for (int64_t pos = pos_start;; ++pos) { @@ -4579,9 +3217,8 @@ int main(int argc, char** argv) { } int next = 0; - // The draw is Philox(seed, position), as in a verify window (row t at pos0 draws pos0 + t): a seed gives - // the same text whether a token comes from this path or from the speculative loop below. - sp.counter = (uint64_t) pos; + // Teacher-forced prompt rows consume no generation draws. + sp.counter = (uint64_t) produced.size(); strata::kernels::sample_tokens(d_logits, 1, (int) n_vocab, nullptr, 0, sp, d_next, token_stream); if (cudaMemcpyAsync(&next, d_next, sizeof(int), cudaMemcpyDeviceToHost, (cudaStream_t) token_stream) != cudaSuccess || @@ -4646,16 +3283,24 @@ int main(int argc, char** argv) { const bool ended = o.stop_eos && !produced.empty() && std::find(o.eos_ids.begin(), o.eos_ids.end(), (int64_t) produced.back()) != o.eos_ids.end(); if (spec_pos > 0 && (int64_t) produced.size() < o.max_new && !ended) { - std::vector oracle; - if (!o.spec_oracle.empty()) { - std::ifstream in(o.spec_oracle); + auto read_ids = [](const std::string& path, std::vector& ids) { + std::ifstream in(path); std::string text((std::istreambuf_iterator(in)), std::istreambuf_iterator()); std::string e; - if (!in || !parse_i64_list(text.c_str(), oracle, e)) { - std::fprintf(stderr, "strata generate: cannot read --spec-oracle %s\n", o.spec_oracle.c_str()); - return 2; - } + return (bool) in && parse_i64_list(text.c_str(), ids, e); + }; + std::vector oracle, follow; + if (!o.spec_oracle.empty() && !read_ids(o.spec_oracle, oracle)) { + std::fprintf(stderr, "strata generate: cannot read --spec-oracle %s\n", o.spec_oracle.c_str()); + return 2; + } + if (!o.spec_follow.empty() && !read_ids(o.spec_follow, follow)) { + std::fprintf(stderr, "strata generate: cannot read --spec-follow %s\n", o.spec_follow.c_str()); + return 2; } + // --spec-follow: the run ends with the continuation + const int64_t max_new = follow.empty() ? o.max_new : std::min(o.max_new, (int64_t) follow.size()); + int64_t follow_differ = 0, follow_emitted = 0; if (thits.d_res == nullptr) { std::fprintf(stderr, "strata generate: --spec needs the device residency table (--expert-profile, " "--expert-cache and the token graph)\n"); @@ -4664,12 +3309,14 @@ int main(int argc, char** argv) { mem_mark("the head and the prompt path"); strata::core::Verifier ver; strata::core::VerifyHits vh; - vh.d_res = thits.d_res; + vh.res = host_res.data(); vh.cache_base = thits.cache_base; + vh.slot_off = xcache.slot_offsets(); + vh.slots = xcache.slots(); vh.blob = thits.blob; - vh.slot_off = xcache.slot_offsets(); // E-6: the device plan's pointers - vh.n_slots = xcache.slots(); - if (!ver.init(wt, g, ss, vh, native_head.loaded() ? &native_head : nullptr, o.spec, err)) { + ver.set_split_head(split_head.get()); + ver.set_logits_wanted(!o.window_logits.empty()); + if (!ver.init(wt, g, ss, vh, native_head.loaded() ? &native_head : nullptr, o.max_window(), err)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } @@ -4679,94 +3326,35 @@ int main(int argc, char** argv) { return 1; } mem_mark("the verifier and the drafter's binding"); - ver.set_sampling(sp); // the CLI's own sampling (until 0.1.19 this loop was always greedy); no penalties here ver.set_split(o.spec_split); - // auto: the copy kernel for every pack. DMA (the native packs' default until 0.1.13) has the host call - // cudaMemcpyAsync + cudaLaunchHostFunc inside a verify window while the GPU spins on the flag they raise; - // issue #31's thread dumps show the host stuck in that cudaMemcpyAsync on a driver lock for good. The copy - // kernel needs no host CUDA call there, and costs ~1-3% decode on IQ3_S (45.3 -> 44.8 tok/s, 8 requests). - ver.set_pcie_mode(o.pcie_mode == "dma" ? 0 : o.pcie_mode == "direct" ? 1 : 2); + ver.set_profile(o.window_profile); + ver.set_host_rows(!o.window_hashes.empty()); + ver.set_pcie_mode(o.pcie_mode == "dma" ? 0 : o.pcie_mode == "direct" ? 1 : o.pcie_mode == "kernel" ? 2 + : native_pack ? 0 : 2); // auto: DMA for the native packs, the copy kernel for Q2_0 + if (drive.d.gpu2 != nullptr && gpu2.prefetch_slots() > 0) + ver.set_predict(&drive_predict, &drive, o.adapt_every > 0 && o.adapt_swaps > 0); + if (drive.d.gpu2 != nullptr) { + ver.set_gpu2(true); + ver.set_watch(&drive_watch, &drive); + } + if (drive.tier1 != nullptr || drive.tier2 != nullptr) { // the tiers' updates start at the tail + ver.set_tail(&drive_tail, &drive); + ver.set_rows_in(&drive_rows_in, &drive); + } drive.d.plan = ver.plan_sink(); + drive.d.host_res = ver.residency(); // the snapshot the windows decide their GPU hits by drive.d.pcie_num = (int) (o.pcie_frac * 256.0 + 0.5); if (drive.d.pcie_num < 0) drive.d.pcie_num = 0; if (drive.d.pcie_num > 256) drive.d.pcie_num = 256; + ver.set_pcie_share(drive.d.pcie_num > 0); const int64_t pcie0 = drive.d.pcie_experts; if (o.adapt_every > 0 && o.adapt_swaps > 0) drive.d.usage.assign((size_t) (g.n_layers * g.n_expert), 0.0f); - int64_t swaps_total = 0; - double ms_adapt = 0; - cudaStream_t adapt_stream = nullptr; - if (!drive.d.usage.empty() && cudaStreamCreateWithFlags(&adapt_stream, cudaStreamNonBlocking) != cudaSuccess) { - std::fprintf(stderr, "strata generate: cannot create the refill stream\n"); - return 1; - } - // plan v0.3 P6: swaps in flight - (residency index, slot) admitted when adapt_ev has completed - std::vector> pending; - cudaEvent_t adapt_ev = nullptr; - cudaEventCreateWithFlags(&adapt_ev, cudaEventDisableTiming); - auto apply_pending = [&](bool wait) { - if (pending.empty()) return; - if (wait) cudaEventSynchronize(adapt_ev); - else if (cudaEventQuery(adapt_ev) != cudaSuccess) return; - for (const auto& [i, slot] : pending) host_res[(size_t) i] = slot; - pending.clear(); - if (d_res != nullptr) - cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice); - }; - // Plan v0.3 P6: the VRAM tier follows the conversation. Candidates are missing experts routed at least - // twice (decayed); each is paired with its layer's least-routed resident expert and swapped when it was - // routed clearly more often. Copies run between rounds, when the GPU is idle. - auto adapt = [&]() -> bool { - const Clock::time_point ta = Clock::now(); - if (!pending.empty()) return true; // the previous swaps are still in flight - struct Swap { float gain; int32_t layer, in, out; }; - std::vector swaps; - std::vector> cand, vict; - for (int64_t l = 0; l < g.n_layers; ++l) { - cand.clear(); - vict.clear(); - const float* u = drive.d.usage.data() + l * g.n_expert; - const int32_t* r = host_res.data() + l * g.n_expert; - for (int32_t e = 0; e < (int32_t) g.n_expert; ++e) { - if (r[e] < 0) { if (u[e] >= 2.0f) cand.emplace_back(u[e], e); } - else vict.emplace_back(u[e], e); - } - if (cand.empty() || vict.empty()) continue; - std::sort(cand.begin(), cand.end(), [](auto& a, auto& b) { return a.first > b.first; }); - const size_t nc = std::min(cand.size(), vict.size()); - std::partial_sort(vict.begin(), vict.begin() + (ptrdiff_t) nc, vict.end(), - [](auto& a, auto& b) { return a.first < b.first; }); - for (size_t i = 0; i < nc; ++i) { - if (cand[i].first < vict[i].first + 1.5f) break; - swaps.push_back({cand[i].first - vict[i].first, (int32_t) l, cand[i].second, vict[i].second}); - } - } - std::sort(swaps.begin(), swaps.end(), [](const Swap& a, const Swap& b) { return a.gain > b.gain; }); - if ((int) swaps.size() > o.adapt_swaps) swaps.resize((size_t) o.adapt_swaps); - for (const Swap& s : swaps) { - const size_t in = (size_t) s.layer * g.n_expert + s.in, out = (size_t) s.layer * g.n_expert + s.out; - const int32_t slot = host_res[out]; - const uint8_t* b = srcp->blob(s.layer, s.in); - // asynchronous: the copies run while the MTP drafts; the next window waits for them - if (slot < 0 || b == nullptr || - cudaMemcpyAsync(xcache.device_slot(slot), b, (size_t) strata::kernels::cpu::expert_layout().blob_bytes(s.layer), - cudaMemcpyHostToDevice, adapt_stream) != cudaSuccess) { - std::fprintf(stderr, "strata generate: an adaptive refill failed\n"); - return false; - } - host_res[out] = strata::core::kNotResident; // evicted now: the CPU computes it meanwhile - pending.emplace_back((int32_t) in, slot); // resident once the copy has landed - } - if (!swaps.empty()) cudaEventRecord(adapt_ev, adapt_stream); - for (float& v : drive.d.usage) v *= 0.7f; - swaps_total += (int64_t) swaps.size(); - ms_adapt += std::chrono::duration(Clock::now() - ta).count(); - return true; - }; int64_t p = spec_pos; int32_t x = (int32_t) tok; - std::vector drafts((size_t) o.spec, 0); - std::vector dprob((size_t) o.spec, 1.0f); - std::vector window_hist((size_t) o.spec + 1, 0); + const int W = o.max_window(); + std::vector drafts((size_t) W, 0); + std::vector dprob((size_t) W, 1.0f); + std::vector window_hist((size_t) W + 1, 0); // plan v0.3 P6: with a native pack the first window is the last prompt token alone (it produces the first // generated token and the MTP's first cell); otherwise the token loop already did that. bool first_window = native_pack; @@ -4775,40 +3363,63 @@ int main(int argc, char** argv) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } - std::vector window((size_t) o.spec), outv((size_t) o.spec); - std::vector accepted_hist((size_t) o.spec, 0); + std::vector window((size_t) W), outv((size_t) W); + std::vector fdiff((size_t) W, 0); + std::FILE* hashes = o.window_hashes.empty() ? nullptr : std::fopen(o.window_hashes.c_str(), "w"); + if (!o.window_hashes.empty() && hashes == nullptr) { + std::fprintf(stderr, "strata generate: cannot write %s\n", o.window_hashes.c_str()); + return 1; + } + std::FILE* wlog = o.window_logits.empty() ? nullptr : std::fopen(o.window_logits.c_str(), "w"); + if (!o.window_logits.empty() && wlog == nullptr) { + std::fprintf(stderr, "strata generate: cannot write %s\n", o.window_logits.c_str()); + return 1; + } + std::vector wl(wlog != nullptr ? (size_t) W * (size_t) n_vocab : 0); + std::vector wl_ids(wlog != nullptr ? (size_t) n_vocab : 0); + auto write_logits = [&](int64_t k, int32_t emitted, const float* l) { // generated token k's line + constexpr int K = 64; + double m = l[0], s = 0; + for (int64_t i = 1; i < n_vocab; ++i) m = std::max(m, (double) l[i]); + for (int64_t i = 0; i < n_vocab; ++i) s += std::exp((double) l[i] - m); + const double lse = m + std::log(s); + for (int64_t i = 0; i < n_vocab; ++i) wl_ids[(size_t) i] = (int32_t) i; + std::partial_sort(wl_ids.begin(), wl_ids.begin() + K, wl_ids.end(), + [&](int32_t a, int32_t b) { return l[a] > l[b] || (l[a] == l[b] && a < b); }); + std::fprintf(wlog, "%lld %d", (long long) k, (int) emitted); + for (int j = 0; j < K; ++j) + std::fprintf(wlog, " %d:%.6f", (int) wl_ids[(size_t) j], (double) l[wl_ids[(size_t) j]] - lse); + std::fprintf(wlog, "\n"); + }; + std::vector accepted_hist((size_t) W, 0); int64_t rounds = 0, drafts_total = 0, drafts_ok = 0, corrupt_counter = 0; - const int S_mtp = o.mtp_max_t > 0 ? std::min(o.mtp_max_t, o.spec) : o.spec; - if (use_mtp && S_mtp < o.spec) mtp.set_max_drafts(S_mtp - 1); - strata::spec::SuffixDrafter sfx(std::max(1, o.suffix_draft), 64, (size_t) o.max_context + 4096); - strata::spec::DraftPolicy policy(o.spec); // MTP or lookup window (see draft_policy.hpp) - std::vector sbuf((size_t) o.spec, 0); - int64_t sfx_windows = 0, sfx_drafts = 0, sfx_ok = 0; - if (o.suffix_draft > 0) { - for (int64_t t : o.tokens) sfx.append((int32_t) t); - for (int64_t t : produced) sfx.append((int32_t) t); - } + // prompt lookup: the text so far, the prompt and then each emitted token + strata::spec::SuffixDrafter lookup(o.spec_lookup, 32, + o.spec_lookup > 0 ? (size_t) (n_prompt + max_new) + 16 : 16); + if (o.spec_lookup > 0) { + for (int64_t t : o.tokens) lookup.append((int32_t) t); + for (int64_t t : produced) lookup.append((int32_t) t); + } + std::vector ldrafts((size_t) W, 0); + int lk = 0; // the next window's drafts from the lookup (0: the draft layer's) + int64_t lookup_rounds = 0, lookup_drafts = 0, lookup_ok = 0; const double pool_ms0 = drive.cpu_ms; const int64_t misses0 = drive.d.multi_misses, entries0 = drive.d.multi_entries; - while ((int64_t) produced.size() < o.max_new) { + const uint32_t sleeps0 = pool.sleeps(); + if (o.stats) drive.d.routed.assign((size_t) (g.n_layers * g.n_expert), 0); + if (use_mtp) + mtp.on_draft = [&](int j, int32_t tok, float prob) { // the drafts the next window will verify + if (lk == 0 && prob >= (float) o.spec_min_p) ver.ple_ahead(j + 1, tok); + }; + while ((int64_t) produced.size() < max_new) { const Clock::time_point t0 = Clock::now(); - int T = S_mtp; + int T = o.spec; if (use_mtp && o.spec_min_p > 0.0) { T = 1; - while (T < S_mtp && dprob[(size_t) T - 1] >= (float) o.spec_min_p) ++T; + while (T < o.spec && dprob[(size_t) T - 1] >= (float) o.spec_min_p) ++T; } + if (lk > 0) T = 1 + lk; if (first_window) T = 1; - bool from_sfx = false; - int sfx_match = 0; - if (o.suffix_draft > 0 && !first_window) { - const int k = sfx.propose(o.spec - 1, sbuf.data()); - sfx_match = sfx.last_match(); - if (k > 0 && (!use_mtp || sbuf[0] == drafts[0])) { - const strata::spec::DraftPolicy::Pick pk = policy.choose(T, k, sfx_match); - if (pk.lookup) { T = pk.t; from_sfx = true; } - } - } - const bool timed_round = !first_window; ++window_hist[(size_t) T]; if (p + T > o.max_context) { std::fprintf(stderr, "strata generate: ran out of context at position %lld\n", (long long) p); @@ -4817,90 +3428,136 @@ int main(int argc, char** argv) { window[0] = x; for (int i = 1; i < T; ++i) { const size_t at = produced.size() - 1 + (size_t) i; - int32_t d = from_sfx ? sbuf[(size_t) i - 1] : use_mtp ? drafts[(size_t) i - 1] - : at < oracle.size() ? (int32_t) oracle[at] : 0; + int32_t d = lk > 0 ? ldrafts[(size_t) i - 1] : use_mtp ? drafts[(size_t) i - 1] + : at < oracle.size() ? (int32_t) oracle[at] : 0; if (o.spec_corrupt > 0 && (++corrupt_counter % o.spec_corrupt) == 0) d = (d + 1) % (int32_t) n_vocab; window[(size_t) i] = d; } drive.d.layers = 0; drive.d.experts = 0; drive.d.failed = false; - apply_pending(false); + drive_window(&drive, T); + drive.adapt_next = !drive.d.usage.empty() && ((rounds + 1) % o.adapt_every) == 0; + drive.flush_next = !drive.d.usage.empty() && ((rounds + 2) % o.adapt_every) == 0; if (!ver.run(T, window.data(), p, &drive_pool_multi, &drive, outv.data(), err)) { + drive.join_adapt(); std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } if (drive.d.failed) { + drive.join_adapt(); std::fprintf(stderr, "strata generate: the expert pool failed at layer %lld expert %lld: %s\n", (long long) drive.d.fail_layer, (long long) drive.d.fail_expert, drive.d.fail ? drive.d.fail : "(no message)"); return 1; } + if (wlog != nullptr && !ver.copy_logits(T, wl.data(), err)) { + drive.join_adapt(); + std::fprintf(stderr, "strata generate: %s\n", err.c_str()); + return 1; + } + if (hashes != nullptr) { // FNV-1a over the window's final residual rows + const size_t n_floats = (size_t) T * (size_t) (g.hc * g.n_embd); + uint64_t h = 1469598103934665603ull; + const auto* bytes = (const uint8_t*) ver.final_R_host(); + for (size_t i = 0; i < n_floats * sizeof(float); ++i) h = (h ^ bytes[i]) * 1099511628211ull; + std::fprintf(hashes, "%lld %lld %d %016llx", (long long) rounds, (long long) p, T, (unsigned long long) h); + for (int i = 0; i < T; ++i) std::fprintf(hashes, " %d", (int) outv[(size_t) i]); + std::fprintf(hashes, "\n"); + } + if (!follow.empty()) { // the continuation stands in for the argmax + const size_t k0 = produced.size(); + for (int i = 0; i < T && k0 + (size_t) i < follow.size(); ++i) { + fdiff[(size_t) i] = outv[(size_t) i] != (int32_t) follow[k0 + (size_t) i]; + outv[(size_t) i] = (int32_t) follow[k0 + (size_t) i]; + } + } int a = 0; while (a < T - 1 && window[(size_t) a + 1] == outv[(size_t) a]) ++a; + if (lk > 0) { + ++lookup_rounds; + lookup_drafts += lk; + lookup_ok += a; + } if (first_window) { first_window = false; ttft_ms = std::chrono::duration(Clock::now() - t_start).count(); } - // plan v0.3 P6: the adaptive tier's host work (ranking, copy submission) runs on its own thread while the - // GPU commits and drafts; it touches only the residency tables, which nothing reads until the next window - std::thread adapt_thr; - bool adapt_ok = true; - if (!drive.d.usage.empty() && ((rounds + 1) % o.adapt_every) == 0) - adapt_thr = std::thread([&] { adapt_ok = adapt(); }); - if (!ver.commit(a + 1, err)) { - if (adapt_thr.joinable()) adapt_thr.join(); + if (!ver.commit(a + 1, err, false)) { // beside the draft + drive.join_adapt(); std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } + ver.ple_ahead(0, outv[(size_t) a]); ++rounds; drafts_total += T - 1; drafts_ok += a; ++accepted_hist[(size_t) a]; - if (from_sfx) { ++sfx_windows; sfx_drafts += T - 1; sfx_ok += a; } bool eos = false; - for (int i = 0; i <= a && (int64_t) produced.size() < o.max_new && !eos; ++i) { + for (int i = 0; i <= a && (int64_t) produced.size() < max_new && !eos; ++i) { + if (wlog != nullptr) + write_logits((int64_t) produced.size(), outv[(size_t) i], wl.data() + (size_t) i * (size_t) n_vocab); produced.push_back(outv[(size_t) i]); - if (o.suffix_draft > 0) sfx.append(outv[(size_t) i]); + if (o.spec_lookup > 0) lookup.append(outv[(size_t) i]); eos = o.stop_eos && std::find(o.eos_ids.begin(), o.eos_ids.end(), (int64_t) outv[(size_t) i]) != o.eos_ids.end(); + if (!follow.empty()) { follow_differ += fdiff[(size_t) i]; ++follow_emitted; } } if (eos) { - if (adapt_thr.joinable()) adapt_thr.join(); + drive.join_adapt(); + if (!ver.wait_commit(err)) { + std::fprintf(stderr, "strata generate: %s\n", err.c_str()); + return 1; + } total_ms += std::chrono::duration(Clock::now() - t0).count(); break; } - const bool drafted = !use_mtp || (int64_t) produced.size() >= o.max_new || - mtp.draft(T, outv.data(), p, a, drafts.data(), err, dprob.data(), (float) o.spec_min_p); - if (adapt_thr.joinable()) adapt_thr.join(); - if (!adapt_ok) return 1; + // the next window's drafts: the lookup's when they are more than the draft layer's chain allows (the + // chain stops at its first draft then) + lk = 0; + if (o.spec_lookup > 0 && (int64_t) produced.size() < max_new) { + lk = lookup.propose(W - 1, ldrafts.data()); + if (lk < o.spec) lk = 0; + for (int j = 0; j < lk; ++j) ver.ple_ahead(j + 1, ldrafts[(size_t) j]); + } + mtp.set_max_drafts(lk > 0 ? 1 : o.spec - 1); + const bool drafted = (!use_mtp || (int64_t) produced.size() >= max_new || + mtp.draft(T, outv.data(), p, a, drafts.data(), err, dprob.data(), (float) o.spec_min_p)) && + ver.wait_commit(err); + if (!drive.join_adapt()) { + std::fprintf(stderr, "strata generate: %s\n", adapt_err.c_str()); + return 1; + } if (!drafted) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; } x = outv[(size_t) a]; p += a + 1; - const double round_ms = std::chrono::duration(Clock::now() - t0).count(); - total_ms += round_ms; - if (timed_round) policy.observe(from_sfx, T, a, sfx_match, round_ms); + total_ms += std::chrono::duration(Clock::now() - t0).count(); if (rounds % 64 == 0) std::fprintf(stderr, "strata generate: position %lld, %lld tokens, %lld rounds\n", (long long) p, (long long) produced.size(), (long long) rounds); } + if (hashes != nullptr) std::fclose(hashes); + if (wlog != nullptr) std::fclose(wlog); std::printf("%-24s %lld rounds of %d, drafts accepted %lld of %lld (%.3f), %.2f tokens per round\n", "speculation", (long long) rounds, o.spec, (long long) drafts_ok, (long long) drafts_total, drafts_total > 0 ? (double) drafts_ok / (double) drafts_total : 0.0, rounds > 0 ? (double) (drafts_ok + rounds) / (double) rounds : 0.0); + if (!follow.empty()) + std::printf("%-24s %lld of %lld emitted tokens differ from the argmax\n", "follow", + (long long) follow_differ, (long long) follow_emitted); if (o.spec_min_p > 0.0) { std::printf("%-24s", "window sizes"); for (size_t i = 1; i < window_hist.size(); ++i) std::printf(" T%zu:%lld", i, (long long) window_hist[i]); std::printf(" (min draft probability %.2f)\n", o.spec_min_p); } - if (o.suffix_draft > 0) - std::printf("%-24s %lld windows, drafts accepted %lld of %lld\n", "suffix drafts", (long long) sfx_windows, - (long long) sfx_ok, (long long) sfx_drafts); std::printf("%-24s", "accepted per round"); for (size_t i = 0; i < accepted_hist.size(); ++i) std::printf(" %zu:%lld", i, (long long) accepted_hist[i]); std::printf("\n"); + if (o.spec_lookup > 0) + std::printf("%-24s %lld windows of its drafts, %lld of %lld accepted\n", "prompt lookup", + (long long) lookup_rounds, (long long) lookup_ok, (long long) lookup_drafts); if (rounds > 0) std::printf("%-24s wait for rings %.3f pool %.3f host %.3f commit %.3f ms/round; CPU experts %.2f " "distinct / %.2f routed per layer\n", @@ -4910,26 +3567,58 @@ int main(int argc, char** argv) { (double) (drive.d.multi_entries - entries0) / (double) (rounds * g.n_layers)); if (rounds > 0) std::printf("%-24s gate/up %.3f quantize %.3f down %.3f ms/round; %.1f GB/s over the rows phases; " - "CPU pool call %.3f ms/round\n", "pool multi", pool.ms_multi_gu / rounds, - pool.ms_multi_q / rounds, pool.ms_multi_down / rounds, + "CPU pool call %.3f ms/round; %u worker sleeps; %lld thread moves\n", "pool multi", + pool.ms_multi_gu / rounds, pool.ms_multi_q / rounds, pool.ms_multi_down / rounds, (double) pool.multi_bytes / 1e6 / std::max(1e-9, pool.ms_multi_gu + pool.ms_multi_down), - (drive.cpu_ms - pool_ms0) / rounds); + (drive.cpu_ms - pool_ms0) / rounds, pool.sleeps() - sleeps0, (long long) placement.moves); if (rounds > 0) std::printf("%-24s plan %.3f activation quantize %.3f jobs %.3f run %.3f ms/round\n", "dispatch", drive.d.ms_plan / rounds, drive.d.ms_actq / rounds, drive.d.ms_jobs / rounds, drive.d.ms_run / rounds); + ver.print_profile(); if (rounds > 0 && !drive.d.usage.empty()) - std::printf("%-24s %lld experts swapped into the VRAM tier (every %d rounds, %.3f ms/round)\n", "adaptive tier", - (long long) swaps_total, o.adapt_every, ms_adapt / rounds); + std::printf("%-24s %lld experts swapped in and %lld into empty slots (every %d rounds, %.3f ms/round), " + "%lld moves dropped; %.1f MB/round copied\n", "adaptive tier", (long long) tier.swaps, + (long long) tier.fills, o.adapt_every, tier.ms / rounds, (long long) tier.dropped, + (double) tier.sent_bytes / 1e6 / (double) rounds); + if (rounds > 0 && drive.d.gpu2 != nullptr) + std::printf("%-24s %.2f routed entries and %.2f experts per round in %.2f layers (%.2f left to the CPU), " + "its rows after the CPU's in %.2f layers, %.3f ms/round; %lld experts swapped in and %lld into " + "empty slots (%lld still empty, %lld moves dropped; its copies between windows, %.2f ms, at " + "%.1f GB/s); %.2f experts per round prefetched, %.2f of them used (%.3f ms/round of host " + "time, %lld predictions waited for)\n", "second GPU", + (double) gpu2.entries_done / (double) rounds, (double) gpu2.experts / (double) rounds, + (double) gpu2.layers / (double) rounds, (double) drive.d.gpu2_skipped / (double) rounds, + (double) ver.late2 / (double) rounds, ver.ms_late2 / (double) rounds, + (long long) tier2.swaps, (long long) tier2.fills, + (long long) tier2.free_slots(), (long long) tier2.dropped, drive.gap_ms, + tier2.copy_rate() / 1e6, (double) gpu2.prefetch_copied / (double) rounds, + (double) gpu2.prefetch_used / (double) rounds, ver.ms_predict / (double) rounds, + (long long) ver.predict_late); + if (rounds > 0 && !drive.d.routed.empty()) { + // the share of the routed entries the N most-routed experts of this run take + std::vector c = drive.d.routed; + std::sort(c.begin(), c.end(), std::greater()); + double total = 0; + for (uint32_t v : c) total += v; + std::printf("%-24s", "routing concentration"); + double acc = 0; + size_t at = 0; + for (size_t n : {1000, 2000, 4000, 6000, 8000, 10000, 12000, 16000, 20000}) { + for (; at < n && at < c.size(); ++at) acc += c[at]; + std::printf(" %zu:%.3f", n, total > 0 ? acc / total : 0.0); + } + std::printf(" (top-N experts' share of %.0f routed entries)\n", total); + } if (rounds > 0 && drive.d.pcie_num > 0) std::printf("%-24s %.2f distinct experts per layer read over PCIe (share %d/256 of the misses)\n", "pcie experts", (double) (drive.d.pcie_experts - pcie0) / (double) (rounds * g.n_layers), drive.d.pcie_num); (void) pool_ms0; if (use_mtp && rounds > 0) - std::printf("%-24s %.3f ms/round drafting (%lld rounds), MTP prompt %.1f ms, %.0f MiB of VRAM\n", "mtp", - mtp.ms_draft / (double) mtp.rounds, (long long) mtp.rounds, mtp.ms_prefill, - (double) mtp.vram_bytes() / 1048576.0); + std::printf("%-24s %.3f ms/round drafting (%lld rounds), %.0f MiB of VRAM\n", "mtp", + mtp.ms_draft / (double) mtp.rounds, (long long) mtp.rounds, (double) mtp.vram_bytes() / 1048576.0); + drive.d.host_res = host_res.data(); // the verifier's snapshot goes with it } if (dump != nullptr && std::fclose(dump) != 0) { @@ -5039,21 +3728,16 @@ int main(int argc, char** argv) { std::printf("%-24s %.3f ms/token over %lld layers (%.0f positions, %lld dispatches)\n", " the CPU expert pool", pool_positions > 0.0 ? drive.cpu_ms / pool_positions : 0.0, (long long) g.n_layers, pool_positions, (long long) drive.calls); - // **AND WHERE INSIDE `run()` IT WENT.** Three phases per layer and they were one number, which cannot - // tell a pool that is slow at the WORK from one that is slow at the SYNCHRONISATION - opposite fixes. - // Wait-for-park is expected to be ~0 (the workers re-parked at the end of the previous layer); the - // question is whether the time is in the drain or in the re-park barrier. + // **AND WHERE INSIDE `run()` IT WENT.** They were one number, which cannot tell a pool that is slow at + // the WORK from one that is slow at the SYNCHRONISATION - opposite fixes: the drain, or the barrier that + // closes the phase. if (pool_positions > 0.0) { - double wp = 0, dr = 0, rp = 0; - pool.phase_ms(wp, dr, rp); + double dr = 0, cl = 0; + pool.phase_ms(dr, cl); const double per = pool_positions; - std::printf("%-24s wait-park %.3f drain %.3f re-park %.3f ms/token\n", - " pool phases", wp / per, dr / per, rp / per); + std::printf("%-24s drain %.3f close %.3f ms/token\n", " pool phases", dr / per, cl / per); } std::printf("%-24s %lld blobs read\n", " expert blobs", (long long) srcp->reads()); - for (int r = 0; r < 3; ++r) if (o.expert_cache_remote[(size_t) r] > 0) - std::printf(" CUDA%d experts %lld routed entries computed\n", - r + 1, (long long) remote_experts[(size_t) r].computed()); // ---- **R4's DISPATCH MEASUREMENT: h, ON THE ENGINE'S OWN ROUTING.** No offline trace, no corpus // question, no k-fold - these are the ids the router actually produced on this run. Reported as // hits/lookups so it can be read directly as the h the cache would deliver, and alongside `refused` diff --git a/src/spec/suffix_drafter.cpp b/src/spec/suffix_drafter.cpp index 6f04f76d595..12611797ff1 100644 --- a/src/spec/suffix_drafter.cpp +++ b/src/spec/suffix_drafter.cpp @@ -17,16 +17,19 @@ uint64_t mix(uint64_t x) { SuffixDrafter::SuffixDrafter(int min_match, int max_match, size_t capacity_tokens) : min_match_(std::max(3, min_match)), max_match_(std::max(min_match, max_match)) { - size_t cap = 1; + size_t cap = 16; while (cap < capacity_tokens * 2) cap <<= 1; // load factor <= 0.5 at the nominal capacity table_.assign(cap, Slot{}); mask_ = cap - 1; hist_.reserve(capacity_tokens); + prev_.reserve(capacity_tokens); } void SuffixDrafter::reset() { hist_.clear(); + prev_.clear(); std::fill(table_.begin(), table_.end(), Slot{}); + used_ = 0; last_match_ = 0; } @@ -36,46 +39,73 @@ uint64_t SuffixDrafter::key_at(size_t end) const { } SuffixDrafter::Slot* SuffixDrafter::find_slot(uint64_t key, bool insert) { - for (size_t i = key & mask_, probes = 0; probes <= mask_; i = (i + 1) & mask_, ++probes) { + for (size_t i = key & mask_;; i = (i + 1) & mask_) { // never full: rebuilt at half load Slot& s = table_[i]; if (s.key == key) return &s; if (s.key == 0) { if (!insert) return nullptr; s.key = key; + ++used_; return &s; } } - return nullptr; // table full: the history outgrew its capacity +} + +void SuffixDrafter::link(size_t end) { + int32_t before = -1; + if (end >= 2) { + Slot* s = find_slot(key_at(end), true); + before = s->last; + s->last = (int32_t) end; + } + prev_.push_back(before); +} + +// The indexed positions again, without the slots of undone ones, in a table that holds them at a quarter load. +void SuffixDrafter::rebuild() { + const size_t n = prev_.size(); + size_t cap = table_.size(); + while (cap < 4 * n) cap <<= 1; + table_.assign(cap, Slot{}); + mask_ = cap - 1; + used_ = 0; + prev_.clear(); + for (size_t e = 0; e < n; ++e) link(e); } void SuffixDrafter::append(const int32_t* tokens, size_t n) { for (size_t i = 0; i < n; ++i) { + if (2 * (used_ + 1) > table_.size()) rebuild(); hist_.push_back(tokens[i]); - const size_t end = hist_.size() - 1; - if (end < 2) continue; - Slot* s = find_slot(key_at(end), true); - if (s == nullptr) continue; - for (int w = WAYS - 1; w > 0; --w) s->pos[w] = s->pos[w - 1]; - s->pos[0] = (uint32_t) end; - if (s->n < WAYS) ++s->n; + link(hist_.size() - 1); } } +void SuffixDrafter::assign(const int32_t* tokens, size_t n) { + size_t keep = 0; + const size_t common = std::min(n, hist_.size()); + while (keep < common && hist_[keep] == tokens[keep]) ++keep; + for (size_t e = hist_.size(); e-- > std::max(keep, 2);) // newest first: each trigram's last goes back + if (Slot* s = find_slot(key_at(e), false)) s->last = prev_[e]; + hist_.resize(keep); + prev_.resize(keep); + last_match_ = 0; + append(tokens + keep, n - keep); +} + int SuffixDrafter::propose(int max_k, int32_t* out) { last_match_ = 0; const size_t n = hist_.size(); if (n < 4 || max_k <= 0) return 0; const size_t cur = n - 1; - const Slot* s = find_slot(key_at(cur), false); - if (s == nullptr) return 0; size_t best_end = 0; int best_len = 0; - for (int w = 0; w < s->n; ++w) { - const size_t p = s->pos[w]; - if (p >= cur) continue; // the current suffix itself + int32_t p = prev_[cur]; // the current trigram's earlier occurrences, newest first + for (int w = 0; w < WAYS && p >= 0; ++w, p = prev_[(size_t) p]) { int len = 0; - while (len < max_match_ && len <= (int) p && hist_[p - len] == hist_[cur - len]) ++len; - if (len > best_len) { best_len = len; best_end = p; } // most recent first, so ties keep the newer + while (len < max_match_ && len <= p && hist_[(size_t) (p - len)] == hist_[cur - len]) ++len; + if (len > best_len) { best_len = len; best_end = (size_t) p; } // most recent first, so ties keep the newer + if (best_len == max_match_) break; } if (best_len < min_match_) return 0; last_match_ = best_len; diff --git a/src/spec/suffix_drafter_test.cpp b/src/spec/suffix_drafter_test.cpp index 9fa8611e7c8..c67723a9642 100644 --- a/src/spec/suffix_drafter_test.cpp +++ b/src/spec/suffix_drafter_test.cpp @@ -12,6 +12,7 @@ // No model runs; this measures the drafter on text, not on the model's own outputs. #include "strata/spec/suffix_drafter.hpp" +#include #include #include #include @@ -103,7 +104,7 @@ void unit_tests() { const int n = d.propose(6, out); check(n >= 3 && out[0] == 1 && out[1] == 2 && out[2] == 3, "periodic continuation"); } - { // bounded memory: appending beyond the nominal capacity does not crash and still works + { // beyond the nominal capacity the table grows SuffixDrafter d(3, 32, 1024); std::vector doc; for (int i = 0; i < 5000; ++i) doc.push_back(i % 997); @@ -111,6 +112,54 @@ void unit_tests() { int32_t out[4]; check(d.propose(4, out) > 0, "propose after overflow of nominal capacity"); } + { // assign: the shared prefix stays, the rest of the old text proposes nothing + SuffixDrafter d; + const std::vector a = {1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3}; + d.append(a.data(), a.size()); + int32_t out[4]; + check(d.propose(4, out) > 0 && out[0] == 4, "before assign"); + const std::vector b = {1, 2, 3, 40, 50, 60, 70, 1, 2, 3}; + d.assign(b.data(), b.size()); + check(d.size() == b.size(), "assign replaces the history"); + const int n = d.propose(4, out); + check(n == 4 && out[0] == 40 && out[3] == 70, "the new text's continuation, not the old one's"); + d.assign(a.data(), a.size()); + check(d.propose(4, out) == 4 && out[0] == 4 && out[3] == 7, "the old text again"); + } + { // many texts: the slots of undone positions are dropped when the table is rebuilt + SuffixDrafter d(3, 32, 256); + std::vector doc(200); + for (int r = 0; r < 50; ++r) { + for (int i = 0; i < 200; ++i) doc[(size_t) i] = r * 1000 + i % 150; + d.assign(doc.data(), doc.size()); + } + int32_t out[4]; + check(d.propose(4, out) == 4 && out[0] == 49000 + 50, "propose after many texts"); + } + { // proposals depend on the history alone: after assign undid a longer text and the table grew, as a fresh + // drafter's (a 6-token alphabet: every trigram recurs, more often than WAYS) + uint32_t rng = 12345; + auto next = [&] { rng = rng * 1664525u + 1013904223u; return (int32_t) ((rng >> 16) % 6); }; + std::vector a(3000), b; + for (int32_t& t : a) t = next(); + b.assign(a.begin(), a.begin() + 1700); + for (int i = 0; i < 1300; ++i) b.push_back(next()); + SuffixDrafter x(3, 32, 16), y(3, 32, 1u << 14); + x.append(a.data(), a.size()); + x.assign(b.data(), 1000); + x.append(b.data() + 1000, b.size() - 1000); + y.append(b.data(), b.size()); + bool same = x.size() == y.size(); + int32_t ox[8], oy[8]; + for (int i = 0; i < 300 && same; ++i) { + const int nx = x.propose(8, ox), ny = y.propose(8, oy); + same = nx == ny && x.last_match() == y.last_match() && std::equal(ox, ox + nx, oy); + const int32_t t = next(); + x.append(t); + y.append(t); + } + check(same, "proposals after assign and growth are a fresh drafter's"); + } std::printf("suffix_drafter unit tests: %s\n", g_fail ? "FAILED" : "OK"); } } // namespace diff --git a/tools/iq_pack.py b/tools/iq_pack.py index d86f53c4a1d..07ac44b7e6a 100644 --- a/tools/iq_pack.py +++ b/tools/iq_pack.py @@ -1,30 +1,24 @@ -"""tools/iq_pack.py - plan v0.3 P6: a native pack for any of the model files (Q2_0, IQ2_XS, IQ3_XXS). +"""tools/iq_pack.py - plan v0.3 P6: a native pack for a model file whose tensors the engine reads in GGUF form +(Q2_0, IQ2_XS, IQ3_XXS, UD-Q4_K_XL). - python tools/iq_pack.py --gguf -00001-of-00002.gguf --out pack/iq3_xxs (standalone) + python tools/iq_pack.py --gguf -00001-of-0000N.gguf --out pack/iq3_xxs (standalone) python tools/iq_pack.py --gguf -00001-of-00002.gguf --base pack/full --out ... (share dense.bin) -The i-quant experts cannot be re-expressed in the Q2_0 pack form, so this pack keeps every quantized tensor in -its GGUF form: +The other shards of a split model are found by name. The experts cannot be re-expressed in the Q2_0 pack form, +so this pack keeps every quantized tensor in its GGUF form: - experts.bin per layer, 512 blobs of [gate rows | up rows | down rows], the raw GGUF slices. Blob - size is per layer (the files mix IQ1_M ... IQ3_S gate/up and Q2_0 / IQ4_NL down). - native_experts.txt one line per layer: layer gu_type d_type offset blob_bytes + native_experts.txt one line per layer: layer gu_type d_type offset blob_bytes. The engine fills its expert + arena from the GGUF (tensors found by name); `offset` is the layer's place in that arena. + experts.bin optional (--experts-bin): the arena as a file, per layer 512 blobs of [gate rows | up rows | + down rows], the raw GGUF slices. Blob size is per layer (the files mix formats). index.txt the table the engine loads. Quantized dense tensors, token_embd and output are served natively from the GGUF by the engine (--native): their rows carry shape only. - dense.bin standalone: the BF16/F16/F32 tensors exactly as the GGUF stores them (index kinds 4/5/2). + dense.bin standalone: every other tensor in the float form the engine takes (FORM below; index kinds + 4/5/2 = BF16/F16/F32), converted when the file stores it otherwise. With --base: the base (Q2_0) pack's dense.bin, hard-linked - the float tensors are - byte-identical in all three model files (checked) - plus extra.bin for tensors that are + byte-identical in all three ISTA model files (checked) - plus extra.bin for tensors that are float here but quantized in the base pack (blk.1.ple_key). tokenizer/ exported from the GGUF (tools/strata_tokenizer.py), with the model's chat template. - -Split files: every shard of the model is read (-0000N-of-0000M.gguf beside --gguf), so the layers may be -split anyhow (Swift 1.5's GGUFs put layers 13-47 in shard 2 and the PLE table in shard 1). A layer whose experts -are not in shard 1 names its shard in native_experts.txt (v3). Router tensors stored as F32 whose values are -exactly BF16 (Swift 1.5) are written as BF16, the form the engine's router takes; anything else is refused. - -For ordinary quants, --compat-bf16 dequantizes the small projections that the engine reads as BF16, using -round-to-nearest-even. This introduces BF16 rounding; it does not reconstruct the original full-precision -weights. Experts, native attention projections, token embeddings and the disk-backed PLE table stay unchanged. """ from __future__ import annotations @@ -32,6 +26,7 @@ import json import os import pathlib +import re import shutil import subprocess import sys @@ -43,74 +38,100 @@ import gguf_reader as G # noqa: E402 FLOAT = {"BF16", "F32", "F16"} -ROUTERS = ("ffn_gate_inp.weight", "ffn_gate_inp_shexp.weight") -NOT_IN_PACK = {"per_layer_token_embd.weight"} # the 28.8 GB PLE table: read from its GGUF by the engine - -# These small projections are read as BF16 by the residual, router, GDN, QSA and PLE kernels. -# GSQ-RCO files already store them that way. Ordinary GGUF quants (including OrcaRouter's IQ3_XXS) -# quantize them too; --compat-bf16 explicitly dequantizes and rounds ONLY these tensors. -BF16_PROJECTIONS = ( - "hc_attn_down.weight", "hc_attn_up.weight", "hc_attn_inject.weight", - "hc_ffn_down.weight", "hc_ffn_up.weight", "hc_ffn_inject.weight", - "ssm_alpha.weight", "ssm_beta.weight", "indexer.q_proj.weight", "indexer.k_proj.weight", - "ple_value.weight", *ROUTERS, -) -BF16_OUTPUT = {"output_hc_down.weight", "output_hc_up.weight"} - - -def needs_bf16(name: str, type_name: str) -> bool: - if name in BF16_OUTPUT: - return True - if not name.startswith("blk."): - return False - # The existing native PLE key supports Q2_0 only. Other key encodings use the BF16 path. - return name.endswith(BF16_PROJECTIONS) or (name == "blk.1.ple_key.weight" and type_name != "Q2_0") - - -def bf16_bytes(raw: np.ndarray, type_name: str) -> bytes: - from _paths import add_gguf_py - add_gguf_py() - from gguf import GGMLQuantizationType as Q, quants - values = quants.dequantize(raw, Q[type_name]) - if not np.isfinite(values).all(): - raise ValueError("cannot convert non-finite weights to BF16") - # ggml's round-to-nearest-even conversion, including correct halfway rounding. - return quants.quantize(values, Q.BF16).tobytes() +ROLES = ("gate", "up", "down") +N_EXPERT = 512 +ALIGN = 64 +KIND = {"BF16": "4", "F16": "5", "F32": "2"} + +# The float form the engine takes for each tensor it does not read from the GGUF (name without `blk.N.`). A file +# that stores one otherwise is converted: UD-Q4_K_XL keeps the router, SSM gates and injections in F32 (their +# values are BF16's, so the conversion is exact), the hyper-connection projections and the PLE value in Q8_0 +# (re-rounded to BF16) and the PLE conv in F32. +FORM = { + "ffn_gate_inp.weight": "BF16", "ffn_gate_inp_shexp.weight": "BF16", + "hc_attn_down.weight": "BF16", "hc_attn_up.weight": "BF16", "hc_attn_inject.weight": "BF16", + "hc_ffn_down.weight": "BF16", "hc_ffn_up.weight": "BF16", "hc_ffn_inject.weight": "BF16", + "output_hc_down.weight": "BF16", "output_hc_up.weight": "BF16", + "indexer.k_proj.weight": "BF16", "indexer.q_proj.weight": "BF16", + "ssm_alpha.weight": "BF16", "ssm_beta.weight": "BF16", + "ple_key.weight": "BF16", "ple_value.weight": "BF16", "ple_conv1d.weight": "F16", + "attn_q_norm.weight": "F32", "attn_k_norm.weight": "F32", "hc_attn_norm.weight": "F32", + "hc_ffn_norm.weight": "F32", "output_hc_norm.weight": "F32", "indexer.q_norm.weight": "F32", + "indexer.k_norm.weight": "F32", "ple_norm_conv.weight": "F32", "ple_norm_key.weight": "F32", + "ple_norm_query.weight": "F32", "ssm_a": "F32", "ssm_conv1d.weight": "F32", "ssm_dt.bias": "F32", + "ssm_norm.weight": "F32", +} +# Served from the GGUF when quantized in a type the engine has kernels for: the dense projections and the PLE key +# (NativeDense and NativeHead: native_mmvq_supported), the token embedding (NativeEmbed: iq_supported). +MMVQ_TYPES = {"Q4_0", "Q5_0", "Q8_0", "Q3_K", "Q4_K", "Q5_K", "Q6_K", "IQ4_NL", "IQ4_XS", "Q2_0", "IQ2_XXS", + "IQ2_XS", "IQ3_XXS", "IQ3_S", "IQ2_S", "IQ1_M"} +EMBED_TYPES = {"IQ2_XXS", "IQ2_XS", "IQ3_XXS", "IQ4_NL", "IQ3_S", "IQ2_S", "IQ4_XS", "IQ1_S", "IQ1_M", "Q2_0", "Q3_K", + "Q4_K", "Q5_K", "Q6_K", "Q5_1", "Q8_0"} +NATIVE = {"attn_qkv.weight", "attn_gate.weight", "ssm_out.weight", "attn_q.weight", "attn_k.weight", + "attn_v.weight", "attn_output.weight", "ffn_gate_shexp.weight", "ffn_up_shexp.weight", + "ffn_down_shexp.weight", "ple_key.weight", "output.weight"} +PLE_TABLE = "per_layer_token_embd.weight" # read from the GGUF by the engine (--ple-gguf) class Model: - """All shards of one model: name -> (GGUFFile, TensorInfo, memmap, shard path).""" + """A model's shards (a split model's are found by name from the first), tensors looked up across them.""" def __init__(self, first: pathlib.Path): - import re - m = re.search(r"-(\d{5})-of-(\d{5})\.gguf$", first.name) - paths = [first] - if m: - total = int(m.group(2)) - paths = [first.with_name(first.name[:m.start()] + "-%05d-of-%05d.gguf" % (i, total)) - for i in range(1, total + 1)] - missing = [str(p) for p in paths if not p.is_file()] - if missing: - raise FileNotFoundError("missing model shards (wait for the download): " + ", ".join(missing)) - self.paths = paths + m = re.fullmatch(r"(.*-)00001(-of-(\d{5})\.gguf)", first.name) + names = [f"{m.group(1)}{i:05d}{m.group(2)}" for i in range(1, int(m.group(3)) + 1)] if m else [first.name] + self.shards = [] self.where = {} - for p in paths: - g = G.GGUFFile(p) - mm = np.memmap(p, dtype=np.uint8, mode="r") + for name in names: + path = first.parent / name + if not path.exists(): + raise FileNotFoundError(f"missing model shard {path}") + g = G.GGUFFile(path) + mm = np.memmap(path, dtype=np.uint8, mode="r") for t in g.tensors: - size = t.expected_bytes() - if size is None or g.data_start + t.offset + size > mm.size: - raise ValueError(f"{p.name}: unsupported or truncated tensor {t.name}") if t.name in self.where: - raise ValueError(f"{p.name}: duplicate tensor {t.name}") - self.where[t.name] = (g, t, mm, p) + raise ValueError(f"tensor {t.name} is in two shards") + self.where[t.name] = (g, mm, t) + self.shards.append(g) - def bytes(self, name) -> np.ndarray: - g, t, mm, _ = self.where[name] - return tensor_bytes(mm, g, t) -ROLES = ("gate", "up", "down") -N_EXPERT = 512 -ALIGN = 64 + def tensors(self): + return [t for g in self.shards for t in g.tensors] + + def get(self, name): + return self.where[name][2] if name in self.where else None + + def bytes(self, t) -> np.ndarray: + g, mm, _ = self.where[t.name] + return mm[g.data_start + t.offset: g.data_start + t.offset + t.expected_bytes()] + + +def suffix(name: str) -> str: + return re.sub(r"^blk\.\d+\.", "", name) + + +def to_f32(raw: np.ndarray, type_name: str) -> np.ndarray: + if type_name == "F32": + return raw.view(" bytes: + if form == "F32": + return x.astype("> 16) & 1)) >> 16).astype(" bool: + return name.startswith("blk.") and name.endswith(("_exps.weight",)) def read_index(path: pathlib.Path): @@ -124,68 +145,53 @@ def read_index(path: pathlib.Path): return header, rows -def tensor_bytes(mm, g, t) -> np.ndarray: - n = t.expected_bytes() - return mm[g.data_start + t.offset: g.data_start + t.offset + n] - +def native_row(t) -> list[str]: + ne0 = int(t.shape[0]) + ne1 = int(t.shape[1]) if len(t.shape) > 1 else 0 + return [t.name, "0", "0", "0", "0", "0", "0", str(ne0), str(ne1), "8", "0", "32"] + ["0"] * 7 -def is_expert(name: str) -> bool: - return name.startswith("blk.") and name.endswith(("_exps.weight",)) - -def index_standalone(src, out, model: Model, compat_bf16: bool = False) -> int: - """Every non-expert tensor of the model: floats into dense.bin as stored (exact-BF16 F32 routers as BF16), - quantized ones native-only.""" - if not compat_bf16: - for name, (_, t, _, _) in model.where.items(): - if needs_bf16(name, t.type_name) and t.type_name != "BF16" and not ( - t.type_name == "F32" and name.endswith(ROUTERS)): - print(f"{name} is {t.type_name}, but the engine requires BF16; use --compat-bf16") - return 1 - rows, at = [], 0 - served = 0 - converted = [] +def index_standalone(src, out, model: Model) -> int: + """Every non-expert tensor: quantized ones the engine serves natively as shape-only rows, the others into + dense.bin in their engine form.""" + rows, at, served, converted = [], 0, 0, 0 with open(out / "dense.bin", "wb") as fo: - for name, (g, t, mm, _) in model.where.items(): - if is_expert(t.name) or t.name in NOT_IN_PACK: + for t in model.tensors(): + if is_expert(t.name) or t.name == PLE_TABLE: continue if len(t.shape) > 2: print("tensor %s has %d dimensions; the index holds two" % (t.name, len(t.shape))) return 1 + sfx = suffix(t.name) + if t.name == "token_embd.weight" or (sfx in NATIVE and t.type_name not in FLOAT): + ok = EMBED_TYPES if t.name == "token_embd.weight" else MMVQ_TYPES + if t.type_name not in ok: + print("%s is %s, a type the engine has no native kernels for" % (t.name, t.type_name)) + return 1 + served += 1 + rows.append(native_row(t)) + continue + form = FORM.get(sfx) + if form is None: + print("tensor %s: no engine form known for it" % t.name) + return 1 + raw = model.bytes(t) + if t.type_name == form: + data = raw.tobytes() + else: + data = encode(to_f32(raw, t.type_name), form) + converted += 1 ne0 = int(t.shape[0]) ne1 = int(t.shape[1]) if len(t.shape) > 1 else 0 - convert = compat_bf16 and needs_bf16(t.name, t.type_name) and t.type_name != "BF16" - if t.type_name in FLOAT or convert: - raw = tensor_bytes(mm, g, t).tobytes() - if convert: - raw = bf16_bytes(np.frombuffer(raw, dtype=np.uint8), t.type_name) - kind = "4" - converted.append({"name": t.name, "source_type": t.type_name, "bytes": len(raw)}) - else: - kind = {"BF16": "4", "F16": "5", "F32": "2"}[t.type_name] - if not convert and t.type_name == "F32" and t.name.endswith(ROUTERS): - u = np.frombuffer(raw, dtype=np.uint32) - if np.count_nonzero(u & 0xFFFF): - print("router %s is F32 with values that are not BF16; the engine's router is BF16" % t.name) - return 1 - raw = (u >> 16).astype(np.uint16).tobytes() # the exact BF16 values - kind = "4" - rows.append([t.name, "0", kind, str(at), str(len(raw)), "0", str(len(raw)), str(ne0), str(ne1), - "0", "0", "1"] + ["0"] * 7) - fo.write(raw) - pad = (-len(raw)) % ALIGN - fo.write(b"\0" * pad) - at += len(raw) + pad - else: - served += 1 - rows.append([t.name, "0", "0", "0", "0", "0", "0", str(ne0), str(ne1), "8", "0", "32"] + ["0"] * 7) + rows.append([t.name, "0", KIND[form], str(at), str(len(data)), "0", str(len(data)), str(ne0), str(ne1), + "0", "0", "1"] + ["0"] * 7) + fo.write(data) + pad = (-len(data)) % ALIGN + fo.write(b"\0" * pad) + at += len(data) + pad + if converted: + print("dense.bin: %d tensors converted to their engine form" % converted) write_index(out, rows, src, served, 0) - if compat_bf16: - (out / "compat-bf16.json").write_text(json.dumps({ - "source": str(src), "rounding": "nearest-even", "tensors": converted, - }, indent=2) + "\n", encoding="utf-8") - print("compat-bf16: %d tensors, %.2f GiB; expert and PLE table bytes unchanged" - % (len(converted), sum(t["bytes"] for t in converted) / 2**30)) return 0 @@ -203,25 +209,23 @@ def write_index(out, rows, src, served, n_extra): % (len(rows), served, n_extra, at / 2**30)) -def index_from_base(a, src, base, out, g, T, mm) -> int: +def index_from_base(src, base, out, model: Model) -> int: base_src = pathlib.Path(json.loads((base / "manifest.json").read_text(encoding="utf-8"))["source"]["shard1"]) if not base_src.exists(): print("cannot find the base pack's shard 1 from its manifest.json") return 1 - bg = G.GGUFFile(base_src) - BT = {t.name: t for t in bg.tensors} - bmm = np.memmap(base_src, dtype=np.uint8, mode="r") + bm = Model(base_src) header, rows = read_index(base / "index.txt") new_rows, extra = [], [] served = 0 for name, f in rows.items(): - t, bt = T.get(name), BT.get(name) + t, bt = model.get(name), bm.get(name) if t is None or bt is None: print("tensor %s missing from one of the models" % name) return 1 if t.type_name in FLOAT and bt.type_name in FLOAT: if t.type_name != bt.type_name or t.shape != bt.shape or \ - not np.array_equal(tensor_bytes(mm, g, t), tensor_bytes(bmm, bg, bt)): + not np.array_equal(model.bytes(t), bm.bytes(bt)): print("float tensor %s differs from the base model; this pack cannot reuse its dense.bin" % name) return 1 new_rows.append(list(f)) @@ -231,7 +235,7 @@ def index_from_base(a, src, base, out, g, T, mm) -> int: return 1 nbytes = t.expected_bytes() off = sum(len(b) + (-len(b)) % ALIGN for b in extra) - extra.append(tensor_bytes(mm, g, t).tobytes()) + extra.append(model.bytes(t).tobytes()) # file 3 = extra.bin, raw BF16 (index kind 4) new_rows.append([name, "3", "4", str(off), str(nbytes), "0", str(nbytes), f[7], f[8]] + ["0"] * 10) else: @@ -255,37 +259,20 @@ def index_from_base(a, src, base, out, g, T, mm) -> int: def main() -> int: ap = argparse.ArgumentParser() - ap.add_argument("--gguf", required=True, help="the model's shard 1") + ap.add_argument("--gguf", required=True, help="the model's first shard (the others are found by name)") ap.add_argument("--base", help="optional: a Q2_0 canonical pack whose dense.bin holds the shared float tensors") ap.add_argument("--out", required=True) ap.add_argument("--skip-experts", action="store_true", help="rewrite the index only") - ap.add_argument("--compat-bf16", action="store_true", - help="dequantize small non-native projections to BF16 for ordinary Qwen4Exp GGUFs " - "(rounds weights; leaves experts and the PLE table unchanged)") ap.add_argument("--experts-bin", action="store_true", help="also write experts.bin (the engine otherwise reads the experts from the GGUF itself)") a = ap.parse_args() - if a.compat_bf16 and a.base: - ap.error("--compat-bf16 cannot reuse --base dense weights") - # HF snapshot files are symlinks to hash-named blobs. Keep the shard filename for discovery. - src = pathlib.Path(a.gguf).absolute() + src = pathlib.Path(a.gguf).resolve() base = pathlib.Path(a.base).resolve() if a.base else None out = pathlib.Path(a.out) out.mkdir(parents=True, exist_ok=True) - g = G.GGUFFile(src) - mm = np.memmap(src, dtype=np.uint8, mode="r") model = Model(src) - T = {n: w[1] for n, w in model.where.items()} - if len(model.paths) > 1: - print("model shards: " + ", ".join(p.name for p in model.paths)) - if a.base: - if any(w[3] != src for w in model.where.values() if not w[1].name in NOT_IN_PACK): - print("--base needs a model whose tensors are all in shard 1") - return 1 - rc = index_from_base(a, src, base, out, g, {t.name: t for t in g.tensors}, mm) - else: - rc = index_standalone(src, out, model, a.compat_bf16) + rc = index_from_base(src, base, out, model) if a.base else index_standalone(src, out, model) if rc: return rc if not (out / "tokenizer" / "vocab.json").exists() or not (out / "tokenizer" / "chat_template.jinja").exists(): @@ -293,33 +280,24 @@ def main() -> int: check=True) # writes /tokenizer/ # ---- the experts - n_layers = 1 + max(int(n.split(".")[1]) for n in T if n.startswith("blk.") and n.endswith("_exps.weight")) - n_expert = int(T["blk.0.ffn_gate_inp.weight"].shape[1]) # router rows = experts kept (pruned models ship < 512) - if any(int(T["blk.%d.ffn_gate_inp.weight" % l].shape[1]) != n_expert for l in range(n_layers)): - print("the routers disagree on the expert count; a per-layer pruned model cannot be packed") - return 1 + names = [t.name for t in model.tensors()] + n_layers = 1 + max(int(n.split(".")[1]) for n in names if n.startswith("blk.") and n.endswith("_exps.weight")) layout, offset = [], 0 for l in range(n_layers): - ts = [T["blk.%d.ffn_%s_exps.weight" % (l, r)] for r in ROLES] - per = [t.expected_bytes() // n_expert for t in ts] + ts = [model.get("blk.%d.ffn_%s_exps.weight" % (l, r)) for r in ROLES] + per = [t.expected_bytes() // N_EXPERT for t in ts] if per[0] != per[1] or ts[0].type_name != ts[1].type_name: print("layer %d: gate and up differ in type" % l) return 1 blob = per[0] + per[1] + per[2] layout.append((l, ts[0].type_id, ts[2].type_id, offset, blob, ts)) - offset += blob * n_expert + offset += blob * N_EXPERT with open(out / "native_experts.txt", "w", encoding="utf-8", newline="\n") as fo: - fo.write("# strata native experts v3: layer gu_type d_type offset blob_bytes gate_off up_off down_off [shard] " - "(n_expert %d, total %d; absolute offsets in %s, or in the named shard beside it)\n" - % (n_expert, offset, src.name)) + fo.write("# strata native experts v3: layer gu_type d_type offset blob_bytes (n_expert %d, total %d; the " + "engine takes the experts from %s by tensor name, or from experts.bin)\n" + % (N_EXPERT, offset, src.name)) for l, gt, dt, off, blob, ts in layout: - ws = [model.where[t.name] for t in ts] - if len({w[3] for w in ws}) != 1: - print("layer %d: its gate/up/down tensors are in different shards" % l) - return 1 - gg, shard = ws[0][0], ws[0][3] - line = "%d %d %d %d %d %d %d %d" % (l, gt, dt, off, blob, *[gg.data_start + t.offset for t in ts]) - fo.write(line + ("" if shard == src else " " + shard.name) + "\n") + fo.write("%d %d %d %d %d\n" % (l, gt, dt, off, blob)) if a.skip_experts or not a.experts_bin: if (out / "experts.bin").exists() and not a.experts_bin: print("note: %s/experts.bin exists; the engine reads it instead of the GGUF" % out) @@ -330,9 +308,9 @@ def main() -> int: return 0 with open(path, "wb") as fo: for l, gt, dt, off, blob, ts in layout: - parts = [model.bytes(t.name).reshape(n_expert, -1) for t in ts] - chunk = np.concatenate(parts, axis=1) # (n_expert, blob): gate | up | down per expert - assert chunk.shape == (n_expert, blob) + parts = [model.bytes(t).reshape(N_EXPERT, -1) for t in ts] + chunk = np.concatenate(parts, axis=1) # (512, blob): gate | up | down per expert + assert chunk.shape == (N_EXPERT, blob) fo.write(chunk.tobytes()) if l % 8 == 0: print(" layer %2d %-8s/%-7s blob %8d at %.2f GiB" % (l, ts[0].type_name, ts[2].type_name, blob, diff --git a/tools/mtp_pack.py b/tools/mtp_pack.py index 022d2ea8ff5..b21c27427d8 100644 --- a/tools/mtp_pack.py +++ b/tools/mtp_pack.py @@ -4,16 +4,18 @@ no longer on this PC, so Strata packs its own from the 31 `mtp.*` tensors fetched by tools/mtp_fetch.py: python tools/mtp_pack.py --src Desktop/Strata/mtp-bf16 --experts q2_0 --out mtp-q2_0.gguf - python tools/mtp_pack.py --src ... --experts q4_0 --out mtp-q4_0.gguf (acceptance comparison arm) + python tools/mtp_pack.py --src ... --experts iq4_nl --out mtp-iq4_nl.gguf Layout: dense tensors (attention, indexer, hyper-connections, shared expert, router, fc/norms) stay BF16 (F32 for 1-D norms), ~0.18 GB. The routed experts keep the checkpoint's fused layout - `gate_up_proj` [512, 1280, 2560] and `down_proj` [512, 2560, 640], quantized along the last (input) axis - in one of: - q2_0 64-element blocks, grid {-1, 0, 1, 2} x d. The scale is chosen per block to MINIMIZE squared error over - that grid (the ggml reference sets d = max|w| and never uses the +2 level). Same format as the main - model's experts, so Strata's CPU VNNI kernel and GPU hit kernel serve it unchanged. ~0.71 GB. - q4_0 ggml reference rounding. ~1.42 GB. q8_0 ggml reference. ~2.67 GB. + q2_0 64-element blocks, grid {-1, 0, 1, 2} x d. The scale is chosen per block to MINIMIZE squared error over + that grid (the ggml reference sets d = max|w| and never uses the +2 level). Same format as the main + model's experts, so Strata's CPU VNNI kernel and GPU hit kernel serve it unchanged. ~0.71 GB. + iq4_nl ggml's quantizer without an importance matrix (relative RMS error ~0.078). ~1.42 GB. + q8_0 ggml reference (~0.006). ~2.67 GB. + iq4_nl and q8_0 run through the main model's native expert kernels (tools/mtp_rt.py writes them as GGUF blocks). This is round-to-nearest, not GSQ: the plan picks the expert format by MEASURED draft acceptance (P0.3/P6), not by this file's reconstruction error, which is reported per tensor only as a sanity check. No model runs here. @@ -23,8 +25,10 @@ import argparse import hashlib import json +import os import sys import time +from concurrent.futures import ThreadPoolExecutor from pathlib import Path import numpy as np @@ -82,26 +86,56 @@ def q2_0(w: np.ndarray) -> np.ndarray: return out.reshape(-1) -def q4_0(w: np.ndarray) -> np.ndarray: - """ggml quantize_row_q4_0_ref: d = max / -8 (signed max), code = clamp(round(x/d + 8.5) truncated, 0, 15).""" +IQ4_NL_GRID = np.array([-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113], dtype=np.float32) +IQ4_NL_MID = (IQ4_NL_GRID[:-1] + IQ4_NL_GRID[1:]) / 2 + + +def iq4_nl(w: np.ndarray) -> np.ndarray: + """ggml quantize_row_iq4_nl without an importance matrix (quantize_row_iq4_nl_impl, block 32, ntry 7): the scale + that maps the block's largest |w| onto the grid's end, then the best of 15 trial scales by weighted least + squares (weights w^2); codes are the nearest grid values (ties to the upper). fp16 d, then 16 bytes of 4-bit + codes: element j in the low nibble of byte j, element j + 16 in its high nibble.""" x = w.reshape(-1, 32).astype(np.float32) - idx = np.abs(x).argmax(axis=1) - mx = x[np.arange(x.shape[0]), idx][:, None] - d = mx / -8.0 - inv = np.where(d != 0, 1.0 / np.where(d != 0, d, 1.0), 0.0) - q = np.minimum(15, np.trunc(x * inv + 8.5)).astype(np.uint8) - packed = (q[:, :16] | (q[:, 16:] << 4)).astype(np.uint8) + wt = x * x + rows = np.arange(x.shape[0]) + mx = x[rows, np.abs(x).argmax(axis=1)] # the signed value of the largest |w| (first one) + live = np.abs(mx) >= 1e-15 + mx = np.where(live, mx, np.float32(1)) + + def fit(inv: np.ndarray): + q = IQ4_NL_GRID[np.searchsorted(IQ4_NL_MID, inv[:, None] * x, side="right")] + wq = wt * q + a, b = wq * x, wq * q + sqx, sq2 = a[:, 0].copy(), b[:, 0].copy() + for j in range(1, 32): # ggml's order of the sums + sqx += a[:, j] + sq2 += b[:, j] + return sqx, sq2 + + sqx, sq2 = fit(1 / (-mx / IQ4_NL_GRID[0])) + d = sqx / np.where(sq2 > 0, sq2, 1) + best = d * sqx + for t in range(-7, 8): + sqx, sq2 = fit((t + IQ4_NL_GRID[0]) / mx) + better = (sq2 > 0) & (sqx * sqx > best * sq2) + d = np.where(better, sqx / np.where(sq2 > 0, sq2, 1), d) + best = np.where(better, d * sqx, best) + d = np.where(live, d, 0).astype(np.float32) + inv = np.where(d != 0, 1 / np.where(d != 0, d, 1), 0).astype(np.float32) + codes = np.searchsorted(IQ4_NL_MID, inv[:, None] * x, side="right").astype(np.uint8) out = np.empty((x.shape[0], 18), dtype=np.uint8) out[:, :2] = d.astype(np.float16).view(np.uint8).reshape(-1, 2) - out[:, 2:] = packed + out[:, 2:] = codes[:, :16] | (codes[:, 16:] << 4) return out.reshape(-1) def q8_0(w: np.ndarray) -> np.ndarray: + """ggml quantize_row_q8_0_ref: d = max|w| / 127, codes rounded half away from zero.""" x = w.reshape(-1, 32).astype(np.float32) d = np.abs(x).max(axis=1, keepdims=True) / 127.0 inv = np.where(d > 0, 1.0 / np.where(d > 0, d, 1.0), 0.0) - q = np.rint(x * inv).astype(np.int8) + v = x * inv + q = (np.sign(v) * np.floor(np.abs(v) + 0.5)).astype(np.int8) out = np.empty((x.shape[0], 34), dtype=np.uint8) out[:, :2] = d.astype(np.float16).view(np.uint8).reshape(-1, 2) out[:, 2:] = q.view(np.uint8) @@ -115,19 +149,18 @@ def dequant(kind: str, blob: np.ndarray, n: int) -> np.ndarray: qs = b[:, 2:] codes = np.stack([(qs >> s) & 3 for s in (0, 2, 4, 6)], axis=2).reshape(-1, 64).astype(np.float32) return ((codes - 1.0) * d).reshape(-1)[:n] - if kind == "q4_0": + if kind == "iq4_nl": b = blob.reshape(-1, 18) d = b[:, :2].copy().view(np.float16).astype(np.float32) qs = b[:, 2:] - q = np.concatenate([qs & 15, qs >> 4], axis=1).astype(np.float32) - return ((q - 8.0) * d).reshape(-1)[:n] + return (IQ4_NL_GRID[np.concatenate([qs & 15, qs >> 4], axis=1)] * d).reshape(-1)[:n] b = blob.reshape(-1, 34) d = b[:, :2].copy().view(np.float16).astype(np.float32) return (b[:, 2:].view(np.int8).astype(np.float32) * d).reshape(-1)[:n] QUANT = {"q2_0": (q2_0, gguf.GGMLQuantizationType.Q2_0, 64, 18), - "q4_0": (q4_0, gguf.GGMLQuantizationType.Q4_0, 32, 18), + "iq4_nl": (iq4_nl, gguf.GGMLQuantizationType.IQ4_NL, 32, 18), "q8_0": (q8_0, gguf.GGMLQuantizationType.Q8_0, 32, 34)} @@ -146,7 +179,8 @@ def main() -> int: w.add_string("strata.mtp.source_sha256", hashlib.sha256( json.dumps({t["name"]: t["sha256"] for t in manifest}, sort_keys=True).encode()).hexdigest()) w.add_string("strata.mtp.expert_format", a.experts) - w.add_string("strata.mtp.expert_quantizer", "per-block MSE scale search" if a.experts == "q2_0" else "ggml reference") + w.add_string("strata.mtp.expert_quantizer", {"q2_0": "per-block MSE scale search", "iq4_nl": "ggml quantize_iq4_nl", + "q8_0": "ggml reference"}[a.experts]) report = [] for t in sorted(manifest, key=lambda t: t["name"]): name, shape = t["name"], t["shape"] @@ -158,7 +192,8 @@ def main() -> int: x = _Experts(raw) if shape[-1] % block: raise ValueError(f"{name}: inner dim {shape[-1]} not a multiple of {block}") - parts = [fn(x[e]) for e in range(shape[0])] # one expert at a time bounds memory + with ThreadPoolExecutor(os.cpu_count() or 4) as ex: # numpy releases the GIL in the quantizers + parts = list(ex.map(lambda e: fn(x[e]), range(shape[0]))) blob = np.concatenate(parts) errs = [] for e in range(min(a.check_experts, shape[0])): diff --git a/tools/mtp_rt.py b/tools/mtp_rt.py index 6276e4c85ce..a1e1e9cf188 100644 --- a/tools/mtp_rt.py +++ b/tools/mtp_rt.py @@ -3,9 +3,12 @@ python tools/mtp_rt.py --gguf /mtp-bf16/mtp-q2_0.gguf --out /mtp-bf16/rt Writes - experts.bin 512 routed experts in the engine's blob layout (`include/strata/kernels/cpu/expert.hpp`): gate/up - rows interleaved (2r = gate r, 2r+1 = up r), then down rows; the Q2_0 codes in one plane and the fp16 - scales in another. A lossless relayout of the GGUF's Q2_0 blocks (same bytes, `cpu_expert_fixture.py`). + experts.bin 512 routed experts. Q2_0 (tools/mtp_pack.py's default) in the engine's blob layout + (`include/strata/kernels/cpu/expert.hpp`): gate/up rows interleaved (2r = gate r, 2r+1 = up r), then + down rows; the Q2_0 codes in one plane and the fp16 scales in another, a lossless relayout of the + GGUF's blocks (same bytes, `cpu_expert_fixture.py`). Other formats (iq4_nl, q8_0) as the GGUF's + blocks, [gate rows | up rows | down rows] per expert, for the native expert kernels. + experts.txt only with the GGUF-block layout: the ggml type ids of gate/up and down. dense.bin every other tensor: the large projections quantized to Q8_0 (ggml's reference rounding) so the engine's multi-column MMVQ can run them; the hyper-connection and router weights kept BF16; the RMSNorm weights as F32 with the Gemma "+1" applied (vLLM GemmaRMSNorm scales by 1 + w). @@ -67,12 +70,23 @@ def main() -> int: out.mkdir(parents=True, exist_ok=True) r = gguf.GGUFReader(a.gguf) tens = {t.name: t for t in r.tensors} - gu = np.asarray(tens["mtp.layers.0.mlp.experts.gate_up_proj"].data) - dn = np.asarray(tens["mtp.layers.0.mlp.experts.down_proj"].data) - assert gu.shape == (NE, 2 * FF, 720) and dn.shape == (NE, H, 180), (gu.shape, dn.shape) + gut, dnt = tens["mtp.layers.0.mlp.experts.gate_up_proj"], tens["mtp.layers.0.mlp.experts.down_proj"] + gu, dn = np.asarray(gut.data), np.asarray(dnt.data) + q2 = gguf.GGMLQuantizationType.Q2_0 + fmt = out / "experts.txt" with open(out / "experts.bin", "wb") as f: - for e in range(NE): - f.write(blob_of(gu[e], dn[e])) + if gut.tensor_type == q2 and dnt.tensor_type == q2: + assert gu.shape == (NE, 2 * FF, 720) and dn.shape == (NE, H, 180), (gu.shape, dn.shape) + for e in range(NE): + f.write(blob_of(gu[e], dn[e])) + fmt.unlink(missing_ok=True) + else: + assert gu.shape[:2] == (NE, 2 * FF) and dn.shape[:2] == (NE, H), (gu.shape, dn.shape) + for e in range(NE): + f.write(gu[e].tobytes()) + f.write(dn[e].tobytes()) + fmt.write_text(f"{int(gut.tensor_type)} {int(dnt.tensor_type)}\n", encoding="utf-8") + blob = (out / "experts.bin").stat().st_size // NE lines = [] off = 0 with open(out / "dense.bin", "wb") as f: @@ -101,7 +115,8 @@ def main() -> int: lines.append(f"{short} {kind} {rows} {cols} {off} {len(raw)}") off += len(raw) (out / "dense.txt").write_text("\n".join(lines) + "\n", encoding="utf-8") - print(f"experts.bin {NE * BLOB} B, dense.bin {off} B, {len(lines)} tensors -> {out}") + print(f"experts.bin {NE * blob} B ({gut.tensor_type.name}/{dnt.tensor_type.name}), dense.bin {off} B, " + f"{len(lines)} tensors -> {out}") for l in lines: print(" " + l) return 0