diff --git a/CMakeLists.txt b/CMakeLists.txt index d1cf5fb..4d0afd9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -118,6 +118,7 @@ target_link_libraries(strata_engine PUBLIC strata_kernels) add_library(strata_models STATIC src/models/deepseek/deepseek_executor.cpp src/models/glm52/glm52_executor.cpp + src/models/glm53/glm53_executor.cpp src/models/gemma4/gemma4_executor.cpp src/models/kimi_k3/kimi_k3_executor.cpp src/models/laguna/laguna_executor.cpp @@ -129,6 +130,10 @@ add_library(strata_models STATIC src/models/gemma4/gemma4_runtime.cpp src/models/glm52/glm52_runtime.cpp src/models/glm52/glm52_manifest.cpp + src/models/glm53/glm53_manifest.cpp + src/models/glm53/glm53_sequence.cpp + src/models/glm53/glm53_checkpoint.cpp + src/models/glm53/glm53_runtime.cpp src/models/kimi_k3/kimi_k3_ops.cpp src/models/kimi_k3/kimi_k3_manifest.cpp src/models/kimi_k3/kimi_k3_checkpoint.cpp @@ -367,6 +372,7 @@ if(BUILD_TESTING) tests/test_deepseek_expert_residency.cpp tests/test_fp4_decode.cpp tests/test_glm52_ops.cpp + tests/test_glm53_manifest.cpp tests/test_gemma4_ops.cpp tests/test_gemma4_checkpoint.cpp tests/test_gemma4_image.cpp diff --git a/README.md b/README.md index 2036575..293eae1 100644 --- a/README.md +++ b/README.md @@ -35,8 +35,8 @@ make check # build and run the test suite ./build/strata-chat --model /path/to/checkpoint --model-type gemma4 ``` -`--model-type` is one of `gemma4`, `deepseek`, `glm`, `laguna`, `inkling`, -`kimi-k3`. Add `--devices 0,1` to pin GPUs, `--admission-only --json` to print +`--model-type` is one of `gemma4`, `deepseek`, `glm`, `glm53`, `laguna`, +`inkling`, `kimi-k3`. Add `--devices 0,1` to pin GPUs, `--admission-only --json` to print the placement plan without loading anything. There is also an OpenAI-compatible server (`strata-server`). @@ -51,6 +51,7 @@ There is also an OpenAI-compatible server (`strata-server`). | **Laguna S 2.1** | 48 layers 1:3 global/sliding, 256 experts + 1 shared, top-10 | NVFP4 or MXFP4 experts, BF16 elsewhere | Spine in VRAM; experts stream from RAM | | **Inkling Small** | 42 layers, 256 experts + 2 sinks, top-6, no rotary | NVFP4 or MXFP4 experts, BF16 elsewhere | Experts stream from RAM | | **GLM-5.2** | 78 layers, 256 experts, top-8 | INT4 group-128, W4A16 | Exceeds combined memory; I/O-dependent | +| **GLM-5.3-Flash** | 45 layers, 3 KDA : 1 sparse MLA, 288 experts + 1 shared, top-8 | FP8 E4M3 block-128 with F32 inverse scales | Text-only; streams checkpoint modules, exact through 2,048 tokens | | **Kimi-K3** | 93 layers, 3 KDA : 1 gated MLA, 896 experts, top-16 | MXFP4 experts, BF16 elsewhere | 1.45 TB; I/O-dependent, 38.6 s/step. Vision not implemented | Each runs its declared semantics as-is — hybrid compressed attention and @@ -135,7 +136,7 @@ record those as negatives with their measurements; 0165 records the positive. | [`docs/server.md`](docs/server.md) | the OpenAI-compatible HTTP API | | [`docs/models/`](docs/models/) | copy-paste build, chat, server, and measured-speed runbooks by model | | [`docs/current-architecture.md`](docs/current-architecture.md) | how the code is organised and what is enforced | -| [`docs/model-bringup-guide.md`](docs/model-bringup-guide.md) | adding a seventh model | +| [`docs/model-bringup-guide.md`](docs/model-bringup-guide.md) | adding another model | | [`docs/architecture.md`](docs/architecture.md) | the target scheduler design, not yet built | | [`docs/README.md`](docs/README.md) | product documentation index | | [`CONTRIBUTING.md`](CONTRIBUTING.md) | repository layout, architecture, and change hygiene rules | diff --git a/apps/strata_chat.cpp b/apps/strata_chat.cpp index aaeb8f1..fdc1934 100644 --- a/apps/strata_chat.cpp +++ b/apps/strata_chat.cpp @@ -168,7 +168,7 @@ R"(strata-chat -- interactive chat against a Strata runtime required: --model DIR checkpoint directory - --model-type TYPE gemma4 | deepseek | glm | laguna | inkling | kimi-k3 + --model-type TYPE gemma4 | deepseek | glm | glm53 | laguna | inkling | kimi-k3 session: --prompt TEXT answer TEXT and exit instead of prompting diff --git a/apps/strata_server.cpp b/apps/strata_server.cpp index 78b9eb9..3b5291d 100644 --- a/apps/strata_server.cpp +++ b/apps/strata_server.cpp @@ -15,9 +15,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -86,7 +88,7 @@ void stop_server(int) { void usage() { std::cerr << "usage: strata-server --model DIR --model-type " - "gemma4|deepseek|glm|laguna|inkling|kimi-k3\n" + "gemma4|deepseek|glm|glm53|laguna|inkling|kimi-k3\n" << " [--model-id ID] [--host ADDRESS] [--port N]\n" << " [--context-size N] [--max-new N]\n" << " [--devices 0,1,2] [--vram-fraction F]\n" @@ -867,7 +869,7 @@ class ApiServer { const Options& options_; strata::RuntimeSession& runtime_; strata::ModelTokenizer tokenizer_; - std::uint64_t request_id_{}; + std::atomic request_id_{}; }; enum class RouterModelStatus : std::uint8_t { @@ -1502,6 +1504,18 @@ int main(int argc, char** argv) { std::signal(SIGPIPE, SIG_IGN); std::cerr << "[ready] http://" << options.host << ':' << options.port << "\n"; ApiServer server(options, runtime, std::move(tokenizer.value)); + const bool concurrent_requests = + registration->model == strata::RuntimeModel::Glm53; + // The GLM scheduler multiplexes generation iterations, so its HTTP front + // end must not serialize independently arriving streams. Bound detached + // connection workers from the CPUs the host actually exposes instead of + // accumulating one joinable thread object for the lifetime of the server. + const auto discovered_cpus = std::max(1U, std::thread::hardware_concurrency()); + const auto maximum_clients = std::min(128U, + std::max(4U, discovered_cpus * 2U)); + std::atomic active_clients{}; + std::mutex clients_mutex; + std::condition_variable clients_drained; while (stop_requested == 0) { const int client = accept4(listening_socket, nullptr, nullptr, SOCK_CLOEXEC); if (client < 0) { @@ -1509,16 +1523,50 @@ int main(int argc, char** argv) { std::cerr << "warning: accept failed: " << std::strerror(errno) << '\n'; continue; } - HttpRequest request; - std::string error; - if (!read_request(client, request, error)) { - send_error(client, 400, "Bad Request", error, "invalid_request_error"); + const auto serve = [&server, &active_clients, &clients_drained]( + int socket) { + HttpRequest request; + std::string error; + if (!read_request(socket, request, error)) { + send_error(socket, 400, "Bad Request", error, + "invalid_request_error"); + } else { + server.handle(socket, request); + } + close(socket); + active_clients.fetch_sub(1U, std::memory_order_acq_rel); + clients_drained.notify_all(); + }; + if (concurrent_requests) { + if (active_clients.load(std::memory_order_acquire) >= + maximum_clients) { + send_error(client, 503, "Service Unavailable", + "GLM-5.3 request admission is full", + "server_overloaded"); + close(client); + continue; + } + active_clients.fetch_add(1U, std::memory_order_acq_rel); + try { + std::thread(serve, client).detach(); + } catch (const std::system_error& error) { + active_clients.fetch_sub(1U, std::memory_order_acq_rel); + send_error(client, 503, "Service Unavailable", + std::string("cannot start request worker: ") + + error.what(), + "server_overloaded"); + close(client); + } } else { - server.handle(client, request); + active_clients.fetch_add(1U, std::memory_order_acq_rel); + serve(client); } - close(client); } listening_socket = -1; + std::unique_lock clients_lock(clients_mutex); + clients_drained.wait(clients_lock, [&] { + return active_clients.load(std::memory_order_acquire) == 0U; + }); std::cerr << "[shutdown] stopped cleanly\n"; return 0; } diff --git a/apps/strata_tokenize.cpp b/apps/strata_tokenize.cpp index 5e17f4b..c98bb8a 100644 --- a/apps/strata_tokenize.cpp +++ b/apps/strata_tokenize.cpp @@ -17,15 +17,16 @@ int main(int argc, char** argv) { model_type = argv[++index]; } else { std::cerr << "usage: strata-tokenize --tokenizer FILE --model-type " - "gemma4|glm|deepseek|laguna|raw --prompt TEXT\n"; + "gemma4|glm|glm53|deepseek|laguna|raw --prompt TEXT\n"; return 2; } } if (tokenizer_path.empty() || prompt.empty() || - (model_type != "gemma4" && model_type != "glm" && model_type != "deepseek" && + (model_type != "gemma4" && model_type != "glm" && model_type != "glm53" && + model_type != "deepseek" && model_type != "laguna" && model_type != "raw")) { std::cerr << "usage: strata-tokenize --tokenizer FILE --model-type " - "gemma4|glm|deepseek|laguna|raw --prompt TEXT\n"; + "gemma4|glm|glm53|deepseek|laguna|raw --prompt TEXT\n"; return 2; } const auto loaded = strata::ModelTokenizer::load(tokenizer_path); @@ -36,6 +37,8 @@ int main(int argc, char** argv) { std::string rendered; if (model_type == "glm") { rendered = strata::render_glm52_user_prompt(prompt); + } else if (model_type == "glm53") { + rendered = strata::render_glm53_user_prompt(prompt); } else if (model_type == "deepseek") { rendered = strata::render_deepseek_v4_user_prompt(prompt); } else if (model_type == "gemma4") { diff --git a/docs/README.md b/docs/README.md index 7791b31..9341562 100644 --- a/docs/README.md +++ b/docs/README.md @@ -23,6 +23,8 @@ live in [`../CONTRIBUTING.md`](../CONTRIBUTING.md). - `kimi-k3-runtime.md` — Kimi-K3's pinned contract, cost model, NVMe write constraint, chat format, and a per-gate status table that says which gates have been measured and which have not. +- `models/glm53.md` — GLM-5.3-Flash's text-only contract, exact context bound, + and chat/server commands. - Gemma 4's pinned checkpoint, tokenizer, text/vision graph, and public runtime contract are described in the root README and `current-architecture.md`. - `../kernels/cuda/README.md` — native CUDA and non-CUDA stub behavior. @@ -39,7 +41,7 @@ a previous “not implemented” statement false. - `cli.md` — every command-line flag and the dry-run placement planner. - `sampling.md` — sampler options, their exact semantics, reproducibility. - `server.md` — `strata-server` and the OpenAI-compatible API. -- `model-bringup-guide.md` — **how to add a seventh model.** The procedure, the +- `model-bringup-guide.md` — **how to add another model.** The procedure, the five shared files, and the rules that are not obvious. Start here for a new architecture. diff --git a/docs/current-architecture.md b/docs/current-architecture.md index 637a925..42a5199 100644 --- a/docs/current-architecture.md +++ b/docs/current-architecture.md @@ -16,7 +16,7 @@ strata_device the CUDA backend, or its error-returning stub strata_kernels CPU reference kernels (Q4, INT4 group-128, attention) strata_engine placement solver and cache, residency, sampling, chat and session support, the model registry -strata_models six models, one directory each, plus the shared checkpoint +strata_models seven models, one directory each, plus the shared checkpoint reader, tokenizer and placement inventories strata_app RuntimeSession and the OpenAI protocol ``` @@ -26,7 +26,7 @@ strata_app RuntimeSession and the OpenAI protocol `strata_models` under `--whole-archive`; see "Model registration" below. `src/` mirrors this: `src/platform/`, `src/engine/`, `src/app/`, and -`src/models/{deepseek,glm52,gemma4,kimi_k3,laguna,inkling,common}/`. Headers mirror +`src/models/{deepseek,glm52,glm53,gemma4,kimi_k3,laguna,inkling,common}/`. Headers mirror the same ownership under `include/strata/{platform,device,kernels,engine,app}` and `include/strata/models//`. They are private application interfaces: their headers and libraries are not installed, and no source or ABI @@ -80,7 +80,7 @@ string-to-enum chain. `--whole-archive` on `strata_models` is load-bearing: a static library drops any member nothing references, and nothing references a self-registering translation unit by definition. `strata_app` links ahead of it because `ld` -resolves in one pass. A test asserts all six models are registered, because +resolves in one pass. A test asserts all seven models are registered, because without the flag `find_model` returns null for every model and the rest of the suite still passes. @@ -106,7 +106,9 @@ prefill then token-at-a-time decode. Gemma 4 performs bounded prefill with whole vision blocks, hybrid local/global attention, and a BF16 local-ring/global-full KV cache. DeepSeek performs bounded layer-major prefill pages with a multi-row router projection and exact row-ordered causal -transitions. Kimi-K3 batches over a token span throughout. Inkling has an +transitions. GLM-5.3 runs token-at-a-time hybrid KDA and sparse MLA, with its +exact text context capped at the sparse top-k of 2,048. Kimi-K3 batches over a +token span throughout. Inkling has an opt-in paged prefill (`prefill_page_tokens`) that runs attention and its four short convolutions row-serial — both carry row-ordered state — and batches the routed MoE expert-major between them; it is bit-identical to its @@ -148,7 +150,7 @@ not. present) runs Gemma 4 against `tests/fixtures/gemma4/layer-hash-trace.json` — a per-layer hidden-state hash plus per-operation hashes over a fixed prompt. The types are model-neutral (`include/strata/platform/diagnostics.hpp`); DeepSeek emits -the same records, and the remaining four models do not yet. +the same records, and the remaining five models do not yet. Its limits, stated because a gate nobody understands is worse than none: it covers **prefill only** — once the device KV path engages, a whole device's @@ -168,7 +170,7 @@ hardware probe), and a **plan cache** keyed by checkpoint identity, GPU identity, context size, device list, VRAM fraction and flags. `plan_model_placement` in `strata_engine` is a thin dispatcher through a -registered `PlacementPlanner`; the implementation that opens six different +registered `PlacementPlanner`; the implementation that opens seven different checkpoints lives in `strata_models` and installs itself at static-init. That inversion is why `strata_engine` names no model symbol. diff --git a/docs/model-bringup-guide.md b/docs/model-bringup-guide.md index dac146a..ff3e2ab 100644 --- a/docs/model-bringup-guide.md +++ b/docs/model-bringup-guide.md @@ -1,6 +1,6 @@ # Adding a model -A procedure derived from the six existing adapters. Follow it in order; the +A procedure derived from the seven existing adapters. Follow it in order; the ordering is part of the correctness contract. ## What it actually costs diff --git a/docs/models/README.md b/docs/models/README.md index 5ed07f5..fe929a3 100644 --- a/docs/models/README.md +++ b/docs/models/README.md @@ -12,4 +12,5 @@ transfers to a different context length or machine. | Inkling Small | [Build, chat, serve, and benchmark](inkling.md) | 9.072 tok/s fresh / 28.010 tok/s same-route warm | | DeepSeek V4 | [Build, serve, and benchmark](deepseek.md) | 9.171 tok/s with the routed-expert tier, 8.571 without | | GLM-5.2 | [Build, preflight, and run](glm52.md) | I/O-dependent on the reference workstation | +| GLM-5.3-Flash | [Build, preflight, chat, and serve](glm53.md) | I/O-dependent; exact text context currently capped at 2,048 tokens | | Kimi-K3 | [Build, preflight, and run](kimi-k3.md) | ~0.02 tok/s; SATA-bound on the reference workstation | diff --git a/docs/models/glm53.md b/docs/models/glm53.md new file mode 100644 index 0000000..3da5724 --- /dev/null +++ b/docs/models/glm53.md @@ -0,0 +1,66 @@ +# GLM-5.3-Flash: text runbook + +GLM-5.3-Flash is a 45-layer, 288-routed-expert model with four mHC streams, +three Kimi Delta Attention layers for every sparse MLA layer, eight selected +experts plus one shared expert, and block-128 FP8 E4M3 weights with F32 inverse +scales. Strata consumes the checkpoint as published and streams modules from +storage; it does not requantize the model or reduce its expert count or top-k. + +This adapter currently supports text only. Image or video content is rejected +before generation. The admitted context is capped at 2,048 tokens. At that +length the model's sparse index `top_k` is 2,048, so every causally visible key +is selected and dense causal MLA is exactly equivalent to running the sparse +indexer. Longer contexts fail admission until the k-pool indexer is implemented; +there is no silent dense or truncated fallback. + +## Build and preflight + +```bash +cmake -S . -B build-release -DCMAKE_BUILD_TYPE=Release +cmake --build build-release --parallel --target strata-chat strata-server + +./build-release/strata-chat \ + --model models/glm53f --model-type glm53 \ + --devices 0,1,2 --context-size 2048 --max-new 256 \ + --vram-fraction 0.85 --dry-run +``` + +The pinned release has 62 safetensors shards, 76,108 indexed tensors, and +328,326,771,576 indexed payload bytes. Admission validates those extents, the +hybrid layer schedule, text tensor roles, representative shapes and dtypes, +and the FP8 block geometry before generation. + +## Chat and server + +```bash +./build-release/strata-chat \ + --model models/glm53f --model-type glm53 \ + --devices 0,1,2 --context-size 2048 --max-new 256 + +./build-release/strata-server \ + --model models/glm53f --model-type glm53 --model-id glm53f \ + --devices 0,1,2 --context-size 2048 --max-new 256 --port 8080 +``` + +The server exposes the same text runtime through its OpenAI-compatible chat +completion endpoint. Do not send image content: multimodal support is outside +this adapter's current contract. + +## Current operating point + +The runtime discovers CPU width, free VRAM, peer topology and storage at +startup. On hosts where the routed checkpoint is much larger than the usable +CUDA cache, it maps canonical FP8 experts once and executes them directly from +host memory while keeping the non-expert spine, fused KDA state and mHC +transitions on CUDA. Prompt pages group rows by expert so an expert is traversed +once for every row that selected it; decode does not reread routed experts from +the checkpoint through Strata's explicit I/O path per token (the OS still owns +page-cache residency for mapped payloads). Systems with a high-speed two-GPU +peer fabric additionally admit the full TP2 route; PCIe systems use a +contiguous pipeline schedule. + +MTP drafting and verification are implemented but opt in because acceptance is +workload-dependent. Set `STRATA_GLM53_MTP=1` for an acceptance campaign. The +exact production latency path leaves it disabled. Resident absorbed MLA remains +an explicit experimental route (`STRATA_GLM53_RESIDENT_MLA=1`) until its +BF16-boundary equivalence gate is closed; it is never selected silently. diff --git a/include/strata/app/runtime.hpp b/include/strata/app/runtime.hpp index 96ccd69..08e928b 100644 --- a/include/strata/app/runtime.hpp +++ b/include/strata/app/runtime.hpp @@ -5,7 +5,7 @@ // RuntimeModel, RuntimeConfig, GenerationResult and friends moved to // model_executor.hpp in Phase 4 and are re-exported here, so every existing // `#include "strata/app/runtime.hpp"` keeps working unchanged. They live a tier -// lower because the six models implement ModelExecutor against them: had they +// lower because the seven models implement ModelExecutor against them: had they // stayed here, every model would depend upward on the application tier, which // is exactly the inversion check-symbols exists to catch. diff --git a/include/strata/device/cuda_backend.hpp b/include/strata/device/cuda_backend.hpp index 2fc999b..a8cb16c 100644 --- a/include/strata/device/cuda_backend.hpp +++ b/include/strata/device/cuda_backend.hpp @@ -16,12 +16,17 @@ namespace strata { class CudaBuffer; +class CudaWeight; enum class CudaWeightEncoding : std::uint8_t { Plain, OffsetPackedInt4, OffsetPackedInt8, Fp8E4m3Block128, + // E4M3 payload with one F32 inverse scale per 128x128 weight block. This + // is the standard dynamic-FP8 layout used by GLM-5.3-Flash; it is not the + // E8M0 scale byte used by DeepSeek and the two must never alias. + Fp8E4m3Block128F32, Fp4E2m1Group32, // compressed-tensors "nvfp4-pack-quantized": E2M1 nibble pairs with FP8 // E4M3 group scales and one FP32 per-tensor global scale. Weights @@ -65,6 +70,20 @@ struct CudaMatmulProfile { std::uint64_t d2h_nanoseconds{}; }; +// Independent same-device projections issued as one host-completion group. +// Inputs and outputs remain ordinary F32 host spans; the backend stages every +// source and destination in disjoint pinned regions, reuses its device +// workspace in stream order, and crosses back to the host once after the last +// projection. This changes scheduling only, never a matmul's numerical route. +struct CudaMatmulBatchItem { + const CudaWeight* weight{}; + std::span input; + std::uint32_t rows{}; + std::span output; + bool round_bf16_output{}; + bool fp8_tensor_page{}; +}; + struct CudaSynchronizationStats { std::uint64_t calls{}; std::uint64_t nanoseconds{}; @@ -647,6 +666,59 @@ struct CudaBufferPatch { std::span bytes; }; +// One GLM-5.3 KDA decode command. `state` owns the exact F32 recurrent matrix, +// causal convolution history, and immutable per-layer coefficients packed by +// the model runtime. Projection boundaries remain host-visible for now, but +// the O(H * D^2) recurrence never crosses PCIe after prompt admission. +struct CudaGlm53KdaRequest { + const CudaBuffer* state{}; + // When `input` is present these nine same-device weights turn the command + // into the complete KDA attention sublayer. All intermediate activations + // stay in the persistent state buffer and only the final O projection is + // published. Empty input preserves the narrower projected-input oracle. + const CudaWeight* query_projection{}; + const CudaWeight* key_projection{}; + const CudaWeight* value_projection{}; + const CudaWeight* forget_a_projection{}; + const CudaWeight* beta_projection{}; + const CudaWeight* gate_a_projection{}; + const CudaWeight* forget_b_projection{}; + const CudaWeight* gate_b_projection{}; + // Optional same-device BF16/FP8 output projection. When present, the normalized + // KDA heads never return to the host; `output` is the projection's BF16 + // row rather than the unprojected head row. + const CudaWeight* output_projection{}; + std::span input; + std::span query; + std::span key; + std::span value; + std::span forget; + std::span beta; + std::span gate; + std::uint32_t heads{}; + std::uint32_t head_dim{}; + std::uint32_t convolution_kernel{}; + // Consume the BF16 layer input from, and publish the BF16 branch into, + // the active resident mHC workspace. `input` and `output` are empty and + // the complete attention command remains stream ordered in this mode. + bool mhc_source_destination{}; +}; + +struct CudaGlm53MlaRequest { + const CudaBuffer* state{}; + const CudaWeight* query_a{}; + const CudaWeight* key_value_a{}; + const CudaWeight* query_b{}; + const CudaWeight* key_value_b{}; + const CudaWeight* output{}; + std::uint32_t position{}; + std::uint32_t maximum_context{}; + std::uint32_t heads{}; + std::uint32_t head_dim{}; + std::uint32_t query_rank{}; + std::uint32_t key_value_rank{}; +}; + // Persistent target-format inputs for one DeepSeek mHC pre boundary. The // projection remains F32 as in the accepted SM86 contract; scale/base and the // BF16 norm weight are packed into one immutable auxiliary allocation. @@ -844,7 +916,9 @@ enum class CudaMatmulRoute : std::uint8_t { PackedOffsetInt, Nvfp4Group16, Fp8TensorPage, + Fp8F32TensorPage, Fp8E4m3Block128, + Fp8E4m3Block128F32, Fp4E2m1Group32, // MIX-2 register-fed skinny kernels. These are the accepted QPN shapes made // model agnostic: any weight whose encoding and extents admit the m16n8k16 @@ -852,6 +926,7 @@ enum class CudaMatmulRoute : std::uint8_t { // shapes the fragment layout cannot express. The census distinguishes them // so a run can show which of the two actually served a dispatch. Fp8RegisterFed, + Fp8F32RegisterFed, Fp4RegisterFed, GemmaMarlin, // MoE expert batches dispatch through CudaBackend::enqueue_moe, which @@ -859,6 +934,8 @@ enum class CudaMatmulRoute : std::uint8_t { // matmul_impl is load-only. These are counted separately so a census can // distinguish a load-time dispatch from a per-token one. MoePlainBf16, + MoeFp8E4m3Block128F32, + MoeFp8F32RegisterFed, MoeNvfp4Group16, MoeFp4E2m1Group32, MoePackedInt4, @@ -911,8 +988,9 @@ void record_cuda_matmul_route(CudaMatmulRoute route) noexcept; class CudaBackend { public: - // Permutes an Fp8E4m3Block128 or Fp4E2m1Group32 weight from its canonical - // layout into m16n8k16 fragment order, in place. The fragment order + // Permutes an Fp8E4m3Block128, Fp8E4m3Block128F32, or Fp4E2m1Group32 + // weight from its canonical layout into m16n8k16 fragment order, in place. + // The fragment order // REPLACES the canonical device layout -- one-copy residency, not a second // buffer -- so every consumer of that weight must expect fragment order // afterwards. matmul_impl and the shared expert call this themselves on @@ -941,6 +1019,14 @@ class CudaBackend { [[nodiscard]] static bool compiled() noexcept; [[nodiscard]] static std::vector available_devices(); [[nodiscard]] static ParseResult device_memory(int device); + // Host NUMA node nearest this CUDA device, discovered through its PCI + // identity. Returns -1 when the platform cannot describe that affinity. + [[nodiscard]] static int device_numa_node(int device) noexcept; + // True only for the CUDA driver's best peer-performance rank. Mere peer + // addressability over a contended PCIe host bridge is not enough to make + // fine-grained cross-device projection barriers profitable. + [[nodiscard]] static bool high_speed_peer_access_supported( + int source, int destination) noexcept; [[nodiscard]] static std::uint64_t weight_storage_bytes( std::uint64_t weight_bytes, std::uint64_t scale_bytes) noexcept; @@ -975,6 +1061,11 @@ class CudaBackend { enum class UploadCompletion : std::uint8_t { Synchronous, Deferred, + // The caller additionally guarantees that any weight which an + // in-flight MoE command references remains leased. Allocation and H2D + // may then use the independent upload stream while that command runs; + // synchronize_uploads() orders the next execution-stream consumer. + DeferredConcurrent, }; // `prepack` asks for the weight to be permuted into m16n8k16 fragment // order as part of the upload, stream-ordered behind the copy so the loader @@ -1007,6 +1098,15 @@ class CudaBackend { std::span output); [[nodiscard]] ValidationResult allocate_buffer( int device, std::uint64_t bytes, CudaBuffer& output); + [[nodiscard]] ValidationResult glm53_kda_decode( + const CudaGlm53KdaRequest& request, std::span output); + [[nodiscard]] ValidationResult glm53_mhc_router( + int device, const CudaWeight& router, std::span logits); + [[nodiscard]] ValidationResult glm53_mhc_swiglu( + int device, const CudaWeight& gate, const CudaWeight& up, + const CudaWeight& down, std::uint32_t intermediate); + [[nodiscard]] ValidationResult glm53_mla_decode_to_mhc( + const CudaGlm53MlaRequest& request); [[nodiscard]] ValidationResult upload_gemma4_kv( const CudaBuffer& cache, std::span keys, std::span values, std::uint32_t start, @@ -1039,6 +1139,8 @@ class CudaBackend { bool round_bf16_output = false, CudaMatmulProfile* profile = nullptr, bool dsv4_fp8_tensor_page = false); + [[nodiscard]] ValidationResult matmul_batch( + std::span items); [[nodiscard]] ValidationResult matmul_softcap( const CudaWeight& weight, std::span input, float softcap, std::span output); @@ -1061,6 +1163,13 @@ class CudaBackend { // projection can use the SM86 BF16-WMMA path. Other capabilities retain // the native FP8 CUDA-core kernel as a numerical fallback. [[nodiscard]] bool dsv4_fp8_tensor_page_supported(int device) const noexcept; + // Model-neutral tensor-core page projection for E4M3 weights with F32 + // block-128 scales and continuous per-row activation scales. Capability is + // discovered from the selected CUDA device; callers still opt in per + // projection so a numerical route change is never implicit. + [[nodiscard]] bool fp8_f32_tensor_page_supported(int device) const noexcept; + [[nodiscard]] bool fp8_f32_register_fed_supported( + int device) const noexcept; [[nodiscard]] ValidationResult validate_dsv4_mhc_device( int device) const; // Executes the model-neutral forward attention primitive under the @@ -1152,6 +1261,8 @@ class CudaBackend { std::span hidden); [[nodiscard]] ValidationResult dsv4_mhc_finish_device( int device, std::span hidden); + [[nodiscard]] ValidationResult dsv4_mhc_download_layer_input( + int device, std::span layer_input); // Device-only rank-local mHC bridges. These preserve the existing state // machine while keeping the attention/FFN boundary on the CUDA stream. [[nodiscard]] ValidationResult dsv4_mhc_device_view( @@ -1339,11 +1450,18 @@ class CudaBackend { [[nodiscard]] ValidationResult enqueue_moe( int device, std::span hidden, std::uint32_t rows, std::span routed, - const CudaMoeExpert* shared = nullptr); + const CudaMoeExpert* shared = nullptr, + float swiglu_limit = 0.0F); + [[nodiscard]] ValidationResult enqueue_glm53_moe_from_mhc( + int device, std::span routed, + const CudaMoeExpert& shared, std::span coefficients, + float swiglu_limit); [[nodiscard]] ValidationResult collect_moe( int device, std::span routed_output, std::span shared_output = {}); [[nodiscard]] ValidationResult synchronize(int device); + [[nodiscard]] ValidationResult profiler_start(); + [[nodiscard]] ValidationResult profiler_stop(); [[nodiscard]] CudaBackendStats stats() const noexcept; @@ -1355,6 +1473,12 @@ class CudaBackend { int device, const CudaDsv4MhcWeights& weights, std::span hidden, std::span weighted, std::span layer_input, bool device_only); + [[nodiscard]] ValidationResult enqueue_moe_impl( + int device, std::span hidden, std::uint32_t rows, + std::span routed, + const CudaMoeExpert* shared, float swiglu_limit, + bool mhc_source_destination, + std::span routed_coefficients); [[nodiscard]] ValidationResult dsv4_mhc_transition_impl( int device, const CudaDsv4MhcWeights& next_weights, std::span branch_output, std::span weighted, @@ -1378,7 +1502,10 @@ class CudaBackend { std::uint64_t rows_per_group, std::span output, float softcap, bool round_output = false, CudaMatmulProfile* profile = nullptr, - bool dsv4_fp8_tensor_page = false); + bool dsv4_fp8_tensor_page = false, + const std::byte* batch_input = nullptr, + std::byte* batch_output = nullptr, + bool defer_completion = false); struct Impl; std::unique_ptr impl_; }; diff --git a/include/strata/engine/model_executor.hpp b/include/strata/engine/model_executor.hpp index b33bbef..fbc3e93 100644 --- a/include/strata/engine/model_executor.hpp +++ b/include/strata/engine/model_executor.hpp @@ -3,7 +3,7 @@ // The model seam. // // Everything above this header (RuntimeSession, the CLIs, the server) speaks -// only ModelExecutor and the registry. Everything below it (the six concrete +// only ModelExecutor and the registry. Everything below it (the seven concrete // runtimes) implements ModelExecutor and registers itself from its own // translation unit. // @@ -38,6 +38,7 @@ namespace strata { enum class RuntimeModel : std::uint8_t { Glm52, + Glm53, DeepSeekV4, Gemma4, KimiK3, diff --git a/include/strata/engine/placement.hpp b/include/strata/engine/placement.hpp index fb2e260..dae4784 100644 --- a/include/strata/engine/placement.hpp +++ b/include/strata/engine/placement.hpp @@ -21,6 +21,7 @@ inline constexpr std::uint32_t kPlacementPlanVersion = 2U; enum class PlacementModel : std::uint8_t { Glm52, + Glm53, DeepSeekV4, Gemma4, KimiK3, @@ -238,8 +239,8 @@ using PlacementPlanResult = ParseResult; // Opens the checkpoint index and shard headers, sizes every component, and // solves. Reads no tensor payload and uploads nothing. // The planner seam. plan_model_placement below is a thin dispatcher in -// strata_engine; the implementation that knows how to open six different -// checkpoints and size six different inventories lives in strata_models and +// strata_engine; the implementation that knows how to open seven different +// checkpoints and size seven different inventories lives in strata_models and // installs itself here at static-init time. // // This exists to invert one specific dependency. Before it, placement_cache.cpp diff --git a/include/strata/engine/runtime_support.hpp b/include/strata/engine/runtime_support.hpp index 7476457..09010f8 100644 --- a/include/strata/engine/runtime_support.hpp +++ b/include/strata/engine/runtime_support.hpp @@ -90,7 +90,7 @@ using RuntimeDevicePlanResult = ParseResult; // Render a chat template, encode it, and make it fit the context. // -// Five of the six runtimes wrote this loop themselves, identically apart from +// Six of the seven runtimes wrote this loop themselves, identically apart from // which template function they called and whether an over-long prompt trims or // reports. The shape is: render, encode, check that the prompt plus the // requested generation fits, and if it does not, drop the oldest turn and diff --git a/include/strata/models/common/tokenizer.hpp b/include/strata/models/common/tokenizer.hpp index 417dd02..19ea8ca 100644 --- a/include/strata/models/common/tokenizer.hpp +++ b/include/strata/models/common/tokenizer.hpp @@ -79,6 +79,16 @@ class ModelTokenizer { std::span messages, std::string_view reasoning_effort = "medium-high", bool enable_thinking = true); +// GLM-5.3 ships a distinct template: no newline after , three reasoning +// budgets with Max as the fallback, and cleared reasoning on prior assistant +// turns in chat mode. +[[nodiscard]] std::string render_glm53_user_prompt( + std::string_view user_text, std::string_view reasoning_effort = "max", + bool clear_thinking = true); +[[nodiscard]] std::string render_glm53_chat_prompt( + std::span messages, + std::string_view reasoning_effort = "max", + bool clear_thinking = true); [[nodiscard]] std::string render_deepseek_v4_user_prompt( std::string_view user_text, bool enable_thinking = false); [[nodiscard]] std::string render_deepseek_v4_chat_prompt( diff --git a/include/strata/models/glm53/glm53_checkpoint.hpp b/include/strata/models/glm53/glm53_checkpoint.hpp new file mode 100644 index 0000000..44b95f5 --- /dev/null +++ b/include/strata/models/glm53/glm53_checkpoint.hpp @@ -0,0 +1,90 @@ +#pragma once + +#include "strata/device/cuda_backend.hpp" +#include "strata/models/glm53/glm53_manifest.hpp" +#include "strata/platform/checkpoint_io.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace strata { + +class Glm53CheckpointReader; + +struct Glm53CheckpointOpenResult { + std::unique_ptr value; + std::vector errors; + + [[nodiscard]] bool ok() const noexcept { + return errors.empty() && value != nullptr; + } +}; + +class Glm53CheckpointReader { +public: + Glm53CheckpointReader(const Glm53CheckpointReader&) = delete; + Glm53CheckpointReader& operator=(const Glm53CheckpointReader&) = delete; + Glm53CheckpointReader(Glm53CheckpointReader&&) = delete; + Glm53CheckpointReader& operator=(Glm53CheckpointReader&&) = delete; + ~Glm53CheckpointReader(); + + [[nodiscard]] static Glm53CheckpointOpenResult open( + std::string model_directory); + [[nodiscard]] const Glm53ManifestTensor* find( + std::string_view name) const noexcept; + [[nodiscard]] ParseResult> read( + std::string_view name, std::uint64_t maximum_bytes) const; + [[nodiscard]] ParseResult> view( + std::string_view name) const; + [[nodiscard]] ParseResult> read_f32( + std::string_view name, std::uint64_t maximum_elements) const; + [[nodiscard]] ParseResult> read_f32_row( + std::string_view name, std::uint64_t row) const; + [[nodiscard]] std::uint64_t cuda_linear_storage_bytes( + std::string_view base_name) const; + [[nodiscard]] std::uint64_t cuda_linear_slice_storage_bytes( + std::string_view base_name, std::uint64_t row_begin, + std::uint64_t row_count) const; + [[nodiscard]] ValidationResult load_cuda_linear( + std::string_view base_name, std::uint64_t rows, + std::uint64_t columns, int device, CudaBackend& backend, + CudaWeight& output, bool concurrent_prefetch = false) const; + [[nodiscard]] ValidationResult load_cuda_linear_slice( + std::string_view base_name, std::uint64_t total_rows, + std::uint64_t columns, std::uint64_t row_begin, + std::uint64_t row_count, int device, CudaBackend& backend, + CudaWeight& output) const; + + [[nodiscard]] const Glm53TextConfig& config() const noexcept { + return config_; + } + [[nodiscard]] const Glm53IndexManifest& manifest() const noexcept { + return manifest_; + } + +private: + Glm53CheckpointReader() = default; + [[nodiscard]] ParseResult> read_slice( + const Glm53ManifestTensor& tensor, std::uint64_t offset, + std::uint64_t bytes) const; + + std::string model_directory_; + Glm53TextConfig config_; + Glm53IndexManifest manifest_; + std::unordered_map by_name_; + CheckpointShardSet shards_; + struct ShardMapping { + std::byte* address{}; + std::uint64_t bytes{}; + }; + mutable std::mutex mapping_mutex_; + mutable std::unordered_map mappings_; +}; + +} // namespace strata diff --git a/include/strata/models/glm53/glm53_manifest.hpp b/include/strata/models/glm53/glm53_manifest.hpp new file mode 100644 index 0000000..678bbf5 --- /dev/null +++ b/include/strata/models/glm53/glm53_manifest.hpp @@ -0,0 +1,173 @@ +#pragma once + +#include "strata/platform/result.hpp" +#include "strata/platform/safetensors.hpp" + +#include +#include +#include +#include +#include +#include + +namespace strata { + +struct Glm53TextConfig { + std::uint32_t hidden_size{}; + std::uint32_t layer_count{}; + std::uint32_t attention_heads{}; + std::uint32_t key_value_heads{}; + std::uint32_t query_lora_rank{}; + std::uint32_t kv_lora_rank{}; + std::uint32_t nope_head_dim{}; + std::uint32_t rope_head_dim{}; + std::uint32_t value_head_dim{}; + std::uint32_t linear_attention_heads{}; + std::uint32_t linear_head_dim{}; + std::uint32_t short_conv_kernel{}; + std::uint32_t dense_intermediate_size{}; + std::uint32_t expert_intermediate_size{}; + std::uint32_t routed_experts{}; + std::uint32_t experts_per_token{}; + std::uint32_t expert_groups{}; + std::uint32_t selected_expert_groups{}; + std::uint32_t shared_experts{}; + std::uint32_t dense_prefix_layers{}; + std::uint32_t vocabulary_size{}; + std::uint32_t maximum_context_tokens{}; + std::uint32_t mhc_multiplier{}; + std::uint32_t mhc_sinkhorn_iterations{}; + std::uint32_t index_heads{}; + std::uint32_t index_head_dim{}; + std::uint32_t index_topk{}; + std::uint32_t index_pool{}; + std::uint32_t fp8_block_rows{}; + std::uint32_t fp8_block_columns{}; + float rms_epsilon{}; + float mhc_epsilon{}; + float routed_scale{}; + float swiglu_limit{}; + float kda_gate_lower_bound{}; + bool normalize_topk{}; + bool mhc{}; + bool mla_use_nope{}; + bool index_pool_compress{}; + bool index_pool_select_tail{}; + bool tie_word_embeddings{}; + std::string architecture; + std::string model_type; + std::string hidden_activation; + std::string router_scoring; + std::string topk_method; + std::string quantization_method; + std::string quantization_format; + std::vector attention_layer_types; + std::vector mlp_layer_types; + std::vector full_attention_layers; + std::vector kda_layers; +}; + +struct Glm53ConfigResult { + Glm53TextConfig value; + std::vector errors; + + [[nodiscard]] bool ok() const noexcept { return errors.empty(); } +}; + +enum class Glm53TensorRole : std::uint8_t { + Embedding, + OutputHead, + Norm, + Mhc, + KdaAttention, + SparseAttention, + AttentionIndexer, + DenseMlp, + Router, + SharedExpert, + RoutedExpert, + Mtp, + Vision, + Count, +}; + +enum class Glm53TensorComponent : std::uint8_t { + Weight, + Scale, + Bias, + State, +}; + +enum class Glm53TensorEncoding : std::uint8_t { + Plain, + Fp8E4m3Block128F32, +}; + +struct Glm53ManifestTensor { + std::string name; + std::string shard; + Glm53TensorRole role{Glm53TensorRole::Norm}; + Glm53TensorComponent component{Glm53TensorComponent::State}; + Glm53TensorEncoding encoding{Glm53TensorEncoding::Plain}; + std::int32_t layer{-1}; + std::int32_t expert{-1}; + SafetensorsDtype source_dtype{SafetensorsDtype::Other}; + std::vector source_shape; + std::uint64_t source_offset{}; + std::uint64_t source_bytes{}; +}; + +struct Glm53IndexManifest { + std::uint64_t indexed_tensor_bytes{}; + std::vector shards; + std::vector tensors; + std::array(Glm53TensorRole::Count)> + role_counts{}; + std::uint64_t fp8_modules{}; + std::uint64_t dense_spine_bytes{}; + std::uint64_t routed_expert_bytes{}; + std::uint64_t vision_bytes{}; + std::uint64_t scanned_shards{}; + std::uint64_t shard_file_bytes{}; + std::uint64_t tensor_payload_bytes{}; +}; + +struct Glm53ManifestResult { + Glm53IndexManifest manifest; + std::vector errors; + + [[nodiscard]] bool ok() const noexcept { return errors.empty(); } +}; + +struct Glm53CheckpointOptions { + bool require_all_shards{true}; + std::size_t maximum_errors{64U}; +}; + +[[nodiscard]] constexpr bool glm53_full_attention_layer( + std::uint32_t layer) noexcept { + return layer < 45U && (layer + 1U) % 4U == 0U; +} + +[[nodiscard]] constexpr bool glm53_kda_layer(std::uint32_t layer) noexcept { + return layer < 45U && !glm53_full_attention_layer(layer); +} + +[[nodiscard]] constexpr bool glm53_moe_layer(std::uint32_t layer) noexcept { + return layer >= 3U && layer < 45U; +} + +[[nodiscard]] Glm53ConfigResult parse_glm53_config(std::string_view json); +[[nodiscard]] ValidationResult validate_glm53_config( + const Glm53TextConfig& config); +[[nodiscard]] Glm53TensorRole classify_glm53_tensor( + std::string_view name, std::int32_t& layer, std::int32_t& expert) noexcept; +[[nodiscard]] Glm53ManifestResult build_glm53_index_manifest( + SafetensorsIndex index); +[[nodiscard]] Glm53ManifestResult validate_glm53_checkpoint( + const std::string& model_directory, Glm53IndexManifest manifest, + const Glm53CheckpointOptions& options = {}); +[[nodiscard]] std::string_view to_string(Glm53TensorRole role) noexcept; +[[nodiscard]] std::string_view to_string(Glm53TensorEncoding encoding) noexcept; + +} // namespace strata diff --git a/include/strata/models/glm53/glm53_runtime.hpp b/include/strata/models/glm53/glm53_runtime.hpp new file mode 100644 index 0000000..14ce758 --- /dev/null +++ b/include/strata/models/glm53/glm53_runtime.hpp @@ -0,0 +1,80 @@ +#pragma once + +#include "strata/engine/chat_protocol.hpp" +#include "strata/engine/sampling.hpp" +#include "strata/platform/result.hpp" +#include "strata/platform/types.hpp" + +#include +#include +#include +#include +#include +#include + +namespace strata { + +// Deterministic weighted assignment used by both warmup and execution. The +// same projection key therefore has exactly one CUDA home, while independent +// projections spread according to discovered cache capacity. +[[nodiscard]] std::vector glm53_projection_slots( + std::span keys, + std::span costs, + std::span capacities, + std::size_t preferred_slot); + +struct Glm53RuntimeConfig { + std::vector devices; + double vram_cache_fraction{0.85}; + std::uint32_t maximum_context_tokens{2048U}; + double sampling_temperature{}; + std::uint64_t sampling_seed{33'377'335U}; + bool verbose{}; + bool load_progress{}; +}; + +struct Glm53RunMetrics { + std::uint64_t prompt_tokens{}; + std::uint64_t prefill_tokens{}; + std::uint64_t reused_prompt_tokens{}; + std::uint64_t decode_tokens{}; + double prefill_seconds{}; + double decode_seconds{}; +}; + +struct Glm53GenerationResult { + std::string text; + std::vector prompt_token_ids; + std::vector generated_token_ids; + std::vector logprobs; + Glm53RunMetrics metrics; + std::vector errors; + bool stopped{}; + + [[nodiscard]] bool ok() const noexcept { return errors.empty(); } +}; + +class Glm53Runtime { +public: + Glm53Runtime(); + ~Glm53Runtime(); + Glm53Runtime(Glm53Runtime&&) noexcept; + Glm53Runtime& operator=(Glm53Runtime&&) noexcept; + Glm53Runtime(const Glm53Runtime&) = delete; + Glm53Runtime& operator=(const Glm53Runtime&) = delete; + + [[nodiscard]] ValidationResult initialize( + const std::string& model_directory, + const Glm53RuntimeConfig& config = {}); + [[nodiscard]] Glm53GenerationResult generate_chat_stream( + std::span messages, + std::uint32_t maximum_new_tokens, const SamplingOptions& sampling, + std::span stop, + const TokenStreamCallback& on_token = {}); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace strata diff --git a/include/strata/models/glm53/glm53_sequence.hpp b/include/strata/models/glm53/glm53_sequence.hpp new file mode 100644 index 0000000..357f6c9 --- /dev/null +++ b/include/strata/models/glm53/glm53_sequence.hpp @@ -0,0 +1,102 @@ +#pragma once + +#include "strata/platform/result.hpp" + +#include +#include +#include +#include +#include +#include + +namespace strata { + +// 45 target layers plus the checkpoint's exact next-token prediction layer. +inline constexpr std::uint32_t kGlm53LayerCount = 46U; +inline constexpr std::uint32_t kGlm53KdaHeads = 64U; +inline constexpr std::uint32_t kGlm53KdaHeadWidth = 128U; +inline constexpr std::uint32_t kGlm53KdaWidth = + kGlm53KdaHeads * kGlm53KdaHeadWidth; +inline constexpr std::uint32_t kGlm53MlaRank = 512U; + +// Physical, copy-on-write rows used by the sparse-MLA cache. A logical +// sequence owns a page table; prefix snapshots share immutable pages and only +// the page receiving the next append is copied. State is F32 exactly as the +// reference runtime -- this is paging, not cache quantization. +class Glm53PagedRows { +public: + Glm53PagedRows() = default; + Glm53PagedRows(std::uint32_t columns, std::uint32_t page_rows); + + [[nodiscard]] ValidationResult reset( + std::uint32_t columns, std::uint32_t page_rows); + [[nodiscard]] ValidationResult append(std::span row); + [[nodiscard]] ValidationResult append_rows( + std::span rows, std::uint32_t row_count); + [[nodiscard]] ValidationResult truncate(std::uint32_t rows); + [[nodiscard]] std::span row(std::uint32_t index) const noexcept; + [[nodiscard]] std::vector materialize() const; + + [[nodiscard]] std::uint32_t columns() const noexcept { return columns_; } + [[nodiscard]] std::uint32_t page_rows() const noexcept { return page_rows_; } + [[nodiscard]] std::uint32_t rows() const noexcept { return rows_; } + [[nodiscard]] std::size_t physical_pages() const noexcept { + return pages_.size(); + } + [[nodiscard]] std::uint64_t private_bytes() const noexcept; + +private: + struct Page { + explicit Page(std::size_t elements) : values(elements) {} + std::vector values; + }; + + [[nodiscard]] ValidationResult ensure_append_page(); + + std::uint32_t columns_{}; + std::uint32_t page_rows_{}; + std::uint32_t rows_{}; + std::vector> pages_; +}; + +// All mutable state for one GLM text sequence. Large KDA matrices and short +// convolution histories are allocated lazily and copied on first write after +// a prefix fork. Sparse MLA uses the physical page table above. +class Glm53SequenceState { +public: + Glm53SequenceState() = default; + + [[nodiscard]] ValidationResult reset( + std::uint32_t maximum_context_tokens, + std::uint32_t mla_page_rows = 64U); + [[nodiscard]] std::span recurrent(std::uint32_t layer); + [[nodiscard]] std::span convolution( + std::uint32_t layer, std::uint32_t projection); + [[nodiscard]] Glm53PagedRows& mla(std::uint32_t layer); + [[nodiscard]] const Glm53PagedRows& mla(std::uint32_t layer) const; + void copy_mla_from(std::uint32_t layer, + const Glm53SequenceState& source); + + [[nodiscard]] std::uint32_t token_count() const noexcept { + return token_count_; + } + void set_token_count(std::uint32_t value) noexcept { token_count_ = value; } + [[nodiscard]] std::uint32_t maximum_context_tokens() const noexcept { + return maximum_context_tokens_; + } + [[nodiscard]] std::uint64_t private_bytes() const noexcept; + +private: + using Buffer = std::shared_ptr>; + [[nodiscard]] static std::span writable( + Buffer& buffer, std::size_t elements); + + std::array recurrent_{}; + std::array, kGlm53LayerCount> convolution_{}; + std::array mla_{}; + std::uint32_t maximum_context_tokens_{}; + std::uint32_t mla_page_rows_{64U}; + std::uint32_t token_count_{}; +}; + +} // namespace strata diff --git a/include/strata/platform/hardware_profile.hpp b/include/strata/platform/hardware_profile.hpp index 8c9e581..16fef4d 100644 --- a/include/strata/platform/hardware_profile.hpp +++ b/include/strata/platform/hardware_profile.hpp @@ -38,6 +38,10 @@ struct HardwareProfile { // the machine's total -- a cgroup or a taskset makes those differ, and the // affinity mask is the one that governs how many threads are useful. std::size_t usable_cpus{}; + // Exact affinity-filtered logical CPU ids. A device-addressed worker pool + // needs identities, not only a count, to stay local without assuming the + // host's CPU numbering. + std::vector usable_cpu_ids; NumaTopology numa; // CPUs on the smallest online node. The relevant figure for anything that // assigns one worker pool per node, because the smallest node is what diff --git a/include/strata/platform/quantization.hpp b/include/strata/platform/quantization.hpp index 0b7e5bf..1aafa82 100644 --- a/include/strata/platform/quantization.hpp +++ b/include/strata/platform/quantization.hpp @@ -6,7 +6,7 @@ // so they belong below the model tier. They lived in model.hpp only because // that is where the first consumer needed them, which forced // compressed_tensors.hpp -- a strata_platform header with no model dependency -// of its own -- to include the header that declares all six models' Spec +// of its own -- to include the header that declared every model's Spec // types. That was the last recorded layering exception. #include "strata/platform/types.hpp" diff --git a/kernels/cuda/backend.cu b/kernels/cuda/backend.cu index 7713b4e..9dcc3bf 100644 --- a/kernels/cuda/backend.cu +++ b/kernels/cuda/backend.cu @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -16,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -38,6 +40,7 @@ namespace strata { #include "detail/backend_core.inc.cuh" #include "detail/backend_dense_page.inc.cuh" #include "detail/backend_indexing.inc.cuh" +#include "detail/backend_hybrid_recurrence.inc.cuh" #include "detail/backend_flash_attention.inc.cuh" #include "detail/backend_prepared_attention.inc.cuh" #include "detail/backend_mhc.inc.cuh" diff --git a/kernels/cuda/detail/backend_core.inc.cuh b/kernels/cuda/detail/backend_core.inc.cuh index b97103d..dc52fd9 100644 --- a/kernels/cuda/detail/backend_core.inc.cuh +++ b/kernels/cuda/detail/backend_core.inc.cuh @@ -68,6 +68,38 @@ ParseResult CudaBackend::device_memory(int device) { return result; } +int CudaBackend::device_numa_node(int device) noexcept { + char bus_id[32]{}; + if (cudaDeviceGetPCIBusId(bus_id, sizeof(bus_id), device) != cudaSuccess) { + return -1; + } + char path[128]{}; + const int written = std::snprintf( + path, sizeof(path), "/sys/bus/pci/devices/%s/numa_node", bus_id); + if (written <= 0 || static_cast(written) >= sizeof(path)) { + return -1; + } + std::ifstream input(path); + int node = -1; + return input >> node && node >= 0 ? node : -1; +} + +bool CudaBackend::high_speed_peer_access_supported( + int source, int destination) noexcept { + if (source == destination) return true; + int supported = 0; + if (cudaDeviceCanAccessPeer(&supported, source, destination) != + cudaSuccess || + supported == 0) { + return false; + } + int rank = -1; + return cudaDeviceGetP2PAttribute( + &rank, cudaDevP2PAttrPerformanceRank, source, destination) == + cudaSuccess && + rank == 0; +} + std::uint64_t CudaBackend::weight_storage_bytes( std::uint64_t weight_bytes, std::uint64_t scale_bytes) noexcept { if (weight_bytes == 0U) return 0U; @@ -117,6 +149,12 @@ ValidationResult CudaBackend::initialize(std::span devices, state.lightning_index_supported = state.flash_attention_supported; state.dsv4_fp8_tensor_page_supported = properties.major == 8 && properties.minor == 6; + // BF16 WMMA is an architectural CUDA capability from Ampere onward. + // Keep DeepSeek's experimentally-bound SM86 flag above exact, while + // allowing the model-neutral F32-scale route to follow the hardware + // discovered at runtime instead of a machine-specific device list. + state.fp8_f32_tensor_page_supported = properties.major >= 8; + state.fp8_f32_register_fed_supported = properties.major >= 8; if (auto status = cudaStreamCreateWithFlags(&state.stream, cudaStreamNonBlocking); status != cudaSuccess) { return cuda_error(status, "create CUDA stream"); @@ -320,7 +358,8 @@ ValidationResult CudaBackend::upload(int device, const CudaWeightDescriptor& des result.errors.emplace_back("weight upload targets an uninitialized CUDA device"); return result; } - if (found->second.moe_in_flight) { + if (found->second.moe_in_flight && + completion != UploadCompletion::DeferredConcurrent) { result.errors.emplace_back( "weight upload cannot overlap an in-flight DeepSeek MoE command"); return result; @@ -424,6 +463,22 @@ ValidationResult CudaBackend::upload(int device, const CudaWeightDescriptor& des result.errors.emplace_back("invalid native FP8 CUDA weight descriptor"); return result; } + } else if (descriptor.encoding == + CudaWeightEncoding::Fp8E4m3Block128F32) { + const auto expected_scale_columns = (descriptor.columns + 127U) / 128U; + const auto expected_scale_rows = (descriptor.rows + 127U) / 128U; + if (descriptor.dtype != SafetensorsDtype::F8E4M3 || + descriptor.packed_columns != descriptor.columns || + descriptor.scale_columns != expected_scale_columns || + descriptor.group_size != 128U || + !checked_bytes(descriptor.rows, descriptor.columns, 1U, + expected_weights) || + !checked_bytes(expected_scale_rows, descriptor.scale_columns, 4U, + expected_scales)) { + result.errors.emplace_back( + "invalid F32-scaled native FP8 CUDA weight descriptor"); + return result; + } } else { result.errors.emplace_back("unsupported CUDA weight encoding"); return result; @@ -489,7 +544,7 @@ ValidationResult CudaBackend::upload(int device, const CudaWeightDescriptor& des // upload_ready before anything reads the weight. A synchronous upload keeps // the execution stream, because its caller's host payload dies at return // and the wait below is what keeps it alive long enough. - const bool deferred = completion == UploadCompletion::Deferred; + const bool deferred = completion != UploadCompletion::Synchronous; auto* const upload_stream = deferred ? state.upload_stream : state.stream; const auto upload_error = [&state, &target, upload_stream]( cudaError_t status, const char* operation) { @@ -498,8 +553,66 @@ ValidationResult CudaBackend::upload(int device, const CudaWeightDescriptor& des } return cuda_error(status, operation); }; + // cudaMemcpyAsync from an mmap-backed checkpoint is only superficially + // asynchronous: the runtime blocks the submitting thread while it copies + // pageable input into an internal pinned buffer. Nsight on GLM-5.3 showed + // 1.98 seconds of host API time and 9.99 GB of H2D traffic in one token. + // Keep a reusable pinned ring instead. Its capacity follows the admitted + // device arena (1/64th), which gives larger GPUs a deeper overlap window + // without baking this host's VRAM size into the runtime. + const bool stage_pageable = deferred && payload_bytes >= (1ULL << 20U) && + state.weight_arena != nullptr; + if (stage_pageable && state.weight_host_staging == nullptr) { + const auto arena = state.weight_arena->occupancy(); + const auto desired = std::max( + payload_bytes, arena.capacity / 64U); + void* staging = nullptr; + if (desired <= std::numeric_limits::max() && + cudaMallocHost(&staging, static_cast(desired)) == + cudaSuccess) { + state.weight_host_staging = static_cast(staging); + state.weight_host_staging_bytes = desired; + } else { + // Pinned staging is an optimization. A host whose lockable-memory + // budget is smaller keeps the exact pageable path. + static_cast(cudaGetLastError()); + } + } + const bool pinned_stage_ready = + stage_pageable && state.weight_host_staging != nullptr && + payload_bytes <= state.weight_host_staging_bytes; + const auto reserve_staging = [&](std::uint64_t bytes, + const std::byte*& source, + const std::byte* original) -> cudaError_t { + if (!pinned_stage_ready || bytes == 0U) { + source = original; + return cudaSuccess; + } + constexpr std::uint64_t alignment = 256U; + auto cursor = (state.weight_host_staging_cursor + alignment - 1U) & + ~(alignment - 1U); + if (cursor > state.weight_host_staging_bytes || + bytes > state.weight_host_staging_bytes - cursor) { + // Every earlier slice is consumed by this one upload stream. Wait + // only when the ring wraps, then recycle the whole arena at once. + const auto drained = cudaStreamSynchronize(upload_stream); + if (drained != cudaSuccess) return drained; + cursor = 0U; + } + auto* destination = state.weight_host_staging + cursor; + std::memcpy(destination, original, static_cast(bytes)); + state.weight_host_staging_cursor = cursor + bytes; + source = destination; + return cudaSuccess; + }; + const std::byte* weight_source = weights.data(); + if (auto status = reserve_staging(weights.size(), weight_source, + weights.data()); + status != cudaSuccess) { + return upload_error(status, "recycle pinned CUDA weight staging"); + } auto copy_started = std::chrono::steady_clock::now(); - if (auto status = cudaMemcpyAsync(target->weights, weights.data(), weights.size(), + if (auto status = cudaMemcpyAsync(target->weights, weight_source, weights.size(), cudaMemcpyHostToDevice, upload_stream); status != cudaSuccess) { return upload_error(status, "upload CUDA weights"); @@ -517,8 +630,14 @@ ValidationResult CudaBackend::upload(int device, const CudaWeightDescriptor& des allocation_nanoseconds += elapsed_nanoseconds_since(scale_allocation_started); } + const std::byte* scale_source = scales.data(); + if (auto status = reserve_staging(scales.size(), scale_source, + scales.data()); + status != cudaSuccess) { + return upload_error(status, "recycle pinned CUDA scale staging"); + } copy_started = std::chrono::steady_clock::now(); - if (auto status = cudaMemcpyAsync(target->scales, scales.data(), scales.size(), + if (auto status = cudaMemcpyAsync(target->scales, scale_source, scales.size(), cudaMemcpyHostToDevice, upload_stream); status != cudaSuccess) { return upload_error(status, "upload CUDA scales"); @@ -530,7 +649,9 @@ ValidationResult CudaBackend::upload(int device, const CudaWeightDescriptor& des // costs one read and one write of what was staged, measured at 0.509 // ms/token against Laguna's 65.05 ms staging term (experiment 0168), so it // does not need to move off the staging path. - if (prepack == FragmentLayout::Prepack && regfed_matmul_enabled()) { + if (prepack == FragmentLayout::Prepack && regfed_matmul_enabled() && + (descriptor.encoding != CudaWeightEncoding::Fp8E4m3Block128F32 || + state.fp8_f32_register_fed_supported)) { const auto scratch_bytes = fragment_prepack_scratch_bytes(descriptor); if (scratch_bytes != 0U) { bool ready = true; @@ -822,3 +943,10 @@ ValidationResult CudaBackend::allocate_buffer( } return result; } +ValidationResult CudaBackend::profiler_start() { + return cuda_error(cudaProfilerStart(), "start CUDA profiler capture"); +} + +ValidationResult CudaBackend::profiler_stop() { + return cuda_error(cudaProfilerStop(), "stop CUDA profiler capture"); +} diff --git a/kernels/cuda/detail/backend_hybrid_recurrence.inc.cuh b/kernels/cuda/detail/backend_hybrid_recurrence.inc.cuh new file mode 100644 index 0000000..fc0d477 --- /dev/null +++ b/kernels/cuda/detail/backend_hybrid_recurrence.inc.cuh @@ -0,0 +1,1152 @@ +namespace { + +__device__ __forceinline__ float glm53_bf16(float value) { + return __bfloat162float(__float2bfloat16_rn(value)); +} + +__device__ __forceinline__ float glm53_sigmoid(float value) { + return 1.0F / (1.0F + expf(-value)); +} + +__global__ void glm53_kda_conv_kernel( + float* activations, float* convolution, const float* taps, + std::uint32_t width, std::uint32_t kernel) { + const auto channel = blockIdx.x * blockDim.x + threadIdx.x; + if (channel >= width) return; + const auto history_width = kernel - 1U; + for (std::uint32_t projection = 0U; projection < 3U; ++projection) { + auto* values = activations + static_cast(projection) * width; + auto* history = convolution + + static_cast(projection) * width * history_width + + static_cast(channel) * history_width; + const auto* weights = taps + + static_cast(projection) * width * kernel + + static_cast(channel) * kernel; + float sum = weights[kernel - 1U] * values[channel]; + for (std::uint32_t offset = 0U; offset < history_width; ++offset) { + sum += weights[offset] * history[offset]; + } + for (std::uint32_t offset = 0U; offset + 1U < history_width; ++offset) { + history[offset] = history[offset + 1U]; + } + history[history_width - 1U] = values[channel]; + values[channel] = glm53_bf16(sum * glm53_sigmoid(sum)); + } +} + +__global__ void glm53_kda_recurrence_kernel( + float* recurrent, const float* a_log, const float* dt_bias, + const float* norm_weight, const float* activations, float* output, + std::uint32_t heads, std::uint32_t head_dim) { + const auto head = blockIdx.x; + const auto lane = threadIdx.x; + if (head >= heads || lane >= head_dim) return; + const auto width = heads * head_dim; + const auto base = head * head_dim; + const auto* query = activations; + const auto* key = activations + width; + const auto* value = activations + 2U * width; + const auto* forget = activations + 3U * width; + const auto* gate = activations + 4U * width; + const auto* beta = activations + 5U * width; + extern __shared__ float scratch[]; + auto* normalized_query = scratch; + auto* normalized_key = normalized_query + head_dim; + auto* decay = normalized_key + head_dim; + auto* raw = decay + head_dim; + __shared__ float query_inverse; + __shared__ float key_inverse; + __shared__ float output_inverse; + if (lane == 0U) { + float query_square = 0.0F; + float key_square = 0.0F; + for (std::uint32_t index = 0U; index < head_dim; ++index) { + const auto q = query[base + index]; + const auto k = key[base + index]; + query_square += q * q; + key_square += k * k; + } + query_inverse = rsqrtf(query_square + 1.0e-6F) / + sqrtf(static_cast(head_dim)); + key_inverse = rsqrtf(key_square + 1.0e-6F); + } + __syncthreads(); + normalized_query[lane] = query[base + lane] * query_inverse; + normalized_key[lane] = key[base + lane] * key_inverse; + decay[lane] = expf(-5.0F * glm53_sigmoid( + expf(a_log[head]) * (forget[base + lane] + dt_bias[base + lane]))); + __syncthreads(); + auto* state_row = recurrent + + (static_cast(head) * head_dim + lane) * head_dim; + float projected = 0.0F; + for (std::uint32_t index = 0U; index < head_dim; ++index) { + state_row[index] *= decay[index]; + projected += state_row[index] * normalized_key[index]; + } + const auto delta = + (value[base + lane] - projected) * beta[head]; + float mixed = 0.0F; + for (std::uint32_t index = 0U; index < head_dim; ++index) { + state_row[index] += delta * normalized_key[index]; + mixed += state_row[index] * normalized_query[index]; + } + raw[lane] = glm53_bf16(mixed); + __syncthreads(); + if (lane == 0U) { + float square = 0.0F; + for (std::uint32_t index = 0U; index < head_dim; ++index) { + square += raw[index] * raw[index]; + } + output_inverse = rsqrtf( + square / static_cast(head_dim) + 1.0e-5F); + } + __syncthreads(); + output[base + lane] = glm53_bf16( + norm_weight[lane] * raw[lane] * output_inverse * + glm53_sigmoid(gate[base + lane])); +} + +__global__ void glm53_kda_beta_kernel(float* beta, std::uint32_t heads) { + const auto index = blockIdx.x * blockDim.x + threadIdx.x; + if (index < heads) beta[index] = glm53_bf16(glm53_sigmoid(beta[index])); +} + +__global__ void glm53_moe_join_mhc_kernel( + const float* expert_output, const float* coefficients, + std::uint32_t routed, __nv_bfloat16* branch, std::uint32_t hidden, + unsigned int* error_flag) { + const auto column = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + if (column >= hidden) return; + float value = expert_output[static_cast(routed) * hidden + + column]; + for (std::uint32_t expert = 0U; expert < routed; ++expert) { + const float weighted = glm53_bf16( + coefficients[expert] * + expert_output[static_cast(expert) * hidden + + column]); + value = glm53_bf16(value + weighted); + } + if (!isfinite(value)) atomicExch(error_flag, 1U); + branch[column] = __float2bfloat16_rn(value); +} + +__global__ void glm53_swiglu_kernel( + float* gate, const float* up, std::uint32_t elements) { + const auto index = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + if (index >= elements) return; + const float g = fminf(glm53_bf16(gate[index]), 10.0F); + const float u = fminf(fmaxf(glm53_bf16(up[index]), -10.0F), 10.0F); + gate[index] = glm53_bf16(g * glm53_sigmoid(g) * u); +} + +__global__ void glm53_rms_norm_bf16_kernel( + float* values, const float* weights, std::uint32_t columns) { + __shared__ float scale; + if (threadIdx.x == 0U) { + float sum = 0.0F; + for (std::uint32_t column = 0U; column < columns; ++column) { + sum = __fadd_rn(sum, __fmul_rn(values[column], values[column])); + } + scale = 1.0F / sqrtf(sum / static_cast(columns) + 1.0e-5F); + } + __syncthreads(); + for (std::uint32_t column = threadIdx.x; column < columns; + column += blockDim.x) { + values[column] = glm53_bf16( + weights[column] * (values[column] * scale)); + } +} + +__global__ void glm53_mla_attention_kernel( + const float* query, const float* expanded, float* attended, + std::uint32_t history, std::uint32_t heads, std::uint32_t head_dim) { + extern __shared__ float scores[]; + const auto head = static_cast(blockIdx.x); + if (head >= heads) return; + if (threadIdx.x == 0U) { + float highest = -INFINITY; + for (std::uint32_t token = 0U; token < history; ++token) { + const auto* key = expanded + + (static_cast(token) * heads + head) * + (2U * head_dim); + const auto* q = query + + static_cast(head) * head_dim; + float score = 0.0F; + for (std::uint32_t column = 0U; column < head_dim; ++column) { + score = __fadd_rn(score, __fmul_rn(q[column], key[column])); + } + score *= rsqrtf(static_cast(head_dim)); + scores[token] = score; + highest = fmaxf(highest, score); + } + float total = 0.0F; + for (std::uint32_t token = 0U; token < history; ++token) { + scores[token] = expf(scores[token] - highest); + total += scores[token]; + } + for (std::uint32_t token = 0U; token < history; ++token) { + scores[token] = glm53_bf16(scores[token] / total); + } + } + __syncthreads(); + for (std::uint32_t column = threadIdx.x; column < head_dim; + column += blockDim.x) { + float value = 0.0F; + for (std::uint32_t token = 0U; token < history; ++token) { + const auto* source = expanded + + (static_cast(token) * heads + head) * + (2U * head_dim) + head_dim; + value += scores[token] * source[column]; + } + attended[static_cast(head) * head_dim + column] = + glm53_bf16(value); + } +} + +__global__ void glm53_mla_absorb_query_kernel( + const float* query, const __nv_bfloat16* key_value_weights, + float* compressed, std::uint32_t heads, std::uint32_t head_dim, + std::uint32_t latent_dim) { + const auto index = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + const auto elements = static_cast(heads) * latent_dim; + if (index >= elements) return; + const auto head = static_cast(index / latent_dim); + const auto latent = static_cast(index % latent_dim); + float value = 0.0F; + for (std::uint32_t column = 0U; column < head_dim; ++column) { + const auto weight_row = + static_cast(head) * 2U * head_dim + column; + value = __fadd_rn( + value, __fmul_rn( + query[static_cast(head) * head_dim + column], + __bfloat162float( + key_value_weights[weight_row * latent_dim + latent]))); + } + compressed[index] = value; +} + +__global__ void glm53_mla_latent_attention_kernel( + const float* compressed_query, const float* latent_cache, + float* weighted_latent, std::uint32_t history, std::uint32_t heads, + std::uint32_t head_dim, std::uint32_t latent_dim) { + extern __shared__ float scores[]; + const auto head = static_cast(blockIdx.x); + if (head >= heads) return; + if (threadIdx.x == 0U) { + float highest = -INFINITY; + for (std::uint32_t token = 0U; token < history; ++token) { + float score = 0.0F; + for (std::uint32_t column = 0U; column < latent_dim; ++column) { + score = __fadd_rn( + score, + __fmul_rn( + compressed_query[ + static_cast(head) * latent_dim + + column], + latent_cache[ + static_cast(token) * latent_dim + + column])); + } + score *= rsqrtf(static_cast(head_dim)); + scores[token] = score; + highest = fmaxf(highest, score); + } + float total = 0.0F; + for (std::uint32_t token = 0U; token < history; ++token) { + scores[token] = expf(scores[token] - highest); + total += scores[token]; + } + for (std::uint32_t token = 0U; token < history; ++token) { + scores[token] = glm53_bf16(scores[token] / total); + } + } + __syncthreads(); + for (std::uint32_t column = threadIdx.x; column < latent_dim; + column += blockDim.x) { + float value = 0.0F; + for (std::uint32_t token = 0U; token < history; ++token) { + value += scores[token] * + latent_cache[static_cast(token) * latent_dim + + column]; + } + weighted_latent[static_cast(head) * latent_dim + + column] = value; + } +} + +__global__ void glm53_mla_expand_value_kernel( + const float* weighted_latent, const __nv_bfloat16* key_value_weights, + float* attended, std::uint32_t heads, std::uint32_t head_dim, + std::uint32_t latent_dim) { + const auto index = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + const auto elements = static_cast(heads) * head_dim; + if (index >= elements) return; + const auto head = static_cast(index / head_dim); + const auto column = static_cast(index % head_dim); + const auto weight_row = static_cast(head) * 2U * head_dim + + head_dim + column; + float value = 0.0F; + for (std::uint32_t latent = 0U; latent < latent_dim; ++latent) { + value = __fadd_rn( + value, + __fmul_rn( + weighted_latent[ + static_cast(head) * latent_dim + latent], + __bfloat162float( + key_value_weights[weight_row * latent_dim + latent]))); + } + attended[index] = glm53_bf16(value); +} + +} // namespace + +ValidationResult CudaBackend::glm53_kda_decode( + const CudaGlm53KdaRequest& request, std::span output) { + ValidationResult result; + if (request.state == nullptr || !request.state->valid() || + request.heads == 0U || request.head_dim == 0U || + request.head_dim > 256U || request.convolution_kernel < 2U) { + return {{"CUDA GLM-5.3 KDA command is invalid"}}; + } + const auto width = static_cast(request.heads) * + request.head_dim; + const bool full_layer = !request.input.empty() || + request.mhc_source_destination; + const auto projected = request.output_projection != nullptr; + const auto projected_rows = projected && request.output_projection->valid() + ? request.output_projection->impl_->descriptor.rows + : 0U; + if ((!full_layer && + (request.query.size() != width || request.key.size() != width || + request.value.size() != width || request.forget.size() != width || + request.gate.size() != width || request.beta.size() != request.heads)) || + ((!request.mhc_source_destination && + output.size() != (projected ? projected_rows : width)) || + (request.mhc_source_destination && !output.empty()))) { + return {{"CUDA GLM-5.3 KDA operands have incompatible shapes"}}; + } + const auto recurrent_floats = width * request.head_dim; + const auto convolution_floats = 3ULL * width * + (request.convolution_kernel - 1U); + const auto tap_floats = 3ULL * width * request.convolution_kernel; + const auto required_state_floats = recurrent_floats + convolution_floats + + tap_floats + request.heads + width + request.head_dim; + if (required_state_floats > + std::numeric_limits::max() / sizeof(float) || + request.state->device_bytes() < required_state_floats * sizeof(float)) { + return {{"CUDA GLM-5.3 KDA persistent state has an invalid extent"}}; + } + const int device = request.state->device(); + const auto projection_encoding = projected + ? request.output_projection->impl_->descriptor.encoding + : CudaWeightEncoding::Plain; + const auto plain_bf16_projection = projected && + projection_encoding == CudaWeightEncoding::Plain && + request.output_projection->impl_->descriptor.dtype == + SafetensorsDtype::Bf16; + const auto fp8_projection = projected && + projection_encoding == CudaWeightEncoding::Fp8E4m3Block128F32 && + request.output_projection->impl_->fragment_prepacked; + if (projected && + (request.output_projection->device() != device || + request.output_projection->impl_->descriptor.columns != width || + (!plain_bf16_projection && !fp8_projection))) { + return {{"CUDA GLM-5.3 KDA output projection is not resident in the " + "required register-fed layout"}}; + } + const auto found = impl_->devices.find(device); + if (found == impl_->devices.end()) { + return {{"CUDA GLM-5.3 KDA state targets an uninitialized device"}}; + } + auto& device_state = found->second; + if (device_state.moe_in_flight) { + return {{"CUDA GLM-5.3 KDA cannot overlap an in-flight MoE command"}}; + } + if (auto status = cudaSetDevice(device); status != cudaSuccess) { + return cuda_error(status, "select CUDA device for GLM-5.3 KDA"); + } + if (full_layer) { + const auto hidden = projected_rows; + const auto workspace_floats = 2ULL * hidden + 6ULL * width + + 2ULL * request.head_dim + request.heads; + if (!projected || + (!request.mhc_source_destination && + request.input.size() != hidden) || + (request.mhc_source_destination && !request.input.empty()) || + required_state_floats + workspace_floats > + request.state->device_bytes() / sizeof(float)) { + return {{"CUDA GLM-5.3 fused KDA layer workspace is invalid"}}; + } + const std::array projections{ + request.query_projection, request.key_projection, + request.value_projection, request.forget_a_projection, + request.beta_projection, request.gate_a_projection, + request.forget_b_projection, request.gate_b_projection, + request.output_projection}; + if (std::any_of(projections.begin(), projections.end(), + [device](const CudaWeight* weight) { + if (weight == nullptr || !weight->valid() || + weight->device() != device) { + return true; + } + const auto& descriptor = weight->impl_->descriptor; + const auto plain_bf16 = + descriptor.encoding == CudaWeightEncoding::Plain && + descriptor.dtype == SafetensorsDtype::Bf16; + const auto regfed_fp8 = + descriptor.encoding == + CudaWeightEncoding::Fp8E4m3Block128F32 && + weight->impl_->fragment_prepacked; + return !plain_bf16 && !regfed_fp8; + })) { + return {{"CUDA GLM-5.3 fused KDA layer weights are not resident " + "BF16 or register-fed FP8 tensors"}}; + } + const auto shape = [](const CudaWeight* weight, std::uint64_t rows, + std::uint64_t columns) { + return weight->impl_->descriptor.rows == rows && + weight->impl_->descriptor.columns == columns; + }; + if (!shape(request.query_projection, width, hidden) || + !shape(request.key_projection, width, hidden) || + !shape(request.value_projection, width, hidden) || + !shape(request.forget_a_projection, request.head_dim, hidden) || + !shape(request.beta_projection, request.heads, hidden) || + !shape(request.gate_a_projection, request.head_dim, hidden) || + !shape(request.forget_b_projection, width, request.head_dim) || + !shape(request.gate_b_projection, width, request.head_dim) || + !shape(request.output_projection, hidden, width)) { + return {{"CUDA GLM-5.3 fused KDA layer weight shapes are invalid"}}; + } + if (request.mhc_source_destination && + (!device_state.dsv4_mhc_supported || + device_state.dsv4_mhc_stage != 1U || + device_state.dsv4_mhc_workspace == nullptr || + device_state.dsv4_mhc_branch_ready || + device_state.dsv4_mhc_failed)) { + return {{"CUDA GLM-5.3 fused KDA mHC command order is invalid"}}; + } + const auto input_bytes = hidden * sizeof(float); + const auto output_bytes = hidden * sizeof(float); + const auto grow_pinned = [](std::byte*& pointer, + std::uint64_t& capacity, + std::uint64_t required) -> cudaError_t { + if (required <= capacity) return cudaSuccess; + void* replacement = nullptr; + if (auto status = cudaMallocHost(&replacement, required); + status != cudaSuccess) return status; + if (pointer != nullptr) static_cast(cudaFreeHost(pointer)); + pointer = static_cast(replacement); + capacity = required; + return cudaSuccess; + }; + if (auto status = grow_pinned( + device_state.matmul_host_input, + device_state.matmul_host_input_bytes, input_bytes); + status != cudaSuccess) { + return cuda_error(status, + "allocate fused GLM-5.3 KDA input staging"); + } + if (auto status = grow_pinned( + device_state.matmul_host_output, + device_state.matmul_host_output_bytes, output_bytes); + status != cudaSuccess) { + return cuda_error(status, + "allocate fused GLM-5.3 KDA output staging"); + } + if (!request.mhc_source_destination) { + std::memcpy(device_state.matmul_host_input, request.input.data(), + input_bytes); + } + auto* packed = static_cast(request.state->impl_->data); + auto* workspace = packed + required_state_floats; + auto* device_input = workspace; + auto* activations = device_input + hidden; + auto* query = activations; + auto* key = query + width; + auto* value = key + width; + auto* forget = value + width; + auto* gate = forget + width; + auto* beta = gate + width; + auto* forget_low = beta + request.heads; + auto* gate_low = forget_low + request.head_dim; + auto* heads_output = gate_low + request.head_dim; + auto* final_output = heads_output + width; + if (request.mhc_source_destination) { + constexpr std::uint32_t threads = 256U; + constexpr std::uint32_t blocks = + (kDsv4MhcHidden + threads - 1U) / threads; + dsv4_bf16_to_fp32<<>>( + device_state.dsv4_mhc_workspace->layer_input, device_input, + kDsv4MhcHidden); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, + "convert resident GLM-5.3 KDA input"); + } + } else if (auto status = cudaMemcpyAsync( + device_input, device_state.matmul_host_input, + input_bytes, cudaMemcpyHostToDevice, + device_state.stream); + status != cudaSuccess) { + return cuda_error(status, + "upload fused GLM-5.3 KDA layer input"); + } + const auto project = [&](const CudaWeight* weight, + float* source, float* destination, + std::uint64_t rows, + const char* operation) -> ValidationResult { + auto& projection = *weight->impl_; + const auto plain_bf16 = + projection.descriptor.encoding == CudaWeightEncoding::Plain && + projection.descriptor.dtype == SafetensorsDtype::Bf16; + if (plain_bf16) { + constexpr unsigned int threads = 256U; + constexpr unsigned int warps_per_block = threads / 32U; + const auto blocks = static_cast( + (rows + warps_per_block - 1U) / warps_per_block); + bf16_matvec_kernel<<>>( + destination, source, + static_cast(projection.weights), + projection.descriptor.columns, rows); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, operation); + } + } else { + if (auto status = launch_regfed_fp8_matvec( + device_state.moe_regfed, projection.descriptor, + projection.weights, projection.scales, + projection.fragment_prepacked, source, destination, + device_state.stream); + status != cudaSuccess) { + return cuda_error(status, operation); + } + } + round_bf16_rows_kernel<<< + static_cast((rows + 255U) / 256U), 256U, 0U, + device_state.stream>>>(destination, rows); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, operation); + } + return {}; + }; + for (const auto& command : std::array{ + std::tuple{request.query_projection, device_input, query, + width, "project fused GLM-5.3 KDA query"}, + std::tuple{request.key_projection, device_input, key, width, + "project fused GLM-5.3 KDA key"}, + std::tuple{request.value_projection, device_input, value, + width, "project fused GLM-5.3 KDA value"}, + std::tuple{request.forget_a_projection, device_input, + forget_low, + static_cast(request.head_dim), + "project fused GLM-5.3 KDA forget A"}, + std::tuple{request.beta_projection, device_input, beta, + static_cast(request.heads), + "project fused GLM-5.3 KDA beta"}, + std::tuple{request.gate_a_projection, device_input, gate_low, + static_cast(request.head_dim), + "project fused GLM-5.3 KDA gate A"}}) { + auto projected_result = std::apply(project, command); + if (!projected_result.ok()) return projected_result; + } + auto projected_result = project( + request.forget_b_projection, forget_low, forget, width, + "project fused GLM-5.3 KDA forget B"); + if (!projected_result.ok()) return projected_result; + projected_result = project( + request.gate_b_projection, gate_low, gate, width, + "project fused GLM-5.3 KDA gate B"); + if (!projected_result.ok()) return projected_result; + glm53_kda_beta_kernel<<<1U, 128U, 0U, device_state.stream>>>( + beta, request.heads); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, "activate fused GLM-5.3 KDA beta"); + } + auto* recurrent = packed; + auto* convolution = recurrent + recurrent_floats; + const auto* taps = convolution + convolution_floats; + const auto* a_log = taps + tap_floats; + const auto* dt_bias = a_log + request.heads; + const auto* norm_weight = dt_bias + width; + constexpr std::uint32_t threads = 256U; + glm53_kda_conv_kernel<<< + static_cast((width + threads - 1U) / threads), + threads, 0U, device_state.stream>>>( + activations, convolution, taps, + static_cast(width), request.convolution_kernel); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, + "launch fused GLM-5.3 KDA convolution"); + } + glm53_kda_recurrence_kernel<<< + request.heads, request.head_dim, + static_cast(4U * request.head_dim * sizeof(float)), + device_state.stream>>>( + recurrent, a_log, dt_bias, norm_weight, activations, + heads_output, request.heads, request.head_dim); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, + "launch fused GLM-5.3 KDA recurrence"); + } + projected_result = project( + request.output_projection, heads_output, final_output, hidden, + "project fused GLM-5.3 KDA output"); + if (!projected_result.ok()) return projected_result; + if (request.mhc_source_destination) { + constexpr std::uint32_t threads = 256U; + constexpr std::uint32_t blocks = + (kDsv4MhcHidden + threads - 1U) / threads; + dsv4_fp32_to_bf16<<>>( + final_output, device_state.dsv4_mhc_workspace->branch, + kDsv4MhcHidden); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, + "publish resident GLM-5.3 KDA branch"); + } + device_state.dsv4_mhc_branch_ready = true; + return {}; + } + if (auto status = cudaMemcpyAsync( + device_state.matmul_host_output, final_output, output_bytes, + cudaMemcpyDeviceToHost, device_state.stream); + status != cudaSuccess) { + return cuda_error(status, + "download fused GLM-5.3 KDA layer output"); + } + const auto wait_started = std::chrono::steady_clock::now(); + if (auto status = cudaStreamSynchronize(device_state.stream); + status != cudaSuccess) { + return cuda_error(status, + "synchronize fused GLM-5.3 KDA layer"); + } + std::memcpy(output.data(), device_state.matmul_host_output, + output_bytes); + const auto wait_nanoseconds = elapsed_nanoseconds_since(wait_started); + { + std::scoped_lock lock(impl_->mutex); + auto& stats = *std::find_if( + impl_->stats.devices.begin(), impl_->stats.devices.end(), + [device](const auto& value) { return value.device == device; }); + stats.activation_h2d_bytes += input_bytes; + stats.activation_d2h_bytes += output_bytes; + record_synchronization(stats, SynchronizationSubsystem::Other, + 1U, wait_nanoseconds); + } + return {}; + } + const auto input_floats = 5ULL * width + request.heads; + const auto input_bytes = input_floats * sizeof(float); + const auto output_bytes = (projected ? projected_rows : width) * + sizeof(float); + if (input_bytes > device_state.input_bytes) { + if (device_state.input != nullptr) { + static_cast(cudaFree(device_state.input)); + } + if (auto status = cudaMalloc(&device_state.input, input_bytes); + status != cudaSuccess) { + device_state.input = nullptr; + device_state.input_bytes = 0U; + return cuda_error(status, "allocate GLM-5.3 KDA input workspace"); + } + device_state.input_bytes = input_bytes; + } + if (output_bytes > device_state.output_bytes) { + if (device_state.output != nullptr) { + static_cast(cudaFree(device_state.output)); + } + if (auto status = cudaMalloc(&device_state.output, output_bytes); + status != cudaSuccess) { + device_state.output = nullptr; + device_state.output_bytes = 0U; + return cuda_error(status, "allocate GLM-5.3 KDA output workspace"); + } + device_state.output_bytes = output_bytes; + } + const auto grow_pinned = [](std::byte*& pointer, std::uint64_t& capacity, + std::uint64_t required) -> cudaError_t { + if (required <= capacity) return cudaSuccess; + void* replacement = nullptr; + if (auto status = cudaMallocHost(&replacement, required); + status != cudaSuccess) return status; + if (pointer != nullptr) static_cast(cudaFreeHost(pointer)); + pointer = static_cast(replacement); + capacity = required; + return cudaSuccess; + }; + if (auto status = grow_pinned(device_state.matmul_host_input, + device_state.matmul_host_input_bytes, + input_bytes); + status != cudaSuccess) { + return cuda_error(status, "allocate GLM-5.3 KDA input staging"); + } + if (auto status = grow_pinned(device_state.matmul_host_output, + device_state.matmul_host_output_bytes, + output_bytes); + status != cudaSuccess) { + return cuda_error(status, "allocate GLM-5.3 KDA output staging"); + } + auto* staged = reinterpret_cast(device_state.matmul_host_input); + std::uint64_t cursor = 0U; + for (const auto values : {request.query, request.key, request.value, + request.forget, request.gate}) { + std::copy(values.begin(), values.end(), staged + cursor); + cursor += width; + } + std::copy(request.beta.begin(), request.beta.end(), staged + cursor); + if (auto status = cudaMemcpyAsync( + device_state.input, staged, input_bytes, cudaMemcpyHostToDevice, + device_state.stream); + status != cudaSuccess) { + return cuda_error(status, "upload GLM-5.3 KDA activations"); + } + auto* packed = static_cast(request.state->impl_->data); + auto* recurrent = packed; + auto* convolution = recurrent + recurrent_floats; + const auto* taps = convolution + convolution_floats; + const auto* a_log = taps + tap_floats; + const auto* dt_bias = a_log + request.heads; + const auto* norm_weight = dt_bias + width; + constexpr std::uint32_t threads = 256U; + glm53_kda_conv_kernel<<< + static_cast((width + threads - 1U) / threads), threads, + 0U, device_state.stream>>>( + device_state.input, convolution, taps, static_cast(width), + request.convolution_kernel); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, "launch GLM-5.3 fused short convolution"); + } + glm53_kda_recurrence_kernel<<< + request.heads, request.head_dim, + static_cast(4U * request.head_dim * sizeof(float)), + device_state.stream>>>( + recurrent, a_log, dt_bias, norm_weight, device_state.input, + device_state.output, request.heads, request.head_dim); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, "launch GLM-5.3 fused KDA recurrence"); + } + const float* device_result = device_state.output; + if (projected) { + auto& projection = *request.output_projection->impl_; + if (plain_bf16_projection) { + constexpr unsigned int threads = 256U; + constexpr unsigned int warps_per_block = threads / 32U; + const auto blocks = static_cast( + (projected_rows + warps_per_block - 1U) / warps_per_block); + bf16_matvec_kernel<<>>( + device_state.input, device_state.output, + static_cast(projection.weights), + width, projected_rows); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error( + status, "launch GLM-5.3 fused BF16 KDA output projection"); + } + } else { + if (auto status = launch_regfed_fp8_matvec( + device_state.moe_regfed, projection.descriptor, + projection.weights, projection.scales, + projection.fragment_prepacked, device_state.output, + device_state.input, device_state.stream); + status != cudaSuccess) { + return cuda_error( + status, "launch GLM-5.3 fused FP8 KDA output projection"); + } + } + round_bf16_rows_kernel<<< + static_cast((projected_rows + 255U) / 256U), 256U, + 0U, device_state.stream>>>(device_state.input, projected_rows); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, + "round GLM-5.3 KDA projected activation"); + } + device_result = device_state.input; + } + if (auto status = cudaMemcpyAsync( + device_state.matmul_host_output, device_result, output_bytes, + cudaMemcpyDeviceToHost, device_state.stream); + status != cudaSuccess) { + return cuda_error(status, "download GLM-5.3 KDA activation"); + } + const auto wait_started = std::chrono::steady_clock::now(); + if (auto status = cudaStreamSynchronize(device_state.stream); + status != cudaSuccess) { + return cuda_error(status, "synchronize GLM-5.3 KDA recurrence"); + } + std::memcpy(output.data(), device_state.matmul_host_output, output_bytes); + const auto wait_nanoseconds = elapsed_nanoseconds_since(wait_started); + { + std::scoped_lock lock(impl_->mutex); + auto& stats = *std::find_if( + impl_->stats.devices.begin(), impl_->stats.devices.end(), + [device](const auto& value) { return value.device == device; }); + stats.activation_h2d_bytes += input_bytes; + stats.activation_d2h_bytes += output_bytes; + record_synchronization(stats, SynchronizationSubsystem::Other, 1U, + wait_nanoseconds); + } + return result; +} + +ValidationResult CudaBackend::glm53_mhc_router( + int device, const CudaWeight& router, std::span logits) { + if (!router.valid() || router.device() != device || logits.size() != 288U) { + return {{"CUDA GLM-5.3 resident router command is invalid"}}; + } + const auto& descriptor = router.impl_->descriptor; + if (descriptor.encoding != CudaWeightEncoding::Plain || + descriptor.dtype != SafetensorsDtype::Bf16 || + descriptor.rows != logits.size() || + descriptor.columns != kDsv4MhcHidden) { + return {{"CUDA GLM-5.3 resident router requires a 288x4096 BF16 " + "projection"}}; + } + const auto found = impl_->devices.find(device); + if (found == impl_->devices.end()) { + return {{"CUDA GLM-5.3 resident router targets an uninitialized " + "device"}}; + } + auto& state = found->second; + if (!state.dsv4_mhc_supported || state.dsv4_mhc_stage != 1U || + state.dsv4_mhc_workspace == nullptr || state.dsv4_mhc_branch_ready || + state.dsv4_mhc_failed || state.moe_in_flight) { + return {{"CUDA GLM-5.3 resident router violates mHC command order"}}; + } + if (auto status = cudaSetDevice(device); status != cudaSuccess) { + return cuda_error(status, + "select CUDA device for GLM-5.3 resident router"); + } + const auto bytes = logits.size_bytes(); + if (bytes > state.matmul_host_output_bytes) { + void* replacement = nullptr; + if (auto status = cudaMallocHost(&replacement, bytes); + status != cudaSuccess) { + return cuda_error(status, + "allocate GLM-5.3 resident router staging"); + } + if (state.matmul_host_output != nullptr) { + static_cast(cudaFreeHost(state.matmul_host_output)); + } + state.matmul_host_output = static_cast(replacement); + state.matmul_host_output_bytes = bytes; + } + constexpr unsigned int threads = 256U; + constexpr unsigned int warps_per_block = threads / 32U; + const auto blocks = static_cast( + (descriptor.rows + warps_per_block - 1U) / warps_per_block); + bf16_input_matvec_kernel<<>>( + state.dsv4_mhc_workspace->glm53_router_logits, + state.dsv4_mhc_workspace->layer_input, + static_cast(router.impl_->weights), + descriptor.columns, descriptor.rows); + if (auto status = cudaMemcpyAsync( + state.matmul_host_output, + state.dsv4_mhc_workspace->glm53_router_logits, bytes, + cudaMemcpyDeviceToHost, state.stream); status != cudaSuccess) { + return cuda_error(status, "download GLM-5.3 resident router logits"); + } + if (auto status = cudaStreamSynchronize(state.stream); + status != cudaSuccess) { + return cuda_error(status, + "synchronize GLM-5.3 resident router logits"); + } + std::memcpy(logits.data(), state.matmul_host_output, bytes); + { + std::scoped_lock lock(impl_->mutex); + auto& stats = *std::find_if( + impl_->stats.devices.begin(), impl_->stats.devices.end(), + [device](const auto& value) { return value.device == device; }); + stats.activation_d2h_bytes += bytes; + } + return {}; +} + +ValidationResult CudaBackend::glm53_mhc_swiglu( + int device, const CudaWeight& gate, const CudaWeight& up, + const CudaWeight& down, std::uint32_t intermediate) { + const auto found = impl_->devices.find(device); + if (found == impl_->devices.end() || intermediate == 0U) { + return {{"CUDA GLM-5.3 resident SwiGLU command is invalid"}}; + } + auto& state = found->second; + const auto valid = [device](const CudaWeight& weight, + std::uint64_t rows, + std::uint64_t columns) { + if (!weight.valid() || weight.device() != device || + weight.impl_->descriptor.rows != rows || + weight.impl_->descriptor.columns != columns) return false; + const auto& descriptor = weight.impl_->descriptor; + return (descriptor.encoding == CudaWeightEncoding::Plain && + descriptor.dtype == SafetensorsDtype::Bf16) || + (descriptor.encoding == + CudaWeightEncoding::Fp8E4m3Block128F32 && + weight.impl_->fragment_prepacked); + }; + if (!valid(gate, intermediate, kDsv4MhcHidden) || + !valid(up, intermediate, kDsv4MhcHidden) || + !valid(down, kDsv4MhcHidden, intermediate) || + !state.dsv4_mhc_supported || state.dsv4_mhc_stage != 1U || + state.dsv4_mhc_workspace == nullptr || state.dsv4_mhc_branch_ready || + state.dsv4_mhc_failed || state.moe_in_flight) { + return {{"CUDA GLM-5.3 resident SwiGLU weights or command order are " + "invalid"}}; + } + if (auto status = cudaSetDevice(device); status != cudaSuccess) { + return cuda_error(status, + "select CUDA device for GLM-5.3 resident SwiGLU"); + } + const auto hidden_bytes = + static_cast(kDsv4MhcHidden) * sizeof(float); + const auto activation_bytes = + static_cast(intermediate) * 2U * sizeof(float); + const auto ensure = [&](float*& pointer, std::uint64_t& capacity, + std::uint64_t required) -> cudaError_t { + if (capacity >= required) return cudaSuccess; + if (pointer != nullptr) static_cast(cudaFree(pointer)); + pointer = nullptr; + capacity = 0U; + if (auto status = cudaMalloc(&pointer, required); + status != cudaSuccess) return status; + capacity = required; + return cudaSuccess; + }; + if (auto status = ensure(state.moe_hidden, state.moe_hidden_bytes, + hidden_bytes); status != cudaSuccess) { + return cuda_error(status, "allocate resident SwiGLU hidden buffer"); + } + if (auto status = ensure(state.moe_activations, + state.moe_activation_bytes, activation_bytes); + status != cudaSuccess) { + return cuda_error(status, + "allocate resident SwiGLU activation buffer"); + } + constexpr unsigned int threads = 256U; + constexpr unsigned int warps = threads / 32U; + constexpr unsigned int hidden_blocks = + (kDsv4MhcHidden + threads - 1U) / threads; + dsv4_bf16_to_fp32<<>>( + state.dsv4_mhc_workspace->layer_input, state.moe_hidden, + kDsv4MhcHidden); + auto* gate_output = state.moe_activations; + auto* up_output = gate_output + intermediate; + const auto project = [&](const CudaWeight& weight, float* source, + float* destination, + std::uint64_t rows) -> cudaError_t { + const auto& descriptor = weight.impl_->descriptor; + if (descriptor.encoding == CudaWeightEncoding::Plain) { + const auto blocks = static_cast( + (rows + warps - 1U) / warps); + bf16_matvec_kernel<<>>( + destination, source, + static_cast(weight.impl_->weights), + descriptor.columns, rows); + return cudaGetLastError(); + } + return launch_regfed_fp8_matvec( + state.moe_regfed, descriptor, weight.impl_->weights, + weight.impl_->scales, weight.impl_->fragment_prepacked, + source, destination, state.stream); + }; + if (auto status = project(gate, state.moe_hidden, gate_output, + intermediate); status != cudaSuccess) { + return cuda_error(status, "project resident GLM-5.3 SwiGLU gate"); + } + if (auto status = project(up, state.moe_hidden, up_output, intermediate); + status != cudaSuccess) { + return cuda_error(status, "project resident GLM-5.3 SwiGLU up"); + } + round_bf16_rows_kernel<<< + static_cast((intermediate + threads - 1U) / threads), + threads, 0U, state.stream>>>(gate_output, intermediate); + round_bf16_rows_kernel<<< + static_cast((intermediate + threads - 1U) / threads), + threads, 0U, state.stream>>>(up_output, intermediate); + glm53_swiglu_kernel<<< + static_cast((intermediate + threads - 1U) / threads), + threads, 0U, state.stream>>>(gate_output, up_output, intermediate); + if (auto status = project(down, gate_output, state.moe_hidden, + kDsv4MhcHidden); status != cudaSuccess) { + return cuda_error(status, "project resident GLM-5.3 SwiGLU down"); + } + round_bf16_rows_kernel<<>>( + state.moe_hidden, kDsv4MhcHidden); + dsv4_fp32_to_bf16<<>>( + state.moe_hidden, state.dsv4_mhc_workspace->branch, + kDsv4MhcHidden); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, "launch resident GLM-5.3 SwiGLU"); + } + state.dsv4_mhc_branch_ready = true; + return {}; +} + +ValidationResult CudaBackend::glm53_mla_decode_to_mhc( + const CudaGlm53MlaRequest& request) { + if (request.state == nullptr || !request.state->valid() || + request.position >= request.maximum_context || request.heads == 0U || + request.head_dim == 0U || request.query_rank == 0U || + request.key_value_rank == 0U) { + return {{"CUDA GLM-5.3 resident MLA command is invalid"}}; + } + const auto device = request.state->device(); + const auto found = impl_->devices.find(device); + if (found == impl_->devices.end()) { + return {{"CUDA GLM-5.3 resident MLA targets an uninitialized device"}}; + } + auto& state = found->second; + if (!state.dsv4_mhc_supported || state.dsv4_mhc_stage != 1U || + state.dsv4_mhc_workspace == nullptr || state.dsv4_mhc_branch_ready || + state.dsv4_mhc_failed || state.moe_in_flight) { + return {{"CUDA GLM-5.3 resident MLA violates mHC command order"}}; + } + const auto width = static_cast(request.heads) * + request.head_dim; + const auto expanded_width = 2ULL * width; + const auto cache_floats = + static_cast(request.maximum_context) * + request.key_value_rank; + const auto state_floats = cache_floats + request.query_rank + + request.key_value_rank; + if (request.state->device_bytes() < state_floats * sizeof(float)) { + return {{"CUDA GLM-5.3 resident MLA state extent is invalid"}}; + } + const auto valid = [device](const CudaWeight* weight, + std::uint64_t rows, + std::uint64_t columns) { + if (weight == nullptr || !weight->valid() || + weight->device() != device || + weight->impl_->descriptor.rows != rows || + weight->impl_->descriptor.columns != columns) return false; + const auto& descriptor = weight->impl_->descriptor; + return (descriptor.encoding == CudaWeightEncoding::Plain && + descriptor.dtype == SafetensorsDtype::Bf16) || + (descriptor.encoding == + CudaWeightEncoding::Fp8E4m3Block128F32 && + weight->impl_->fragment_prepacked); + }; + if (!valid(request.query_a, request.query_rank, kDsv4MhcHidden) || + !valid(request.key_value_a, request.key_value_rank, + kDsv4MhcHidden) || + !valid(request.query_b, width, request.query_rank) || + !valid(request.key_value_b, expanded_width, + request.key_value_rank) || + !valid(request.output, kDsv4MhcHidden, width)) { + return {{"CUDA GLM-5.3 resident MLA projection shapes are invalid"}}; + } + const auto history = static_cast(request.position) + 1U; + const auto compressed_width = + static_cast(request.heads) * request.key_value_rank; + const auto workspace_floats = kDsv4MhcHidden + request.query_rank + width + + 2U * compressed_width + width + kDsv4MhcHidden; + const auto workspace_bytes = workspace_floats * sizeof(float); + if (auto status = cudaSetDevice(device); status != cudaSuccess) { + return cuda_error(status, + "select CUDA device for GLM-5.3 resident MLA"); + } + if (workspace_bytes > state.glm53_mla_workspace_bytes) { + if (state.glm53_mla_workspace != nullptr) { + static_cast(cudaFree(state.glm53_mla_workspace)); + } + state.glm53_mla_workspace = nullptr; + state.glm53_mla_workspace_bytes = 0U; + if (auto status = cudaMalloc(&state.glm53_mla_workspace, + workspace_bytes); + status != cudaSuccess) { + return cuda_error(status, + "allocate GLM-5.3 resident MLA workspace"); + } + state.glm53_mla_workspace_bytes = workspace_bytes; + } + auto* packed = static_cast(request.state->impl_->data); + auto* q_norm = packed + cache_floats; + auto* kv_norm = q_norm + request.query_rank; + auto* workspace = reinterpret_cast(state.glm53_mla_workspace); + auto* input = workspace; + auto* q_rank = input + kDsv4MhcHidden; + auto* query = q_rank + request.query_rank; + auto* compressed_query = query + width; + auto* weighted_latent = compressed_query + compressed_width; + auto* attended = weighted_latent + compressed_width; + auto* output = attended + width; + auto* latent = packed + static_cast(request.position) * + request.key_value_rank; + constexpr unsigned int threads = 256U; + constexpr unsigned int warps = threads / 32U; + constexpr unsigned int hidden_blocks = + (kDsv4MhcHidden + threads - 1U) / threads; + dsv4_bf16_to_fp32<<>>( + state.dsv4_mhc_workspace->layer_input, input, kDsv4MhcHidden); + const auto project_one = [&](const CudaWeight* weight, float* source, + float* destination, + std::uint64_t rows) -> cudaError_t { + const auto& descriptor = weight->impl_->descriptor; + if (descriptor.encoding == CudaWeightEncoding::Plain) { + const auto blocks = static_cast( + (rows + warps - 1U) / warps); + bf16_matvec_kernel<<>>( + destination, source, + static_cast(weight->impl_->weights), + descriptor.columns, rows); + } else if (auto status = launch_regfed_fp8_matvec( + state.moe_regfed, descriptor, weight->impl_->weights, + weight->impl_->scales, + weight->impl_->fragment_prepacked, source, destination, + state.stream); status != cudaSuccess) { + return status; + } + round_bf16_rows_kernel<<< + static_cast((rows + threads - 1U) / threads), + threads, 0U, state.stream>>>(destination, rows); + return cudaGetLastError(); + }; + if (auto status = project_one(request.query_a, input, q_rank, + request.query_rank); status != cudaSuccess) { + return cuda_error(status, "project resident GLM-5.3 MLA query A"); + } + if (auto status = project_one(request.key_value_a, input, latent, + request.key_value_rank); + status != cudaSuccess) { + return cuda_error(status, "project resident GLM-5.3 MLA KV A"); + } + glm53_rms_norm_bf16_kernel<<<1U, threads, 0U, state.stream>>>( + q_rank, q_norm, request.query_rank); + glm53_rms_norm_bf16_kernel<<<1U, threads, 0U, state.stream>>>( + latent, kv_norm, request.key_value_rank); + if (auto status = project_one(request.query_b, q_rank, query, width); + status != cudaSuccess) { + return cuda_error(status, "project resident GLM-5.3 MLA query B"); + } + const auto* kv_weights = static_cast( + request.key_value_b->impl_->weights); + glm53_mla_absorb_query_kernel<<< + static_cast((compressed_width + threads - 1U) / threads), + threads, 0U, state.stream>>>( + query, kv_weights, compressed_query, request.heads, request.head_dim, + request.key_value_rank); + glm53_mla_latent_attention_kernel<<< + request.heads, threads, + static_cast(history * sizeof(float)), + state.stream>>>(compressed_query, packed, weighted_latent, + static_cast(history), request.heads, + request.head_dim, request.key_value_rank); + glm53_mla_expand_value_kernel<<< + static_cast((width + threads - 1U) / threads), threads, + 0U, state.stream>>>(weighted_latent, kv_weights, attended, + request.heads, request.head_dim, + request.key_value_rank); + if (auto status = project_one(request.output, attended, output, + kDsv4MhcHidden); + status != cudaSuccess) { + return cuda_error(status, "project resident GLM-5.3 MLA output"); + } + dsv4_fp32_to_bf16<<>>( + output, state.dsv4_mhc_workspace->branch, kDsv4MhcHidden); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + return cuda_error(status, "launch resident GLM-5.3 MLA"); + } + state.dsv4_mhc_branch_ready = true; + return {}; +} diff --git a/kernels/cuda/detail/backend_indexing.inc.cuh b/kernels/cuda/detail/backend_indexing.inc.cuh index 13bb8f6..23a5124 100644 --- a/kernels/cuda/detail/backend_indexing.inc.cuh +++ b/kernels/cuda/detail/backend_indexing.inc.cuh @@ -10,6 +10,113 @@ ValidationResult CudaBackend::matmul(const CudaWeight& weight, dsv4_fp8_tensor_page); } +ValidationResult CudaBackend::matmul_batch( + std::span items) { + ValidationResult result; + if (items.empty() || items.front().weight == nullptr || + !items.front().weight->valid()) { + result.errors.emplace_back("CUDA matmul batch is empty or invalid"); + return result; + } + const int device = items.front().weight->device(); + if (impl_->detailed_timing) { + for (const auto& item : items) { + if (item.weight == nullptr || item.weight->device() != device) { + return {{"CUDA matmul batch spans invalid or different devices"}}; + } + auto completed = matmul(*item.weight, item.input, item.rows, + item.output, item.round_bf16_output, nullptr, + item.fp8_tensor_page); + if (!completed.ok()) return completed; + } + return result; + } + + std::vector input_offsets(items.size()); + std::vector output_offsets(items.size()); + std::uint64_t input_bytes = 0U; + std::uint64_t output_bytes = 0U; + for (std::size_t index = 0U; index < items.size(); ++index) { + const auto& item = items[index]; + if (item.weight == nullptr || !item.weight->valid() || + item.weight->device() != device || item.input.empty() || + item.output.empty() || item.rows == 0U || + item.input.size_bytes() > + std::numeric_limits::max() - input_bytes || + item.output.size_bytes() > + std::numeric_limits::max() - output_bytes) { + result.errors.emplace_back( + "CUDA matmul batch spans invalid or different devices"); + return result; + } + input_offsets[index] = input_bytes; + output_offsets[index] = output_bytes; + input_bytes += item.input.size_bytes(); + output_bytes += item.output.size_bytes(); + } + auto& state = impl_->devices.at(device); + const auto grow_pinned = [](std::byte*& pointer, std::uint64_t& capacity, + std::uint64_t required) -> cudaError_t { + if (required <= capacity) return cudaSuccess; + void* replacement = nullptr; + if (auto status = cudaMallocHost(&replacement, + static_cast(required)); + status != cudaSuccess) { + return status; + } + if (pointer != nullptr) static_cast(cudaFreeHost(pointer)); + pointer = static_cast(replacement); + capacity = required; + return cudaSuccess; + }; + if (auto status = grow_pinned(state.matmul_host_input, + state.matmul_host_input_bytes, input_bytes); + status != cudaSuccess) { + return cuda_error(status, "allocate CUDA matmul batch input staging"); + } + if (auto status = grow_pinned(state.matmul_host_output, + state.matmul_host_output_bytes, output_bytes); + status != cudaSuccess) { + return cuda_error(status, "allocate CUDA matmul batch output staging"); + } + for (std::size_t index = 0U; index < items.size(); ++index) { + std::memcpy(state.matmul_host_input + input_offsets[index], + items[index].input.data(), items[index].input.size_bytes()); + } + for (std::size_t index = 0U; index < items.size(); ++index) { + const auto& item = items[index]; + auto issued = matmul_impl( + *item.weight, item.input, item.rows, 0U, 0U, item.output, 0.0F, + item.round_bf16_output, nullptr, item.fp8_tensor_page, + state.matmul_host_input + input_offsets[index], + state.matmul_host_output + output_offsets[index], true); + if (!issued.ok()) { + static_cast(cudaStreamSynchronize(state.stream)); + return issued; + } + } + const auto wait_started = std::chrono::steady_clock::now(); + if (auto status = cudaStreamSynchronize(state.stream); status != cudaSuccess) { + return cuda_error(status, "synchronize CUDA matmul batch"); + } + for (std::size_t index = 0U; index < items.size(); ++index) { + std::memcpy(items[index].output.data(), + state.matmul_host_output + output_offsets[index], + items[index].output.size_bytes()); + } + const auto wait_nanoseconds = elapsed_nanoseconds_since(wait_started); + { + std::scoped_lock lock(impl_->mutex); + auto& device_stats = *std::find_if( + impl_->stats.devices.begin(), impl_->stats.devices.end(), + [device](const auto& value) { return value.device == device; }); + record_synchronization(device_stats, + SynchronizationSubsystem::Projection, 1U, + wait_nanoseconds); + } + return result; +} + ValidationResult CudaBackend::matmul_softcap( const CudaWeight& weight, std::span input, float softcap, std::span output) { @@ -64,6 +171,18 @@ bool CudaBackend::dsv4_fp8_tensor_page_supported(int device) const noexcept { found->second.dsv4_fp8_tensor_page_supported; } +bool CudaBackend::fp8_f32_tensor_page_supported(int device) const noexcept { + const auto found = impl_->devices.find(device); + return found != impl_->devices.end() && + found->second.fp8_f32_tensor_page_supported; +} + +bool CudaBackend::fp8_f32_register_fed_supported(int device) const noexcept { + const auto found = impl_->devices.find(device); + return found != impl_->devices.end() && + found->second.fp8_f32_register_fed_supported; +} + ValidationResult CudaBackend::validate_dsv4_mhc_device(int device) const { ValidationResult result; const auto found = impl_->devices.find(device); @@ -1292,4 +1411,3 @@ ValidationResult CudaBackend::dsv4_physical_lightning_index( } return result; } - diff --git a/kernels/cuda/detail/backend_kernels.cuh b/kernels/cuda/detail/backend_kernels.cuh index ee15807..0a30e5a 100644 --- a/kernels/cuda/detail/backend_kernels.cuh +++ b/kernels/cuda/detail/backend_kernels.cuh @@ -245,6 +245,36 @@ __global__ void quantize_activation_e4m3_kernel(float* values, value = quantize_e4m3_value(value / scale) * scale; } +// GLM-5.3's compressed-tensors dynamic activation contract stores ordinary +// F32 inverse scales, not E8M0 powers of two. Simulate its per-token K128 +// quantization in place so the scalar matmul consumes the same values as the +// fused FP8 kernel while retaining the existing F32 activation workspace. +__global__ void quantize_activation_e4m3_f32_scale_kernel( + float* values, std::uint64_t columns, std::uint32_t rows) { + const std::uint32_t row = blockIdx.y; + const std::uint64_t group_begin = + static_cast(blockIdx.x) * 128U; + if (row >= rows || group_begin >= columns) return; + const std::uint64_t index = group_begin + threadIdx.x; + const float magnitude = index < columns + ? fabsf(values[static_cast(row) * columns + index]) + : 0.0F; + __shared__ float maximum[128]; + maximum[threadIdx.x] = magnitude; + __syncthreads(); + for (unsigned int stride = 64U; stride != 0U; stride >>= 1U) { + if (threadIdx.x < stride) { + maximum[threadIdx.x] = fmaxf(maximum[threadIdx.x], + maximum[threadIdx.x + stride]); + } + __syncthreads(); + } + if (index >= columns) return; + const float scale = maximum[0] > 0.0F ? maximum[0] / 448.0F : 1.0F; + auto& value = values[static_cast(row) * columns + index]; + value = quantize_e4m3_value(value / scale) * scale; +} + __global__ void round_bf16_rows_kernel(float* values, std::uint64_t elements) { for (std::uint64_t index = blockIdx.x * blockDim.x + threadIdx.x; index < elements; index += gridDim.x * blockDim.x) { @@ -295,6 +325,45 @@ __global__ void quantize_activation_e4m3_bytes_kernel( } } +// Compact continuous-scale activation used by GLM-5.3. This is byte-for-byte +// the same E4M3 quantization simulated by +// quantize_activation_e4m3_f32_scale_kernel, but keeps one raw code per value +// and one F32 scale per row/K128 block. Keeping the scale separate lets tensor +// cores dot exactly representable raw E4M3 values before the two continuous +// scales are applied in F32. +__global__ void quantize_activation_e4m3_f32_bytes_kernel( + unsigned char* values, float* scales, const float* source, + std::uint64_t columns, std::uint32_t rows) { + const std::uint32_t row = blockIdx.y; + const std::uint64_t group_begin = + static_cast(blockIdx.x) * 128U; + if (row >= rows || group_begin >= columns) return; + const std::uint64_t index = group_begin + threadIdx.x; + const float value = index < columns + ? source[static_cast(row) * columns + + index] + : 0.0F; + __shared__ float maximum[128]; + maximum[threadIdx.x] = fabsf(value); + __syncthreads(); + for (unsigned int stride = 64U; stride != 0U; stride >>= 1U) { + if (threadIdx.x < stride) { + maximum[threadIdx.x] = fmaxf(maximum[threadIdx.x], + maximum[threadIdx.x + stride]); + } + __syncthreads(); + } + const float scale = maximum[0] > 0.0F ? maximum[0] / 448.0F : 1.0F; + if (threadIdx.x == 0U) { + scales[static_cast(row) * gridDim.x + blockIdx.x] = + scale; + } + if (index < columns) { + values[static_cast(row) * columns + index] = + encode_e4m3_value(quantize_e4m3_value(value / scale)); + } +} + // The persistent mHC workspace stores its layer input as BF16. Decode and // apply the same 128-column FP8 activation simulation in one launch so the // shared expert sees exactly the values produced by the existing host bridge @@ -1473,9 +1542,38 @@ __global__ void native_fp8_matmul_kernel( } } +__global__ void native_fp8_f32_scale_matmul_kernel( + float* output, const float* input, const unsigned char* weights, + const float* scales, std::uint64_t scale_columns, + std::uint32_t batch, std::uint64_t columns, std::uint64_t rows, + std::uint32_t groups, std::uint64_t rows_per_group) { + const std::uint64_t output_row = blockIdx.x; + const std::uint32_t batch_row = blockIdx.y; + if (output_row >= rows || batch_row >= batch) return; + const std::uint64_t input_row = groups == 0U + ? batch_row + : static_cast(batch_row) * groups + + output_row / rows_per_group; + const std::uint64_t input_base = input_row * columns; + const std::uint64_t weight_base = output_row * columns; + float sum = 0.0F; + for (std::uint64_t column = threadIdx.x; column < columns; + column += blockDim.x) { + const float weight = fp8_e4m3_value(weights[weight_base + column]); + const float scale = scales[(output_row / 128U) * scale_columns + + column / 128U]; + sum += input[input_base + column] * weight * scale; + } + sum = reduce_block(sum); + if (threadIdx.x == 0U) { + output[static_cast(batch_row) * rows + output_row] = sum; + } +} + constexpr std::uint32_t kDsv4Fp8TensorBlockM = 64U; constexpr std::uint32_t kDsv4Fp8TensorBlockN = 128U; constexpr std::uint32_t kDsv4Fp8TensorBlockK = 128U; +constexpr std::uint32_t kFp8F32TensorBlockN = 64U; // SM86 page-projection path. Both operands remain byte FP8 in global memory; // each tile widens them exactly to BF16 in shared memory and uses BF16 WMMA. @@ -1597,6 +1695,174 @@ __global__ void dsv4_fp8_decode_bf16_tensor_kernel( } } +// QPN-derived continuous-scale page projection. Unlike DeepSeek's E8M0 route, +// neither GLM scale is a power of two, so widening a scaled activation to BF16 +// would discard part of the checkpoint's declared arithmetic. Instead each +// K128 tile performs a tensor-core dot over the raw E4M3 values (all exactly +// representable in BF16), publishes that partial, and applies the activation +// and weight scales in F32. A 64x64 output tile leaves the raw A/B tiles, the +// published partial and the accumulated result within the 48 KiB SM86 shared +// memory budget without assuming a particular installed GPU. +__device__ __forceinline__ unsigned char fp8_fragment_prepacked_code( + const unsigned char* codes, std::uint32_t row, std::uint32_t column, + std::uint32_t columns) { + const std::uint32_t pair = column / 32U; + const std::uint32_t within16 = column & 15U; + const std::uint32_t lane = (row & 7U) * 4U + + ((within16 & 7U) / 2U); + const auto packed = reinterpret_cast(codes)[ + (static_cast(row / 16U) * (columns / 32U) + pair) * 32U + + lane]; + const std::uint32_t words[4] = {packed.x, packed.y, packed.z, packed.w}; + const bool upper_columns = within16 >= 8U; + const std::uint32_t word = (column & 31U) / 16U * 2U + + (upper_columns ? 1U : 0U); + const std::uint32_t i = (upper_columns ? 4U : 0U) + + ((row & 15U) >= 8U ? 2U : 0U) + + (within16 & 1U); + return static_cast((words[word] >> ((i & 3U) * 8U)) & + 0xFFU); +} + +template +__global__ void fp8_f32_decode_bf16_tensor_kernel( + float* output, const unsigned char* input, const float* input_scales, + const unsigned char* weights, const float* weight_scales, + std::uint32_t batch, std::uint32_t columns, std::uint32_t rows) { + using namespace nvcuda; + union SharedAOrPartial { + __nv_bfloat16 a[kDsv4Fp8TensorBlockM * kDsv4Fp8TensorBlockK]; + float partial[kDsv4Fp8TensorBlockM * kFp8F32TensorBlockN]; + }; + __shared__ SharedAOrPartial shared_a_or_partial; + __shared__ __nv_bfloat16 shared_b[ + kDsv4Fp8TensorBlockK * kFp8F32TensorBlockN]; + __shared__ float totals[ + kDsv4Fp8TensorBlockM * kFp8F32TensorBlockN]; + + const std::uint32_t tile_m = blockIdx.y * kDsv4Fp8TensorBlockM; + const std::uint32_t tile_n = blockIdx.x * kFp8F32TensorBlockN; + const std::uint32_t warp = threadIdx.x / warpSize; + const std::uint32_t warp_m = warp & 3U; + const std::uint32_t warp_n_group = warp >> 2U; + constexpr std::uint32_t fragments_per_warp = 2U; + const std::uint32_t scale_columns = columns / kDsv4Fp8TensorBlockK; + + for (std::uint32_t index = threadIdx.x; + index < kDsv4Fp8TensorBlockM * kFp8F32TensorBlockN; + index += blockDim.x) { + totals[index] = 0.0F; + } + __syncthreads(); + + for (std::uint32_t tile_k = 0U; tile_k < columns; + tile_k += kDsv4Fp8TensorBlockK) { + for (std::uint32_t index = threadIdx.x; + index < kDsv4Fp8TensorBlockM * kDsv4Fp8TensorBlockK; + index += blockDim.x) { + const std::uint32_t local_m = index / kDsv4Fp8TensorBlockK; + const std::uint32_t local_k = index % kDsv4Fp8TensorBlockK; + const std::uint32_t global_m = tile_m + local_m; + const auto encoded = global_m < batch + ? input[static_cast(global_m) * columns + + tile_k + local_k] + : 0U; + shared_a_or_partial.a[index] = + __float2bfloat16_rn(fp8_e4m3_value(encoded)); + } + for (std::uint32_t index = threadIdx.x; + index < kDsv4Fp8TensorBlockK * kFp8F32TensorBlockN; + index += blockDim.x) { + const std::uint32_t local_k = index / kFp8F32TensorBlockN; + const std::uint32_t local_n = index % kFp8F32TensorBlockN; + const std::uint32_t global_n = tile_n + local_n; + const auto encoded = kFragmentPrepacked + ? fp8_fragment_prepacked_code(weights, global_n, + tile_k + local_k, columns) + : weights[static_cast(global_n) * columns + + tile_k + local_k]; + shared_b[index] = + __float2bfloat16_rn(fp8_e4m3_value(encoded)); + } + __syncthreads(); + + wmma::fragment a_fragment; + wmma::fragment b_fragment; + wmma::fragment + accumulators[fragments_per_warp]; + for (std::uint32_t fragment = 0U; fragment < fragments_per_warp; + ++fragment) { + wmma::fill_fragment(accumulators[fragment], 0.0F); + } + for (std::uint32_t local_k = 0U; + local_k < kDsv4Fp8TensorBlockK; local_k += 16U) { + wmma::load_matrix_sync( + a_fragment, + shared_a_or_partial.a + + warp_m * 16U * kDsv4Fp8TensorBlockK + local_k, + kDsv4Fp8TensorBlockK); + for (std::uint32_t fragment = 0U; + fragment < fragments_per_warp; ++fragment) { + const std::uint32_t fragment_n = + warp_n_group * fragments_per_warp + fragment; + wmma::load_matrix_sync( + b_fragment, + shared_b + local_k * kFp8F32TensorBlockN + + fragment_n * 16U, + kFp8F32TensorBlockN); + wmma::mma_sync(accumulators[fragment], a_fragment, b_fragment, + accumulators[fragment]); + } + } + __syncthreads(); + for (std::uint32_t fragment = 0U; fragment < fragments_per_warp; + ++fragment) { + const std::uint32_t fragment_n = + warp_n_group * fragments_per_warp + fragment; + float* destination = shared_a_or_partial.partial + + warp_m * 16U * kFp8F32TensorBlockN + fragment_n * 16U; + wmma::store_matrix_sync(destination, accumulators[fragment], + kFp8F32TensorBlockN, + wmma::mem_row_major); + } + __syncthreads(); + + for (std::uint32_t index = threadIdx.x; + index < kDsv4Fp8TensorBlockM * kFp8F32TensorBlockN; + index += blockDim.x) { + const std::uint32_t local_m = index / kFp8F32TensorBlockN; + const std::uint32_t local_n = index % kFp8F32TensorBlockN; + const std::uint32_t global_m = tile_m + local_m; + const std::uint32_t global_n = tile_n + local_n; + if (global_m < batch) { + const float activation_scale = input_scales[ + static_cast(global_m) * scale_columns + + tile_k / kDsv4Fp8TensorBlockK]; + const float weight_scale = weight_scales[ + static_cast(global_n / 128U) * + scale_columns + tile_k / kDsv4Fp8TensorBlockK]; + totals[index] += shared_a_or_partial.partial[index] * + activation_scale * weight_scale; + } + } + __syncthreads(); + } + + for (std::uint32_t index = threadIdx.x; + index < kDsv4Fp8TensorBlockM * kFp8F32TensorBlockN; + index += blockDim.x) { + const std::uint32_t local_m = index / kFp8F32TensorBlockN; + const std::uint32_t local_n = index % kFp8F32TensorBlockN; + const std::uint32_t global_m = tile_m + local_m; + if (global_m < batch) { + output[static_cast(global_m) * rows + tile_n + + local_n] = totals[index]; + } + } +} + __global__ void native_fp4_matmul_kernel( float* output, const float* input, const unsigned char* weights, const unsigned char* scales, std::uint64_t packed_columns, @@ -1725,6 +1991,21 @@ struct Mxfp4MoeBatch { std::uint32_t rows{}; }; +// Standard compressed-tensors FP8: E4M3 payloads with ordinary F32 inverse +// scales per 128x128 weight block. Kept separate from the E8M0-scaled FP8 +// batches so their scale bytes can never be reinterpreted silently. +struct Fp8F32MoeBatch { + const unsigned char* gate_weights[kMaxMoeExperts]{}; + const float* gate_scales[kMaxMoeExperts]{}; + const unsigned char* up_weights[kMaxMoeExperts]{}; + const float* up_scales[kMaxMoeExperts]{}; + const unsigned char* down_weights[kMaxMoeExperts]{}; + const float* down_scales[kMaxMoeExperts]{}; + float coefficients[kMaxMoeExperts]{}; + std::uint32_t count{}; + std::uint32_t rows{}; +}; + // Laguna carries the routed experts of layers 40-47 as plain BF16. struct PlainBf16MoeBatch { const __nv_bfloat16* gate_weights[kMaxMoeExperts]{}; @@ -2633,6 +2914,92 @@ __global__ void mxfp4_moe_down_kernel( } } +__global__ void fp8_f32_moe_gate_up_kernel( + float* activations, const float* hidden, Fp8F32MoeBatch batch, + std::uint64_t columns, std::uint64_t intermediate, + std::uint64_t scale_columns, float swiglu_limit, + unsigned int* error_flag) { + const std::uint64_t output_row = blockIdx.x; + const std::uint32_t batch_row = blockIdx.y; + const auto expert = batch_row / batch.rows; + const auto row = batch_row % batch.rows; + if (output_row >= intermediate || expert >= batch.count) return; + + const auto* gate_weights = batch.gate_weights[expert]; + const auto* gate_scales = batch.gate_scales[expert]; + const auto* up_weights = batch.up_weights[expert]; + const auto* up_scales = batch.up_scales[expert]; + const auto weight_base = output_row * columns; + const auto scale_base = (output_row / 128U) * scale_columns; + const auto input_base = static_cast(row) * columns; + float gate = 0.0F; + float up = 0.0F; + for (std::uint64_t column = threadIdx.x; column < columns; + column += blockDim.x) { + const float input = hidden[input_base + column]; + gate += input * fp8_e4m3_value(gate_weights[weight_base + column]) * + gate_scales[scale_base + column / 128U]; + up += input * fp8_e4m3_value(up_weights[weight_base + column]) * + up_scales[scale_base + column / 128U]; + } + gate = reduce_block(gate); + __syncthreads(); + up = reduce_block(up); + if (threadIdx.x == 0U) { + gate = bf16_round(gate); + up = bf16_round(up); + if (!isfinite(gate) || !isfinite(up)) { + atomicExch(error_flag, 1U); + return; + } + const float limited_gate = swiglu_limit > 0.0F + ? fminf(gate, swiglu_limit) : gate; + const float limited_up = swiglu_limit > 0.0F + ? fminf(fmaxf(up, -swiglu_limit), swiglu_limit) : up; + const float exponential = limited_gate >= 0.0F + ? expf(-limited_gate) : expf(limited_gate); + const float sigmoid = limited_gate >= 0.0F + ? 1.0F / (1.0F + exponential) + : exponential / (1.0F + exponential); + const auto activation = + (static_cast(expert) * batch.rows + row) * + intermediate + output_row; + activations[activation] = bf16_round( + limited_gate * sigmoid * limited_up); + } +} + +__global__ void fp8_f32_moe_down_kernel( + float* output, const float* activations, Fp8F32MoeBatch batch, + std::uint64_t columns, std::uint64_t rows, + std::uint64_t scale_columns, unsigned int* error_flag) { + const std::uint64_t output_row = blockIdx.x; + const std::uint32_t batch_row = blockIdx.y; + const auto expert = batch_row / batch.rows; + const auto row = batch_row % batch.rows; + if (output_row >= rows || expert >= batch.count) return; + + const auto* weights = batch.down_weights[expert]; + const auto* scales = batch.down_scales[expert]; + const auto weight_base = output_row * columns; + const auto scale_base = (output_row / 128U) * scale_columns; + const auto input_base = + (static_cast(expert) * batch.rows + row) * columns; + float sum = 0.0F; + for (std::uint64_t column = threadIdx.x; column < columns; + column += blockDim.x) { + sum += activations[input_base + column] * + fp8_e4m3_value(weights[weight_base + column]) * + scales[scale_base + column / 128U]; + } + sum = reduce_block(sum); + if (threadIdx.x == 0U) { + if (!isfinite(sum)) atomicExch(error_flag, 1U); + output[(static_cast(expert) * batch.rows + row) * rows + + output_row] = bf16_round(sum); + } +} + // Plain BF16 counterparts. These keep bf16_matvec_kernel's one-warp-per-output- // row layout and its __fadd_rn/__shfl_down_sync reduction rather than the // block reduction the quantized batches use, so the dot product is summed in @@ -3678,6 +4045,47 @@ __global__ void regfed_activation_fragment_kernel( } } +// Compact E4M3 activation codes to MMA B-fragment order. GLM keeps its +// continuous activation scale beside these codes; the tensor dot sees only +// the raw representable values and the scale is applied to each K128 partial +// in F32 by regfed_fp8_f32_matmul_kernel. +__global__ void regfed_fp8_activation_fragment_kernel( + uint2* __restrict__ destination, + const unsigned char* __restrict__ source, std::uint32_t m, + std::uint32_t columns, std::uint32_t column_blocks, + std::uint32_t groups_per_block) { + const std::uint32_t k_tiles = columns / kRegfedTileK; + const std::uint32_t total = + k_tiles * column_blocks * groups_per_block * 4U; + constexpr std::uint32_t unit_factor = + ((127U + 120U) << 7U) * 0x0001'0001U; + for (std::uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + index < total; index += gridDim.x * blockDim.x) { + const std::uint32_t thread = index % 4U; + const std::uint32_t group = (index / 4U) % groups_per_block; + const std::uint32_t block = + (index / (4U * groups_per_block)) % column_blocks; + const std::uint32_t k_tile = + index / (4U * groups_per_block * column_blocks); + const std::uint32_t column = block * kRegfedTileM + group; + uint2 value = make_uint2(0U, 0U); + if (column < m) { + const auto* row = source + + static_cast(column) * columns + + k_tile * kRegfedTileK; + const std::uint32_t first = + static_cast(row[thread * 2U]) | + (static_cast(row[thread * 2U + 1U]) << 8U); + const std::uint32_t second = + static_cast(row[thread * 2U + 8U]) | + (static_cast(row[thread * 2U + 9U]) << 8U); + value = make_uint2(dsv4_fp8_decode_pair(first, unit_factor), + dsv4_fp8_decode_pair(second, unit_factor)); + } + destination[index] = value; + } +} + // ---- the kernels ----------------------------------------------------------- // One warp owns one (N-tile, K-slice). The last slice of a tile folds the @@ -3917,6 +4325,143 @@ __global__ __launch_bounds__(128) void regfed_fp8_matmul_kernel( } } +// Register-fed GLM W8A8. Weight and activation codes remain one byte through +// HBM and are decoded directly into MMA registers. Each K128 dot is completed +// before its two arbitrary F32 scales are applied, preserving continuous +// dynamic activation scaling without widening scaled operands to BF16. +template +__global__ __launch_bounds__(128) void regfed_fp8_f32_matmul_kernel( + float* __restrict__ output, const uint4* __restrict__ codes, + const float* __restrict__ weight_scales, + const uint2* __restrict__ activations, + const float* __restrict__ activation_scales, std::uint32_t columns, + std::uint32_t rows, std::uint32_t scale_columns, std::uint32_t split, + std::uint32_t m, std::uint32_t groups_per_block, + float* __restrict__ partials, std::uint32_t* __restrict__ counters) { + const std::uint32_t lane = threadIdx.x & 31U; + const std::uint32_t warp = threadIdx.x >> 5U; + const std::uint32_t n_tiles = rows / kRegfedTileN; + const std::uint32_t pairs = columns / 32U; + const std::uint32_t pairs_per_slice = pairs / split; + const std::uint32_t group = lane >> 2U; + const std::uint32_t thread = lane & 3U; + constexpr std::uint32_t unit_factor = + ((127U + 120U) << 7U) * 0x0001'0001U; + __shared__ std::uint32_t arrived[kRegfedWarpsPerBlock]; + + bool live[kColBlocks]; + std::size_t activation_offset[kColBlocks]; +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { + live[c] = group < groups_per_block && + c * kRegfedTileM + group < m; + activation_offset[c] = + (static_cast(c) * groups_per_block + group) * 4U + + thread; + } + + for (std::uint32_t work = blockIdx.x * kRegfedWarpsPerBlock + warp; + work < n_tiles * split; work += gridDim.x * kRegfedWarpsPerBlock) { + const std::uint32_t n_tile = work / split; + const std::uint32_t slice = work % split; + float totals_for_slice[kColBlocks][4]{}; + const std::uint32_t begin = slice * pairs_per_slice; + const std::uint32_t end = begin + pairs_per_slice; + for (std::uint32_t pair_block = begin; pair_block < end; + pair_block += 4U) { + float block_acc[kColBlocks][4]{}; +#pragma unroll + for (std::uint32_t within = 0U; within < 4U; ++within) { + const std::uint32_t pair = pair_block + within; + const std::uint32_t k_tile = pair * 2U; + const uint4 packed = codes[ + (static_cast(n_tile) * pairs + pair) * 32U + + lane]; + const std::uint32_t word[4] = {packed.x, packed.y, packed.z, + packed.w}; +#pragma unroll + for (std::uint32_t half = 0U; half < 2U; ++half) { + const std::uint32_t low = word[half * 2U]; + const std::uint32_t high = word[half * 2U + 1U]; + const std::uint32_t a0 = + dsv4_fp8_decode_pair(low & 0xFFFFU, unit_factor); + const std::uint32_t a1 = + dsv4_fp8_decode_pair(low >> 16U, unit_factor); + const std::uint32_t a2 = + dsv4_fp8_decode_pair(high & 0xFFFFU, unit_factor); + const std::uint32_t a3 = + dsv4_fp8_decode_pair(high >> 16U, unit_factor); + const std::size_t tile_base = + (static_cast(k_tile) + half) * + kColBlocks * groups_per_block * 4U; +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { + const uint2 b = live[c] + ? activations[tile_base + activation_offset[c]] + : make_uint2(0U, 0U); + dsv4_mma_m16n8k16( + block_acc[c][0], block_acc[c][1], + block_acc[c][2], block_acc[c][3], a0, a1, a2, a3, + b.x, b.y); + } + } + } + const std::uint32_t scale_column = pair_block / 4U; + const float weight_scale = weight_scales[ + static_cast(n_tile / 8U) * scale_columns + + scale_column]; +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { +#pragma unroll + for (std::uint32_t i = 0U; i < 4U; ++i) { + const std::uint32_t input_row = + c * kRegfedTileM + thread * 2U + (i & 1U); + if (input_row < m) { + totals_for_slice[c][i] += block_acc[c][i] * + activation_scales[ + static_cast(input_row) * + scale_columns + scale_column] * + weight_scale; + } + } + } + } + + float* slot = partials + static_cast(work) * + kRegfedTileN * kRegfedMaxM; +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { +#pragma unroll + for (std::uint32_t i = 0U; i < 4U; ++i) { + const std::uint32_t row = group + ((i >= 2U) ? 8U : 0U); + const std::uint32_t column = + c * kRegfedTileM + thread * 2U + (i & 1U); + if (column < m) slot[row * m + column] = totals_for_slice[c][i]; + } + } + + __threadfence(); + __syncwarp(); + if (lane == 0U) arrived[warp] = atomicAdd(&counters[n_tile], 1U); + __syncwarp(); + if (arrived[warp] == split - 1U) { + if (lane < kRegfedTileN) { + for (std::uint32_t column = 0U; column < m; ++column) { + float sum = 0.0F; + for (std::uint32_t s = 0U; s < split; ++s) { + sum += partials[(static_cast(n_tile) * + split + s) * kRegfedTileN * + kRegfedMaxM + lane * m + column]; + } + output[static_cast(column) * rows + + n_tile * kRegfedTileN + lane] = sum; + } + } + if (lane == 0U) counters[n_tile] = 0U; + } + } +} + // Shape admission for the register-fed routes. Stated once, used by both the // load-time prepack and the dispatch, so a weight can never be prepacked into a // layout the kernel will not read. @@ -4019,6 +4564,364 @@ __global__ void regfed_moe_activation_fragment_kernel( } } +__global__ void regfed_fp8_moe_activation_fragment_kernel( + uint2* __restrict__ destination, + const unsigned char* __restrict__ source, std::uint32_t experts, + std::uint32_t m, std::uint32_t columns, std::uint32_t column_blocks, + std::uint32_t groups_per_block) { + const std::uint32_t k_tiles = columns / kRegfedTileK; + const std::uint32_t per_expert = + k_tiles * column_blocks * groups_per_block * 4U; + const std::uint64_t total = + static_cast(experts) * per_expert; + constexpr std::uint32_t unit_factor = + ((127U + 120U) << 7U) * 0x0001'0001U; + for (std::uint64_t index = blockIdx.x * blockDim.x + threadIdx.x; + index < total; index += gridDim.x * blockDim.x) { + const auto local = static_cast(index % per_expert); + const auto expert = static_cast(index / per_expert); + const std::uint32_t thread = local % 4U; + const std::uint32_t group = (local / 4U) % groups_per_block; + const std::uint32_t block = + (local / (4U * groups_per_block)) % column_blocks; + const std::uint32_t k_tile = + local / (4U * groups_per_block * column_blocks); + const std::uint32_t column = block * kRegfedTileM + group; + uint2 value = make_uint2(0U, 0U); + if (column < m) { + const auto* row = source + + (static_cast(expert) * m + column) * columns + + k_tile * kRegfedTileK; + const std::uint32_t first = + static_cast(row[thread * 2U]) | + (static_cast(row[thread * 2U + 1U]) << 8U); + const std::uint32_t second = + static_cast(row[thread * 2U + 8U]) | + (static_cast(row[thread * 2U + 9U]) << 8U); + value = make_uint2(dsv4_fp8_decode_pair(first, unit_factor), + dsv4_fp8_decode_pair(second, unit_factor)); + } + destination[index] = value; + } +} + +template +__global__ __launch_bounds__(128) void regfed_fp8_f32_moe_gate_up_kernel( + float* __restrict__ gate_partials, float* __restrict__ up_partials, + const uint2* __restrict__ activations, + const float* __restrict__ activation_scales, Fp8F32MoeBatch batch, + std::uint32_t columns, std::uint32_t intermediate, std::uint32_t split, + std::uint32_t m, std::uint32_t groups_per_block) { + const std::uint32_t lane = threadIdx.x & 31U; + const std::uint32_t warp = threadIdx.x >> 5U; + const std::uint32_t n_tiles = intermediate / kRegfedTileN; + const std::uint32_t pairs = columns / 32U; + const std::uint32_t pairs_per_slice = pairs / split; + const std::uint32_t scale_columns = columns / 128U; + const std::uint32_t group = lane >> 2U; + const std::uint32_t thread = lane & 3U; + const std::uint32_t total = batch.count * n_tiles * split; + constexpr std::uint32_t unit_factor = + ((127U + 120U) << 7U) * 0x0001'0001U; + bool live[kColBlocks]; + std::size_t offset[kColBlocks]; +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { + live[c] = group < groups_per_block && + c * kRegfedTileM + group < m; + offset[c] = + (static_cast(c) * groups_per_block + group) * 4U + + thread; + } + + for (std::uint32_t work = blockIdx.x * kRegfedWarpsPerBlock + warp; + work < total; work += gridDim.x * kRegfedWarpsPerBlock) { + const std::uint32_t slice = work % split; + const std::uint32_t flat = work / split; + const std::uint32_t n_tile = flat % n_tiles; + const std::uint32_t expert = flat / n_tiles; + const auto* gate4 = + reinterpret_cast(batch.gate_weights[expert]); + const auto* up4 = + reinterpret_cast(batch.up_weights[expert]); + const auto* gate_scales = batch.gate_scales[expert]; + const auto* up_scales = batch.up_scales[expert]; + float gate_totals[kColBlocks][4]{}; + float up_totals[kColBlocks][4]{}; + const std::uint32_t begin = slice * pairs_per_slice; + const std::uint32_t end = begin + pairs_per_slice; + for (std::uint32_t pair_block = begin; pair_block < end; + pair_block += 4U) { + float gate_block[kColBlocks][4]{}; + float up_block[kColBlocks][4]{}; +#pragma unroll + for (std::uint32_t within = 0U; within < 4U; ++within) { + const std::uint32_t pair = pair_block + within; + const std::size_t code_index = + (static_cast(n_tile) * pairs + pair) * 32U + + lane; + const uint4 gate_packed = gate4[code_index]; + const uint4 up_packed = up4[code_index]; + const std::uint32_t gate_words[4] = { + gate_packed.x, gate_packed.y, gate_packed.z, gate_packed.w}; + const std::uint32_t up_words[4] = { + up_packed.x, up_packed.y, up_packed.z, up_packed.w}; +#pragma unroll + for (std::uint32_t half = 0U; half < 2U; ++half) { + const auto decode = [&](const std::uint32_t* words, + std::uint32_t (&decoded)[4]) { + const auto low = words[half * 2U]; + const auto high = words[half * 2U + 1U]; + decoded[0] = dsv4_fp8_decode_pair( + low & 0xFFFFU, unit_factor); + decoded[1] = dsv4_fp8_decode_pair( + low >> 16U, unit_factor); + decoded[2] = dsv4_fp8_decode_pair( + high & 0xFFFFU, unit_factor); + decoded[3] = dsv4_fp8_decode_pair( + high >> 16U, unit_factor); + }; + std::uint32_t gate_a[4]; + std::uint32_t up_a[4]; + decode(gate_words, gate_a); + decode(up_words, up_a); + const std::size_t base = + (static_cast(pair * 2U) + half) * + kColBlocks * groups_per_block * 4U; +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { + const uint2 b = live[c] + ? activations[base + offset[c]] + : make_uint2(0U, 0U); + dsv4_mma_m16n8k16( + gate_block[c][0], gate_block[c][1], + gate_block[c][2], gate_block[c][3], gate_a[0], + gate_a[1], gate_a[2], gate_a[3], b.x, b.y); + dsv4_mma_m16n8k16( + up_block[c][0], up_block[c][1], up_block[c][2], + up_block[c][3], up_a[0], up_a[1], up_a[2], up_a[3], + b.x, b.y); + } + } + } + const std::uint32_t scale_column = pair_block / 4U; + const auto weight_scale_index = + static_cast(n_tile / 8U) * scale_columns + + scale_column; +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { +#pragma unroll + for (std::uint32_t i = 0U; i < 4U; ++i) { + const std::uint32_t input_row = + c * kRegfedTileM + thread * 2U + (i & 1U); + if (input_row < m) { + const float activation_scale = activation_scales[ + static_cast(input_row) * + scale_columns + scale_column]; + gate_totals[c][i] += gate_block[c][i] * + activation_scale * + gate_scales[weight_scale_index]; + up_totals[c][i] += up_block[c][i] * activation_scale * + up_scales[weight_scale_index]; + } + } + } + } +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { +#pragma unroll + for (std::uint32_t i = 0U; i < 4U; ++i) { + const std::uint32_t row = group + ((i >= 2U) ? 8U : 0U); + const std::uint32_t column = + c * kRegfedTileM + thread * 2U + (i & 1U); + if (column >= m) continue; + const std::size_t slot = + ((static_cast(expert) * intermediate + + n_tile * kRegfedTileN + row) * m + column) * split + + slice; + gate_partials[slot] = gate_totals[c][i]; + up_partials[slot] = up_totals[c][i]; + } + } + } +} + +__global__ void regfed_fp8_f32_moe_swiglu_kernel( + float* __restrict__ activations, const float* __restrict__ gate_partials, + const float* __restrict__ up_partials, std::uint32_t experts, + std::uint32_t intermediate, std::uint32_t m, std::uint32_t split, + float swiglu_limit, unsigned int* __restrict__ error_flag) { + const std::uint64_t total = + static_cast(experts) * intermediate * m; + for (std::uint64_t index = blockIdx.x * blockDim.x + threadIdx.x; + index < total; index += gridDim.x * blockDim.x) { + const std::uint32_t column = static_cast(index % m); + float gate = 0.0F; + float up = 0.0F; + for (std::uint32_t slice = 0U; slice < split; ++slice) { + gate += gate_partials[index * split + slice]; + up += up_partials[index * split + slice]; + } + gate = bf16_round(gate); + up = bf16_round(up); + if (!isfinite(gate) || !isfinite(up)) { + atomicExch(error_flag, 1U); + continue; + } + const float limited_gate = fminf(gate, swiglu_limit); + const float limited_up = fminf(fmaxf(up, -swiglu_limit), swiglu_limit); + const float exponential = limited_gate >= 0.0F + ? expf(-limited_gate) : expf(limited_gate); + const float sigmoid = limited_gate >= 0.0F + ? 1.0F / (1.0F + exponential) + : exponential / (1.0F + exponential); + const std::uint32_t rest = static_cast(index / m); + const std::uint32_t output_row = rest % intermediate; + const std::uint32_t expert = rest / intermediate; + activations[(static_cast(expert) * m + column) * + intermediate + output_row] = + bf16_round(limited_gate * sigmoid * limited_up); + } +} + +template +__global__ __launch_bounds__(128) void regfed_fp8_f32_moe_down_kernel( + float* __restrict__ partials, const uint2* __restrict__ activations, + const float* __restrict__ activation_scales, Fp8F32MoeBatch batch, + std::uint32_t columns, std::uint32_t rows, std::uint32_t split, + std::uint32_t m, std::uint32_t groups_per_block) { + const std::uint32_t lane = threadIdx.x & 31U; + const std::uint32_t warp = threadIdx.x >> 5U; + const std::uint32_t n_tiles = rows / kRegfedTileN; + const std::uint32_t pairs = columns / 32U; + const std::uint32_t pairs_per_slice = pairs / split; + const std::uint32_t scale_columns = columns / 128U; + const std::uint32_t group = lane >> 2U; + const std::uint32_t thread = lane & 3U; + const std::uint32_t total = batch.count * n_tiles * split; + const std::size_t per_expert_activation = + static_cast(columns / kRegfedTileK) * kColBlocks * + groups_per_block * 4U; + constexpr std::uint32_t unit_factor = + ((127U + 120U) << 7U) * 0x0001'0001U; + bool live[kColBlocks]; + std::size_t offset[kColBlocks]; +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { + live[c] = group < groups_per_block && + c * kRegfedTileM + group < m; + offset[c] = + (static_cast(c) * groups_per_block + group) * 4U + + thread; + } + for (std::uint32_t work = blockIdx.x * kRegfedWarpsPerBlock + warp; + work < total; work += gridDim.x * kRegfedWarpsPerBlock) { + const std::uint32_t slice = work % split; + const std::uint32_t flat = work / split; + const std::uint32_t n_tile = flat % n_tiles; + const std::uint32_t expert = flat / n_tiles; + const auto* codes4 = + reinterpret_cast(batch.down_weights[expert]); + const auto* weight_scales = batch.down_scales[expert]; + float totals_for_slice[kColBlocks][4]{}; + const std::uint32_t begin = slice * pairs_per_slice; + const std::uint32_t end = begin + pairs_per_slice; + for (std::uint32_t pair_block = begin; pair_block < end; + pair_block += 4U) { + float block_acc[kColBlocks][4]{}; +#pragma unroll + for (std::uint32_t within = 0U; within < 4U; ++within) { + const std::uint32_t pair = pair_block + within; + const uint4 packed = codes4[ + (static_cast(n_tile) * pairs + pair) * 32U + + lane]; + const std::uint32_t words[4] = { + packed.x, packed.y, packed.z, packed.w}; +#pragma unroll + for (std::uint32_t half = 0U; half < 2U; ++half) { + const auto low = words[half * 2U]; + const auto high = words[half * 2U + 1U]; + const std::uint32_t a[4] = { + dsv4_fp8_decode_pair(low & 0xFFFFU, unit_factor), + dsv4_fp8_decode_pair(low >> 16U, unit_factor), + dsv4_fp8_decode_pair(high & 0xFFFFU, unit_factor), + dsv4_fp8_decode_pair(high >> 16U, unit_factor)}; + const std::size_t base = + static_cast(expert) * + per_expert_activation + + (static_cast(pair * 2U) + half) * + kColBlocks * groups_per_block * 4U; +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { + const uint2 b = live[c] + ? activations[base + offset[c]] + : make_uint2(0U, 0U); + dsv4_mma_m16n8k16( + block_acc[c][0], block_acc[c][1], + block_acc[c][2], block_acc[c][3], a[0], a[1], a[2], + a[3], b.x, b.y); + } + } + } + const std::uint32_t scale_column = pair_block / 4U; + const float weight_scale = weight_scales[ + static_cast(n_tile / 8U) * scale_columns + + scale_column]; +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { +#pragma unroll + for (std::uint32_t i = 0U; i < 4U; ++i) { + const std::uint32_t input_row = + c * kRegfedTileM + thread * 2U + (i & 1U); + if (input_row < m) { + totals_for_slice[c][i] += block_acc[c][i] * + activation_scales[ + (static_cast(expert) * m + + input_row) * scale_columns + scale_column] * + weight_scale; + } + } + } + } +#pragma unroll + for (std::uint32_t c = 0U; c < kColBlocks; ++c) { +#pragma unroll + for (std::uint32_t i = 0U; i < 4U; ++i) { + const std::uint32_t row = group + ((i >= 2U) ? 8U : 0U); + const std::uint32_t column = + c * kRegfedTileM + thread * 2U + (i & 1U); + if (column >= m) continue; + const std::size_t slot = + ((static_cast(expert) * rows + + n_tile * kRegfedTileN + row) * m + column) * split + + slice; + partials[slot] = totals_for_slice[c][i]; + } + } + } +} + +__global__ void regfed_fp8_f32_moe_reduce_kernel( + float* __restrict__ output, const float* __restrict__ partials, + std::uint32_t experts, std::uint32_t rows, std::uint32_t m, + std::uint32_t split, unsigned int* __restrict__ error_flag) { + const std::uint64_t total = static_cast(experts) * rows * m; + for (std::uint64_t index = blockIdx.x * blockDim.x + threadIdx.x; + index < total; index += gridDim.x * blockDim.x) { + const std::uint32_t column = static_cast(index % m); + const std::uint32_t rest = static_cast(index / m); + const std::uint32_t output_row = rest % rows; + const std::uint32_t expert = rest / rows; + float sum = 0.0F; + for (std::uint32_t slice = 0U; slice < split; ++slice) { + sum += partials[index * split + slice]; + } + if (!isfinite(sum)) atomicExch(error_flag, 1U); + output[(static_cast(expert) * m + column) * rows + + output_row] = bf16_round(sum); + } +} + // One warp owns one (expert, N-tile, K-slice). Gate and up share the activation // fragment, so a single pass over the hidden vector feeds both weight streams. template @@ -4267,7 +5170,8 @@ __global__ void regfed_mxfp4_moe_reduce_kernel( // device, so no second persistent copy of any weight exists. [[nodiscard]] std::uint64_t fragment_prepack_scratch_bytes( const CudaWeightDescriptor& descriptor) noexcept { - if (descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128) { + if (descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128 || + descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128F32) { if (!regfed_fp8_shape_admissible(descriptor.rows, descriptor.columns)) { return 0U; } @@ -4295,7 +5199,8 @@ cudaError_t launch_fragment_prepack(const CudaWeightDescriptor& descriptor, return static_cast( std::min((total + threads - 1U) / threads, 65535U)); }; - if (descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128) { + if (descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128 || + descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128F32) { const std::uint64_t bytes = descriptor.rows * descriptor.columns; if (auto status = cudaMemcpyAsync(scratch, weights, static_cast(bytes), diff --git a/kernels/cuda/detail/backend_matmul.inc.cuh b/kernels/cuda/detail/backend_matmul.inc.cuh index 50dd2be..9d3cc4d 100644 --- a/kernels/cuda/detail/backend_matmul.inc.cuh +++ b/kernels/cuda/detail/backend_matmul.inc.cuh @@ -71,7 +71,7 @@ ValidationResult CudaBackend::prepack_marlin(int device, ValidationResult CudaBackend::prepack_fragment(int device, const CudaWeight& weight) { ValidationResult result; - if (!weight.valid()) { + if (!weight.valid() || weight.device() != device) { result.errors.emplace_back("fragment prepack received an invalid weight"); return result; } @@ -81,6 +81,14 @@ ValidationResult CudaBackend::prepack_fragment(int device, return result; } const auto& descriptor = weight.impl_->descriptor; + const auto found = impl_->devices.find(device); + if (descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128F32 && + (found == impl_->devices.end() || + !found->second.fp8_f32_register_fed_supported)) { + result.errors.emplace_back( + "F32-scaled FP8 fragment prepack requires BF16 tensor cores"); + return result; + } const auto scratch_bytes = fragment_prepack_scratch_bytes(descriptor); if (scratch_bytes == 0U) { result.errors.emplace_back( @@ -134,12 +142,22 @@ const char* cuda_matmul_route_name(CudaMatmulRoute route) noexcept { case CudaMatmulRoute::PackedOffsetInt: return "packed_offset_int"; case CudaMatmulRoute::Nvfp4Group16: return "nvfp4_group16"; case CudaMatmulRoute::Fp8TensorPage: return "fp8_tensor_page"; + case CudaMatmulRoute::Fp8F32TensorPage: + return "fp8_f32_tensor_page"; case CudaMatmulRoute::Fp8E4m3Block128: return "fp8_e4m3_block128"; + case CudaMatmulRoute::Fp8E4m3Block128F32: + return "fp8_e4m3_block128_f32"; case CudaMatmulRoute::Fp4E2m1Group32: return "fp4_e2m1_group32"; case CudaMatmulRoute::Fp8RegisterFed: return "fp8_register_fed"; + case CudaMatmulRoute::Fp8F32RegisterFed: + return "fp8_f32_register_fed"; case CudaMatmulRoute::Fp4RegisterFed: return "fp4_register_fed"; case CudaMatmulRoute::GemmaMarlin: return "gemma_marlin"; case CudaMatmulRoute::MoePlainBf16: return "moe_plain_bf16"; + case CudaMatmulRoute::MoeFp8E4m3Block128F32: + return "moe_fp8_e4m3_block128_f32"; + case CudaMatmulRoute::MoeFp8F32RegisterFed: + return "moe_fp8_f32_register_fed"; case CudaMatmulRoute::MoeNvfp4Group16: return "moe_nvfp4_group16"; case CudaMatmulRoute::MoeFp4E2m1Group32: return "moe_fp4_e2m1_group32"; @@ -160,7 +178,8 @@ ValidationResult CudaBackend::matmul_impl( std::uint32_t rows, std::uint32_t groups, std::uint64_t rows_per_group, std::span output, float softcap, bool round_output, CudaMatmulProfile* profile, - bool dsv4_fp8_tensor_page) { + bool dsv4_fp8_tensor_page, const std::byte* batch_input, + std::byte* batch_output, bool defer_completion) { ValidationResult result; if (profile != nullptr) *profile = {}; if (!weight.valid()) { @@ -177,7 +196,10 @@ ValidationResult CudaBackend::matmul_impl( output.size() == descriptor.rows * rows; if (rows == 0U || (!regular_shape && !grouped_shape) || !std::isfinite(softcap) || softcap < 0.0F || - (softcap != 0.0F && (rows != 1U || groups != 0U))) { + (softcap != 0.0F && (rows != 1U || groups != 0U)) || + (defer_completion && + (batch_input == nullptr || batch_output == nullptr || + profile != nullptr || impl_->detailed_timing))) { result.errors.emplace_back("CUDA matmul activation shapes are incompatible"); return result; } @@ -205,10 +227,12 @@ ValidationResult CudaBackend::matmul_impl( // kernel, which is recorded as its own census route rather than hidden. const bool regfed_encoding = descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128 || + descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128F32 || descriptor.encoding == CudaWeightEncoding::Fp4E2m1Group32; const bool regfed_shape = regfed_encoding && groups == 0U && softcap == 0.0F && - (descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128 + (descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128 || + descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128F32 ? regfed_fp8_shape_admissible(descriptor.rows, descriptor.columns) : regfed_fp4_shape_admissible(descriptor.rows, descriptor.columns)); // The register-fed route requires a weight already permuted by an explicit @@ -217,8 +241,22 @@ ValidationResult CudaBackend::matmul_impl( // weight's other consumers. Deciding it here corrupted the DeepSeek V4 // attention output projection, which matmul_impl touches 129 times a run // and the attention path then reads canonically. + const bool f32_fragment_page_candidate = + descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128F32 && + weight.impl_->fragment_prepacked && dsv4_fp8_tensor_page && rows > 16U && + state.fp8_f32_tensor_page_supported && groups == 0U && softcap == 0.0F && + descriptor.columns % kDsv4Fp8TensorBlockK == 0U && + descriptor.rows % 128U == 0U && + descriptor.columns <= std::numeric_limits::max() && + descriptor.rows <= std::numeric_limits::max(); const bool regfed = regfed_shape && regfed_matmul_enabled() && - weight.impl_->fragment_prepacked; + weight.impl_->fragment_prepacked && + (descriptor.encoding != + CudaWeightEncoding::Fp8E4m3Block128F32 || + state.fp8_f32_register_fed_supported) && + !f32_fragment_page_candidate; + const bool f32_regfed = regfed && descriptor.encoding == + CudaWeightEncoding::Fp8E4m3Block128F32; const bool marlin = weight.impl_->marlin_prepacked && groups == 0U && rows <= 128U && descriptor.encoding == @@ -228,7 +266,8 @@ ValidationResult CudaBackend::matmul_impl( // No hidden fallback. Fragment order replaces the canonical layout, so a // permuted weight reaching a canonical kernel does not degrade -- it // decodes a permutation as if it were weights. Refuse instead. - if (weight.impl_->fragment_prepacked && !regfed) { + if (weight.impl_->fragment_prepacked && !regfed && + !f32_fragment_page_candidate) { result.errors.emplace_back( "CUDA matmul received a fragment-prepacked weight but has no " "register-fed route for this call; refusing to read fragment order " @@ -241,16 +280,28 @@ ValidationResult CudaBackend::matmul_impl( "no admissible Marlin route"); return result; } - const bool tensor_page = + const bool dsv4_tensor_page = dsv4_fp8_tensor_page && !regfed && state.dsv4_fp8_tensor_page_supported && rows > 1U && groups == 0U && descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128 && descriptor.columns % kDsv4Fp8TensorBlockK == 0U && descriptor.rows % kDsv4Fp8TensorBlockN == 0U && descriptor.columns <= std::numeric_limits::max() && - descriptor.rows <= std::numeric_limits::max(); + descriptor.rows <= std::numeric_limits::max() && + (!weight.impl_->fragment_prepacked || f32_fragment_page_candidate); + const bool f32_tensor_page = + dsv4_fp8_tensor_page && !regfed && + state.fp8_f32_tensor_page_supported && rows > 1U && groups == 0U && + descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128F32 && + descriptor.columns % kDsv4Fp8TensorBlockK == 0U && + descriptor.rows % 128U == 0U && + descriptor.columns <= std::numeric_limits::max() && + descriptor.rows <= std::numeric_limits::max() && + (!weight.impl_->fragment_prepacked || f32_fragment_page_candidate); + const bool tensor_page = dsv4_tensor_page || f32_tensor_page; const auto input_scale_bytes = tensor_page - ? static_cast(rows) * descriptor.scale_columns + ? static_cast(rows) * descriptor.scale_columns * + (f32_tensor_page ? sizeof(float) : sizeof(unsigned char)) : 0U; const auto compact_input_bytes = tensor_page ? static_cast(input.size()) + input_scale_bytes @@ -263,7 +314,7 @@ ValidationResult CudaBackend::matmul_impl( descriptor.rows > std::numeric_limits::max() / padded_rows / sizeof(float)) { result.errors.emplace_back( - "DeepSeek FP8 tensor page output workspace overflows"); + "FP8 tensor page output workspace overflows"); return result; } const auto tensor_output_bytes = tensor_page @@ -277,7 +328,7 @@ ValidationResult CudaBackend::matmul_impl( const auto marlin_output_bytes = marlin && rows > 1U ? 128U * descriptor.rows * sizeof(float) : output_bytes; - const auto required_output_bytes = tensor_page + const auto required_output_bytes = tensor_page || f32_regfed ? std::max(input_bytes, tensor_output_bytes) : marlin_output_bytes; std::uint64_t workspace_allocation_calls = 0U; @@ -325,11 +376,11 @@ ValidationResult CudaBackend::matmul_impl( capacity = target; return true; }; - const bool stage_input = ensure_host_staging( + const bool stage_input = batch_input != nullptr || ensure_host_staging( state.matmul_host_input, state.matmul_host_input_bytes, input_bytes); - const bool stage_output = ensure_host_staging( + const bool stage_output = batch_output != nullptr || ensure_host_staging( state.matmul_host_output, state.matmul_host_output_bytes, output_bytes); - if (stage_input) { + if (stage_input && batch_input == nullptr) { std::memcpy(state.matmul_host_input, input.data(), input.size_bytes()); } if (impl_->detailed_timing) { @@ -339,9 +390,10 @@ ValidationResult CudaBackend::matmul_impl( } } if (auto status = cudaMemcpyAsync( - tensor_page ? static_cast(state.output) + (tensor_page || f32_regfed) ? static_cast(state.output) : static_cast(state.input), - stage_input ? static_cast(state.matmul_host_input) + batch_input != nullptr ? static_cast(batch_input) + : stage_input ? static_cast(state.matmul_host_input) : static_cast(input.data()), input.size_bytes(), cudaMemcpyHostToDevice, state.stream); status != cudaSuccess) { @@ -354,20 +406,40 @@ ValidationResult CudaBackend::matmul_impl( } } const bool native = descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128 || + descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128F32 || descriptor.encoding == CudaWeightEncoding::Fp4E2m1Group32; const bool w8_group32 = descriptor.encoding == CudaWeightEncoding::OffsetPackedInt8 && rows == 1U && groups == 0U && descriptor.group_size == 32U && descriptor.columns % 32U == 0U; - if (tensor_page) { + if (tensor_page || f32_regfed) { const dim3 quantize_grid( static_cast(descriptor.scale_columns), rows, 1U); auto* compact_values = reinterpret_cast(state.input); - auto* compact_scales = compact_values + input.size(); - quantize_activation_e4m3_bytes_kernel<<< + if (f32_tensor_page || f32_regfed) { + auto* compact_scales = reinterpret_cast( + compact_values + input.size()); + quantize_activation_e4m3_f32_bytes_kernel<<< + quantize_grid, 128U, 0U, state.stream>>>( + compact_values, compact_scales, state.output, + descriptor.columns, rows); + } else { + auto* compact_scales = compact_values + input.size(); + quantize_activation_e4m3_bytes_kernel<<< + quantize_grid, 128U, 0U, state.stream>>>( + compact_values, compact_scales, state.output, + descriptor.columns, rows); + } + } else if (descriptor.encoding == + CudaWeightEncoding::Fp8E4m3Block128F32 && + !marlin && !regfed) { + const auto input_rows = groups == 0U ? rows : rows * groups; + const dim3 quantize_grid( + static_cast((descriptor.columns + 127U) / 128U), + input_rows, 1U); + quantize_activation_e4m3_f32_scale_kernel<<< quantize_grid, 128U, 0U, state.stream>>>( - compact_values, compact_scales, state.output, - descriptor.columns, rows); + state.input, descriptor.columns, input_rows); } else if (native && !marlin) { const auto input_rows = groups == 0U ? rows : rows * groups; const dim3 quantize_grid( @@ -476,10 +548,20 @@ ValidationResult CudaBackend::matmul_impl( const auto n_tiles = static_cast(descriptor.rows / kRegfedTileN); const std::uint32_t units = - descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128 + (descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128 || + descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128F32) ? static_cast(descriptor.columns / 32U) : k_tiles / kRegfedKPerLoad; - const std::uint32_t split = regfed_split_k(units, n_tiles); + std::uint32_t split = regfed_split_k(units, n_tiles); + if (f32_regfed) { + // A continuous scale belongs to a complete K128 block (four + // 32-column pairs), so split-K boundaries may never bisect one. + split = 1U; + while (split < 16U && units % ((split * 2U) * 4U) == 0U && + n_tiles * split * 2U <= 4096U) { + split *= 2U; + } + } const std::uint64_t activation_bytes = static_cast(k_tiles) * column_blocks * groups_per_block * 4U * sizeof(uint2); @@ -545,15 +627,29 @@ ValidationResult CudaBackend::matmul_impl( static_cast(chunk_activation_bytes); const auto fragment_total = static_cast(k_tiles) * chunk_blocks * chunk_groups * 4U; - regfed_activation_fragment_kernel<<< - static_cast(std::min( - (fragment_total + 255U) / 256U, 65535U)), - 256U, 0U, state.stream>>>( - static_cast(state.regfed_activation), - state.input + static_cast(start) * - descriptor.columns, - chunk, static_cast(descriptor.columns), - chunk_blocks, chunk_groups); + if (f32_regfed) { + const auto* compact_values = + reinterpret_cast(state.input); + regfed_fp8_activation_fragment_kernel<<< + static_cast(std::min( + (fragment_total + 255U) / 256U, 65535U)), + 256U, 0U, state.stream>>>( + static_cast(state.regfed_activation), + compact_values + static_cast(start) * + descriptor.columns, + chunk, static_cast(descriptor.columns), + chunk_blocks, chunk_groups); + } else { + regfed_activation_fragment_kernel<<< + static_cast(std::min( + (fragment_total + 255U) / 256U, 65535U)), + 256U, 0U, state.stream>>>( + static_cast(state.regfed_activation), + state.input + static_cast(start) * + descriptor.columns, + chunk, static_cast(descriptor.columns), + chunk_blocks, chunk_groups); + } float* chunk_output = state.output + static_cast(start) * descriptor.rows; if (descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128) { @@ -583,6 +679,35 @@ ValidationResult CudaBackend::matmul_impl( split, chunk, chunk_groups, state.regfed_partials, state.regfed_counters); } + } else if (descriptor.encoding == + CudaWeightEncoding::Fp8E4m3Block128F32) { + record_cuda_matmul_route(CudaMatmulRoute::Fp8F32RegisterFed); + const auto* compact_values = + reinterpret_cast(state.input); + const auto* compact_scales = reinterpret_cast( + compact_values + input.size()); + const auto launch = [&](auto tag) { + constexpr std::uint32_t kBlocks = decltype(tag)::value; + regfed_fp8_f32_matmul_kernel<<< + blocks, kRegfedWarpsPerBlock * 32U, 0U, + state.stream>>>( + chunk_output, + static_cast(weight.impl_->weights), + static_cast(weight.impl_->scales), + static_cast(state.regfed_activation), + compact_scales + static_cast(start) * + descriptor.scale_columns, + static_cast(descriptor.columns), + static_cast(descriptor.rows), + static_cast(descriptor.scale_columns), + split, chunk, chunk_groups, state.regfed_partials, + state.regfed_counters); + }; + if (chunk_blocks == 1U) { + launch(std::integral_constant{}); + } else { + launch(std::integral_constant{}); + } } else { record_cuda_matmul_route(CudaMatmulRoute::Fp4RegisterFed); if (chunk_blocks == 1U) { @@ -610,7 +735,7 @@ ValidationResult CudaBackend::matmul_impl( } } } - } else if (tensor_page) { + } else if (dsv4_tensor_page) { record_cuda_matmul_route(CudaMatmulRoute::Fp8TensorPage); const dim3 tensor_grid( static_cast( @@ -627,6 +752,32 @@ ValidationResult CudaBackend::matmul_impl( static_cast(weight.impl_->scales), rows, static_cast(descriptor.columns), static_cast(descriptor.rows)); + } else if (f32_tensor_page) { + record_cuda_matmul_route(CudaMatmulRoute::Fp8F32TensorPage); + const dim3 tensor_grid( + static_cast( + descriptor.rows / kFp8F32TensorBlockN), + static_cast( + padded_rows / kDsv4Fp8TensorBlockM), 1U); + const auto* compact_values = + reinterpret_cast(state.input); + const auto* compact_scales = reinterpret_cast( + compact_values + input.size()); + const auto launch = [&](auto tag) { + constexpr bool kPrepacked = decltype(tag)::value; + fp8_f32_decode_bf16_tensor_kernel<<< + tensor_grid, threads, 0U, state.stream>>>( + state.output, compact_values, compact_scales, + static_cast(weight.impl_->weights), + static_cast(weight.impl_->scales), rows, + static_cast(descriptor.columns), + static_cast(descriptor.rows)); + }; + if (weight.impl_->fragment_prepacked) { + launch(std::true_type{}); + } else { + launch(std::false_type{}); + } } else if (descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128) { record_cuda_matmul_route(CudaMatmulRoute::Fp8E4m3Block128); native_fp8_matmul_kernel<<>>( @@ -635,6 +786,15 @@ ValidationResult CudaBackend::matmul_impl( static_cast(weight.impl_->scales), descriptor.scale_columns, rows, descriptor.columns, descriptor.rows, groups, rows_per_group); + } else if (descriptor.encoding == + CudaWeightEncoding::Fp8E4m3Block128F32) { + record_cuda_matmul_route(CudaMatmulRoute::Fp8E4m3Block128F32); + native_fp8_f32_scale_matmul_kernel<<>>( + state.output, state.input, + static_cast(weight.impl_->weights), + static_cast(weight.impl_->scales), + descriptor.scale_columns, rows, descriptor.columns, descriptor.rows, + groups, rows_per_group); } else if (descriptor.encoding == CudaWeightEncoding::Fp4E2m1Group32) { record_cuda_matmul_route(CudaMatmulRoute::Fp4E2m1Group32); native_fp4_matmul_kernel<<>>( @@ -685,7 +845,8 @@ ValidationResult CudaBackend::matmul_impl( } } if (auto status = cudaMemcpyAsync( - stage_output ? static_cast(state.matmul_host_output) + batch_output != nullptr ? static_cast(batch_output) + : stage_output ? static_cast(state.matmul_host_output) : static_cast(output.data()), state.output, output.size_bytes(), cudaMemcpyDeviceToHost, state.stream); @@ -702,6 +863,21 @@ ValidationResult CudaBackend::matmul_impl( const auto issue_nanoseconds = static_cast( std::chrono::duration_cast( wait_started - issue_started).count()); + if (defer_completion) { + std::scoped_lock lock(impl_->mutex); + auto& device_stats = *std::find_if( + impl_->stats.devices.begin(), impl_->stats.devices.end(), + [&weight](const auto& value) { + return value.device == weight.impl_->device; + }); + device_stats.activation_h2d_bytes += input_bytes; + device_stats.activation_d2h_bytes += output_bytes; + ++device_stats.matmul_calls; + device_stats.workspace_allocation_calls += workspace_allocation_calls; + device_stats.workspace_allocation_bytes += workspace_allocation_bytes; + device_stats.matmul_issue_nanoseconds += issue_nanoseconds; + return result; + } if (auto status = cudaStreamSynchronize(state.stream); status != cudaSuccess) { return cuda_error(status, "synchronize CUDA matmul"); } diff --git a/kernels/cuda/detail/backend_mhc.inc.cuh b/kernels/cuda/detail/backend_mhc.inc.cuh index 79cac14..4b9a5d1 100644 --- a/kernels/cuda/detail/backend_mhc.inc.cuh +++ b/kernels/cuda/detail/backend_mhc.inc.cuh @@ -769,6 +769,50 @@ ValidationResult CudaBackend::dsv4_mhc_finish_device( return dsv4_mhc_finish_impl(device, {}, hidden, true); } +ValidationResult CudaBackend::dsv4_mhc_download_layer_input( + int device, std::span layer_input) { + ValidationResult result; + const auto found = impl_->devices.find(device); + if (found == impl_->devices.end()) { + return {{"mHC layer-input download targets an uninitialized device"}}; + } + auto& state = found->second; + if (!state.dsv4_mhc_supported || state.dsv4_mhc_stage != 1U || + state.dsv4_mhc_workspace == nullptr || state.dsv4_mhc_branch_ready || + state.moe_in_flight || state.dsv4_mhc_failed || + layer_input.size() != kDsv4MhcHidden || + state.dsv4_mhc_host_staging == nullptr || + state.dsv4_mhc_host_staging_bytes < + kDsv4MhcHidden * sizeof(std::uint16_t)) { + return {{"mHC layer-input download violates command order"}}; + } + if (auto status = cudaSetDevice(device); status != cudaSuccess) { + return cuda_error(status, "select CUDA device for mHC input download"); + } + const auto bytes = kDsv4MhcHidden * sizeof(std::uint16_t); + if (auto status = cudaMemcpyAsync( + state.dsv4_mhc_host_staging, + state.dsv4_mhc_workspace->layer_input, bytes, + cudaMemcpyDeviceToHost, state.stream); + status != cudaSuccess) { + return cuda_error(status, "download mHC normalized layer input"); + } + if (auto status = cudaStreamSynchronize(state.stream); + status != cudaSuccess) { + return cuda_error(status, "synchronize mHC layer-input download"); + } + const auto* encoded = reinterpret_cast( + state.dsv4_mhc_host_staging); + for (std::size_t index = 0U; index < layer_input.size(); ++index) { + layer_input[index] = std::bit_cast( + static_cast(encoded[index]) << 16U); + if (!std::isfinite(layer_input[index])) { + return {{"mHC normalized layer input is non-finite"}}; + } + } + return result; +} + ValidationResult CudaBackend::dsv4_mhc_device_view( int device, CudaDsv4MhcDeviceView& view) { view = {}; diff --git a/kernels/cuda/detail/backend_model_kernels.cuh b/kernels/cuda/detail/backend_model_kernels.cuh index 1e99e9a..f75dcf7 100644 --- a/kernels/cuda/detail/backend_model_kernels.cuh +++ b/kernels/cuda/detail/backend_model_kernels.cuh @@ -2376,6 +2376,7 @@ struct alignas(256) Dsv4MhcWorkspace { float post[kDsv4MhcMultiplier]; float combination[kDsv4MhcMultiplier * kDsv4MhcMultiplier]; float router_logits[kDsv4MhcRouterLogits]; + float glm53_router_logits[288U]; unsigned int failure{}; }; diff --git a/kernels/cuda/detail/backend_moe.inc.cuh b/kernels/cuda/detail/backend_moe.inc.cuh index 26cf0ab..7110a5f 100644 --- a/kernels/cuda/detail/backend_moe.inc.cuh +++ b/kernels/cuda/detail/backend_moe.inc.cuh @@ -2036,7 +2036,25 @@ ValidationResult CudaBackend::collect_deepseek_moe_rows( ValidationResult CudaBackend::enqueue_moe( int device, std::span hidden, std::uint32_t rows, - std::span routed, const CudaMoeExpert* shared) { + std::span routed, const CudaMoeExpert* shared, + float swiglu_limit) { + return enqueue_moe_impl(device, hidden, rows, routed, shared, + swiglu_limit, false, {}); +} + +ValidationResult CudaBackend::enqueue_glm53_moe_from_mhc( + int device, std::span routed, + const CudaMoeExpert& shared, std::span coefficients, + float swiglu_limit) { + return enqueue_moe_impl(device, {}, 1U, routed, &shared, swiglu_limit, + true, coefficients); +} + +ValidationResult CudaBackend::enqueue_moe_impl( + int device, std::span hidden, std::uint32_t rows, + std::span routed, const CudaMoeExpert* shared, + float swiglu_limit, bool mhc_source_destination, + std::span routed_coefficients) { ValidationResult result; const auto found = impl_->devices.find(device); if (found == impl_->devices.end()) { @@ -2044,13 +2062,21 @@ ValidationResult CudaBackend::enqueue_moe( return result; } auto& state = found->second; - if (state.moe_in_flight) { + if (state.moe_in_flight || + (mhc_source_destination && + (!state.dsv4_mhc_supported || state.dsv4_mhc_stage != 1U || + state.dsv4_mhc_workspace == nullptr || + state.dsv4_mhc_branch_ready || state.dsv4_mhc_failed))) { result.errors.emplace_back("MoE workspace already has an in-flight command"); return result; } const auto expert_count = routed.size() + (shared == nullptr ? 0U : 1U); if (rows == 0U || expert_count == 0U || expert_count > kMaxMoeExperts || - routed.size() > kMaxRoutedMoeExperts) { + routed.size() > kMaxRoutedMoeExperts || + (mhc_source_destination && + (rows != 1U || shared == nullptr || + routed_coefficients.size() != routed.size())) || + (!mhc_source_destination && !routed_coefficients.empty())) { result.errors.emplace_back("MoE command has an unsupported row or expert count"); return result; } @@ -2069,12 +2095,20 @@ ValidationResult CudaBackend::enqueue_moe( const bool nvfp4_batch = batch_encoding == CudaWeightEncoding::Nvfp4Group16; const bool mxfp4_batch = batch_encoding == CudaWeightEncoding::Fp4E2m1Group32; + const bool fp8_f32_batch = + batch_encoding == CudaWeightEncoding::Fp8E4m3Block128F32; const bool plain_batch = batch_encoding == CudaWeightEncoding::Plain; - if (!nvfp4_batch && !mxfp4_batch && !plain_batch && + if (!nvfp4_batch && !mxfp4_batch && !fp8_f32_batch && !plain_batch && batch_encoding != CudaWeightEncoding::OffsetPackedInt4) { result.errors.emplace_back("MoE command has an unsupported weight encoding"); return result; } + if (fp8_f32_batch && + (!std::isfinite(swiglu_limit) || swiglu_limit <= 0.0F)) { + result.errors.emplace_back( + "F32-scaled FP8 MoE requires a positive finite SwiGLU limit"); + return result; + } std::uint64_t hidden_columns = 0U; std::uint64_t intermediate_columns = 0U; @@ -2095,6 +2129,10 @@ ValidationResult CudaBackend::enqueue_moe( : mxfp4_batch ? (weight->impl_->descriptor.dtype == SafetensorsDtype::I8 && weight->impl_->descriptor.group_size == 32U) + : fp8_f32_batch + ? (weight->impl_->descriptor.dtype == + SafetensorsDtype::F8E4M3 && + weight->impl_->descriptor.group_size == 128U) : plain_batch ? weight->impl_->descriptor.dtype == SafetensorsDtype::Bf16 : (weight->impl_->descriptor.dtype == SafetensorsDtype::I32 && @@ -2108,8 +2146,10 @@ ValidationResult CudaBackend::enqueue_moe( const auto& gate = expert.gate->impl_->descriptor; const auto& up = expert.up->impl_->descriptor; const auto& down = expert.down->impl_->descriptor; - const auto expected_down_packed = (nvfp4_batch || mxfp4_batch) - ? (down.columns + 1U) / 2U : (down.columns + 7U) / 8U; + const auto expected_down_packed = fp8_f32_batch + ? down.columns + : (nvfp4_batch || mxfp4_batch) + ? (down.columns + 1U) / 2U : (down.columns + 7U) / 8U; const auto expected_down_scales = nvfp4_batch ? (down.columns + 15U) / 16U : mxfp4_batch ? (down.columns + 31U) / 32U @@ -2130,7 +2170,8 @@ ValidationResult CudaBackend::enqueue_moe( // because scaling before the down projection is not float-equal to // scaling after it and Laguna's reference scales after. if (!std::isfinite(expert.coefficient) || - ((shared_expert || nvfp4_batch || mxfp4_batch || plain_batch) && + ((shared_expert || nvfp4_batch || mxfp4_batch || fp8_f32_batch || + plain_batch) && expert.coefficient != 1.0F)) { result.errors.emplace_back("MoE expert coefficient is invalid"); return false; @@ -2153,8 +2194,11 @@ ValidationResult CudaBackend::enqueue_moe( std::uint64_t hidden_elements = 0U; if (!checked_bytes(rows, hidden_columns, 1U, hidden_elements) || - hidden.size() != hidden_elements || + (mhc_source_destination ? !hidden.empty() + : hidden.size() != hidden_elements) || std::any_of(hidden.begin(), hidden.end(), + [](float value) { return !std::isfinite(value); }) || + std::any_of(routed_coefficients.begin(), routed_coefficients.end(), [](float value) { return !std::isfinite(value); })) { result.errors.emplace_back("MoE hidden rows are incompatible"); return result; @@ -2243,6 +2287,7 @@ ValidationResult CudaBackend::enqueue_moe( PackedInt4MoeBatch batch; Nvfp4MoeBatch nvfp4_batch_data; Mxfp4MoeBatch mxfp4_batch_data; + Fp8F32MoeBatch fp8_f32_batch_data; PlainBf16MoeBatch plain_batch_data; state.moe_weights.clear(); state.moe_weights.reserve(expert_count * 3U); @@ -2255,6 +2300,22 @@ ValidationResult CudaBackend::enqueue_moe( static_cast(expert.up->impl_->weights); plain_batch_data.down_weights[index] = static_cast(expert.down->impl_->weights); + } else if (fp8_f32_batch) { + fp8_f32_batch_data.gate_weights[index] = + static_cast(expert.gate->impl_->weights); + fp8_f32_batch_data.gate_scales[index] = + static_cast(expert.gate->impl_->scales); + fp8_f32_batch_data.up_weights[index] = + static_cast(expert.up->impl_->weights); + fp8_f32_batch_data.up_scales[index] = + static_cast(expert.up->impl_->scales); + fp8_f32_batch_data.down_weights[index] = + static_cast(expert.down->impl_->weights); + fp8_f32_batch_data.down_scales[index] = + static_cast(expert.down->impl_->scales); + fp8_f32_batch_data.coefficients[index] = + index < routed_coefficients.size() + ? routed_coefficients[index] : 1.0F; } else if (nvfp4_batch) { nvfp4_batch_data.gate_weights[index] = static_cast(expert.gate->impl_->weights); @@ -2316,9 +2377,51 @@ ValidationResult CudaBackend::enqueue_moe( nvfp4_batch_data.rows = rows; mxfp4_batch_data.count = batch.count; mxfp4_batch_data.rows = rows; + fp8_f32_batch_data.count = batch.count; + fp8_f32_batch_data.rows = rows; plain_batch_data.count = batch.count; plain_batch_data.rows = rows; + bool fp8_f32_prepacked = fp8_f32_batch; + bool fp8_f32_any_prepacked = false; + if (fp8_f32_batch) { + for (std::uint32_t index = 0U; index < batch.count; ++index) { + const auto* expert = index < routed.size() ? &routed[index] : shared; + const bool ready = expert->gate->impl_->fragment_prepacked && + expert->up->impl_->fragment_prepacked && + expert->down->impl_->fragment_prepacked && + regfed_fp8_shape_admissible( + expert->gate->impl_->descriptor.rows, + expert->gate->impl_->descriptor.columns) && + regfed_fp8_shape_admissible( + expert->up->impl_->descriptor.rows, + expert->up->impl_->descriptor.columns) && + regfed_fp8_shape_admissible( + expert->down->impl_->descriptor.rows, + expert->down->impl_->descriptor.columns); + fp8_f32_prepacked = fp8_f32_prepacked && ready; + fp8_f32_any_prepacked = fp8_f32_any_prepacked || + expert->gate->impl_->fragment_prepacked || + expert->up->impl_->fragment_prepacked || + expert->down->impl_->fragment_prepacked; + } + if (fp8_f32_any_prepacked && !fp8_f32_prepacked) { + result.errors.emplace_back( + "F32-scaled FP8 MoE batch mixes fragment-prepacked and " + "canonical experts"); + return result; + } + if (fp8_f32_prepacked && rows > kRegfedMaxM) { + result.errors.emplace_back( + "F32-scaled FP8 MoE batch is fragment-prepacked but exceeds " + "the register-fed row width"); + return result; + } + } + const bool fp8_f32_regfed = + fp8_f32_prepacked && regfed_matmul_enabled() && + state.fp8_f32_register_fed_supported; + state.moe_hidden_columns = hidden_columns; state.moe_intermediate_columns = intermediate_columns; state.moe_rows = rows; @@ -2329,7 +2432,7 @@ ValidationResult CudaBackend::enqueue_moe( state.moe_has_shared = shared != nullptr; state.moe_shared_phase_timing_valid = false; state.moe_host_join = false; - state.moe_output_to_mhc = false; + state.moe_output_to_mhc = mhc_source_destination; state.moe_host_callback = {}; state.moe_kernel_launches = 0U; state.moe_in_flight = true; @@ -2348,6 +2451,25 @@ ValidationResult CudaBackend::enqueue_moe( } }; + if (fp8_f32_regfed) { + const std::uint64_t hidden_compact = hidden_elements + + static_cast(rows) * + (hidden_columns / 128U) * sizeof(float); + const std::uint64_t activation_elements = + activation_rows * intermediate_columns; + const std::uint64_t activation_compact = activation_elements + + activation_rows * (intermediate_columns / 128U) * sizeof(float); + const auto compact_bytes = std::max(hidden_compact, + activation_compact); + if (regfed_grow(state.moe_regfed_compact, + state.moe_regfed_compact_bytes, compact_bytes, false, + state.stream) != cudaSuccess) { + abort_enqueue(cudaErrorMemoryAllocation, + "allocate F32-scaled FP8 MoE compact workspace"); + return result; + } + } + if (auto status = cudaEventRecord(state.moe_start, state.stream); status != cudaSuccess) { abort_enqueue(status, "record MoE start"); @@ -2359,10 +2481,23 @@ ValidationResult CudaBackend::enqueue_moe( abort_enqueue(status, "reset MoE error flag"); return result; } - if (auto status = cudaMemcpyAsync( - state.moe_hidden, hidden.data(), static_cast(hidden_bytes), - cudaMemcpyHostToDevice, state.stream); - status != cudaSuccess) { + if (mhc_source_destination) { + constexpr std::uint32_t convert_threads = 256U; + constexpr std::uint32_t convert_blocks = + (kDsv4MhcHidden + convert_threads - 1U) / convert_threads; + dsv4_bf16_to_fp32<<>>( + state.dsv4_mhc_workspace->layer_input, state.moe_hidden, + kDsv4MhcHidden); + if (auto status = cudaGetLastError(); status != cudaSuccess) { + abort_enqueue(status, "convert resident GLM-5.3 MoE input"); + return result; + } + } else if (auto status = cudaMemcpyAsync( + state.moe_hidden, hidden.data(), + static_cast(hidden_bytes), + cudaMemcpyHostToDevice, state.stream); + status != cudaSuccess) { abort_enqueue(status, "upload MoE hidden rows"); return result; } @@ -2371,7 +2506,36 @@ ValidationResult CudaBackend::enqueue_moe( abort_enqueue(status, "record MoE hidden upload"); return result; } - if (mxfp4_batch) { + if (fp8_f32_regfed) { + auto* compact_values = + static_cast(state.moe_regfed_compact); + auto* compact_scales = reinterpret_cast( + compact_values + hidden_elements); + const dim3 quantize_grid( + static_cast(hidden_columns / 128U), rows, 1U); + quantize_activation_e4m3_f32_bytes_kernel<<< + quantize_grid, 128U, 0U, state.stream>>>( + compact_values, compact_scales, state.moe_hidden, + hidden_columns, rows); + ++state.moe_kernel_launches; + if (auto status = cudaGetLastError(); status != cudaSuccess) { + abort_enqueue(status, + "compact F32-scaled FP8 MoE hidden activation"); + return result; + } + } else if (fp8_f32_batch) { + const dim3 quantize_grid( + static_cast((hidden_columns + 127U) / 128U), rows, + 1U); + quantize_activation_e4m3_f32_scale_kernel<<< + quantize_grid, 128U, 0U, state.stream>>>( + state.moe_hidden, hidden_columns, rows); + ++state.moe_kernel_launches; + if (auto status = cudaGetLastError(); status != cudaSuccess) { + abort_enqueue(status, "launch F32-scaled FP8 MoE hidden quantization"); + return result; + } + } else if (mxfp4_batch) { const dim3 quantize_grid( static_cast((hidden_columns + 127U) / 128U), rows, 1U); @@ -2434,6 +2598,8 @@ ValidationResult CudaBackend::enqueue_moe( const bool mxfp4_regfed = mxfp4_prepacked && regfed_matmul_enabled(); record_cuda_matmul_route( plain_batch ? CudaMatmulRoute::MoePlainBf16 + : fp8_f32_regfed ? CudaMatmulRoute::MoeFp8F32RegisterFed + : fp8_f32_batch ? CudaMatmulRoute::MoeFp8E4m3Block128F32 : nvfp4_batch ? CudaMatmulRoute::MoeNvfp4Group16 : mxfp4_regfed ? CudaMatmulRoute::MoeFp4RegisterFed : mxfp4_batch ? CudaMatmulRoute::MoeFp4E2m1Group32 @@ -2442,6 +2608,96 @@ ValidationResult CudaBackend::enqueue_moe( plain_bf16_moe_gate_up_kernel<<>>( state.moe_activations, state.moe_hidden, plain_batch_data, hidden_columns, intermediate_columns, state.moe_error); + } else if (fp8_f32_regfed) { + const auto experts = batch.count; + const auto column_blocks = static_cast( + (rows + kRegfedTileM - 1U) / kRegfedTileM); + const auto groups = static_cast( + std::min(rows, kRegfedTileM)); + const auto n_tiles = + static_cast(intermediate_columns / kRegfedTileN); + const auto pairs = + static_cast(hidden_columns / 32U); + std::uint32_t split = 1U; + while (split < 16U && pairs % ((split * 2U) * 4U) == 0U && + static_cast(experts) * n_tiles * split * 2U <= + 4096U) { + split *= 2U; + } + const std::uint64_t partial_bytes = + static_cast(experts) * intermediate_columns * rows * + split * sizeof(float); + const std::uint64_t fragment_total = + (hidden_columns / kRegfedTileK) * column_blocks * groups * 4U; + const auto grow = [&](void*& pointer, std::uint64_t& capacity, + std::uint64_t required) { + return regfed_grow(pointer, capacity, required, false, + state.stream); + }; + if (grow(state.moe_regfed_gate_partials, + state.moe_regfed_gate_partial_bytes, partial_bytes) != + cudaSuccess || + grow(state.moe_regfed_up_partials, + state.moe_regfed_up_partial_bytes, partial_bytes) != + cudaSuccess || + grow(state.moe_regfed_hidden_fragment, + state.moe_regfed_hidden_fragment_bytes, + fragment_total * sizeof(uint2)) != cudaSuccess) { + abort_enqueue(cudaErrorMemoryAllocation, + "allocate register-fed FP8 MoE gate/up workspaces"); + return result; + } + const auto* compact_values = + static_cast(state.moe_regfed_compact); + const auto* compact_scales = reinterpret_cast( + compact_values + hidden_elements); + regfed_fp8_moe_activation_fragment_kernel<<< + static_cast(std::min( + (fragment_total + 255U) / 256U, 65535U)), + 256U, 0U, state.stream>>>( + static_cast(state.moe_regfed_hidden_fragment), + compact_values, 1U, rows, + static_cast(hidden_columns), column_blocks, groups); + const auto blocks = static_cast( + std::min( + (static_cast(experts) * n_tiles * split + + kRegfedWarpsPerBlock - 1U) / + kRegfedWarpsPerBlock, + 65535U)); + const auto launch = [&](auto tag) { + constexpr std::uint32_t kBlocks = decltype(tag)::value; + regfed_fp8_f32_moe_gate_up_kernel<<< + blocks, kRegfedWarpsPerBlock * 32U, 0U, state.stream>>>( + static_cast(state.moe_regfed_gate_partials), + static_cast(state.moe_regfed_up_partials), + static_cast(state.moe_regfed_hidden_fragment), + compact_scales, fp8_f32_batch_data, + static_cast(hidden_columns), + static_cast(intermediate_columns), split, rows, + groups); + }; + if (column_blocks == 1U) { + launch(std::integral_constant{}); + } else { + launch(std::integral_constant{}); + } + const std::uint64_t swiglu_total = + static_cast(experts) * intermediate_columns * rows; + regfed_fp8_f32_moe_swiglu_kernel<<< + static_cast(std::min( + (swiglu_total + 255U) / 256U, 65535U)), + 256U, 0U, state.stream>>>( + state.moe_activations, + static_cast(state.moe_regfed_gate_partials), + static_cast(state.moe_regfed_up_partials), experts, + static_cast(intermediate_columns), rows, split, + swiglu_limit, state.moe_error); + state.moe_kernel_launches += 2U; + } else if (fp8_f32_batch) { + fp8_f32_moe_gate_up_kernel<<>>( + state.moe_activations, state.moe_hidden, fp8_f32_batch_data, + hidden_columns, intermediate_columns, gate.scale_columns, + swiglu_limit, state.moe_error); } else if (nvfp4_batch) { nvfp4_moe_gate_up_kernel<<>>( state.moe_activations, state.moe_hidden, nvfp4_batch_data, @@ -2537,7 +2793,41 @@ ValidationResult CudaBackend::enqueue_moe( abort_enqueue(status, "launch MoE gate/up SwiGLU"); return result; } - if (mxfp4_batch) { + if (fp8_f32_regfed) { + const auto activation_elements = + activation_rows * intermediate_columns; + auto* compact_values = + static_cast(state.moe_regfed_compact); + auto* compact_scales = reinterpret_cast( + compact_values + activation_elements); + const dim3 quantize_grid( + static_cast(intermediate_columns / 128U), + static_cast(activation_rows), 1U); + quantize_activation_e4m3_f32_bytes_kernel<<< + quantize_grid, 128U, 0U, state.stream>>>( + compact_values, compact_scales, state.moe_activations, + intermediate_columns, static_cast(activation_rows)); + ++state.moe_kernel_launches; + if (auto status = cudaGetLastError(); status != cudaSuccess) { + abort_enqueue(status, + "compact F32-scaled FP8 MoE down activation"); + return result; + } + } else if (fp8_f32_batch) { + const dim3 quantize_grid( + static_cast((intermediate_columns + 127U) / 128U), + static_cast(activation_rows), 1U); + quantize_activation_e4m3_f32_scale_kernel<<< + quantize_grid, 128U, 0U, state.stream>>>( + state.moe_activations, intermediate_columns, + static_cast(activation_rows)); + ++state.moe_kernel_launches; + if (auto status = cudaGetLastError(); status != cudaSuccess) { + abort_enqueue(status, + "launch F32-scaled FP8 MoE activation quantization"); + return result; + } + } else if (mxfp4_batch) { // Runs for both routes: it is what makes the activation E4M3, and an // E4M3 value's BF16 image is exact, which is why the register-fed // tensor op multiplies the same numbers the scalar kernel does. @@ -2564,6 +2854,95 @@ ValidationResult CudaBackend::enqueue_moe( plain_bf16_moe_down_kernel<<>>( state.moe_output, state.moe_activations, plain_batch_data, intermediate_columns, hidden_columns, state.moe_error); + } else if (fp8_f32_regfed) { + const auto experts = batch.count; + const auto column_blocks = static_cast( + (rows + kRegfedTileM - 1U) / kRegfedTileM); + const auto groups = static_cast( + std::min(rows, kRegfedTileM)); + const auto n_tiles = + static_cast(hidden_columns / kRegfedTileN); + const auto pairs = + static_cast(intermediate_columns / 32U); + std::uint32_t split = 1U; + while (split < 16U && pairs % ((split * 2U) * 4U) == 0U && + static_cast(experts) * n_tiles * split * 2U <= + 4096U) { + split *= 2U; + } + const std::uint64_t partial_bytes = + static_cast(experts) * hidden_columns * rows * split * + sizeof(float); + const std::uint64_t fragment_total = + static_cast(experts) * + (intermediate_columns / kRegfedTileK) * column_blocks * groups * 4U; + const auto grow = [&](void*& pointer, std::uint64_t& capacity, + std::uint64_t required) { + return regfed_grow(pointer, capacity, required, false, + state.stream); + }; + if (grow(state.moe_regfed_down_partials, + state.moe_regfed_down_partial_bytes, partial_bytes) != + cudaSuccess || + grow(state.moe_regfed_activation_fragment, + state.moe_regfed_activation_fragment_bytes, + fragment_total * sizeof(uint2)) != cudaSuccess) { + abort_enqueue(cudaErrorMemoryAllocation, + "allocate register-fed FP8 MoE down workspaces"); + return result; + } + const auto activation_elements = + activation_rows * intermediate_columns; + const auto* compact_values = + static_cast(state.moe_regfed_compact); + const auto* compact_scales = reinterpret_cast( + compact_values + activation_elements); + regfed_fp8_moe_activation_fragment_kernel<<< + static_cast(std::min( + (fragment_total + 255U) / 256U, 65535U)), + 256U, 0U, state.stream>>>( + static_cast(state.moe_regfed_activation_fragment), + compact_values, experts, rows, + static_cast(intermediate_columns), column_blocks, + groups); + const auto blocks = static_cast( + std::min( + (static_cast(experts) * n_tiles * split + + kRegfedWarpsPerBlock - 1U) / + kRegfedWarpsPerBlock, + 65535U)); + const auto launch = [&](auto tag) { + constexpr std::uint32_t kBlocks = decltype(tag)::value; + regfed_fp8_f32_moe_down_kernel<<< + blocks, kRegfedWarpsPerBlock * 32U, 0U, state.stream>>>( + static_cast(state.moe_regfed_down_partials), + static_cast(state.moe_regfed_activation_fragment), + compact_scales, fp8_f32_batch_data, + static_cast(intermediate_columns), + static_cast(hidden_columns), split, rows, + groups); + }; + if (column_blocks == 1U) { + launch(std::integral_constant{}); + } else { + launch(std::integral_constant{}); + } + const std::uint64_t reduce_total = + static_cast(experts) * hidden_columns * rows; + regfed_fp8_f32_moe_reduce_kernel<<< + static_cast(std::min( + (reduce_total + 255U) / 256U, 65535U)), + 256U, 0U, state.stream>>>( + state.moe_output, + static_cast(state.moe_regfed_down_partials), experts, + static_cast(hidden_columns), rows, split, + state.moe_error); + state.moe_kernel_launches += 2U; + } else if (fp8_f32_batch) { + fp8_f32_moe_down_kernel<<>>( + state.moe_output, state.moe_activations, fp8_f32_batch_data, + intermediate_columns, hidden_columns, down.scale_columns, + state.moe_error); } else if (nvfp4_batch) { nvfp4_moe_down_kernel<<>>( state.moe_output, state.moe_activations, nvfp4_batch_data, @@ -2654,6 +3033,23 @@ ValidationResult CudaBackend::enqueue_moe( abort_enqueue(status, "launch MoE down projection"); return result; } + if (mhc_source_destination) { + constexpr unsigned int join_threads = 256U; + const auto join_blocks = static_cast( + (hidden_columns + join_threads - 1U) / join_threads); + glm53_moe_join_mhc_kernel<<>>( + state.moe_output, fp8_f32_batch_data.coefficients, + static_cast(routed.size()), + state.dsv4_mhc_workspace->branch, + static_cast(hidden_columns), state.moe_error); + ++state.moe_kernel_launches; + if (auto status = cudaGetLastError(); status != cudaSuccess) { + abort_enqueue(status, "join resident GLM-5.3 MoE branch"); + return result; + } + state.dsv4_mhc_branch_ready = true; + } if (auto status = cudaEventRecord(state.moe_kernel_finished, state.stream); status != cudaSuccess) { abort_enqueue(status, "record MoE kernel completion"); @@ -2664,14 +3060,17 @@ ValidationResult CudaBackend::enqueue_moe( auto& device_stats = *std::find_if( impl_->stats.devices.begin(), impl_->stats.devices.end(), [device](const auto& value) { return value.device == device; }); - device_stats.activation_h2d_bytes += hidden_bytes; + device_stats.activation_h2d_bytes += + mhc_source_destination ? 0U : hidden_bytes; device_stats.matmul_calls += 3U * expert_count; device_stats.workspace_allocation_calls += allocation_calls; device_stats.workspace_allocation_bytes += allocation_bytes; ++device_stats.deepseek_moe_calls; device_stats.deepseek_moe_kernel_launches += state.moe_kernel_launches; - ++device_stats.deepseek_moe_h2d_transfers; - device_stats.deepseek_moe_h2d_bytes += hidden_bytes; + device_stats.deepseek_moe_h2d_transfers += + mhc_source_destination ? 0U : 1U; + device_stats.deepseek_moe_h2d_bytes += + mhc_source_destination ? 0U : hidden_bytes; } return result; } diff --git a/kernels/cuda/detail/backend_state.cuh b/kernels/cuda/detail/backend_state.cuh index 848761e..a4977ec 100644 --- a/kernels/cuda/detail/backend_state.cuh +++ b/kernels/cuda/detail/backend_state.cuh @@ -97,6 +97,15 @@ struct CudaBackend::Impl { // the lock. void* upload_prepack_scratch{}; std::uint64_t upload_prepack_scratch_bytes{}; + // Capacity-bounded pinned ring for streamed weight uploads. Mapped + // checkpoint pages are pageable, so cudaMemcpyAsync otherwise performs + // its own blocking host staging for every cache miss. The ring makes + // that staging explicit, reusable, and large enough for many expert + // projections so CPU copies overlap the copy engine. It is sized from + // the device weight arena rather than from a particular GPU model. + std::byte* weight_host_staging{}; + std::uint64_t weight_host_staging_bytes{}; + std::uint64_t weight_host_staging_cursor{}; // Register-fed fused MoE workspaces: split-K partials for gate, up and // down, plus the two B-fragment activation buffers. Grown geometrically // and kept, so a decode step that repeats the same shapes allocates @@ -106,11 +115,13 @@ struct CudaBackend::Impl { void* moe_regfed_down_partials{}; void* moe_regfed_hidden_fragment{}; void* moe_regfed_activation_fragment{}; + void* moe_regfed_compact{}; std::uint64_t moe_regfed_gate_partial_bytes{}; std::uint64_t moe_regfed_up_partial_bytes{}; std::uint64_t moe_regfed_down_partial_bytes{}; std::uint64_t moe_regfed_hidden_fragment_bytes{}; std::uint64_t moe_regfed_activation_fragment_bytes{}; + std::uint64_t moe_regfed_compact_bytes{}; RegfedWorkspace gemma_regfed{}; GemmaMarlinWorkspace gemma_marlin{}; float* moe_regfed_gate{}; @@ -180,6 +191,8 @@ struct CudaBackend::Impl { int dsv4_deferred_attention_source_device{-1}; bool dsv4_deferred_attention_cross_transition{}; Dsv4MhcWorkspace* dsv4_mhc_workspace{}; + std::byte* glm53_mla_workspace{}; + std::uint64_t glm53_mla_workspace_bytes{}; std::byte* dsv4_mhc_host_staging{}; std::uint64_t dsv4_mhc_workspace_bytes{}; std::uint64_t dsv4_mhc_host_staging_bytes{}; @@ -307,6 +320,8 @@ struct CudaBackend::Impl { bool dsv4_mhc_supported{}; bool lightning_index_supported{}; bool dsv4_fp8_tensor_page_supported{}; + bool fp8_f32_tensor_page_supported{}; + bool fp8_f32_register_fed_supported{}; }; std::unordered_map devices; @@ -343,6 +358,9 @@ struct CudaBackend::Impl { } else if (state.dsv4_mhc_workspace != nullptr) { static_cast(cudaFree(state.dsv4_mhc_workspace)); } + if (state.glm53_mla_workspace != nullptr) { + static_cast(cudaFree(state.glm53_mla_workspace)); + } if (state.gemma_workspace != nullptr) { static_cast(cudaFree(state.gemma_workspace)); } @@ -361,12 +379,16 @@ struct CudaBackend::Impl { state.moe_regfed_up_partials, state.moe_regfed_down_partials, state.moe_regfed_hidden_fragment, - state.moe_regfed_activation_fragment}) { + state.moe_regfed_activation_fragment, + state.moe_regfed_compact}) { if (pointer != nullptr) static_cast(cudaFree(pointer)); } if (state.upload_prepack_scratch != nullptr) { static_cast(cudaFree(state.upload_prepack_scratch)); } + if (state.weight_host_staging != nullptr) { + static_cast(cudaFreeHost(state.weight_host_staging)); + } for (void* pointer : {state.moe_regfed.activation, state.moe_regfed.partials, state.moe_regfed.counters, diff --git a/src/engine/placement.cpp b/src/engine/placement.cpp index 987a32c..a4c6331 100644 --- a/src/engine/placement.cpp +++ b/src/engine/placement.cpp @@ -193,6 +193,7 @@ struct ClassAccumulator { std::string_view to_string(PlacementModel model) noexcept { switch (model) { case PlacementModel::Glm52: return "glm"; + case PlacementModel::Glm53: return "glm53"; case PlacementModel::DeepSeekV4: return "deepseek"; case PlacementModel::Gemma4: return "gemma4"; case PlacementModel::KimiK3: return "kimi-k3"; @@ -231,6 +232,7 @@ std::string_view to_string(PlacementClass component) noexcept { bool parse_placement_model(std::string_view text, PlacementModel& model) noexcept { if (text == "glm") { model = PlacementModel::Glm52; return true; } + if (text == "glm53") { model = PlacementModel::Glm53; return true; } if (text == "deepseek") { model = PlacementModel::DeepSeekV4; return true; } if (text == "gemma4") { model = PlacementModel::Gemma4; return true; } if (text == "kimi-k3") { model = PlacementModel::KimiK3; return true; } @@ -482,8 +484,13 @@ PlacementPlanResult solve_placement(const PlacementInventory& inventory, for (const auto& item : inventory.items) { if (item.spillable) continue; if (item.preferred_tier != PlacementTier::Device) { - if (!add(plan.host_resident_bytes, item.host_bytes)) { - result.errors.emplace_back("placement host byte total overflows"); + auto* resident = item.preferred_tier == PlacementTier::Host + ? &plan.host_resident_bytes + : &plan.storage_resident_bytes; + if (!add(*resident, item.host_bytes)) { + result.errors.push_back( + std::string(to_string(item.preferred_tier)) + + " placement byte total overflows"); return result; } account(item.component, item.preferred_tier, item.host_bytes, @@ -701,6 +708,7 @@ PlacementPlanResult solve_placement(const PlacementInventory& inventory, "host tier withholds " + format_bytes(inventory.host_reserve_bytes) + " for activations, worker stacks, and page cache"); } + const bool glm53_dynamic_cache = inventory.model == PlacementModel::Glm53; if (plan.io_dependent) { std::string where = "the checkpoint"; if (hardware.storage.resolved) { @@ -712,7 +720,8 @@ PlacementPlanResult solve_placement(const PlacementInventory& inventory, ')'; } plan.notes.emplace_back( - "steady-state decode reads " + + std::string(glm53_dynamic_cache ? "cold-cache decode can read " + : "steady-state decode reads ") + format_bytes(plan.decode_storage_read_bytes) + " per step from " + where + ": this configuration is I/O dependent"); } @@ -732,7 +741,12 @@ PlacementPlanResult solve_placement(const PlacementInventory& inventory, " holds no layers at this operating point"); } } - if (!inventory.prescriptive) { + if (glm53_dynamic_cache) { + plan.notes.emplace_back( + "GLM dynamically pins its non-expert spine and demand-caches experts " + "within live free VRAM; storage residency and reads above are " + "checkpoint-backing cold-cache bounds, not steady-state cache misses"); + } else if (!inventory.prescriptive) { plan.notes.emplace_back( "descriptive plan: it reports and admits the placement this runtime " "already performs and does not change it"); diff --git a/src/models/common/executor_support.hpp b/src/models/common/executor_support.hpp index fcdfc5f..57b170b 100644 --- a/src/models/common/executor_support.hpp +++ b/src/models/common/executor_support.hpp @@ -1,9 +1,9 @@ #pragma once -// Shared result mapping for the six ModelExecutor implementations. +// Shared result mapping for the seven ModelExecutor implementations. // // Every concrete runtime returns its own *GenerationResult type. The fields -// below exist on all six with identical meaning, so the copy is written once +// below exist on all seven with identical meaning, so the copy is written once // here instead of once per model -- which is what it was before Phase 4, and // where two models quietly lost fields nobody noticed for a month. // diff --git a/src/models/common/placement_model.cpp b/src/models/common/placement_model.cpp index a43e4af..85e199b 100644 --- a/src/models/common/placement_model.cpp +++ b/src/models/common/placement_model.cpp @@ -7,6 +7,7 @@ #include "strata/models/deepseek/deepseek_admission.hpp" #include "strata/models/deepseek/deepseek_checkpoint.hpp" #include "strata/models/gemma4/gemma4_checkpoint.hpp" +#include "strata/models/glm53/glm53_checkpoint.hpp" #include "strata/models/kimi_k3/kimi_k3_checkpoint.hpp" #include "strata/models/inkling/inkling_checkpoint.hpp" #include "strata/models/laguna/laguna_checkpoint.hpp" @@ -336,6 +337,110 @@ struct Gemma4Linear { } +// ------------------------------------------------------------- GLM-5.3 + +[[nodiscard]] PlacementClass glm53_component( + Glm53TensorRole role) noexcept { + switch (role) { + case Glm53TensorRole::Embedding: return PlacementClass::Embedding; + case Glm53TensorRole::OutputHead: return PlacementClass::OutputHead; + case Glm53TensorRole::Norm: + case Glm53TensorRole::Mhc: return PlacementClass::Norm; + case Glm53TensorRole::KdaAttention: + case Glm53TensorRole::SparseAttention: + case Glm53TensorRole::AttentionIndexer: + return PlacementClass::Attention; + case Glm53TensorRole::DenseMlp: return PlacementClass::FeedForward; + case Glm53TensorRole::Router: return PlacementClass::Router; + case Glm53TensorRole::SharedExpert: return PlacementClass::SharedExpert; + case Glm53TensorRole::RoutedExpert: return PlacementClass::RoutedExpert; + case Glm53TensorRole::Vision: return PlacementClass::Vision; + case Glm53TensorRole::Mtp: + case Glm53TensorRole::Count: break; + } + return PlacementClass::Norm; +} + +[[nodiscard]] std::string glm53_module_base(std::string_view name) { + for (const auto suffix : {std::string_view{".weight_scale_inv"}, + std::string_view{".weight"}, + std::string_view{".bias"}}) { + if (name.size() >= suffix.size() && + name.substr(name.size() - suffix.size()) == suffix) { + return std::string(name.substr(0U, name.size() - suffix.size())); + } + } + return std::string(name); +} + +[[nodiscard]] ParseResult build_glm53_inventory( + const Glm53IndexManifest& manifest, std::uint32_t context_tokens) { + ParseResult result; + auto& inventory = result.value; + inventory.model = PlacementModel::Glm53; + inventory.model_name = "GLM-5.3-Flash"; + inventory.layer_count = 45U; + inventory.maximum_context_tokens = context_tokens; + inventory.per_device_workspace_bytes = 2ULL << 30U; + inventory.minimum_device_budget_bytes = kMinimumDeviceBudget; + inventory.contiguous_layer_blocks = false; + // The checkpoint remains the canonical backing store. At runtime the + // non-expert spine is pinned opportunistically and experts enter a + // free-VRAM-sized demand cache, so this inventory is a cold-cache bound; + // it does not prescribe or claim a fixed resident layout. + inventory.prescriptive = false; + + std::map modules; + std::map classes; + for (const auto& tensor : manifest.tensors) { + if (tensor.role == Glm53TensorRole::Vision || + tensor.role == Glm53TensorRole::Mtp) { + continue; + } + const auto base = glm53_module_base(tensor.name); + auto& sizes = modules[base]; + sizes.layer = tensor.layer; + classes[base] = glm53_component(tensor.role); + if (tensor.component == Glm53TensorComponent::Scale) { + sizes.scale_bytes += tensor.source_bytes; + } else if (tensor.component == Glm53TensorComponent::Weight) { + sizes.weight_bytes += tensor.source_bytes; + } else { + sizes.host_bytes += tensor.source_bytes; + } + } + for (const auto& [base, sizes] : modules) { + auto source = sizes; + source.host_bytes = sizes.source_bytes(); + source.weight_bytes = 0U; + source.scale_bytes = 0U; + auto reads = source.host_bytes; + const auto component = classes.at(base); + if (component == PlacementClass::Embedding) { + reads = 4096ULL * sizeof(std::uint16_t); + } else if (component == PlacementClass::RoutedExpert) { + // The manifest has every expert projection as its own module, but + // exact decode reads only the eight selected experts in each MoE + // layer. Summing this fraction over all 288 experts gives eight + // complete gate/up/down triplets per sparse layer. + reads = source.host_bytes * 8ULL / 288ULL; + } + inventory.items.push_back(make_item( + component, source, reads, PlacementTier::Storage, false)); + } + + ModuleSizes state; + constexpr std::uint64_t kKdaLayers = 34U; + constexpr std::uint64_t kSparseLayers = 11U; + state.host_bytes = kKdaLayers * 64ULL * 128ULL * 128ULL * sizeof(float); + state.host_bytes += kKdaLayers * 3ULL * 8192ULL * 3ULL * sizeof(float); + state.host_bytes += kSparseLayers * context_tokens * 512ULL * sizeof(float); + inventory.items.push_back(make_item(PlacementClass::KvCache, state, 0U, + PlacementTier::Host, false)); + return result; +} + + // --------------------------------------------------------------- Inkling [[nodiscard]] ParseResult build_inkling_inventory( @@ -1002,6 +1107,7 @@ struct Gemma4Linear { struct OpenCheckpoints { std::unique_ptr gemma4; std::unique_ptr glm; + std::unique_ptr glm53; std::unique_ptr deepseek; std::unique_ptr kimi; std::unique_ptr laguna; @@ -1020,6 +1126,9 @@ struct OpenCheckpoints { case PlacementModel::Glm52: return build_glm_inventory(checkpoints.glm->manifest(), context_tokens, request.flash_attention); + case PlacementModel::Glm53: + return build_glm53_inventory(checkpoints.glm53->manifest(), + context_tokens); case PlacementModel::Laguna: return build_laguna_inventory(*checkpoints.laguna, context_tokens); case PlacementModel::Inkling: @@ -1046,6 +1155,8 @@ struct OpenCheckpoints { return kGemma4ExecutionContract.maximum_context_tokens; case PlacementModel::Glm52: return kGlm52ExecutionContract.maximum_context_tokens; + case PlacementModel::Glm53: + return 2048U; case PlacementModel::Laguna: return kLagunaExecutionContract.maximum_context_tokens; case PlacementModel::Inkling: @@ -1094,6 +1205,15 @@ PlacementPlanResult plan_model_placement_impl(const PlacementRequest& request, checkpoints.glm = std::move(opened.value); break; } + case PlacementModel::Glm53: { + auto opened = Glm53CheckpointReader::open(request.model_directory); + if (!opened.ok()) { + result.errors = std::move(opened.errors); + return result; + } + checkpoints.glm53 = std::move(opened.value); + break; + } case PlacementModel::Laguna: { auto opened = LagunaCheckpointReader::open(request.model_directory); if (!opened.ok()) { diff --git a/src/models/common/tokenizer.cpp b/src/models/common/tokenizer.cpp index 3933920..baba125 100644 --- a/src/models/common/tokenizer.cpp +++ b/src/models/common/tokenizer.cpp @@ -1623,6 +1623,56 @@ std::string render_glm52_chat_prompt(std::span messages, return output; } +std::string render_glm53_chat_prompt(std::span messages, + std::string_view reasoning_effort, + bool clear_thinking) { + const auto effort = reasoning_effort == "low" ? "Low" + : reasoning_effort == "high" ? "High" : "Max"; + std::string output = + std::string("[gMASK]<|system|>Reasoning Effort: ") + effort; + bool observing = false; + for (const auto& message : messages) { + switch (message.role) { + case ChatRole::System: + output += "<|system|>" + message.content; + observing = false; + break; + case ChatRole::User: + output += "<|user|>" + message.content; + observing = false; + break; + case ChatRole::Assistant: { + output += "<|assistant|>"; + if (clear_thinking) output += ""; + auto begin = message.content.find_first_not_of(" \t\r\n"); + if (begin != std::string::npos) { + const auto end = message.content.find_last_not_of(" \t\r\n"); + output.append(message.content, begin, end - begin + 1U); + } + observing = false; + break; + } + case ChatRole::Tool: + if (!observing) output += "<|observation|>"; + output += "" + message.content + + ""; + observing = true; + break; + } + } + output += "<|assistant|>"; + return output; +} + +std::string render_glm53_user_prompt(std::string_view user_text, + std::string_view reasoning_effort, + bool clear_thinking) { + const std::array messages{ChatMessage{ChatRole::User, + std::string(user_text)}}; + return render_glm53_chat_prompt(messages, reasoning_effort, + clear_thinking); +} + std::string render_deepseek_v4_user_prompt(std::string_view user_text, bool enable_thinking) { const std::array messages{ChatMessage{ChatRole::User, diff --git a/src/models/glm53/glm53_checkpoint.cpp b/src/models/glm53/glm53_checkpoint.cpp new file mode 100644 index 0000000..8870300 --- /dev/null +++ b/src/models/glm53/glm53_checkpoint.cpp @@ -0,0 +1,492 @@ +#include "strata/models/glm53/glm53_checkpoint.hpp" + +#include "../common/checkpoint_common.hpp" + +#include +#include +#include +#include +#include +#include + +namespace strata { +namespace { + +constexpr std::uint64_t kMaximumConfigBytes = 4ULL << 20U; +constexpr std::uint64_t kMaximumIndexBytes = 64ULL << 20U; + +[[nodiscard]] bool deferred_weights_enabled() noexcept { + const char* value = std::getenv("STRATA_GLM53_DEFERRED_WEIGHTS"); + return value == nullptr || + (std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off"); +} + +[[nodiscard]] bool phase_scheduled_reads_enabled() noexcept { + const char* value = std::getenv("STRATA_GLM53_PHASE_SCHEDULER"); + return value == nullptr || + (std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off"); +} + +} // namespace + +Glm53CheckpointReader::~Glm53CheckpointReader() { + std::scoped_lock lock(mapping_mutex_); + for (const auto& [name, mapping] : mappings_) { + static_cast(name); + if (mapping.address != nullptr && mapping.bytes != 0U) { + static_cast(munmap(mapping.address, + static_cast(mapping.bytes))); + } + } +} + +Glm53CheckpointOpenResult Glm53CheckpointReader::open( + std::string model_directory) { + Glm53CheckpointOpenResult result; + const auto root = std::filesystem::path(model_directory); + auto config_text = load_bounded_text_file( + (root / "config.json").string(), kMaximumConfigBytes); + if (!config_text.ok()) { + result.errors = std::move(config_text.errors); + return result; + } + auto config = parse_glm53_config(config_text.value); + if (!config.ok()) { + result.errors = std::move(config.errors); + return result; + } + auto config_gate = validate_glm53_config(config.value); + if (!config_gate.ok()) { + result.errors = std::move(config_gate.errors); + return result; + } + auto index = load_safetensors_index(model_directory, kMaximumIndexBytes); + if (!index.ok()) { + result.errors = std::move(index.errors); + return result; + } + auto built = build_glm53_index_manifest(std::move(index.value)); + if (!built.ok()) { + result.errors = std::move(built.errors); + return result; + } + auto validated = validate_glm53_checkpoint( + model_directory, std::move(built.manifest)); + if (!validated.ok()) { + result.errors = std::move(validated.errors); + return result; + } + auto reader = std::unique_ptr( + new Glm53CheckpointReader()); + reader->model_directory_ = std::move(model_directory); + reader->config_ = std::move(config.value); + reader->manifest_ = std::move(validated.manifest); + reader->by_name_.reserve(reader->manifest_.tensors.size()); + for (std::size_t index_of = 0U; + index_of < reader->manifest_.tensors.size(); ++index_of) { + reader->by_name_.emplace(reader->manifest_.tensors[index_of].name, + index_of); + } + auto opened = reader->shards_.open(reader->model_directory_, + reader->manifest_.shards, + "GLM-5.3-Flash"); + if (!opened.ok()) { + result.errors = std::move(opened.errors); + return result; + } + result.value = std::move(reader); + return result; +} + +const Glm53ManifestTensor* Glm53CheckpointReader::find( + std::string_view name) const noexcept { + const auto found = by_name_.find(name); + return found == by_name_.end() ? nullptr + : &manifest_.tensors[found->second]; +} + +ParseResult> Glm53CheckpointReader::read_slice( + const Glm53ManifestTensor& tensor, std::uint64_t offset, + std::uint64_t bytes) const { + ParseResult> result; + if (offset > tensor.source_bytes || bytes > tensor.source_bytes - offset) { + result.errors.push_back("GLM-5.3 tensor slice is out of bounds: " + + tensor.name); + return result; + } + result.value.resize(static_cast(bytes)); + auto status = shards_.read(tensor.shard, tensor.source_offset + offset, + result.value, tensor.name); + if (!status.ok()) { + result.errors = std::move(status.errors); + result.value.clear(); + } + return result; +} + +ParseResult> Glm53CheckpointReader::read( + std::string_view name, std::uint64_t maximum_bytes) const { + ParseResult> result; + const auto* tensor = find(name); + if (tensor == nullptr) { + result.errors.push_back("unknown GLM-5.3 tensor " + std::string(name)); + return result; + } + if (tensor->source_bytes > maximum_bytes) { + result.errors.push_back("GLM-5.3 tensor exceeds the caller byte budget: " + + std::string(name)); + return result; + } + return read_slice(*tensor, 0U, tensor->source_bytes); +} + +ParseResult> Glm53CheckpointReader::view( + std::string_view name) const { + ParseResult> result; + const auto* tensor = find(name); + if (tensor == nullptr) { + result.errors.push_back("unknown GLM-5.3 tensor " + std::string(name)); + return result; + } + std::scoped_lock lock(mapping_mutex_); + auto mapping = mappings_.find(tensor->shard); + if (mapping == mappings_.end()) { + const int descriptor = shards_.descriptor(tensor->shard); + struct stat status {}; + if (descriptor < 0 || fstat(descriptor, &status) != 0 || + status.st_size <= 0) { + result.errors.push_back("cannot size GLM-5.3 checkpoint shard " + + tensor->shard + ": " + + std::strerror(errno)); + return result; + } + const auto bytes = static_cast(status.st_size); + void* address = mmap(nullptr, static_cast(bytes), PROT_READ, + MAP_SHARED, descriptor, 0); + if (address == MAP_FAILED) { + result.errors.push_back("cannot map GLM-5.3 checkpoint shard " + + tensor->shard + ": " + + std::strerror(errno)); + return result; + } + mapping = mappings_.emplace( + tensor->shard, + ShardMapping{static_cast(address), bytes}).first; + } + if (tensor->source_offset > mapping->second.bytes || + tensor->source_bytes > mapping->second.bytes - tensor->source_offset) { + result.errors.push_back("mapped GLM-5.3 tensor exceeds its shard: " + + tensor->name); + return result; + } + result.value = std::span( + mapping->second.address + tensor->source_offset, + static_cast(tensor->source_bytes)); + return result; +} + +ParseResult> Glm53CheckpointReader::read_f32( + std::string_view name, std::uint64_t maximum_elements) const { + ParseResult> result; + const auto* tensor = find(name); + if (tensor == nullptr) { + result.errors.push_back("unknown GLM-5.3 tensor " + std::string(name)); + return result; + } + const auto width = safetensors_dtype_bytes(tensor->source_dtype); + if ((tensor->source_dtype != SafetensorsDtype::Bf16 && + tensor->source_dtype != SafetensorsDtype::F16 && + tensor->source_dtype != SafetensorsDtype::F32) || width == 0U || + tensor->source_bytes % width != 0U || + tensor->source_bytes / width != maximum_elements) { + result.errors.push_back("GLM-5.3 tensor does not match the required F32 extent: " + + std::string(name)); + return result; + } + auto encoded = read_slice(*tensor, 0U, tensor->source_bytes); + if (!encoded.ok()) { + result.errors = std::move(encoded.errors); + return result; + } + const auto elements = tensor->source_bytes / width; + result.value.resize(static_cast(elements)); + for (std::size_t index = 0U; index < result.value.size(); ++index) { + result.value[index] = detail::decode_plain_scalar( + encoded.value.data() + index * width, tensor->source_dtype); + } + return result; +} + +ParseResult> Glm53CheckpointReader::read_f32_row( + std::string_view name, std::uint64_t row) const { + ParseResult> result; + const auto* tensor = find(name); + if (tensor == nullptr || tensor->source_shape.size() != 2U || + row >= tensor->source_shape[0]) { + result.errors.push_back("GLM-5.3 tensor row is invalid: " + + std::string(name)); + return result; + } + const auto width = safetensors_dtype_bytes(tensor->source_dtype); + if ((tensor->source_dtype != SafetensorsDtype::Bf16 && + tensor->source_dtype != SafetensorsDtype::F16 && + tensor->source_dtype != SafetensorsDtype::F32) || width == 0U) { + result.errors.push_back("GLM-5.3 row tensor is not plain floating point: " + + std::string(name)); + return result; + } + const auto columns = tensor->source_shape[1]; + const auto offset = row * columns * width; + const auto bytes = columns * width; + std::vector owned; + std::span encoded; + if (phase_scheduled_reads_enabled()) { + auto mapped = view(name); + if (!mapped.ok()) { + result.errors = std::move(mapped.errors); + return result; + } + if (offset > mapped.value.size_bytes() || + bytes > mapped.value.size_bytes() - offset) { + result.errors.push_back( + "GLM-5.3 mapped row exceeds its tensor: " + + std::string(name)); + return result; + } + encoded = mapped.value.subspan(static_cast(offset), + static_cast(bytes)); + } else { + auto loaded = read_slice(*tensor, offset, bytes); + if (!loaded.ok()) { + result.errors = std::move(loaded.errors); + return result; + } + owned = std::move(loaded.value); + encoded = owned; + } + result.value.resize(static_cast(columns)); + for (std::size_t column = 0U; column < result.value.size(); ++column) { + result.value[column] = detail::decode_plain_scalar( + encoded.data() + column * width, tensor->source_dtype); + } + return result; +} + +std::uint64_t Glm53CheckpointReader::cuda_linear_storage_bytes( + std::string_view base_name) const { + const auto weight_name = std::string(base_name) + ".weight"; + const auto* weight = find(weight_name); + if (weight == nullptr) return 0U; + std::uint64_t scale_bytes = 0U; + if (weight->source_dtype == SafetensorsDtype::F8E4M3) { + const auto* scale = find(std::string(base_name) + ".weight_scale_inv"); + if (scale == nullptr) return 0U; + scale_bytes = scale->source_bytes; + } + return CudaBackend::weight_storage_bytes(weight->source_bytes, + scale_bytes); +} + +std::uint64_t Glm53CheckpointReader::cuda_linear_slice_storage_bytes( + std::string_view base_name, std::uint64_t row_begin, + std::uint64_t row_count) const { + const auto* weight = find(std::string(base_name) + ".weight"); + if (weight == nullptr || weight->source_shape.size() != 2U || + row_count == 0U || row_begin > weight->source_shape[0] || + row_count > weight->source_shape[0] - row_begin) { + return 0U; + } + if (weight->source_dtype == SafetensorsDtype::F8E4M3) { + if (row_begin % 128U != 0U) return 0U; + const auto* scale = find( + std::string(base_name) + ".weight_scale_inv"); + if (scale == nullptr || scale->source_dtype != SafetensorsDtype::F32) { + return 0U; + } + const auto scale_columns = + (weight->source_shape[1] + 127U) / 128U; + const auto scale_rows = (row_count + 127U) / 128U; + return CudaBackend::weight_storage_bytes( + row_count * weight->source_shape[1], + scale_rows * scale_columns * sizeof(float)); + } + if (weight->source_dtype != SafetensorsDtype::Bf16 && + weight->source_dtype != SafetensorsDtype::F16 && + weight->source_dtype != SafetensorsDtype::F32) { + return 0U; + } + const auto width = safetensors_dtype_bytes(weight->source_dtype); + return CudaBackend::weight_storage_bytes( + row_count * weight->source_shape[1] * width, 0U); +} + +ValidationResult Glm53CheckpointReader::load_cuda_linear( + std::string_view base_name, std::uint64_t rows, std::uint64_t columns, + int device, CudaBackend& backend, CudaWeight& output, + bool concurrent_prefetch) const { + ValidationResult result; + const auto weight_name = std::string(base_name) + ".weight"; + const auto* weight = find(weight_name); + if (weight == nullptr || weight->source_shape != + std::vector{rows, columns}) { + result.errors.push_back("GLM-5.3 linear has an unexpected or missing weight: " + + weight_name); + return result; + } + CudaWeightDescriptor descriptor; + descriptor.dtype = weight->source_dtype; + descriptor.rows = rows; + descriptor.columns = columns; + std::string scale_name; + if (weight->source_dtype == SafetensorsDtype::F8E4M3) { + scale_name = std::string(base_name) + ".weight_scale_inv"; + const auto* scale = find(scale_name); + const auto scale_rows = (rows + 127U) / 128U; + const auto scale_columns = (columns + 127U) / 128U; + if (scale == nullptr || scale->source_dtype != SafetensorsDtype::F32 || + scale->source_shape != + std::vector{scale_rows, scale_columns}) { + result.errors.push_back("GLM-5.3 FP8 linear has an invalid scale: " + + scale_name); + return result; + } + descriptor.encoding = CudaWeightEncoding::Fp8E4m3Block128F32; + descriptor.packed_columns = columns; + descriptor.scale_columns = scale_columns; + descriptor.group_size = 128U; + } else if (weight->source_dtype == SafetensorsDtype::Bf16 || + weight->source_dtype == SafetensorsDtype::F16 || + weight->source_dtype == SafetensorsDtype::F32) { + descriptor.encoding = CudaWeightEncoding::Plain; + } else { + result.errors.push_back("GLM-5.3 linear uses an unsupported dtype: " + + weight_name); + return result; + } + const auto fragment_layout = + descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128F32 && + backend.fp8_f32_register_fed_supported(device) + ? CudaBackend::FragmentLayout::Prepack + : CudaBackend::FragmentLayout::Canonical; + if (deferred_weights_enabled()) { + auto weights = view(weight_name); + if (!weights.ok()) return {std::move(weights.errors)}; + std::span scales; + if (!scale_name.empty()) { + auto mapped_scales = view(scale_name); + if (!mapped_scales.ok()) return {std::move(mapped_scales.errors)}; + scales = mapped_scales.value; + } + // Both mapped views outlive every upload, so the copy stream may retain + // them until its device-side completion event without a heap copy. + return backend.upload(device, descriptor, weights.value, scales, output, + concurrent_prefetch + ? CudaBackend::UploadCompletion::DeferredConcurrent + : CudaBackend::UploadCompletion::Deferred, + fragment_layout); + } + + // Same-binary control route for performance campaigns. This is the former + // production contract: heap-copy each tensor and block each upload. + auto weights = read(weight_name, rows * columns * 4U); + if (!weights.ok()) return {std::move(weights.errors)}; + std::vector scales; + if (!scale_name.empty()) { + auto loaded_scales = read(scale_name, + descriptor.scale_columns * + ((rows + 127U) / 128U) * sizeof(float)); + if (!loaded_scales.ok()) return {std::move(loaded_scales.errors)}; + scales = std::move(loaded_scales.value); + } + return backend.upload(device, descriptor, weights.value, scales, output, + CudaBackend::UploadCompletion::Synchronous, + fragment_layout); +} + +ValidationResult Glm53CheckpointReader::load_cuda_linear_slice( + std::string_view base_name, std::uint64_t total_rows, + std::uint64_t columns, std::uint64_t row_begin, + std::uint64_t row_count, int device, CudaBackend& backend, + CudaWeight& output) const { + ValidationResult result; + const auto weight_name = std::string(base_name) + ".weight"; + const auto* weight = find(weight_name); + if (weight == nullptr || + weight->source_shape != std::vector{total_rows, columns} || + row_count == 0U || row_begin > total_rows || + row_count > total_rows - row_begin || + (weight->source_dtype != SafetensorsDtype::F8E4M3 && + weight->source_dtype != SafetensorsDtype::Bf16 && + weight->source_dtype != SafetensorsDtype::F16 && + weight->source_dtype != SafetensorsDtype::F32)) { + return {{"GLM-5.3 sliced linear has an invalid weight or row range: " + + weight_name}}; + } + const auto width = safetensors_dtype_bytes(weight->source_dtype); + const auto row_bytes = columns * width; + const auto offset = row_begin * row_bytes; + const auto bytes = row_count * row_bytes; + CudaWeightDescriptor descriptor; + descriptor.dtype = weight->source_dtype; + descriptor.rows = row_count; + descriptor.columns = columns; + std::span scales; + if (weight->source_dtype == SafetensorsDtype::F8E4M3) { + if (row_begin % 128U != 0U) { + return {{"GLM-5.3 FP8 slice must begin on a scale-block row"}}; + } + const auto scale_name = std::string(base_name) + ".weight_scale_inv"; + auto mapped_scales = view(scale_name); + const auto* scale = find(scale_name); + const auto scale_columns = (columns + 127U) / 128U; + const auto scale_row_begin = row_begin / 128U; + const auto scale_rows = (row_count + 127U) / 128U; + const auto scale_offset = + scale_row_begin * scale_columns * sizeof(float); + const auto scale_bytes = scale_rows * scale_columns * sizeof(float); + if (!mapped_scales.ok()) return {std::move(mapped_scales.errors)}; + if (scale == nullptr || scale->source_dtype != SafetensorsDtype::F32 || + scale_offset > mapped_scales.value.size_bytes() || + scale_bytes > mapped_scales.value.size_bytes() - scale_offset) { + return {{"GLM-5.3 FP8 slice has an invalid scale extent"}}; + } + scales = mapped_scales.value.subspan( + static_cast(scale_offset), + static_cast(scale_bytes)); + descriptor.encoding = CudaWeightEncoding::Fp8E4m3Block128F32; + descriptor.packed_columns = columns; + descriptor.scale_columns = scale_columns; + descriptor.group_size = 128U; + } else { + descriptor.encoding = CudaWeightEncoding::Plain; + } + if (deferred_weights_enabled()) { + auto mapped = view(weight_name); + if (!mapped.ok()) return {std::move(mapped.errors)}; + if (offset > mapped.value.size_bytes() || + bytes > mapped.value.size_bytes() - offset) { + return {{"GLM-5.3 sliced linear exceeds its mapped weight"}}; + } + return backend.upload( + device, descriptor, + mapped.value.subspan(static_cast(offset), + static_cast(bytes)), + scales, output, CudaBackend::UploadCompletion::Deferred, + descriptor.encoding == CudaWeightEncoding::Fp8E4m3Block128F32 && + backend.fp8_f32_register_fed_supported(device) + ? CudaBackend::FragmentLayout::Prepack + : CudaBackend::FragmentLayout::Canonical); + } + if (weight->source_dtype == SafetensorsDtype::F8E4M3) { + return {{"GLM-5.3 FP8 slices require stable deferred mappings"}}; + } + auto loaded = read_slice(*weight, offset, bytes); + if (!loaded.ok()) return {std::move(loaded.errors)}; + return backend.upload(device, descriptor, loaded.value, {}, output); +} + +} // namespace strata diff --git a/src/models/glm53/glm53_executor.cpp b/src/models/glm53/glm53_executor.cpp new file mode 100644 index 0000000..b6738aa --- /dev/null +++ b/src/models/glm53/glm53_executor.cpp @@ -0,0 +1,48 @@ +#include "../common/executor_support.hpp" +#include "strata/engine/runtime_support.hpp" +#include "strata/models/glm53/glm53_runtime.hpp" + +namespace strata { +namespace { + +class Glm53Executor final : public ModelExecutor { +public: + ValidationResult initialize(const std::string& model_directory, + const RuntimeConfig& config, + const PlacementPlan*) override { + Glm53RuntimeConfig concrete; + concrete.devices = resolve_runtime_devices(config.devices); + concrete.vram_cache_fraction = config.vram_cache_fraction; + concrete.maximum_context_tokens = config.maximum_context_tokens; + concrete.sampling_temperature = config.sampling.temperature; + concrete.sampling_seed = config.sampling.seed; + concrete.verbose = config.verbose; + concrete.load_progress = config.load_progress; + return runtime_.initialize(model_directory, concrete); + } + + GenerationResult generate_chat_stream( + std::span messages, const GenerationOptions& options, + const TokenStreamCallback& on_token) override { + GenerationResult result; + auto concrete = runtime_.generate_chat_stream( + messages, options.maximum_new_tokens, options.sampling, + options.stop, on_token); + detail::copy_common_generation(result, concrete); + result.metrics.reused_prompt_tokens = + concrete.metrics.reused_prompt_tokens; + result.metrics.incremental_kv_continuation = true; + return result; + } + +private: + Glm53Runtime runtime_; +}; + +const ModelRegistrar registrar{{ + RuntimeModel::Glm53, "GLM-5.3-Flash", "glm53", PlacementModel::Glm53, + false, true, false, false, + [] { return std::unique_ptr(new Glm53Executor()); }}}; + +} // namespace +} // namespace strata diff --git a/src/models/glm53/glm53_manifest.cpp b/src/models/glm53/glm53_manifest.cpp new file mode 100644 index 0000000..aebe8ba --- /dev/null +++ b/src/models/glm53/glm53_manifest.cpp @@ -0,0 +1,513 @@ +#include "strata/models/glm53/glm53_manifest.hpp" + +#include "../../platform/json_cursor.hpp" + +#include +#include +#include +#include +#include + +namespace strata { +namespace { + +using detail::JsonCursor; + +constexpr std::string_view kLayerPrefix = "model.language_model.layers."; +constexpr std::uint64_t kIndexedBytes = 328'326'771'576ULL; +constexpr std::uint64_t kShardFileBytes = 328'337'455'672ULL; +constexpr std::uint64_t kTensorCount = 76'108ULL; +constexpr std::uint32_t kShardCount = 62U; + +bool take_index(std::string_view name, std::size_t& cursor, + std::uint32_t& value) noexcept { + const auto begin = cursor; + while (cursor < name.size() && name[cursor] >= '0' && name[cursor] <= '9') { + ++cursor; + } + if (cursor == begin) return false; + const auto parsed = std::from_chars(name.data() + begin, + name.data() + cursor, value); + return parsed.ec == std::errc{}; +} + +bool has_suffix(std::string_view value, std::string_view suffix) noexcept { + return value.size() >= suffix.size() && value.ends_with(suffix); +} + +void parse_uint_list(JsonCursor& cursor, std::vector& output) { + cursor.expect('['); + if (cursor.consume(']')) return; + for (;;) { + output.push_back(static_cast(cursor.parse_uint64())); + if (cursor.consume(']')) return; + cursor.expect(','); + } +} + +void parse_string_list(JsonCursor& cursor, std::vector& output) { + cursor.expect('['); + if (cursor.consume(']')) return; + for (;;) { + output.push_back(cursor.parse_string()); + if (cursor.consume(']')) return; + cursor.expect(','); + } +} + +void parse_linear_attention(JsonCursor& cursor, Glm53TextConfig& config) { + cursor.expect('{'); + if (cursor.consume('}')) return; + for (;;) { + const auto key = cursor.parse_string(); + cursor.expect(':'); + if (key == "num_heads") { + config.linear_attention_heads = + static_cast(cursor.parse_uint64()); + } else if (key == "head_dim") { + config.linear_head_dim = + static_cast(cursor.parse_uint64()); + } else if (key == "short_conv_kernel_size") { + config.short_conv_kernel = + static_cast(cursor.parse_uint64()); + } else if (key == "gate_lower_bound") { + config.kda_gate_lower_bound = static_cast(cursor.parse_number()); + } else if (key == "full_attn_layers") { + parse_uint_list(cursor, config.full_attention_layers); + } else if (key == "kda_layers") { + parse_uint_list(cursor, config.kda_layers); + } else { + cursor.skip_value(); + } + if (cursor.consume('}')) return; + cursor.expect(','); + } +} + +void parse_text(JsonCursor& cursor, Glm53TextConfig& config) { + cursor.expect('{'); + if (cursor.consume('}')) return; + for (;;) { + const auto key = cursor.parse_string(); + cursor.expect(':'); + if (key == "model_type") config.model_type = cursor.parse_string(); + else if (key == "hidden_size") config.hidden_size = static_cast(cursor.parse_uint64()); + else if (key == "num_hidden_layers") config.layer_count = static_cast(cursor.parse_uint64()); + else if (key == "num_attention_heads") config.attention_heads = static_cast(cursor.parse_uint64()); + else if (key == "num_key_value_heads") config.key_value_heads = static_cast(cursor.parse_uint64()); + else if (key == "q_lora_rank") config.query_lora_rank = static_cast(cursor.parse_uint64()); + else if (key == "kv_lora_rank") config.kv_lora_rank = static_cast(cursor.parse_uint64()); + else if (key == "qk_nope_head_dim") config.nope_head_dim = static_cast(cursor.parse_uint64()); + else if (key == "qk_rope_head_dim") config.rope_head_dim = static_cast(cursor.parse_uint64()); + else if (key == "v_head_dim") config.value_head_dim = static_cast(cursor.parse_uint64()); + else if (key == "intermediate_size") config.dense_intermediate_size = static_cast(cursor.parse_uint64()); + else if (key == "moe_intermediate_size") config.expert_intermediate_size = static_cast(cursor.parse_uint64()); + else if (key == "n_routed_experts") config.routed_experts = static_cast(cursor.parse_uint64()); + else if (key == "num_experts_per_tok") config.experts_per_token = static_cast(cursor.parse_uint64()); + else if (key == "n_group") config.expert_groups = static_cast(cursor.parse_uint64()); + else if (key == "topk_group") config.selected_expert_groups = static_cast(cursor.parse_uint64()); + else if (key == "n_shared_experts") config.shared_experts = static_cast(cursor.parse_uint64()); + else if (key == "first_k_dense_replace") config.dense_prefix_layers = static_cast(cursor.parse_uint64()); + else if (key == "vocab_size") config.vocabulary_size = static_cast(cursor.parse_uint64()); + else if (key == "max_position_embeddings") config.maximum_context_tokens = static_cast(cursor.parse_uint64()); + else if (key == "hc_mult") config.mhc_multiplier = static_cast(cursor.parse_uint64()); + else if (key == "hc_sinkhorn_iters") config.mhc_sinkhorn_iterations = static_cast(cursor.parse_uint64()); + else if (key == "index_n_heads") config.index_heads = static_cast(cursor.parse_uint64()); + else if (key == "index_head_dim") config.index_head_dim = static_cast(cursor.parse_uint64()); + else if (key == "index_topk") config.index_topk = static_cast(cursor.parse_uint64()); + else if (key == "index_kpool") config.index_pool = static_cast(cursor.parse_uint64()); + else if (key == "rms_norm_eps") config.rms_epsilon = static_cast(cursor.parse_number()); + else if (key == "hc_eps") config.mhc_epsilon = static_cast(cursor.parse_number()); + else if (key == "routed_scaling_factor") config.routed_scale = static_cast(cursor.parse_number()); + else if (key == "swiglu_limit") config.swiglu_limit = static_cast(cursor.parse_number()); + else if (key == "norm_topk_prob") config.normalize_topk = cursor.parse_bool(); + else if (key == "mhc") config.mhc = cursor.parse_bool(); + else if (key == "mla_use_nope") config.mla_use_nope = cursor.parse_bool(); + else if (key == "index_kpool_compress") config.index_pool_compress = cursor.parse_bool(); + else if (key == "index_kpool_always_select_tail") config.index_pool_select_tail = cursor.parse_bool(); + else if (key == "tie_word_embeddings") config.tie_word_embeddings = cursor.parse_bool(); + else if (key == "hidden_act") config.hidden_activation = cursor.parse_string(); + else if (key == "scoring_func") config.router_scoring = cursor.parse_string(); + else if (key == "topk_method") config.topk_method = cursor.parse_string(); + else if (key == "layer_types") parse_string_list(cursor, config.attention_layer_types); + else if (key == "mlp_layer_types") parse_string_list(cursor, config.mlp_layer_types); + else if (key == "linear_attn_config") parse_linear_attention(cursor, config); + else cursor.skip_value(); + if (cursor.consume('}')) return; + cursor.expect(','); + } +} + +void parse_quantization(JsonCursor& cursor, Glm53TextConfig& config) { + cursor.expect('{'); + if (cursor.consume('}')) return; + for (;;) { + const auto key = cursor.parse_string(); + cursor.expect(':'); + if (key == "quant_method") config.quantization_method = cursor.parse_string(); + else if (key == "fmt") config.quantization_format = cursor.parse_string(); + else if (key == "weight_block_size") { + std::vector dimensions; + parse_uint_list(cursor, dimensions); + if (dimensions.size() == 2U) { + config.fp8_block_rows = dimensions[0]; + config.fp8_block_columns = dimensions[1]; + } + } else cursor.skip_value(); + if (cursor.consume('}')) return; + cursor.expect(','); + } +} + +Glm53TensorComponent component_of(std::string_view name) noexcept { + if (has_suffix(name, ".weight_scale_inv")) return Glm53TensorComponent::Scale; + if (has_suffix(name, ".bias") || has_suffix(name, "_bias")) { + return Glm53TensorComponent::Bias; + } + if (has_suffix(name, ".weight") || has_suffix(name, "_fn")) { + return Glm53TensorComponent::Weight; + } + return Glm53TensorComponent::State; +} + +} // namespace + +Glm53ConfigResult parse_glm53_config(std::string_view json) { + Glm53ConfigResult result; + try { + JsonCursor cursor(json); + cursor.expect('{'); + if (!cursor.consume('}')) { + for (;;) { + const auto key = cursor.parse_string(); + cursor.expect(':'); + if (key == "architectures") { + std::vector architectures; + parse_string_list(cursor, architectures); + if (architectures.size() == 1U) { + result.value.architecture = std::move(architectures[0]); + } + } else if (key == "text_config") { + parse_text(cursor, result.value); + } else if (key == "quantization_config") { + parse_quantization(cursor, result.value); + } else { + cursor.skip_value(); + } + if (cursor.consume('}')) break; + cursor.expect(','); + } + } + } catch (const detail::JsonError& error) { + result.errors.push_back("GLM-5.3 config.json is malformed at offset " + + std::to_string(error.offset()) + ": " + + error.what()); + } + return result; +} + +ValidationResult validate_glm53_config(const Glm53TextConfig& config) { + ValidationResult result; + const auto require = [&result](bool condition, std::string message) { + if (!condition) result.errors.push_back(std::move(message)); + }; + const auto equal = [&require](auto actual, auto wanted, std::string_view name) { + require(actual == wanted, "GLM-5.3 " + std::string(name) + " is " + + std::to_string(actual) + ", expected " + + std::to_string(wanted)); + }; + require(config.architecture == "Glm5NextForConditionalGeneration", + "GLM-5.3 architecture must be Glm5NextForConditionalGeneration"); + require(config.model_type == "glm5_next_text", + "GLM-5.3 text model_type must be glm5_next_text"); + equal(config.hidden_size, 4096U, "hidden_size"); + equal(config.layer_count, 45U, "num_hidden_layers"); + equal(config.attention_heads, 64U, "num_attention_heads"); + equal(config.key_value_heads, 64U, "num_key_value_heads"); + equal(config.query_lora_rank, 1536U, "q_lora_rank"); + equal(config.kv_lora_rank, 512U, "kv_lora_rank"); + equal(config.nope_head_dim, 256U, "qk_nope_head_dim"); + equal(config.rope_head_dim, 0U, "qk_rope_head_dim"); + equal(config.value_head_dim, 256U, "v_head_dim"); + equal(config.linear_attention_heads, 64U, "linear attention heads"); + equal(config.linear_head_dim, 128U, "linear attention head_dim"); + equal(config.short_conv_kernel, 4U, "linear convolution kernel"); + equal(config.dense_intermediate_size, 12288U, "intermediate_size"); + equal(config.expert_intermediate_size, 2048U, "moe_intermediate_size"); + equal(config.routed_experts, 288U, "n_routed_experts"); + equal(config.experts_per_token, 8U, "num_experts_per_tok"); + equal(config.expert_groups, 1U, "n_group"); + equal(config.selected_expert_groups, 1U, "topk_group"); + equal(config.shared_experts, 1U, "n_shared_experts"); + equal(config.dense_prefix_layers, 3U, "first_k_dense_replace"); + equal(config.vocabulary_size, 154880U, "vocab_size"); + equal(config.maximum_context_tokens, 1'048'576U, "max_position_embeddings"); + equal(config.mhc_multiplier, 4U, "hc_mult"); + equal(config.mhc_sinkhorn_iterations, 20U, "hc_sinkhorn_iters"); + equal(config.index_heads, 32U, "index_n_heads"); + equal(config.index_head_dim, 128U, "index_head_dim"); + equal(config.index_topk, 2048U, "index_topk"); + equal(config.index_pool, 4U, "index_kpool"); + equal(config.fp8_block_rows, 128U, "FP8 block rows"); + equal(config.fp8_block_columns, 128U, "FP8 block columns"); + require(std::abs(config.rms_epsilon - 1.0e-5F) <= 1.0e-12F, + "GLM-5.3 rms_norm_eps must be 1e-5"); + require(std::abs(config.mhc_epsilon - 1.0e-6F) <= 1.0e-12F, + "GLM-5.3 hc_eps must be 1e-6"); + require(config.routed_scale == 2.5F, + "GLM-5.3 routed_scaling_factor must be 2.5"); + require(config.swiglu_limit == 10.0F, "GLM-5.3 swiglu_limit must be 10"); + require(config.kda_gate_lower_bound == -5.0F, + "GLM-5.3 KDA gate lower bound must be -5"); + require(config.normalize_topk, "GLM-5.3 must normalize routed top-k weights"); + require(config.mhc, "GLM-5.3 must enable mHC"); + require(config.mla_use_nope, "GLM-5.3 sparse MLA must use NoPE"); + require(config.index_pool_compress && config.index_pool_select_tail, + "GLM-5.3 index k-pool compression and visible tail are required"); + require(!config.tie_word_embeddings, + "GLM-5.3 output head must not be tied to embeddings"); + require(config.hidden_activation == "silu", "GLM-5.3 activation must be silu"); + require(config.router_scoring == "sigmoid", "GLM-5.3 router must use sigmoid scores"); + require(config.topk_method == "noaux_tc", "GLM-5.3 router must use noaux_tc selection"); + require(config.quantization_method == "fp8" && + config.quantization_format == "e4m3", + "GLM-5.3 quantization must be FP8 E4M3 block-128"); + require(config.attention_layer_types.size() == config.layer_count, + "GLM-5.3 attention layer_types must cover all 45 layers"); + require(config.mlp_layer_types.size() == config.layer_count, + "GLM-5.3 mlp_layer_types must cover all 45 layers"); + for (std::uint32_t layer = 0U; layer < config.layer_count; ++layer) { + const auto attention = glm53_kda_layer(layer) ? "linear_attention" + : "deepseek_sparse_attention"; + const auto mlp = glm53_moe_layer(layer) ? "sparse" : "dense"; + require(layer < config.attention_layer_types.size() && + config.attention_layer_types[layer] == attention, + "GLM-5.3 attention schedule differs at layer " + + std::to_string(layer)); + require(layer < config.mlp_layer_types.size() && + config.mlp_layer_types[layer] == mlp, + "GLM-5.3 MLP schedule differs at layer " + + std::to_string(layer)); + } + std::vector expected_full_attention; + std::vector expected_kda; + for (std::uint32_t layer = 0U; layer < 45U; ++layer) { + (glm53_kda_layer(layer) ? expected_kda : expected_full_attention) + .push_back(layer); + } + require(config.full_attention_layers == expected_full_attention && + config.kda_layers == expected_kda, + "GLM-5.3 linear attention lists must exactly partition the pinned " + "34 KDA and 11 sparse layers"); + return result; +} + +Glm53TensorRole classify_glm53_tensor(std::string_view name, + std::int32_t& layer, + std::int32_t& expert) noexcept { + layer = -1; + expert = -1; + if (name.starts_with("model.visual.")) return Glm53TensorRole::Vision; + if (name == "model.language_model.embed_tokens.weight") return Glm53TensorRole::Embedding; + if (name == "model.language_model.norm.weight") return Glm53TensorRole::Norm; + if (name == "lm_head.weight") return Glm53TensorRole::OutputHead; + if (!name.starts_with(kLayerPrefix)) return Glm53TensorRole::Count; + std::size_t cursor = kLayerPrefix.size(); + std::uint32_t ordinal = 0U; + if (!take_index(name, cursor, ordinal) || cursor >= name.size() || + name[cursor] != '.') { + return Glm53TensorRole::Count; + } + layer = static_cast(ordinal); + if (ordinal >= 45U) return Glm53TensorRole::Mtp; + const auto leaf = name.substr(cursor + 1U); + if (leaf.starts_with("hc_")) return Glm53TensorRole::Mhc; + if (leaf.starts_with("input_layernorm") || + leaf.starts_with("post_attention_layernorm")) return Glm53TensorRole::Norm; + if (leaf.starts_with("self_attn.indexer.")) return Glm53TensorRole::AttentionIndexer; + if (leaf.starts_with("self_attn.")) { + return glm53_kda_layer(ordinal) ? Glm53TensorRole::KdaAttention + : Glm53TensorRole::SparseAttention; + } + if (!leaf.starts_with("mlp.")) return Glm53TensorRole::Count; + const auto mlp = leaf.substr(4U); + if (ordinal < 3U) return Glm53TensorRole::DenseMlp; + if (mlp.starts_with("gate.")) return Glm53TensorRole::Router; + if (mlp.starts_with("shared_experts.")) return Glm53TensorRole::SharedExpert; + if (mlp.starts_with("experts.")) { + std::size_t inner = 8U; + std::uint32_t selected = 0U; + if (!take_index(mlp, inner, selected)) return Glm53TensorRole::Count; + expert = static_cast(selected); + return Glm53TensorRole::RoutedExpert; + } + return Glm53TensorRole::Count; +} + +Glm53ManifestResult build_glm53_index_manifest(SafetensorsIndex index) { + Glm53ManifestResult result; + auto& manifest = result.manifest; + manifest.indexed_tensor_bytes = index.total_size; + manifest.shards = std::move(index.shards); + manifest.tensors.reserve(index.entries.size()); + if (manifest.indexed_tensor_bytes != kIndexedBytes || + index.entries.size() != kTensorCount || + manifest.shards.size() != kShardCount) { + result.errors.emplace_back( + "GLM-5.3 checkpoint index extent does not match the pinned Flash release"); + return result; + } + for (auto& entry : index.entries) { + std::int32_t layer = -1; + std::int32_t expert = -1; + const auto role = classify_glm53_tensor(entry.name, layer, expert); + if (role == Glm53TensorRole::Count) { + if (result.errors.size() < 64U) { + result.errors.push_back("unclassified GLM-5.3 tensor " + entry.name); + } + continue; + } + Glm53ManifestTensor tensor; + tensor.name = std::move(entry.name); + tensor.shard = std::move(entry.shard); + tensor.role = role; + tensor.component = component_of(tensor.name); + tensor.encoding = tensor.component == Glm53TensorComponent::Scale + ? Glm53TensorEncoding::Fp8E4m3Block128F32 + : Glm53TensorEncoding::Plain; + tensor.layer = layer; + tensor.expert = expert; + ++manifest.role_counts[static_cast(role)]; + if (tensor.component == Glm53TensorComponent::Scale) ++manifest.fp8_modules; + manifest.tensors.push_back(std::move(tensor)); + } + return result; +} + +Glm53ManifestResult validate_glm53_checkpoint( + const std::string& model_directory, Glm53IndexManifest manifest, + const Glm53CheckpointOptions& options) { + Glm53ManifestResult result; + result.manifest = std::move(manifest); + auto& out = result.manifest; + std::unordered_map indexed; + indexed.reserve(out.tensors.size()); + for (std::size_t index = 0U; index < out.tensors.size(); ++index) { + if (!indexed.emplace(out.tensors[index].name, index).second) { + result.errors.push_back("duplicate GLM-5.3 tensor " + out.tensors[index].name); + } + } + if (!result.errors.empty()) return result; + std::uint64_t resolved = 0U; + for (const auto& shard_name : out.shards) { + auto shard = load_safetensors_shard( + (std::filesystem::path(model_directory) / shard_name).string()); + if (!shard.ok()) { + if (options.require_all_shards) { + for (auto& error : shard.errors) { + if (result.errors.size() < options.maximum_errors) { + result.errors.push_back(std::move(error)); + } + } + } + continue; + } + ++out.scanned_shards; + out.shard_file_bytes += shard.value.file_size; + for (const auto& source : shard.value.tensors) { + const auto found = indexed.find(source.name); + if (found == indexed.end()) { + if (result.errors.size() < options.maximum_errors) { + result.errors.push_back("unindexed GLM-5.3 tensor " + source.name); + } + continue; + } + auto& target = out.tensors[found->second]; + if (target.shard != shard_name) { + result.errors.push_back("misplaced GLM-5.3 tensor " + source.name); + continue; + } + target.source_dtype = source.dtype; + target.source_shape = source.shape; + target.source_offset = source.absolute_begin; + target.source_bytes = source.bytes(); + if (target.role == Glm53TensorRole::Vision) out.vision_bytes += target.source_bytes; + else if (target.role == Glm53TensorRole::RoutedExpert) out.routed_expert_bytes += target.source_bytes; + else out.dense_spine_bytes += target.source_bytes; + out.tensor_payload_bytes += target.source_bytes; + ++resolved; + } + } + if (!result.errors.empty()) return result; + if (out.scanned_shards != kShardCount || out.shard_file_bytes != kShardFileBytes || + resolved != out.tensors.size() || out.tensor_payload_bytes != kIndexedBytes) { + result.errors.emplace_back( + "GLM-5.3 shard or tensor extent does not match the pinned Flash release"); + return result; + } + const auto expect = [&](std::string_view name, + std::vector shape, + SafetensorsDtype dtype) { + const auto found = indexed.find(name); + if (found == indexed.end()) { + result.errors.push_back("GLM-5.3 checkpoint is missing " + std::string(name)); + return; + } + const auto& tensor = out.tensors[found->second]; + if (tensor.source_shape != shape || tensor.source_dtype != dtype) { + result.errors.push_back("GLM-5.3 tensor " + std::string(name) + + " has an unexpected shape or dtype"); + } + }; + expect("model.language_model.embed_tokens.weight", {154880U, 4096U}, + SafetensorsDtype::Bf16); + expect("model.language_model.norm.weight", {4096U}, SafetensorsDtype::Bf16); + expect("lm_head.weight", {154880U, 4096U}, SafetensorsDtype::Bf16); + for (const auto layer : {0U, 3U, 44U}) { + const auto prefix = "model.language_model.layers." + + std::to_string(layer) + "."; + expect(prefix + "hc_attn_fn", {24U, 16384U}, SafetensorsDtype::Bf16); + expect(prefix + "hc_ffn_fn", {24U, 16384U}, SafetensorsDtype::Bf16); + expect(prefix + "input_layernorm.weight", {4096U}, SafetensorsDtype::Bf16); + if (glm53_kda_layer(layer)) { + expect(prefix + "self_attn.q_proj.weight", {8192U, 4096U}, SafetensorsDtype::Bf16); + expect(prefix + "self_attn.A_log", {64U}, SafetensorsDtype::F32); + } else { + expect(prefix + "self_attn.q_a_proj.weight", {1536U, 4096U}, SafetensorsDtype::F8E4M3); + expect(prefix + "self_attn.kv_b_proj.weight", {32768U, 512U}, SafetensorsDtype::Bf16); + } + if (result.errors.size() >= options.maximum_errors) break; + } + return result; +} + +std::string_view to_string(Glm53TensorRole role) noexcept { + switch (role) { + case Glm53TensorRole::Embedding: return "embedding"; + case Glm53TensorRole::OutputHead: return "output-head"; + case Glm53TensorRole::Norm: return "norm"; + case Glm53TensorRole::Mhc: return "mhc"; + case Glm53TensorRole::KdaAttention: return "kda"; + case Glm53TensorRole::SparseAttention: return "sparse-attention"; + case Glm53TensorRole::AttentionIndexer: return "attention-indexer"; + case Glm53TensorRole::DenseMlp: return "dense-mlp"; + case Glm53TensorRole::Router: return "router"; + case Glm53TensorRole::SharedExpert: return "shared-expert"; + case Glm53TensorRole::RoutedExpert: return "routed-expert"; + case Glm53TensorRole::Mtp: return "mtp"; + case Glm53TensorRole::Vision: return "vision"; + case Glm53TensorRole::Count: break; + } + return "unknown"; +} + +std::string_view to_string(Glm53TensorEncoding encoding) noexcept { + switch (encoding) { + case Glm53TensorEncoding::Plain: return "plain"; + case Glm53TensorEncoding::Fp8E4m3Block128F32: + return "fp8-e4m3-block128-f32-scale"; + } + return "unknown"; +} + +} // namespace strata diff --git a/src/models/glm53/glm53_runtime.cpp b/src/models/glm53/glm53_runtime.cpp new file mode 100644 index 0000000..cda71d2 --- /dev/null +++ b/src/models/glm53/glm53_runtime.cpp @@ -0,0 +1,4932 @@ +#include "strata/models/glm53/glm53_runtime.hpp" +#include "strata/models/glm53/glm53_sequence.hpp" + +#include "strata/engine/runtime_support.hpp" +#include "strata/engine/route_predictor.hpp" +#include "strata/models/common/tokenizer.hpp" +#include "strata/models/deepseek/deepseek_ops.hpp" +#include "strata/models/glm53/glm53_checkpoint.hpp" +#include "strata/models/kimi_k3/kimi_k3_ops.hpp" +#include "strata/platform/hardware_profile.hpp" +#include "strata/platform/numerics.hpp" +#include "strata/platform/worker_pool.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__x86_64__) || defined(_M_X64) +#include +#define STRATA_GLM53_HOST_AVX2 1 +#else +#define STRATA_GLM53_HOST_AVX2 0 +#endif + +namespace strata { + +std::vector glm53_projection_slots( + std::span keys, + std::span costs, + std::span capacities, + std::size_t preferred_slot) { + if (keys.empty() || keys.size() != costs.size() || capacities.empty() || + preferred_slot >= capacities.size() || + std::any_of(capacities.begin(), capacities.end(), + [](std::uint64_t value) { return value == 0U; })) { + return {}; + } + std::vector order(keys.size()); + std::iota(order.begin(), order.end(), 0U); + std::stable_sort(order.begin(), order.end(), [&](std::size_t left, + std::size_t right) { + if (costs[left] != costs[right]) return costs[left] > costs[right]; + if (keys[left] != keys[right]) return keys[left] < keys[right]; + return left < right; + }); + std::vector loads(capacities.size(), 0.0L); + std::vector slots(keys.size()); + for (const auto index : order) { + std::size_t best = 0U; + for (std::size_t slot = 1U; slot < capacities.size(); ++slot) { + const auto candidate = loads[slot] / + static_cast(capacities[slot]); + const auto incumbent = loads[best] / + static_cast(capacities[best]); + const auto candidate_distance = + (slot + capacities.size() - preferred_slot) % capacities.size(); + const auto incumbent_distance = + (best + capacities.size() - preferred_slot) % capacities.size(); + if (candidate < incumbent || + (candidate == incumbent && + candidate_distance < incumbent_distance)) { + best = slot; + } + } + slots[index] = best; + loads[best] += static_cast(std::max( + costs[index], 1U)); + } + return slots; +} + +namespace { + +constexpr std::uint32_t kHidden = 4096U; +constexpr std::uint32_t kLayers = 45U; +constexpr std::uint32_t kMtpLayer = 45U; +constexpr std::uint32_t kHeads = 64U; +constexpr std::uint32_t kLinearHead = 128U; +constexpr std::uint32_t kLinearWidth = kHeads * kLinearHead; + +struct Glm53HostFp8Linear { + std::span weights; + std::span scales; + std::uint32_t rows{}; + std::uint32_t columns{}; +}; + +[[nodiscard]] float glm53_quantize_e4m3(float value) noexcept { + const float magnitude = std::min(std::abs(value), 448.0F); + float quantized = 0.0F; + if (magnitude < 0.015625F) { + quantized = std::rint(std::ldexp(magnitude, 9)) * + std::ldexp(1.0F, -9); + } else { + int exponent = 0; + static_cast(std::frexp(magnitude, &exponent)); + exponent = std::clamp(exponent - 1, -6, 8); + const float step = std::ldexp(1.0F, exponent - 3); + quantized = std::min(std::rint(magnitude / step) * step, 448.0F); + } + return std::copysign(quantized, value); +} + +void glm53_quantize_activation(std::span values) noexcept { + constexpr std::size_t block = 128U; + for (std::size_t begin = 0U; begin < values.size(); begin += block) { + const auto end = std::min(begin + block, values.size()); + float maximum = 0.0F; + for (auto index = begin; index < end; ++index) { + maximum = std::max(maximum, std::abs(values[index])); + } + const float scale = maximum > 0.0F ? maximum / 448.0F : 1.0F; + for (auto index = begin; index < end; ++index) { + values[index] = glm53_quantize_e4m3(values[index] / scale) * scale; + } + } +} + +[[nodiscard]] const std::array& glm53_fp8_values() noexcept { + static const auto values = [] { + std::array result{}; + for (std::size_t index = 0U; index < result.size(); ++index) { + result[index] = fp8_e4m3_f32(static_cast(index)); + } + return result; + }(); + return values; +} + +[[nodiscard]] float glm53_host_fp8_dot_scalar( + const std::byte* weights, const float* scales, + std::span input) noexcept { + const auto& values = glm53_fp8_values(); + float sum = 0.0F; + for (std::size_t column = 0U; column < input.size(); ++column) { + const auto code = std::to_integer(weights[column]); + sum = std::fma(input[column] * values[code], scales[column / 128U], sum); + } + return sum; +} + +#if STRATA_GLM53_HOST_AVX2 +__attribute__((target("avx2,fma"))) +[[nodiscard]] float glm53_host_fp8_dot_avx2( + const std::byte* weights, const float* scales, + std::span input) noexcept { + const auto& values = glm53_fp8_values(); + __m256 accumulators[8]{ + _mm256_setzero_ps(), _mm256_setzero_ps(), _mm256_setzero_ps(), + _mm256_setzero_ps(), _mm256_setzero_ps(), _mm256_setzero_ps(), + _mm256_setzero_ps(), _mm256_setzero_ps()}; + std::size_t column = 0U; + for (; column + 64U <= input.size(); column += 64U) { + const auto scale = _mm256_set1_ps(scales[column / 128U]); + for (std::size_t group = 0U; group < 8U; ++group) { + const auto offset = column + group * 8U; + const auto bytes = _mm_loadl_epi64(reinterpret_cast( + weights + offset)); + const auto indices = _mm256_cvtepu8_epi32(bytes); + const auto decoded = _mm256_i32gather_ps(values.data(), indices, 4); + const auto activation = _mm256_loadu_ps(input.data() + offset); + accumulators[group] = _mm256_fmadd_ps( + _mm256_mul_ps(decoded, scale), activation, + accumulators[group]); + } + } + // Keep eight independent dependency chains through the matrix and combine + // only once at the end. This is the host analogue of DeepSeek's tiled + // executor: the checkpoint byte is decoded in-register and never expanded + // into a second resident copy. + for (std::size_t width = 4U; width != 0U; width >>= 1U) { + for (std::size_t index = 0U; index < width; ++index) { + accumulators[index] = _mm256_add_ps( + accumulators[index], accumulators[index + width]); + } + } + const __m128 low = _mm256_castps256_ps128(accumulators[0]); + const __m128 high = _mm256_extractf128_ps(accumulators[0], 1); + __m128 total = _mm_add_ps(low, high); + total = _mm_hadd_ps(total, total); + total = _mm_hadd_ps(total, total); + float sum = _mm_cvtss_f32(total); + for (; column < input.size(); ++column) { + const auto code = std::to_integer(weights[column]); + sum = std::fma(input[column] * values[code], scales[column / 128U], sum); + } + return sum; +} +#endif + +[[nodiscard]] float glm53_host_fp8_dot( + const std::byte* weights, const float* scales, + std::span input) noexcept { +#if STRATA_GLM53_HOST_AVX2 + if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) { + return glm53_host_fp8_dot_avx2(weights, scales, input); + } +#endif + return glm53_host_fp8_dot_scalar(weights, scales, input); +} +constexpr std::uint32_t kMlaHead = 256U; +constexpr std::uint32_t kMlaWidth = kHeads * kMlaHead; +constexpr std::uint32_t kQueryRank = 1536U; +constexpr std::uint32_t kKvRank = 512U; +constexpr std::uint32_t kMhc = 4U; +constexpr std::uint32_t kVocabulary = 154880U; +constexpr std::uint32_t kExactSparseContext = 2048U; +constexpr std::uint64_t kKdaWorkspaceFloats = + 2ULL * kHidden + 6ULL * kLinearWidth + 2ULL * kLinearHead + kHeads; +constexpr std::uint64_t kDeviceWorkspaceReserve = 2ULL << 30U; +constexpr std::uint64_t kMinimumDeviceBudget = 2ULL << 30U; + +[[nodiscard]] std::size_t prefix_cache_entries( + std::uint32_t maximum_context_tokens) noexcept { + const auto kda_state = 34ULL * kHeads * kLinearHead * kLinearHead * + sizeof(float); + const auto convolution_state = + 34ULL * 3ULL * kLinearWidth * 3ULL * sizeof(float); + const auto mla_state = 11ULL * maximum_context_tokens * kKvRank * + sizeof(float); + const auto state_bytes = std::max( + kda_state + convolution_state + mla_state, 1U); + const auto budget = host_hardware_profile().host_usable_bytes(0.05); + if (budget == 0U) return 1U; + return std::clamp( + static_cast(budget / state_bytes), 1U, 64U); +} + +[[nodiscard]] bool batched_projections_enabled() noexcept { + static const bool enabled = [] { + const char* value = std::getenv("STRATA_GLM53_BATCHED_PROJECTIONS"); + return value == nullptr || + (std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off"); + }(); + return enabled; +} + +[[nodiscard]] bool cross_gpu_projections_enabled( + std::span devices) noexcept { + static const int policy = [] { + const char* value = std::getenv("STRATA_GLM53_CROSS_GPU_PROJECTIONS"); + if (value == nullptr) return -1; + return std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off" + ? 1 + : 0; + }(); + if (policy >= 0) return policy != 0; + for (std::size_t source = 0U; source < devices.size(); ++source) { + for (std::size_t destination = source + 1U; + destination < devices.size(); ++destination) { + if (!CudaBackend::high_speed_peer_access_supported( + devices[source], devices[destination]) || + !CudaBackend::high_speed_peer_access_supported( + devices[destination], devices[source])) { + return false; + } + } + } + return devices.size() > 1U; +} + +[[nodiscard]] bool tensor_parallel_head_enabled() noexcept { + static const bool enabled = [] { + const char* value = std::getenv("STRATA_GLM53_TENSOR_PARALLEL_HEAD"); + return value == nullptr || + (std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off"); + }(); + return enabled; +} + +[[nodiscard]] bool full_tensor_parallel_enabled() noexcept { + static const bool enabled = [] { + const char* value = std::getenv("STRATA_GLM53_FULL_TP"); + return value == nullptr || + (std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off"); + }(); + return enabled; +} + +[[nodiscard]] bool replay_ssm_enabled() noexcept { + static const bool enabled = [] { + const char* value = std::getenv("STRATA_GLM53_REPLAY_SSM"); + return value == nullptr || + (std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off"); + }(); + return enabled; +} + +[[nodiscard]] bool phase_scheduler_enabled() noexcept { + static const bool enabled = [] { + const char* value = std::getenv("STRATA_GLM53_PHASE_SCHEDULER"); + return value == nullptr || + (std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off"); + }(); + return enabled; +} + +[[nodiscard]] bool fused_kda_enabled() noexcept { + static const bool enabled = [] { + const char* value = std::getenv("STRATA_GLM53_FUSED_KDA"); + return value == nullptr || + (std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off"); + }(); + return enabled; +} + +[[nodiscard]] bool resident_mla_enabled() noexcept { + static const bool enabled = [] { + const char* value = std::getenv("STRATA_GLM53_RESIDENT_MLA"); + // The absorbed resident MLA route remains a profiling candidate until + // its layer-by-layer exactness gate is closed. Never make an + // experimental arithmetic path the production default. + return value != nullptr && std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off"; + }(); + return enabled; +} + +[[nodiscard]] bool profiler_capture_enabled() noexcept { + const char* value = std::getenv("STRATA_GLM53_NSYS_CAPTURE"); + return value != nullptr && std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off"; +} + +// -1 selects from the discovered CPU width and the admitted CUDA residency; +// 0/1 are explicit campaign overrides. +[[nodiscard]] int host_moe_override() noexcept { + const char* value = std::getenv("STRATA_GLM53_HOST_MOE"); + if (value == nullptr || std::string_view(value) == "auto") return -1; + return std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off" + ? 1 : 0; +} + +[[nodiscard]] bool mtp_enabled() noexcept { + static const bool enabled = [] { + const char* value = std::getenv("STRATA_GLM53_MTP"); + // The checkpoint MTP layer is fully wired, but verification only pays + // when its measured acceptance rate amortizes the extra draft pass. + // Keep the production latency route deterministic and opt in to an + // MTP campaign explicitly until that gate has been established for a + // workload. + return value != nullptr && std::string_view(value) != "0" && + std::string_view(value) != "false" && + std::string_view(value) != "off"; + }(); + return enabled; +} + +struct Glm53RowRange { + std::uint64_t begin{}; + std::uint64_t count{}; +}; + +[[nodiscard]] std::vector weighted_row_ranges( + std::uint64_t rows, std::span capacities, + std::uint64_t alignment) { + if (rows == 0U || capacities.empty() || alignment == 0U || + std::any_of(capacities.begin(), capacities.end(), + [](std::uint64_t value) { return value == 0U; })) { + return {}; + } + if (rows < capacities.size() * alignment) { + alignment = 1U; + } + long double total_capacity = 0.0L; + for (const auto capacity : capacities) { + total_capacity += static_cast(capacity); + } + std::vector ranges; + ranges.reserve(capacities.size()); + std::uint64_t begin = 0U; + long double cumulative = 0.0L; + for (std::size_t slot = 0U; slot < capacities.size(); ++slot) { + std::uint64_t end = rows; + if (slot + 1U != capacities.size()) { + cumulative += static_cast(capacities[slot]); + const auto target = static_cast( + static_cast(rows) * cumulative / total_capacity); + end = target - target % alignment; + const auto minimum = begin + alignment; + const auto remaining = static_cast( + capacities.size() - slot - 1U) * alignment; + end = std::clamp(end, minimum, rows - remaining); + } + ranges.push_back({begin, end - begin}); + begin = end; + } + return ranges; +} + +[[nodiscard]] std::vector contiguous_layer_schedule( + std::uint32_t layers, std::span capacities) { + const auto ranges = weighted_row_ranges(layers, capacities, 1U); + if (ranges.size() != capacities.size()) return {}; + std::vector schedule(layers); + for (std::size_t slot = 0U; slot < ranges.size(); ++slot) { + const auto range = ranges[slot]; + for (std::uint64_t layer = range.begin; + layer < range.begin + range.count; ++layer) { + schedule[static_cast(layer)] = slot; + } + } + return schedule; +} + +[[nodiscard]] std::vector projection_worker_cpus( + std::span devices) { + const auto& hardware = host_hardware_profile(); + std::vector chosen; + chosen.reserve(devices.size()); + const auto usable = [&](int cpu) { + return std::find(hardware.usable_cpu_ids.begin(), + hardware.usable_cpu_ids.end(), cpu) != + hardware.usable_cpu_ids.end(); + }; + const auto available = [&](int cpu) { + return usable(cpu) && + std::find(chosen.begin(), chosen.end(), cpu) == chosen.end(); + }; + for (const int device : devices) { + const int node = CudaBackend::device_numa_node(device); + const std::vector* local = nullptr; + if (node >= 0 && static_cast(node) < + hardware.numa.node_primary_cpus.size() && + !hardware.numa.node_primary_cpus[static_cast(node)] + .empty()) { + local = &hardware.numa.node_primary_cpus[ + static_cast(node)]; + } else if (node >= 0 && static_cast(node) < + hardware.numa.node_cpus.size()) { + local = &hardware.numa.node_cpus[static_cast(node)]; + } + auto selected = hardware.usable_cpu_ids.end(); + if (local != nullptr) { + const auto candidate = std::find_if( + local->begin(), local->end(), available); + if (candidate != local->end()) { + selected = std::find(hardware.usable_cpu_ids.begin(), + hardware.usable_cpu_ids.end(), *candidate); + } + } + if (selected == hardware.usable_cpu_ids.end()) { + selected = std::find_if(hardware.usable_cpu_ids.begin(), + hardware.usable_cpu_ids.end(), available); + } + if (selected == hardware.usable_cpu_ids.end()) return {}; + chosen.push_back(*selected); + } + return chosen; +} + +[[nodiscard]] std::vector compute_worker_cpus() { + const auto& hardware = host_hardware_profile(); + std::vector cpus; + const auto usable = [&](int cpu) { + return std::find(hardware.usable_cpu_ids.begin(), + hardware.usable_cpu_ids.end(), cpu) != + hardware.usable_cpu_ids.end(); + }; + for (const auto& node : hardware.numa.node_primary_cpus) { + for (const int cpu : node) { + if (usable(cpu)) cpus.push_back(cpu); + } + } + if (cpus.empty()) cpus = hardware.usable_cpu_ids; + return cpus; +} + +double now_seconds() { + return std::chrono::duration( + std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +float sigmoid(float value) noexcept { + return value >= 0.0F ? 1.0F / (1.0F + std::exp(-value)) + : std::exp(value) / (1.0F + std::exp(value)); +} + +void round_bf16(std::span values) noexcept { + for (auto& value : values) value = bf16_round_f32(value); +} + +void append(std::vector& destination, + std::vector source) { + for (auto& error : source) destination.push_back(std::move(error)); +} + +[[nodiscard]] std::string projection_group_key( + std::string_view base, Glm53TensorRole role) { + const auto separator = base.find_last_of('.'); + const auto prefix = base.substr(0U, separator + 1U); + const auto leaf = base.substr(separator + 1U); + if (role == Glm53TensorRole::KdaAttention) { + if (leaf == "q_proj" || leaf == "k_proj" || leaf == "v_proj" || + leaf == "f_a_proj" || leaf == "b_proj" || leaf == "g_a_proj") { + return std::string(prefix) + "#kda-input"; + } + if (leaf == "f_b_proj" || leaf == "g_b_proj") { + return std::string(prefix) + "#kda-low-rank"; + } + } else if (role == Glm53TensorRole::SparseAttention) { + if (leaf == "q_a_proj" || leaf == "kv_a_proj_with_mqa") { + return std::string(prefix) + "#mla-input"; + } + if (leaf == "q_b_proj" || leaf == "kv_b_proj") { + return std::string(prefix) + "#mla-expanded"; + } + } else if (role == Glm53TensorRole::DenseMlp && + (leaf == "gate_proj" || leaf == "up_proj")) { + return std::string(prefix) + "#dense-gate-up"; + } + return std::string(base); +} + +class Glm53WeightCache { + struct Entry { + CudaWeight weight; + bool pinned{}; + bool prefetched{}; + std::uint32_t leases{}; + std::list::iterator recency; + }; + + struct State { + std::mutex mutex; + std::unordered_map entries; + std::list recency; + std::uint64_t capacity{}; + std::uint64_t used{}; + std::uint64_t pinned{}; + std::uint64_t hits{}; + std::uint64_t misses{}; + std::uint64_t evictions{}; + std::uint64_t prefetches{}; + std::uint64_t useful_prefetches{}; + std::uint64_t failed_prefetches{}; + }; + +public: + struct LinearRequest { + std::string_view base; + std::uint64_t output_columns{}; + std::uint64_t input_columns{}; + std::span input; + std::uint32_t rows{}; + std::span output; + bool bf16_output{}; + std::uint64_t weight_rows{}; + std::uint64_t weight_row_begin{}; + }; + + struct Stats { + std::vector capacity; + std::vector used; + std::vector pinned; + std::uint64_t hits{}; + std::uint64_t misses{}; + std::uint64_t evictions{}; + std::uint64_t prefetches{}; + std::uint64_t useful_prefetches{}; + std::uint64_t failed_prefetches{}; + }; + + Glm53WeightCache(Glm53CheckpointReader& checkpoint, CudaBackend& backend, + std::vector devices, + std::vector capacities) + : checkpoint_(checkpoint), backend_(backend), + devices_(std::move(devices)) { + std::uint64_t largest_linear = 0U; + for (const auto& tensor : checkpoint_.manifest().tensors) { + if ((tensor.role != Glm53TensorRole::RoutedExpert && + tensor.role != Glm53TensorRole::SharedExpert) || + !tensor.name.ends_with(".weight") || + tensor.source_shape.size() != 2U) { + continue; + } + largest_linear = std::max( + largest_linear, + checkpoint_.cuda_linear_storage_bytes( + tensor.name.substr(0U, tensor.name.size() - 7U))); + } + const auto fragmentation_reserve = + largest_linear <= std::numeric_limits::max() / 2U + ? 2U * largest_linear + : largest_linear; + states_.reserve(capacities.size()); + for (const auto capacity : capacities) { + auto state = std::make_unique(); + state->capacity = capacity > fragmentation_reserve + ? capacity - fragmentation_reserve : capacity; + states_.push_back(std::move(state)); + } + } + + [[nodiscard]] ValidationResult preload( + std::size_t slot, std::string_view base, std::uint64_t rows, + std::uint64_t columns, bool& admitted) { + admitted = false; + const auto bytes = checkpoint_.cuda_linear_storage_bytes(base); + if (slot >= states_.size() || bytes == 0U) { + return {{"GLM-5.3 preload references an invalid CUDA linear"}}; + } + auto& state = *states_[slot]; + std::scoped_lock lock(state.mutex); + auto found = state.entries.find(std::string(base)); + if (found != state.entries.end()) { + if (!found->second.pinned) { + found->second.pinned = true; + state.pinned += found->second.weight.device_bytes(); + state.recency.erase(found->second.recency); + } + admitted = true; + ++state.hits; + return {}; + } + // A smaller or busier GPU may not fit its complete share of the + // resident spine. Skipping residency changes only performance: the + // exact weight is admitted through the demand/LRU path when needed. + if (bytes > state.capacity - state.used) return {}; + Entry entry; + auto loaded = checkpoint_.load_cuda_linear( + base, rows, columns, devices_[slot], backend_, entry.weight); + if (!loaded.ok()) return loaded; + entry.pinned = true; + const auto actual = entry.weight.device_bytes(); + if (actual > state.capacity - state.used) { + return {{"GLM-5.3 resident linear exceeded its admitted CUDA cache"}}; + } + state.used += actual; + state.pinned += actual; + state.entries.emplace(std::string(base), std::move(entry)); + admitted = true; + ++state.misses; + return {}; + } + + [[nodiscard]] ValidationResult preload_slice( + std::size_t slot, std::string_view base, std::uint64_t total_rows, + std::uint64_t columns, std::uint64_t row_begin, + std::uint64_t row_count, bool& admitted) { + admitted = false; + const auto key = slice_key(base, row_begin, row_count); + const auto bytes = checkpoint_.cuda_linear_slice_storage_bytes( + base, row_begin, row_count); + if (slot >= states_.size() || bytes == 0U) { + return {{"GLM-5.3 preload references an invalid CUDA slice"}}; + } + auto& state = *states_[slot]; + std::scoped_lock lock(state.mutex); + auto found = state.entries.find(key); + if (found != state.entries.end()) { + admitted = true; + ++state.hits; + return {}; + } + if (bytes > state.capacity - state.used) return {}; + Entry entry; + auto loaded = checkpoint_.load_cuda_linear_slice( + base, total_rows, columns, row_begin, row_count, devices_[slot], + backend_, entry.weight); + if (!loaded.ok()) return loaded; + entry.pinned = true; + const auto actual = entry.weight.device_bytes(); + if (actual > state.capacity - state.used) { + return {{"GLM-5.3 resident slice exceeded its admitted CUDA cache"}}; + } + state.used += actual; + state.pinned += actual; + state.entries.emplace(key, std::move(entry)); + admitted = true; + ++state.misses; + return {}; + } + + [[nodiscard]] ValidationResult matmul( + std::size_t slot, std::string_view base, std::uint64_t output_columns, + std::uint64_t input_columns, std::span input, + std::uint32_t rows, std::span output, bool bf16_output) { + const LinearRequest request{base, output_columns, input_columns, input, + rows, output, bf16_output, 0U, 0U}; + return matmul_batch(slot, std::span(&request, 1U)); + } + + [[nodiscard]] ValidationResult matmul_batch( + std::size_t slot, std::span requests) { + if (slot >= states_.size()) { + return {{"GLM-5.3 linear targets an invalid CUDA cache slot"}}; + } + if (requests.empty()) { + return {{"GLM-5.3 linear batch is empty"}}; + } + auto& state = *states_[slot]; + std::scoped_lock lock(state.mutex); + struct BatchLeases { + State& state; + std::vector keys; + ~BatchLeases() { + for (const auto& key : keys) { + const auto found = state.entries.find(key); + if (found != state.entries.end() && + found->second.leases != 0U) { + --found->second.leases; + } + } + } + } leases{state, {}}; + leases.keys.reserve(requests.size()); + std::vector batch; + batch.reserve(requests.size()); + for (const auto& request : requests) { + const bool sliced = request.weight_rows != 0U; + const std::string key = sliced + ? slice_key(request.base, request.weight_row_begin, + request.output_columns) + : std::string(request.base); + auto found = state.entries.find(key); + if (found == state.entries.end()) { + const auto bytes = sliced + ? checkpoint_.cuda_linear_slice_storage_bytes( + request.base, request.weight_row_begin, + request.output_columns) + : checkpoint_.cuda_linear_storage_bytes(request.base); + if (bytes == 0U || bytes > state.capacity) { + return {{"GLM-5.3 linear is absent or exceeds its CUDA cache: " + + key}}; + } + while (state.used + bytes > state.capacity) { + auto victim_position = state.recency.end(); + for (auto candidate = state.recency.begin(); + candidate != state.recency.end(); ++candidate) { + const auto entry = state.entries.find(*candidate); + if (entry != state.entries.end() && + entry->second.leases == 0U) { + victim_position = candidate; + break; + } + } + if (victim_position == state.recency.end()) { + return {{"GLM-5.3 pinned spine leaves insufficient CUDA " + "cache for an exact demand weight"}}; + } + const auto victim_key = *victim_position; + state.recency.erase(victim_position); + auto victim = state.entries.find(victim_key); + if (victim == state.entries.end() || victim->second.pinned) { + return {{"GLM-5.3 CUDA cache recency bookkeeping is invalid"}}; + } + state.used -= victim->second.weight.device_bytes(); + state.entries.erase(victim); + ++state.evictions; + } + Entry entry; + const auto load = [&] { + return sliced + ? checkpoint_.load_cuda_linear_slice( + request.base, request.weight_rows, + request.input_columns, request.weight_row_begin, + request.output_columns, devices_[slot], backend_, + entry.weight) + : checkpoint_.load_cuda_linear( + request.base, request.output_columns, + request.input_columns, devices_[slot], backend_, + entry.weight); + }; + auto loaded = load(); + while (!loaded.ok() && arena_exhausted(loaded) && + evict_one(state)) { + entry.weight = CudaWeight{}; + loaded = load(); + } + if (!loaded.ok()) return loaded; + const auto actual = entry.weight.device_bytes(); + if (actual > state.capacity - state.used) { + return {{"GLM-5.3 demand linear exceeded its admitted CUDA cache"}}; + } + state.recency.push_back(key); + entry.recency = std::prev(state.recency.end()); + state.used += actual; + found = state.entries.emplace(key, std::move(entry)).first; + ++state.misses; + } else { + ++state.hits; + if (found->second.prefetched) { + found->second.prefetched = false; + ++state.useful_prefetches; + } + if (!found->second.pinned) { + state.recency.splice(state.recency.end(), state.recency, + found->second.recency); + } + } + ++found->second.leases; + leases.keys.push_back(key); + batch.push_back({&found->second.weight, request.input, request.rows, + request.output, request.bf16_output, + request.rows > 1U}); + } + // One device-side event orders all deferred cache-miss uploads before + // the consumer. This never blocks the host and is a no-op on a hit-only + // path; matmul's output completion still protects the LRU entry. + if (auto ordered = backend_.synchronize_uploads(devices_[slot]); + !ordered.ok()) { + return ordered; + } + return backend_.matmul_batch(batch); + } + + // Admit one routed expert without consuming it. The worker holds the + // device cache lock through the copy-stream completion, so a demand can + // never observe a half-uploaded entry and an eviction cannot recycle its + // arena storage early. Storage faults and H2D copies happen on the worker, + // concurrently with the preceding layer's compute stream. + [[nodiscard]] ValidationResult prefetch_expert( + std::size_t slot, std::uint32_t layer, std::uint32_t expert) { + if (slot >= states_.size() || layer >= kLayers || expert >= 288U || + !glm53_moe_layer(layer)) { + return { {"GLM-5.3 expert prefetch has an invalid target"} }; + } + const auto prefix = "model.language_model.layers." + + std::to_string(layer) + ".mlp.experts." + + std::to_string(expert) + "."; + struct Projection { + std::string key; + std::uint64_t rows{}; + std::uint64_t columns{}; + }; + const std::array projections{{ + {prefix + "gate_proj", 2048U, kHidden}, + {prefix + "up_proj", 2048U, kHidden}, + {prefix + "down_proj", kHidden, 2048U}}}; + auto& state = *states_[slot]; + std::scoped_lock lock(state.mutex); + bool admitted = false; + for (const auto& projection : projections) { + if (state.entries.contains(projection.key)) continue; + const auto bytes = checkpoint_.cuda_linear_storage_bytes( + projection.key); + if (bytes == 0U || bytes > state.capacity) { + ++state.failed_prefetches; + return {}; + } + while (state.used + bytes > state.capacity) { + auto victim = state.recency.end(); + for (auto candidate = state.recency.begin(); + candidate != state.recency.end(); ++candidate) { + const auto found = state.entries.find(*candidate); + if (found != state.entries.end() && + !found->second.pinned && found->second.leases == 0U) { + victim = candidate; + break; + } + } + if (victim == state.recency.end()) { + ++state.failed_prefetches; + return {}; + } + auto found = state.entries.find(*victim); + state.used -= found->second.weight.device_bytes(); + state.entries.erase(found); + state.recency.erase(victim); + ++state.evictions; + } + Entry entry; + const auto load = [&] { + return checkpoint_.load_cuda_linear( + projection.key, projection.rows, projection.columns, + devices_[slot], backend_, entry.weight, true); + }; + auto loaded = load(); + while (!loaded.ok() && arena_exhausted(loaded) && + evict_one(state)) { + entry.weight = CudaWeight{}; + loaded = load(); + } + if (!loaded.ok()) { + ++state.failed_prefetches; + return loaded; + } + const auto actual = entry.weight.device_bytes(); + if (actual > state.capacity - state.used) { + ++state.failed_prefetches; + return {}; + } + state.recency.push_back(projection.key); + entry.recency = std::prev(state.recency.end()); + entry.prefetched = true; + state.used += actual; + state.entries.emplace(projection.key, std::move(entry)); + admitted = true; + } + if (!admitted) return {}; + auto ordered = backend_.synchronize_uploads(devices_[slot]); + if (!ordered.ok()) { + ++state.failed_prefetches; + return ordered; + } + ++state.prefetches; + return {}; + } + + [[nodiscard]] bool contains_expert( + std::size_t slot, std::uint32_t layer, std::uint32_t expert) const { + if (slot >= states_.size()) return false; + const auto prefix = "model.language_model.layers." + + std::to_string(layer) + ".mlp.experts." + + std::to_string(expert) + "."; + auto& state = *states_[slot]; + std::scoped_lock lock(state.mutex); + return state.entries.contains(prefix + "gate_proj") && + state.entries.contains(prefix + "up_proj") && + state.entries.contains(prefix + "down_proj"); + } + + [[nodiscard]] ValidationResult kda_decode( + std::size_t slot, std::string_view attention, + CudaGlm53KdaRequest request, std::span output) { + if (slot >= states_.size() || request.state == nullptr) { + return {{"GLM-5.3 fused KDA targets an invalid CUDA cache slot"}}; + } + auto& state = *states_[slot]; + std::scoped_lock lock(state.mutex); + const std::array keys{ + std::string(attention) + "q_proj", + std::string(attention) + "k_proj", + std::string(attention) + "v_proj", + std::string(attention) + "f_a_proj", + std::string(attention) + "b_proj", + std::string(attention) + "g_a_proj", + std::string(attention) + "f_b_proj", + std::string(attention) + "g_b_proj", + std::string(attention) + "o_proj"}; + std::array entries{}; + const auto first = request.input.empty() && + !request.mhc_source_destination + ? keys.size() - 1U + : 0U; + for (std::size_t index = first; index < keys.size(); ++index) { + const auto found = state.entries.find(keys[index]); + if (found == state.entries.end() || + found->second.weight.device() != request.state->device()) { + return {{"GLM-5.3 fused KDA projection was not admitted on " + "its layer device: " + keys[index]}}; + } + entries[index] = &found->second; + } + if (entries.back() == nullptr) { + return {{"GLM-5.3 fused KDA output projection was not admitted " + "on its layer device"}}; + } + struct Lease { + std::span entries; + ~Lease() { + for (auto* entry : entries) { + if (entry != nullptr) --entry->leases; + } + } + } lease{entries}; + for (auto* entry : entries) { + if (entry != nullptr) ++entry->leases; + } + request.query_projection = entries[0] == nullptr + ? nullptr : &entries[0]->weight; + request.key_projection = entries[1] == nullptr + ? nullptr : &entries[1]->weight; + request.value_projection = entries[2] == nullptr + ? nullptr : &entries[2]->weight; + request.forget_a_projection = entries[3] == nullptr + ? nullptr : &entries[3]->weight; + request.beta_projection = entries[4] == nullptr + ? nullptr : &entries[4]->weight; + request.gate_a_projection = entries[5] == nullptr + ? nullptr : &entries[5]->weight; + request.forget_b_projection = entries[6] == nullptr + ? nullptr : &entries[6]->weight; + request.gate_b_projection = entries[7] == nullptr + ? nullptr : &entries[7]->weight; + request.output_projection = &entries[8]->weight; + return backend_.glm53_kda_decode(request, output); + } + + [[nodiscard]] ValidationResult router_mhc( + std::size_t slot, std::string_view key, std::span logits) { + if (slot >= states_.size()) { + return {{"GLM-5.3 resident router targets an invalid cache slot"}}; + } + auto& state = *states_[slot]; + std::scoped_lock lock(state.mutex); + const auto found = state.entries.find(std::string(key)); + if (found == state.entries.end() || + found->second.weight.device() != devices_[slot]) { + return {{"GLM-5.3 resident router was not admitted on its layer " + "device: " + std::string(key)}}; + } + ++found->second.leases; + struct Lease { + Entry& entry; + ~Lease() { --entry.leases; } + } lease{found->second}; + return backend_.glm53_mhc_router( + devices_[slot], found->second.weight, logits); + } + + [[nodiscard]] ValidationResult mla_decode_mhc( + std::size_t slot, std::string_view attention, + CudaGlm53MlaRequest request) { + if (slot >= states_.size() || request.state == nullptr) { + return {{"GLM-5.3 resident MLA targets an invalid cache slot"}}; + } + auto& state = *states_[slot]; + std::scoped_lock lock(state.mutex); + const std::array keys{ + std::string(attention) + "q_a_proj", + std::string(attention) + "kv_a_proj_with_mqa", + std::string(attention) + "q_b_proj", + std::string(attention) + "kv_b_proj", + std::string(attention) + "o_proj"}; + std::array entries{}; + for (std::size_t index = 0U; index < keys.size(); ++index) { + const auto found = state.entries.find(keys[index]); + if (found == state.entries.end() || + found->second.weight.device() != request.state->device()) { + return {{"GLM-5.3 resident MLA projection was not admitted: " + + keys[index]}}; + } + entries[index] = &found->second; + } + for (auto* entry : entries) ++entry->leases; + struct Lease { + std::array& entries; + ~Lease() { + for (auto* entry : entries) --entry->leases; + } + } lease{entries}; + request.query_a = &entries[0]->weight; + request.key_value_a = &entries[1]->weight; + request.query_b = &entries[2]->weight; + request.key_value_b = &entries[3]->weight; + request.output = &entries[4]->weight; + return backend_.glm53_mla_decode_to_mhc(request); + } + + [[nodiscard]] ValidationResult swiglu_mhc( + std::size_t slot, std::string_view prefix, + std::uint32_t intermediate) { + if (slot >= states_.size()) { + return {{"GLM-5.3 resident SwiGLU targets an invalid cache slot"}}; + } + auto& state = *states_[slot]; + std::scoped_lock lock(state.mutex); + const std::array keys{ + std::string(prefix) + "gate_proj", + std::string(prefix) + "up_proj", + std::string(prefix) + "down_proj"}; + std::array entries{}; + for (std::size_t index = 0U; index < keys.size(); ++index) { + const auto found = state.entries.find(keys[index]); + if (found == state.entries.end() || + found->second.weight.device() != devices_[slot]) { + return {{"GLM-5.3 resident SwiGLU projection was not " + "admitted: " + keys[index]}}; + } + entries[index] = &found->second; + } + for (auto* entry : entries) ++entry->leases; + struct Lease { + std::array& entries; + ~Lease() { + for (auto* entry : entries) --entry->leases; + } + } lease{entries}; + return backend_.glm53_mhc_swiglu( + devices_[slot], entries[0]->weight, entries[1]->weight, + entries[2]->weight, intermediate); + } + + [[nodiscard]] ValidationResult moe( + std::size_t slot, std::string_view prefix, + std::span routed, + std::span input, std::span output, + bool mhc_source_destination = false) { + ValidationResult result; + if (slot >= states_.size() || routed.size() != 8U || + (mhc_source_destination + ? (!input.empty() || !output.empty()) + : (input.size() != kHidden || output.size() != kHidden))) { + result.errors.emplace_back("GLM-5.3 MoE command has an invalid shape"); + return result; + } + struct Projection { + std::string key; + std::uint64_t rows{}; + std::uint64_t columns{}; + }; + const auto make_modules = [](const std::string& base) { + return std::array{ + Projection{base + "gate_proj", 2048U, kHidden}, + Projection{base + "up_proj", 2048U, kHidden}, + Projection{base + "down_proj", kHidden, 2048U}}; + }; + struct DeviceGroup { + std::vector routes; + bool has_shared{}; + bool enqueued{}; + std::vector leased; + std::vector descriptors; + CudaMoeExpert shared_descriptor; + std::vector routed_output; + std::vector shared_output; + }; + std::vector groups(states_.size()); + groups[slot].has_shared = true; + + // A best-rank peer fabric makes expert parallelism profitable: split + // the eight independent routes capacity-proportionally and join their + // exact host-visible outputs in original router order. PHB/PCIe keeps + // every route with the layer owner, avoiding duplicate cache traffic. + if (!mhc_source_destination && devices_.size() == 2U && + full_tensor_parallel_enabled() && + cross_gpu_projections_enabled(devices_)) { + std::vector capacities; + capacities.reserve(states_.size()); + for (const auto& state : states_) { + capacities.push_back(state->capacity); + } + const auto ranges = weighted_row_ranges( + routed.size(), capacities, 1U); + if (ranges.size() != groups.size()) { + return {{"GLM-5.3 expert-parallel assignment is invalid"}}; + } + for (std::size_t group_slot = 0U; group_slot < ranges.size(); + ++group_slot) { + for (std::uint64_t route = ranges[group_slot].begin; + route < ranges[group_slot].begin + + ranges[group_slot].count; + ++route) { + groups[group_slot].routes.push_back( + static_cast(route)); + } + } + } else { + for (std::size_t route_index = 0U; route_index < routed.size(); + ++route_index) { + groups[slot].routes.push_back(route_index); + } + } + + const auto release = [&](std::size_t group_slot) { + auto& state = *states_[group_slot]; + std::scoped_lock lock(state.mutex); + for (const auto& key : groups[group_slot].leased) { + const auto found = state.entries.find(key); + if (found != state.entries.end() && found->second.leases != 0U) { + --found->second.leases; + } + } + }; + const auto ensure = [&](State& state, std::size_t target_slot, + const Projection& projection, + std::vector& leased) + -> ValidationResult { + auto found = state.entries.find(projection.key); + if (found == state.entries.end()) { + const auto bytes = + checkpoint_.cuda_linear_storage_bytes(projection.key); + if (bytes == 0U || bytes > state.capacity) { + return {{"GLM-5.3 MoE projection is absent or exceeds its " + "CUDA cache: " + projection.key}}; + } + while (state.used + bytes > state.capacity) { + auto victim = state.recency.end(); + for (auto candidate = state.recency.begin(); + candidate != state.recency.end(); ++candidate) { + const auto entry = state.entries.find(*candidate); + if (entry != state.entries.end() && + !entry->second.pinned && entry->second.leases == 0U) { + victim = candidate; + break; + } + } + if (victim == state.recency.end()) { + return {{"GLM-5.3 exact MoE expert set exceeds the " + "available CUDA cache"}}; + } + auto entry = state.entries.find(*victim); + state.used -= entry->second.weight.device_bytes(); + state.entries.erase(entry); + state.recency.erase(victim); + ++state.evictions; + } + Entry entry; + const auto load = [&] { + return checkpoint_.load_cuda_linear( + projection.key, projection.rows, projection.columns, + devices_[target_slot], backend_, entry.weight); + }; + auto loaded = load(); + while (!loaded.ok() && arena_exhausted(loaded) && + evict_one(state)) { + entry.weight = CudaWeight{}; + loaded = load(); + } + if (!loaded.ok()) return loaded; + const auto actual = entry.weight.device_bytes(); + if (actual > state.capacity - state.used) { + return {{"GLM-5.3 MoE projection exceeded its admitted " + "CUDA cache"}}; + } + state.recency.push_back(projection.key); + entry.recency = std::prev(state.recency.end()); + state.used += actual; + found = state.entries.emplace(projection.key, + std::move(entry)).first; + ++state.misses; + } else { + ++state.hits; + if (found->second.prefetched) { + found->second.prefetched = false; + ++state.useful_prefetches; + } + if (!found->second.pinned) { + state.recency.splice(state.recency.end(), state.recency, + found->second.recency); + } + } + ++found->second.leases; + leased.push_back(projection.key); + return {}; + }; + + std::vector routed_output(routed.size() * kHidden); + std::vector shared_output(kHidden); + for (std::size_t group_slot = 0U; group_slot < groups.size(); + ++group_slot) { + auto& group = groups[group_slot]; + if (group.routes.empty() && !group.has_shared) continue; + auto& state = *states_[group_slot]; + std::scoped_lock lock(state.mutex); + std::vector> modules; + modules.reserve(group.routes.size()); + for (const auto route_index : group.routes) { + modules.push_back(make_modules( + std::string(prefix) + "experts." + + std::to_string(routed[route_index].expert) + ".")); + } + const auto shared_modules = make_modules( + std::string(prefix) + "shared_experts."); + group.leased.reserve( + (modules.size() + (group.has_shared ? 1U : 0U)) * 3U); + for (const auto& expert : modules) { + for (const auto& projection : expert) { + auto loaded = ensure(state, group_slot, projection, + group.leased); + if (!loaded.ok()) { + append(result.errors, std::move(loaded.errors)); + break; + } + } + if (!result.ok()) break; + } + if (result.ok() && group.has_shared) { + for (const auto& projection : shared_modules) { + auto loaded = ensure(state, group_slot, projection, + group.leased); + if (!loaded.ok()) { + append(result.errors, std::move(loaded.errors)); + break; + } + } + } + if (!result.ok()) break; + group.descriptors.resize(modules.size()); + for (std::size_t index = 0U; index < modules.size(); ++index) { + const auto& expert = modules[index]; + group.descriptors[index] = { + &state.entries.at(expert[0].key).weight, + &state.entries.at(expert[1].key).weight, + &state.entries.at(expert[2].key).weight, 1.0F}; + } + if (group.has_shared) { + group.shared_descriptor = { + &state.entries.at(shared_modules[0].key).weight, + &state.entries.at(shared_modules[1].key).weight, + &state.entries.at(shared_modules[2].key).weight, 1.0F}; + group.shared_output.resize(kHidden); + } + // The whole routed-plus-shared set was admitted with deferred + // copies. Order it once instead of synchronizing all 27 projection + // uploads independently. + auto ordered = backend_.synchronize_uploads(devices_[group_slot]); + if (!ordered.ok()) { + append(result.errors, std::move(ordered.errors)); + break; + } + group.routed_output.resize(group.routes.size() * kHidden); + ValidationResult enqueued; + if (mhc_source_destination) { + std::vector coefficients; + coefficients.reserve(group.routes.size()); + for (const auto route : group.routes) { + coefficients.push_back(routed[route].weight); + } + enqueued = backend_.enqueue_glm53_moe_from_mhc( + devices_[group_slot], group.descriptors, + group.shared_descriptor, coefficients, 10.0F); + } else { + enqueued = backend_.enqueue_moe( + devices_[group_slot], input, 1U, group.descriptors, + group.has_shared ? &group.shared_descriptor : nullptr, + 10.0F); + } + if (!enqueued.ok()) { + append(result.errors, std::move(enqueued.errors)); + break; + } + group.enqueued = true; + } + + // Every active device has been enqueued before the first completion + // boundary, so their expert projections and transfers overlap. + for (std::size_t group_slot = 0U; group_slot < groups.size(); + ++group_slot) { + auto& group = groups[group_slot]; + if (group.enqueued) { + auto collected = mhc_source_destination + ? backend_.finish_deepseek_moe_chain(devices_[group_slot]) + : backend_.collect_moe( + devices_[group_slot], group.routed_output, + group.has_shared + ? std::span(group.shared_output) + : std::span{}); + if (!collected.ok()) { + append(result.errors, std::move(collected.errors)); + } else if (!mhc_source_destination) { + for (std::size_t local = 0U; local < group.routes.size(); + ++local) { + std::copy_n( + group.routed_output.begin() + + static_cast(local * kHidden), + kHidden, + routed_output.begin() + static_cast( + group.routes[local] * kHidden)); + } + if (group.has_shared) { + std::copy(group.shared_output.begin(), + group.shared_output.end(), + shared_output.begin()); + } + } + } + if (!group.leased.empty()) release(group_slot); + } + if (!result.ok()) return result; + if (mhc_source_destination) return result; + std::copy(shared_output.begin(), shared_output.end(), output.begin()); + for (std::size_t expert = 0U; expert < routed.size(); ++expert) { + const auto begin = expert * kHidden; + for (std::size_t column = 0U; column < kHidden; ++column) { + output[column] = bf16_round_f32( + output[column] + bf16_round_f32( + routed[expert].weight * + routed_output[begin + column])); + } + } + return result; + } + + [[nodiscard]] Stats stats() const { + Stats result; + for (const auto& state_ptr : states_) { + auto& state = *state_ptr; + std::scoped_lock lock(state.mutex); + result.capacity.push_back(state.capacity); + result.used.push_back(state.used); + result.pinned.push_back(state.pinned); + result.hits += state.hits; + result.misses += state.misses; + result.evictions += state.evictions; + result.prefetches += state.prefetches; + result.useful_prefetches += state.useful_prefetches; + result.failed_prefetches += state.failed_prefetches; + } + return result; + } + +private: + [[nodiscard]] static bool arena_exhausted( + const ValidationResult& result) noexcept { + return std::any_of( + result.errors.begin(), result.errors.end(), + [](const std::string& error) { + return error.starts_with("CUDA weight arena is exhausted"); + }); + } + + [[nodiscard]] static bool evict_one(State& state) { + for (auto candidate = state.recency.begin(); + candidate != state.recency.end(); ++candidate) { + auto found = state.entries.find(*candidate); + if (found == state.entries.end() || found->second.pinned || + found->second.leases != 0U) { + continue; + } + state.used -= found->second.weight.device_bytes(); + state.entries.erase(found); + state.recency.erase(candidate); + ++state.evictions; + return true; + } + return false; + } + + [[nodiscard]] static std::string slice_key( + std::string_view base, std::uint64_t row_begin, + std::uint64_t row_count) { + return std::string(base) + "#rows=" + std::to_string(row_begin) + "+" + + std::to_string(row_count); + } + + Glm53CheckpointReader& checkpoint_; + CudaBackend& backend_; + std::vector devices_; + std::vector> states_; +}; + +} // namespace + +struct Glm53Runtime::Impl { + struct PrefetchJob { + ExpertKey key; + std::size_t slot{}; + }; + + struct DeviceSequenceState { + std::array kda; + std::array mla; + bool ready{}; + }; + + struct ResidentLayerWeights { + CudaDsv4MhcWeights attention; + CudaDsv4MhcWeights feedforward; + }; + + struct PrefixEntry { + std::vector tokens; + Glm53SequenceState state; + std::vector logits; + std::vector base_hidden; + std::uint64_t recency{}; + }; + + struct ScheduledRequest { + std::vector prompt; + std::uint32_t maximum_new_tokens{}; + SamplingOptions sampling; + std::vector stop; + TokenStreamCallback on_token; + Glm53GenerationResult result; + Glm53SequenceState sequence; + DeviceSequenceState device_sequence; + std::vector logits; + std::vector base_hidden; + std::vector counts; + std::vector sampled; + Glm53WeightCache::Stats decode_cache_start; + std::mt19937_64 generator; + std::unique_ptr streamed; + std::size_t prefill_cursor{}; + double prefill_started{}; + std::uint32_t position{}; + std::uint32_t iteration{}; + double decode_started{}; + std::mutex completion_mutex; + std::condition_variable completion; + bool prepared{}; + bool decoding{}; + bool mtp_ready{}; + bool done{}; + }; + + Glm53RuntimeConfig config; + std::unique_ptr checkpoint; + ModelTokenizer tokenizer; + CudaBackend cuda; + std::vector devices; + std::vector device_schedule; + std::vector weight_capacities; + std::vector lm_head_ranges; + std::unique_ptr weights; + std::array resident_layers; + bool resident_execution_active{}; + RoutePredictor route_predictor; + std::mutex prefetch_mutex; + std::condition_variable prefetch_ready; + std::deque prefetch_queue; + std::unordered_set pending_prefetch; + std::vector prefetch_threads; + std::size_t prefetch_queue_limit{}; + std::size_t prefetch_prediction_limit{}; + double prefetch_minimum_confidence{1.0}; + bool prefetch_stopping{}; + std::atomic prefetch_requests{}; + std::atomic prefetch_completed{}; + std::atomic prefetch_dropped{}; + std::atomic prefetch_errors{}; + std::unique_ptr projection_workers; + std::unique_ptr host_moe_workers; + bool host_moe_active{}; + std::atomic host_moe_calls{}; + std::atomic host_moe_nanoseconds{}; + bool full_tensor_parallel_active{}; + std::unique_ptr kda_workers; + std::atomic parallel_projection_batches{}; + std::atomic parallel_projection_requests{}; + std::atomic tensor_parallel_head_batches{}; + std::atomic parallel_encode_pages{}; + std::atomic prefix_cache_hits{}; + std::atomic prefix_cache_tokens{}; + std::mutex prefix_mutex; + std::vector prefix_cache; + std::size_t prefix_cache_limit{1U}; + std::uint64_t prefix_clock{}; + std::mutex host_tensor_mutex; + std::unordered_map>> host_tensors; + ValidationResult warmup_result; + bool ready{}; + std::thread warmup_thread; + std::mutex warmup_mutex; + std::mutex scheduler_mutex; + std::condition_variable scheduler_ready; + std::deque> pending_requests; + std::vector> active_requests; + std::thread scheduler_thread; + std::size_t scheduler_capacity{1U}; + bool scheduler_stopping{}; + std::atomic scheduler_iterations{}; + std::atomic scheduler_batched_iterations{}; + std::atomic mtp_drafts{}; + std::atomic mtp_accepted{}; + std::atomic profiler_captured{}; + + ~Impl() { + { + std::scoped_lock lock(scheduler_mutex); + scheduler_stopping = true; + } + scheduler_ready.notify_all(); + if (scheduler_thread.joinable()) scheduler_thread.join(); + if (warmup_thread.joinable()) warmup_thread.join(); + { + std::scoped_lock lock(prefetch_mutex); + prefetch_stopping = true; + } + prefetch_ready.notify_all(); + for (auto& worker : prefetch_threads) { + if (worker.joinable()) worker.join(); + } + } + + void prefetch_loop() { + for (;;) { + PrefetchJob job; + { + std::unique_lock lock(prefetch_mutex); + prefetch_ready.wait(lock, [&] { + return prefetch_stopping || !prefetch_queue.empty(); + }); + if (prefetch_stopping && prefetch_queue.empty()) return; + job = prefetch_queue.front(); + prefetch_queue.pop_front(); + } + auto status = weights->prefetch_expert( + job.slot, job.key.layer, job.key.expert); + if (status.ok()) { + prefetch_completed.fetch_add(1U, std::memory_order_relaxed); + } else if (prefetch_errors.fetch_add( + 1U, std::memory_order_relaxed) == 0U) { + std::cerr << "[glm53-residency] first_prefetch_error=" + << status.errors.front() << '\n'; + } + { + std::scoped_lock lock(prefetch_mutex); + pending_prefetch.erase(job.key); + } + } + } + + void request_prefetch(const RoutePrediction& prediction) { + // Nsight measured the fused resident chain at 12 useful predictions + // out of 90, with speculative uploads adding 2.3 GB/token and holding + // the demand cache mutex for 2.55 seconds. Resident decode already + // overlaps a layer's admitted set as one command, so cache pollution + // is more expensive than the predictor's occasional hit. Keep the + // predictor available to the host-bound path where it was originally + // validated, but never let it contend with the fused demand chain. + if (resident_execution_active || host_moe_active || + prefetch_queue_limit == 0U || + prediction.key.layer >= kLayers || + !glm53_moe_layer(prediction.key.layer)) { + return; + } + const auto slot = slot_for(prediction.key.layer); + if (weights->contains_expert(slot, prediction.key.layer, + prediction.key.expert)) { + return; + } + prefetch_requests.fetch_add(1U, std::memory_order_relaxed); + std::scoped_lock lock(prefetch_mutex); + if (prefetch_stopping || pending_prefetch.contains(prediction.key)) { + return; + } + if (prefetch_queue.size() >= prefetch_queue_limit) { + prefetch_dropped.fetch_add(1U, std::memory_order_relaxed); + return; + } + pending_prefetch.insert(prediction.key); + prefetch_queue.push_back({prediction.key, slot}); + prefetch_ready.notify_one(); + } + + [[nodiscard]] static std::uint64_t route_request_key( + const Glm53SequenceState* sequence, std::uint32_t position) noexcept { + // Each logical token owns one transition chain. Prompt execution is + // layer-major, so using only the sequence address would connect rows + // in execution order instead of connecting adjacent layers. + auto value = static_cast( + reinterpret_cast(sequence)); + value ^= static_cast(position) + + 0x9e3779b97f4a7c15ULL + (value << 6U) + (value >> 2U); + return value == 0U ? 1U : value; + } + + void observe_route(std::uint32_t layer, + std::span selected, + std::uint64_t request, std::uint32_t position, + bool schedule_prefetch) { + if (prefetch_prediction_limit == 0U || request == 0U || + layer >= kLayers) { + return; + } + RouteEvent event; + event.request = request; + event.token_position = position; + event.layer = layer; + event.phase = RoutePhase::Decode; + event.experts.reserve(selected.size()); + event.coefficients.reserve(selected.size()); + for (const auto& route : selected) { + event.experts.push_back(route.expert); + event.coefficients.push_back(route.weight); + } + route_predictor.observe(event); + if (!schedule_prefetch) return; + for (const auto& prediction : route_predictor.predict( + event, prefetch_prediction_limit, + prefetch_minimum_confidence)) { + // One-layer lookahead is the useful overlap window: farther + // predictions consume cache before their demand and are more + // likely to evict a nearer expert. + if (prediction.key.layer == layer + 1U) { + request_prefetch(prediction); + } + } + } + + [[nodiscard]] std::size_t slot_for(std::uint32_t layer) const noexcept { + const auto target_layer = std::min(layer, kLayers - 1U); + return device_schedule[target_layer % device_schedule.size()]; + } + + [[nodiscard]] int device_for(std::uint32_t layer) const noexcept { + return devices[slot_for(layer)]; + } + + [[nodiscard]] ParseResult>> + host_tensor(std::string_view name, std::uint64_t elements) { + ParseResult>> result; + const std::string key(name); + { + std::scoped_lock lock(host_tensor_mutex); + const auto found = host_tensors.find(key); + if (found != host_tensors.end()) { + if (found->second->size() != elements) { + result.errors.push_back( + "GLM-5.3 cached host tensor has an invalid extent: " + key); + } else { + result.value = found->second; + } + return result; + } + } + auto loaded = checkpoint->read_f32(name, elements); + if (!loaded.ok()) { + result.errors = std::move(loaded.errors); + return result; + } + auto value = std::make_shared>( + std::move(loaded.value)); + { + std::scoped_lock lock(host_tensor_mutex); + const auto [found, inserted] = host_tensors.emplace(key, value); + result.value = inserted ? std::move(value) : found->second; + } + return result; + } + + [[nodiscard]] ParseResult host_fp8_linear( + std::string_view base, std::uint32_t rows, + std::uint32_t columns) const { + ParseResult result; + const auto weight_name = std::string(base) + ".weight"; + const auto scale_name = std::string(base) + ".weight_scale_inv"; + const auto* descriptor = checkpoint->find(weight_name); + const auto* scale_descriptor = checkpoint->find(scale_name); + const auto scale_rows = (rows + 127U) / 128U; + const auto scale_columns = (columns + 127U) / 128U; + if (descriptor == nullptr || scale_descriptor == nullptr || + descriptor->source_dtype != SafetensorsDtype::F8E4M3 || + descriptor->source_shape != + std::vector{rows, columns} || + scale_descriptor->source_dtype != SafetensorsDtype::F32 || + scale_descriptor->source_shape != + std::vector{scale_rows, scale_columns}) { + result.errors.push_back( + "GLM-5.3 host expert has an invalid FP8 linear: " + + std::string(base)); + return result; + } + auto weight_payload = checkpoint->view(weight_name); + auto scales = checkpoint->view(scale_name); + if (!weight_payload.ok()) { + result.errors = std::move(weight_payload.errors); + return result; + } + if (!scales.ok()) { + result.errors = std::move(scales.errors); + return result; + } + if (weight_payload.value.size_bytes() != + static_cast(rows) * columns || + scales.value.size_bytes() != + static_cast(scale_rows) * scale_columns * + sizeof(float) || + reinterpret_cast(scales.value.data()) % + alignof(float) != 0U) { + result.errors.push_back( + "GLM-5.3 host expert mapped payload is mis-sized"); + return result; + } + result.value = { + weight_payload.value, + std::span( + reinterpret_cast(scales.value.data()), + static_cast(scale_rows) * scale_columns), + rows, columns}; + return result; + } + + [[nodiscard]] ValidationResult host_moe( + std::string_view prefix, std::span routed, + std::span input, std::span output) { + ValidationResult result; + constexpr std::uint32_t intermediate = 2048U; + constexpr std::size_t expert_count = 9U; + if (!host_moe_active || host_moe_workers == nullptr || + routed.size() != 8U || input.size() != kHidden || + output.size() != kHidden) { + return {{"GLM-5.3 host MoE command has an invalid shape"}}; + } + const auto started = std::chrono::steady_clock::now(); + struct Expert { + Glm53HostFp8Linear gate; + Glm53HostFp8Linear up; + Glm53HostFp8Linear down; + }; + std::array experts; + for (std::size_t index = 0U; index < expert_count; ++index) { + const auto module = index < routed.size() + ? std::string(prefix) + "experts." + + std::to_string(routed[index].expert) + "." + : std::string(prefix) + "shared_experts."; + auto gate = host_fp8_linear(module + "gate_proj", intermediate, + kHidden); + auto up = host_fp8_linear(module + "up_proj", intermediate, + kHidden); + auto down = host_fp8_linear(module + "down_proj", kHidden, + intermediate); + if (!gate.ok() || !up.ok() || !down.ok()) { + if (!gate.ok()) append(result.errors, std::move(gate.errors)); + if (!up.ok()) append(result.errors, std::move(up.errors)); + if (!down.ok()) append(result.errors, std::move(down.errors)); + return result; + } + experts[index] = {gate.value, up.value, down.value}; + } + std::vector quantized_input(input.begin(), input.end()); + glm53_quantize_activation(quantized_input); + std::vector activations(expert_count * intermediate); + const auto gate_up = host_moe_workers->parallel_for( + expert_count * intermediate, [&](std::size_t task) { + const auto expert = task / intermediate; + const auto row = task % intermediate; + const auto& module = experts[expert]; + const auto scale_columns = kHidden / 128U; + const auto* gate_weights = module.gate.weights.data() + + row * kHidden; + const auto* up_weights = module.up.weights.data() + + row * kHidden; + const auto* gate_scales = module.gate.scales.data() + + (row / 128U) * scale_columns; + const auto* up_scales = module.up.scales.data() + + (row / 128U) * scale_columns; + auto gate = bf16_round_f32(glm53_host_fp8_dot( + gate_weights, gate_scales, quantized_input)); + auto up = bf16_round_f32(glm53_host_fp8_dot( + up_weights, up_scales, quantized_input)); + gate = std::min(gate, 10.0F); + up = std::clamp(up, -10.0F, 10.0F); + activations[expert * intermediate + row] = + bf16_round_f32(gate * sigmoid(gate) * up); + }); + if (!gate_up.ok()) return gate_up; + for (std::size_t expert = 0U; expert < expert_count; ++expert) { + glm53_quantize_activation(std::span(activations).subspan( + expert * intermediate, intermediate)); + } + std::vector expert_outputs(expert_count * kHidden); + const auto down = host_moe_workers->parallel_for( + expert_count * kHidden, [&](std::size_t task) { + const auto expert = task / kHidden; + const auto row = task % kHidden; + const auto& module = experts[expert].down; + const auto scale_columns = intermediate / 128U; + const auto* weight_row = module.weights.data() + + row * intermediate; + const auto* scales = module.scales.data() + + (row / 128U) * scale_columns; + expert_outputs[expert * kHidden + row] = bf16_round_f32( + glm53_host_fp8_dot( + weight_row, scales, + std::span(activations).subspan( + expert * intermediate, intermediate))); + }); + if (!down.ok()) return down; + std::copy_n(expert_outputs.begin() + 8U * kHidden, kHidden, + output.begin()); + for (std::size_t expert = 0U; expert < routed.size(); ++expert) { + for (std::size_t column = 0U; column < kHidden; ++column) { + output[column] = bf16_round_f32( + output[column] + bf16_round_f32( + routed[expert].weight * + expert_outputs[expert * kHidden + column])); + } + } + host_moe_calls.fetch_add(1U, std::memory_order_relaxed); + host_moe_nanoseconds.fetch_add( + static_cast(std::chrono::duration_cast< + std::chrono::nanoseconds>(std::chrono::steady_clock::now() - + started).count()), + std::memory_order_relaxed); + return result; + } + + [[nodiscard]] ValidationResult host_moe_page( + std::string_view prefix, + std::span> routes, + std::span input, std::span output) { + ValidationResult result; + constexpr std::uint32_t intermediate = 2048U; + constexpr std::size_t routes_per_row = 8U; + constexpr std::size_t outputs_per_row = routes_per_row + 1U; + const auto rows = routes.size(); + if (!host_moe_active || host_moe_workers == nullptr || rows == 0U || + input.size() != rows * kHidden || + output.size() != rows * kHidden) { + return {{"GLM-5.3 host page MoE command has an invalid shape"}}; + } + const auto started = std::chrono::steady_clock::now(); + struct Expert { + Glm53HostFp8Linear gate; + Glm53HostFp8Linear up; + Glm53HostFp8Linear down; + }; + struct Assignment { + std::size_t input_row{}; + std::size_t output_slot{}; + }; + struct Group { + std::uint32_t expert{}; + bool shared{}; + Expert module; + std::vector assignments; + }; + + std::array group_for_expert; + group_for_expert.fill(std::numeric_limits::max()); + std::vector groups; + groups.reserve(std::min(288U, rows * routes_per_row) + 1U); + for (std::size_t row = 0U; row < rows; ++row) { + for (std::size_t route = 0U; route < routes_per_row; ++route) { + const auto expert = routes[row][route].expert; + if (expert >= group_for_expert.size()) { + return {{"GLM-5.3 host page route is out of range"}}; + } + auto& group_index = group_for_expert[expert]; + if (group_index == std::numeric_limits::max()) { + group_index = groups.size(); + groups.push_back({expert, false, {}, {}}); + } + groups[group_index].assignments.push_back( + {row, row * outputs_per_row + route}); + } + } + groups.push_back({0U, true, {}, {}}); + auto& shared = groups.back(); + shared.assignments.reserve(rows); + for (std::size_t row = 0U; row < rows; ++row) { + shared.assignments.push_back( + {row, row * outputs_per_row + routes_per_row}); + } + + for (auto& group : groups) { + const auto module = group.shared + ? std::string(prefix) + "shared_experts." + : std::string(prefix) + "experts." + + std::to_string(group.expert) + "."; + auto gate = host_fp8_linear(module + "gate_proj", intermediate, + kHidden); + auto up = host_fp8_linear(module + "up_proj", intermediate, + kHidden); + auto down = host_fp8_linear(module + "down_proj", kHidden, + intermediate); + if (!gate.ok() || !up.ok() || !down.ok()) { + if (!gate.ok()) append(result.errors, std::move(gate.errors)); + if (!up.ok()) append(result.errors, std::move(up.errors)); + if (!down.ok()) append(result.errors, std::move(down.errors)); + return result; + } + group.module = {gate.value, up.value, down.value}; + } + + std::vector quantized_input(input.begin(), input.end()); + result = host_moe_workers->parallel_for(rows, [&](std::size_t row) { + glm53_quantize_activation( + std::span(quantized_input) + .subspan(row * kHidden, kHidden)); + }); + if (!result.ok()) return result; + + const auto output_slots = rows * outputs_per_row; + std::vector activations(output_slots * intermediate); + result = host_moe_workers->parallel_for( + groups.size() * intermediate, [&](std::size_t task) { + const auto group_index = task / intermediate; + const auto projection_row = task % intermediate; + const auto& group = groups[group_index]; + const auto scale_columns = kHidden / 128U; + const auto* gate_weights = group.module.gate.weights.data() + + projection_row * kHidden; + const auto* up_weights = group.module.up.weights.data() + + projection_row * kHidden; + const auto* gate_scales = group.module.gate.scales.data() + + (projection_row / 128U) * scale_columns; + const auto* up_scales = group.module.up.scales.data() + + (projection_row / 128U) * scale_columns; + for (const auto& assignment : group.assignments) { + const auto source = std::span(quantized_input) + .subspan(assignment.input_row * kHidden, kHidden); + auto gate = bf16_round_f32(glm53_host_fp8_dot( + gate_weights, gate_scales, source)); + auto up = bf16_round_f32(glm53_host_fp8_dot( + up_weights, up_scales, source)); + gate = std::min(gate, 10.0F); + up = std::clamp(up, -10.0F, 10.0F); + activations[assignment.output_slot * intermediate + + projection_row] = + bf16_round_f32(gate * sigmoid(gate) * up); + } + }); + if (!result.ok()) return result; + result = host_moe_workers->parallel_for( + output_slots, [&](std::size_t slot) { + glm53_quantize_activation( + std::span(activations) + .subspan(slot * intermediate, intermediate)); + }); + if (!result.ok()) return result; + + std::vector expert_outputs(output_slots * kHidden); + result = host_moe_workers->parallel_for( + groups.size() * kHidden, [&](std::size_t task) { + const auto group_index = task / kHidden; + const auto projection_row = task % kHidden; + const auto& group = groups[group_index]; + const auto& projection = group.module.down; + const auto scale_columns = intermediate / 128U; + const auto* projection_weights = projection.weights.data() + + projection_row * intermediate; + const auto* projection_scales = projection.scales.data() + + (projection_row / 128U) * scale_columns; + for (const auto& assignment : group.assignments) { + expert_outputs[assignment.output_slot * kHidden + + projection_row] = bf16_round_f32( + glm53_host_fp8_dot( + projection_weights, projection_scales, + std::span(activations).subspan( + assignment.output_slot * intermediate, + intermediate))); + } + }); + if (!result.ok()) return result; + result = host_moe_workers->parallel_for( + rows * kHidden, [&](std::size_t task) { + const auto row = task / kHidden; + const auto column = task % kHidden; + auto value = expert_outputs[ + (row * outputs_per_row + routes_per_row) * kHidden + + column]; + for (std::size_t route = 0U; route < routes_per_row; ++route) { + value = bf16_round_f32( + value + bf16_round_f32( + routes[row][route].weight * + expert_outputs[ + (row * outputs_per_row + route) * kHidden + + column])); + } + output[row * kHidden + column] = value; + }); + if (!result.ok()) return result; + host_moe_calls.fetch_add(rows, std::memory_order_relaxed); + host_moe_nanoseconds.fetch_add( + static_cast(std::chrono::duration_cast< + std::chrono::nanoseconds>(std::chrono::steady_clock::now() - + started).count()), + std::memory_order_relaxed); + return result; + } + + [[nodiscard]] ValidationResult wait_for_warmup() { + std::scoped_lock lock(warmup_mutex); + if (warmup_thread.joinable()) warmup_thread.join(); + return warmup_result; + } + + [[nodiscard]] std::size_t restore_prefix( + std::span tokens, Glm53SequenceState& state, + std::span logits, std::vector& base_hidden) { + std::scoped_lock lock(prefix_mutex); + PrefixEntry* best = nullptr; + for (auto& entry : prefix_cache) { + if (entry.tokens.size() > tokens.size() || + entry.logits.size() != logits.size() || + (best != nullptr && + entry.tokens.size() <= best->tokens.size()) || + !std::equal(entry.tokens.begin(), entry.tokens.end(), + tokens.begin())) { + continue; + } + best = &entry; + } + if (best == nullptr) return 0U; + state = best->state; + std::copy(best->logits.begin(), best->logits.end(), logits.begin()); + base_hidden = best->base_hidden; + best->recency = ++prefix_clock; + prefix_cache_hits.fetch_add(1U, std::memory_order_relaxed); + prefix_cache_tokens.fetch_add(best->tokens.size(), + std::memory_order_relaxed); + return best->tokens.size(); + } + + void store_prefix(std::span tokens, + const Glm53SequenceState& state, + std::span logits, + std::span base_hidden) { + if (tokens.empty() || state.token_count() != tokens.size()) return; + std::scoped_lock lock(prefix_mutex); + for (auto& entry : prefix_cache) { + if (entry.tokens.size() == tokens.size() && + std::equal(entry.tokens.begin(), entry.tokens.end(), + tokens.begin())) { + entry.state = state; + entry.logits.assign(logits.begin(), logits.end()); + entry.base_hidden.assign(base_hidden.begin(), base_hidden.end()); + entry.recency = ++prefix_clock; + return; + } + } + if (prefix_cache.size() >= prefix_cache_limit) { + const auto victim = std::min_element( + prefix_cache.begin(), prefix_cache.end(), + [](const PrefixEntry& left, const PrefixEntry& right) { + return left.recency < right.recency; + }); + if (victim != prefix_cache.end()) prefix_cache.erase(victim); + } + PrefixEntry entry; + entry.tokens.assign(tokens.begin(), tokens.end()); + entry.state = state; + entry.logits.assign(logits.begin(), logits.end()); + entry.base_hidden.assign(base_hidden.begin(), base_hidden.end()); + entry.recency = ++prefix_clock; + prefix_cache.push_back(std::move(entry)); + } + + [[nodiscard]] ValidationResult warmup() { + struct LinearTask { + std::string base; + std::string group; + std::uint64_t rows{}; + std::uint64_t columns{}; + std::uint32_t layer{}; + std::uint64_t weight_rows{}; + std::uint64_t weight_row_begin{}; + }; + struct HostTask { + std::string name; + std::uint64_t elements{}; + }; + ValidationResult result; + std::vector linear_tasks; + std::vector> device_tasks(devices.size()); + std::vector host_tasks; + for (const auto& tensor : checkpoint->manifest().tensors) { + if (tensor.role == Glm53TensorRole::Vision || + tensor.role == Glm53TensorRole::AttentionIndexer || + tensor.role == Glm53TensorRole::RoutedExpert || + tensor.role == Glm53TensorRole::Embedding || + tensor.name.find(".layers.45.mlp.experts.") != + std::string::npos || + tensor.component == Glm53TensorComponent::Scale) { + continue; + } + const bool linear = tensor.name.ends_with(".weight") && + tensor.source_shape.size() == 2U; + if (linear) { + const auto layer = tensor.layer >= 0 + ? static_cast(tensor.layer) + : kLayers - 1U; + const auto base = tensor.name.substr( + 0U, tensor.name.size() - 7U); + if (base == "lm_head" && lm_head_ranges.size() > 1U) { + for (std::size_t slot = 0U; + slot < lm_head_ranges.size(); ++slot) { + const auto range = lm_head_ranges[slot]; + device_tasks[slot].push_back({ + base, base, range.count, tensor.source_shape[1], + layer, tensor.source_shape[0], range.begin}); + } + continue; + } + if (full_tensor_parallel_active) { + const auto ranges = weighted_row_ranges( + tensor.source_shape[0], weight_capacities, 128U); + if (ranges.size() == devices.size() && + std::all_of(ranges.begin(), ranges.end(), + [](const auto& range) { + return range.count != 0U && + range.begin % 128U == 0U; + })) { + for (std::size_t slot = 0U; slot < ranges.size(); + ++slot) { + device_tasks[slot].push_back({ + base, base, ranges[slot].count, + tensor.source_shape[1], layer, + tensor.source_shape[0], ranges[slot].begin}); + } + continue; + } + } + linear_tasks.push_back({ + base, projection_group_key(base, tensor.role), + tensor.source_shape[0], tensor.source_shape[1], layer}); + continue; + } + if (tensor.source_dtype != SafetensorsDtype::Bf16 && + tensor.source_dtype != SafetensorsDtype::F16 && + tensor.source_dtype != SafetensorsDtype::F32) { + continue; + } + std::uint64_t elements = 1U; + bool valid = !tensor.source_shape.empty(); + for (const auto dimension : tensor.source_shape) { + if (dimension == 0U || + elements > std::numeric_limits::max() / + dimension) { + valid = false; + break; + } + elements *= dimension; + } + if (valid) host_tasks.push_back({tensor.name, elements}); + } + + std::map> linear_groups; + for (auto& task : linear_tasks) { + linear_groups[task.group].push_back(std::move(task)); + } + const bool parallel = projection_workers != nullptr && + batched_projections_enabled() && + cross_gpu_projections_enabled(devices) && + devices.size() > 1U; + for (auto& [group, tasks] : linear_groups) { + static_cast(group); + if (!parallel || tasks.size() == 1U) { + for (auto& task : tasks) { + device_tasks[slot_for(task.layer)].push_back(std::move(task)); + } + continue; + } + std::vector keys; + std::vector costs; + keys.reserve(tasks.size()); + costs.reserve(tasks.size()); + for (const auto& task : tasks) { + keys.push_back(task.base); + costs.push_back(checkpoint->cuda_linear_storage_bytes(task.base)); + } + const auto slots = glm53_projection_slots( + keys, costs, weight_capacities, slot_for(tasks.front().layer)); + if (slots.size() != tasks.size()) { + return {{"GLM-5.3 projection warmup assignment is invalid"}}; + } + for (std::size_t index = 0U; index < tasks.size(); ++index) { + device_tasks[slots[index]].push_back(std::move(tasks[index])); + } + } + + std::vector device_results(devices.size()); + std::vector admitted(devices.size()); + std::vector skipped(devices.size()); + std::atomic next_slot{}; + const auto load_devices = [&] { + for (;;) { + const auto slot = next_slot.fetch_add(1U, + std::memory_order_relaxed); + if (slot >= devices.size()) return; + for (const auto& task : device_tasks[slot]) { + bool kept = false; + auto loaded = task.weight_rows == 0U + ? weights->preload(slot, task.base, task.rows, + task.columns, kept) + : weights->preload_slice( + slot, task.base, task.weight_rows, task.columns, + task.weight_row_begin, task.rows, kept); + if (!loaded.ok()) { + append(device_results[slot].errors, + std::move(loaded.errors)); + break; + } + kept ? ++admitted[slot] : ++skipped[slot]; + } + if (device_results[slot].ok()) { + auto ordered = cuda.synchronize_uploads(devices[slot]); + if (!ordered.ok()) { + append(device_results[slot].errors, + std::move(ordered.errors)); + } + } + } + }; + const auto workers = std::min( + devices.size(), host_hardware_profile().worker_threads(0.1)); + std::vector loaders; + loaders.reserve(workers); + for (std::size_t worker = 0U; worker < workers; ++worker) { + loaders.emplace_back(load_devices); + } + // Host-resident norms, convolution taps and mHC projections are only + // 67 MiB for this checkpoint. Load them while independent PCIe links + // receive their layer-split spine weights. + for (const auto& task : host_tasks) { + auto loaded = host_tensor(task.name, task.elements); + if (!loaded.ok()) { + append(result.errors, std::move(loaded.errors)); + break; + } + } + for (auto& loader : loaders) loader.join(); + for (auto& device_result : device_results) { + append(result.errors, std::move(device_result.errors)); + } + if (result.ok() && resident_execution_active) { + for (std::uint32_t layer = 0U; layer < kLayers; ++layer) { + const auto prefix = "model.language_model.layers." + + std::to_string(layer) + "."; + const auto upload_mhc = [&](const std::string& mhc, + const std::string& norm, + CudaDsv4MhcWeights& destination) + -> ValidationResult { + auto projection = host_tensor(mhc + "_fn", 24U * 16384U); + auto base = host_tensor(mhc + "_base", 24U); + auto scale = host_tensor(mhc + "_scale", 3U); + auto norm_weight = host_tensor(norm, kHidden); + ValidationResult status; + if (!projection.ok() || !base.ok() || !scale.ok() || + !norm_weight.ok()) { + append(status.errors, std::move(projection.errors)); + append(status.errors, std::move(base.errors)); + append(status.errors, std::move(scale.errors)); + append(status.errors, std::move(norm_weight.errors)); + return status; + } + return cuda.upload_dsv4_mhc_weights( + device_for(layer), *projection.value, *scale.value, + *base.value, *norm_weight.value, destination); + }; + auto status = upload_mhc( + prefix + "hc_attn", prefix + "input_layernorm.weight", + resident_layers[layer].attention); + if (!status.ok()) return status; + status = upload_mhc( + prefix + "hc_ffn", + prefix + "post_attention_layernorm.weight", + resident_layers[layer].feedforward); + if (!status.ok()) return status; + } + } + if (config.verbose) { + const auto stats = weights->stats(); + for (std::size_t slot = 0U; slot < devices.size(); ++slot) { + std::cerr << "[glm53-load] cuda=" << devices[slot] + << " resident_linears=" << admitted[slot] + << " streamed_linears=" << skipped[slot] + << " pinned_bytes=" << stats.pinned[slot] + << " cache_capacity_bytes=" << stats.capacity[slot] + << '\n'; + } + } + return result; + } + + [[nodiscard]] ValidationResult reset_sequence( + Glm53SequenceState& sequence) const { + return sequence.reset(config.maximum_context_tokens, 64U); + } + + [[nodiscard]] ValidationResult linear( + std::string_view base, std::span input, + std::uint32_t rows, std::uint32_t columns, + std::span output, std::uint32_t layer, + bool bf16_output = true) { + ValidationResult result; + if (input.size() != static_cast(columns) * rows || + output.empty()) { + result.errors.push_back("GLM-5.3 linear activation shape is invalid for " + + std::string(base)); + return result; + } + const auto output_columns = output.size() / rows; + if (output_columns * rows != output.size()) { + result.errors.push_back("GLM-5.3 linear output shape is invalid for " + + std::string(base)); + return result; + } + // PyTorch returns every published BF16/FP8 linear at the model's BF16 + // activation dtype. Keep that boundary even though the host-facing + // CUDA API transports activations as float. + const Glm53WeightCache::LinearRequest request{ + base, output_columns, columns, input, rows, output, bf16_output}; + return linear_batch( + std::span(&request, 1U), + layer); + } + + [[nodiscard]] ValidationResult linear_batch( + std::span requests, + std::uint32_t layer) { + if (requests.empty()) { + return {{"GLM-5.3 linear projection batch is empty"}}; + } + for (const auto& request : requests) { + if (request.rows == 0U || + request.input.size() != + static_cast(request.input_columns) * + request.rows || + request.output.size() != + static_cast(request.output_columns) * + request.rows) { + return {{"GLM-5.3 linear projection batch has an invalid shape"}}; + } + } + if (full_tensor_parallel_active) { + std::vector> ranges(requests.size()); + bool eligible = true; + for (std::size_t index = 0U; index < requests.size(); ++index) { + const auto& request = requests[index]; + ranges[index] = weighted_row_ranges( + request.output_columns, weight_capacities, 128U); + if (request.weight_rows != 0U || + ranges[index].size() != devices.size()) { + eligible = false; + break; + } + for (const auto range : ranges[index]) { + if (range.count == 0U || range.begin % 128U != 0U || + checkpoint->cuda_linear_slice_storage_bytes( + request.base, range.begin, range.count) == 0U) { + eligible = false; + break; + } + } + if (!eligible) break; + } + if (eligible) { + std::vector>> shards( + requests.size(), + std::vector>(devices.size())); + std::vector> + groups(devices.size()); + for (std::size_t index = 0U; index < requests.size(); ++index) { + for (std::size_t slot = 0U; slot < devices.size(); ++slot) { + const auto range = ranges[index][slot]; + shards[index][slot].resize( + static_cast(requests[index].rows) * + range.count); + groups[slot].push_back({ + requests[index].base, range.count, + requests[index].input_columns, + requests[index].input, requests[index].rows, + shards[index][slot], requests[index].bf16_output, + requests[index].output_columns, range.begin}); + } + } + std::vector device_results(devices.size()); + auto dispatched = projection_workers->parallel_for_addressed( + devices.size(), [&](std::size_t slot) { + device_results[slot] = + weights->matmul_batch(slot, groups[slot]); + }); + if (!dispatched.ok()) return dispatched; + for (auto& device_result : device_results) { + if (!device_result.ok()) return device_result; + } + for (std::size_t index = 0U; index < requests.size(); ++index) { + for (std::size_t slot = 0U; slot < devices.size(); ++slot) { + const auto range = ranges[index][slot]; + for (std::uint32_t row = 0U; + row < requests[index].rows; ++row) { + std::copy_n( + shards[index][slot].begin() + + static_cast( + static_cast(row) * + range.count), + range.count, + requests[index].output.begin() + + static_cast( + static_cast(row) * + requests[index].output_columns + + range.begin)); + } + } + } + parallel_projection_batches.fetch_add( + 1U, std::memory_order_relaxed); + parallel_projection_requests.fetch_add( + requests.size(), std::memory_order_relaxed); + return {}; + } + } + if (!batched_projections_enabled()) { + for (const auto& request : requests) { + auto projected = weights->matmul( + slot_for(layer), request.base, request.output_columns, + request.input_columns, request.input, request.rows, + request.output, request.bf16_output); + if (!projected.ok()) return projected; + } + return {}; + } + if (cross_gpu_projections_enabled(devices) && + projection_workers != nullptr && + devices.size() > 1U && requests.size() > 1U) { + std::vector keys; + std::vector costs; + keys.reserve(requests.size()); + costs.reserve(requests.size()); + for (const auto& request : requests) { + keys.push_back(request.base); + costs.push_back( + checkpoint->cuda_linear_storage_bytes(request.base)); + } + const auto slots = glm53_projection_slots( + keys, costs, weight_capacities, slot_for(layer)); + if (slots.size() != requests.size()) { + return {{"GLM-5.3 parallel projection assignment is invalid"}}; + } + std::vector> groups( + devices.size()); + for (std::size_t index = 0U; index < requests.size(); ++index) { + groups[slots[index]].push_back(requests[index]); + } + std::vector device_results(devices.size()); + auto dispatched = projection_workers->parallel_for_addressed( + devices.size(), [&](std::size_t slot) { + if (!groups[slot].empty()) { + device_results[slot] = + weights->matmul_batch(slot, groups[slot]); + } + }); + if (!dispatched.ok()) return dispatched; + ValidationResult joined; + std::uint64_t active_slots = 0U; + for (std::size_t slot = 0U; slot < device_results.size(); ++slot) { + if (!groups[slot].empty()) ++active_slots; + append(joined.errors, std::move(device_results[slot].errors)); + } + if (joined.ok() && active_slots > 1U) { + parallel_projection_batches.fetch_add( + 1U, std::memory_order_relaxed); + parallel_projection_requests.fetch_add( + static_cast(requests.size()), + std::memory_order_relaxed); + } + return joined; + } + return weights->matmul_batch(slot_for(layer), requests); + } + + [[nodiscard]] ValidationResult norm( + std::span output, std::span input, + std::string_view weight_name) { + auto weight = host_tensor(weight_name, input.size()); + if (!weight.ok()) return {std::move(weight.errors)}; + auto result = kimi_rms_norm(output, input, *weight.value, 1.0e-5F); + if (result.ok()) round_bf16(output); + return result; + } + + [[nodiscard]] ValidationResult norm_rows( + std::span output, std::span input, + std::uint32_t rows, std::uint32_t columns, + std::string_view weight_name) { + ValidationResult result; + if (rows == 0U || input.size() != output.size() || + input.size() != static_cast(rows) * columns) { + result.errors.emplace_back("GLM-5.3 RMSNorm page shape is invalid"); + return result; + } + auto weight = host_tensor(weight_name, columns); + if (!weight.ok()) return {std::move(weight.errors)}; + for (std::uint32_t row = 0U; row < rows; ++row) { + const auto begin = static_cast(row) * columns; + auto normalized = kimi_rms_norm( + output.subspan(begin, columns), input.subspan(begin, columns), + *weight.value, 1.0e-5F); + if (!normalized.ok()) return normalized; + round_bf16(output.subspan(begin, columns)); + } + return result; + } + + [[nodiscard]] ValidationResult mhc_pre( + std::span collapsed, Dsv4MhcMix& mix, + std::span streams, const std::string& prefix) { + ValidationResult result; + auto projection = host_tensor(prefix + "_fn", 24U * 16384U); + auto base = host_tensor(prefix + "_base", 24U); + auto scale = host_tensor(prefix + "_scale", 3U); + if (!projection.ok() || !base.ok() || !scale.ok()) { + append(result.errors, std::move(projection.errors)); + append(result.errors, std::move(base.errors)); + append(result.errors, std::move(scale.errors)); + return result; + } + double square_sum = 0.0; + for (const auto value : streams) square_sum += static_cast(value) * value; + const auto reciprocal = 1.0F / std::sqrt( + static_cast(square_sum / + static_cast(streams.size())) + 1.0e-5F); + std::vector projected(24U, 0.0F); + for (std::size_t row = 0U; row < projected.size(); ++row) { + double sum = 0.0; + for (std::size_t column = 0U; column < streams.size(); ++column) { + sum += static_cast((*projection.value)[row * streams.size() + column]) * + streams[column]; + } + projected[row] = static_cast(sum) * reciprocal; + } + auto split = dsv4_mhc_split_sinkhorn_f32( + projected, *scale.value, *base.value, kMhc, 20U, 1.0e-6F); + if (!split.ok()) return {std::move(split.errors)}; + mix = std::move(split.value); + round_bf16(mix.post); + round_bf16(mix.combination); + std::fill(collapsed.begin(), collapsed.end(), 0.0F); + for (std::size_t stream = 0U; stream < kMhc; ++stream) { + for (std::size_t column = 0U; column < kHidden; ++column) { + collapsed[column] += mix.pre[stream] * + streams[stream * kHidden + column]; + } + } + round_bf16(collapsed); + return result; + } + + [[nodiscard]] ValidationResult attention_kda( + std::span output, std::span input, + std::uint32_t layer, const std::string& attention, + Glm53SequenceState& sequence, CudaBuffer* device_state = nullptr) { + ValidationResult result; + if (device_state != nullptr && fused_kda_enabled()) { + CudaGlm53KdaRequest request; + request.state = device_state; + request.input = input; + request.heads = kHeads; + request.head_dim = kLinearHead; + request.convolution_kernel = 4U; + return weights->kda_decode( + slot_for(layer), attention, request, output); + } + std::vector query(kLinearWidth), key(kLinearWidth), + value(kLinearWidth), low(kLinearHead), beta(kHeads), + gate_low(kLinearHead); + const std::array first_bases{ + attention + "q_proj", attention + "k_proj", attention + "v_proj", + attention + "f_a_proj", attention + "b_proj", + attention + "g_a_proj"}; + const std::array first{ + {{first_bases[0], kLinearWidth, kHidden, input, 1U, query, true}, + {first_bases[1], kLinearWidth, kHidden, input, 1U, key, true}, + {first_bases[2], kLinearWidth, kHidden, input, 1U, value, true}, + {first_bases[3], kLinearHead, kHidden, input, 1U, low, true}, + {first_bases[4], kHeads, kHidden, input, 1U, beta, true}, + {first_bases[5], kLinearHead, kHidden, input, 1U, gate_low, true}}}; + result = linear_batch(first, layer); + if (!result.ok()) return result; + if (device_state == nullptr) { + for (std::uint32_t projection = 0U; projection < 3U; ++projection) { + auto taps = host_tensor( + attention + (projection == 0U ? "q_conv1d.weight" + : projection == 1U ? "k_conv1d.weight" + : "v_conv1d.weight"), + static_cast(kLinearWidth) * 4U); + if (!taps.ok()) return {std::move(taps.errors)}; + auto& values = projection == 0U ? query : projection == 1U ? key : value; + auto convolved = values; + result = kimi_short_conv_step( + convolved, values, *taps.value, + sequence.convolution(layer, projection), 4U); + if (!result.ok()) return result; + values = std::move(convolved); + round_bf16(values); + } + } + std::vector forget(kLinearWidth), gate(kLinearWidth); + const std::array second_bases{ + attention + "f_b_proj", attention + "g_b_proj"}; + const std::array second{ + {{second_bases[0], kLinearWidth, kLinearHead, low, 1U, forget, true}, + {second_bases[1], kLinearWidth, kLinearHead, gate_low, 1U, gate, + true}}}; + result = linear_batch(second, layer); + if (!result.ok()) return result; + for (auto& element : beta) { + element = bf16_round_f32(sigmoid(element)); + } + auto a_log = host_tensor(attention + "A_log", kHeads); + auto dt_bias = host_tensor(attention + "dt_bias", kLinearWidth); + auto o_norm = host_tensor(attention + "o_norm.weight", kLinearHead); + if (!a_log.ok() || !dt_bias.ok() || !o_norm.ok()) { + append(result.errors, std::move(a_log.errors)); + append(result.errors, std::move(dt_bias.errors)); + append(result.errors, std::move(o_norm.errors)); + return result; + } + if (device_state != nullptr) { + CudaGlm53KdaRequest request; + request.state = device_state; + request.query = query; + request.key = key; + request.value = value; + request.forget = forget; + request.beta = beta; + request.gate = gate; + request.heads = kHeads; + request.head_dim = kLinearHead; + request.convolution_kernel = 4U; + return weights->kda_decode(slot_for(layer), attention, request, + output); + } + std::vector heads_out(kLinearWidth); + const auto query_scale = 1.0F / std::sqrt(static_cast(kLinearHead)); + for (std::uint32_t head = 0U; head < kHeads; ++head) { + const auto begin = static_cast(head) * kLinearHead; + auto q = std::span(query).subspan(begin, kLinearHead); + auto k = std::span(key).subspan(begin, kLinearHead); + result = kimi_l2_normalize(q, 1.0e-6F); + if (!result.ok()) return result; + result = kimi_l2_normalize(k, 1.0e-6F); + if (!result.ok()) return result; + for (auto& element : q) element *= query_scale; + std::vector decay(kLinearHead); + result = kimi_kda_log_decay( + decay, std::span(forget).subspan(begin, kLinearHead), + std::span(*dt_bias.value).subspan(begin, kLinearHead), + (*a_log.value)[head], -5.0F); + if (!result.ok()) return result; + for (auto& element : decay) element = std::exp(element); + std::vector raw(kLinearHead); + auto state = sequence.recurrent(layer).subspan( + static_cast(head) * kLinearHead * kLinearHead, + static_cast(kLinearHead) * kLinearHead); + result = kimi_kda_step( + raw, state, q, k, + std::span(value).subspan(begin, kLinearHead), + decay, beta[head], kLinearHead, kLinearHead); + if (!result.ok()) return result; + round_bf16(raw); + result = kimi_kda_output_norm( + std::span(heads_out).subspan(begin, kLinearHead), raw, + std::span(gate).subspan(begin, kLinearHead), + *o_norm.value, 1.0e-5F); + if (!result.ok()) return result; + round_bf16(std::span(heads_out).subspan(begin, kLinearHead)); + } + return linear(attention + "o_proj", heads_out, 1U, kLinearWidth, + output, layer); + } + + [[nodiscard]] ValidationResult attention_kda_page( + std::span output, std::span input, + std::uint32_t rows, std::uint32_t layer, + const std::string& attention, Glm53SequenceState& sequence) { + ValidationResult result; + const auto wide_elements = static_cast(rows) * kLinearWidth; + std::vector query(wide_elements), key(wide_elements), + value(wide_elements), + low(static_cast(rows) * kLinearHead), + beta(static_cast(rows) * kHeads), + gate_low(static_cast(rows) * kLinearHead); + const std::array first_bases{ + attention + "q_proj", attention + "k_proj", attention + "v_proj", + attention + "f_a_proj", attention + "b_proj", + attention + "g_a_proj"}; + const std::array first{ + {{first_bases[0], kLinearWidth, kHidden, input, rows, query, true}, + {first_bases[1], kLinearWidth, kHidden, input, rows, key, true}, + {first_bases[2], kLinearWidth, kHidden, input, rows, value, true}, + {first_bases[3], kLinearHead, kHidden, input, rows, low, true}, + {first_bases[4], kHeads, kHidden, input, rows, beta, true}, + {first_bases[5], kLinearHead, kHidden, input, rows, gate_low, true}}}; + result = linear_batch(first, layer); + if (!result.ok()) return result; + for (std::uint32_t projection = 0U; projection < 3U; ++projection) { + auto taps = host_tensor( + attention + (projection == 0U ? "q_conv1d.weight" + : projection == 1U ? "k_conv1d.weight" + : "v_conv1d.weight"), + static_cast(kLinearWidth) * 4U); + if (!taps.ok()) return {std::move(taps.errors)}; + auto& values = projection == 0U ? query : projection == 1U ? key : value; + for (std::uint32_t row = 0U; row < rows; ++row) { + const auto begin = static_cast(row) * kLinearWidth; + std::vector convolved(kLinearWidth); + result = kimi_short_conv_step( + convolved, + std::span(values).subspan(begin, kLinearWidth), + *taps.value, sequence.convolution(layer, projection), 4U); + if (!result.ok()) return result; + round_bf16(convolved); + std::copy(convolved.begin(), convolved.end(), + values.begin() + static_cast(begin)); + } + } + std::vector forget(wide_elements); + std::vector gate(wide_elements); + const std::array second_bases{ + attention + "f_b_proj", attention + "g_b_proj"}; + const std::array second{ + {{second_bases[0], kLinearWidth, kLinearHead, low, rows, forget, true}, + {second_bases[1], kLinearWidth, kLinearHead, gate_low, rows, gate, + true}}}; + result = linear_batch(second, layer); + if (!result.ok()) return result; + for (auto& element : beta) element = bf16_round_f32(sigmoid(element)); + auto a_log = host_tensor(attention + "A_log", kHeads); + auto dt_bias = host_tensor(attention + "dt_bias", kLinearWidth); + auto o_norm = host_tensor(attention + "o_norm.weight", kLinearHead); + if (!a_log.ok() || !dt_bias.ok() || !o_norm.ok()) { + append(result.errors, std::move(a_log.errors)); + append(result.errors, std::move(dt_bias.errors)); + append(result.errors, std::move(o_norm.errors)); + return result; + } + std::vector heads_out(wide_elements); + const auto query_scale = 1.0F / std::sqrt(static_cast(kLinearHead)); + // The chunk form exposes heads to the physical-core pool, but a page + // narrower than the runner count cannot amortize waking that pool. + // Derive the crossover from the discovered pool width rather than a + // token constant measured on one host. + if (kda_workers != nullptr && + rows >= std::min(kHeads, kda_workers->size())) { + std::vector failures(kHeads); + auto replayed = kda_workers->parallel_for( + kHeads, [&](std::size_t head) { + std::vector q(static_cast(rows) * + kLinearHead); + std::vector k(q.size()), v(q.size()), decay(q.size()); + std::vector head_beta(rows); + for (std::uint32_t row = 0U; row < rows; ++row) { + const auto source = static_cast(row) * + kLinearWidth + + head * kLinearHead; + const auto target = static_cast(row) * + kLinearHead; + auto q_row = std::span(q).subspan( + target, kLinearHead); + auto k_row = std::span(k).subspan( + target, kLinearHead); + std::copy_n(query.begin() + + static_cast(source), + kLinearHead, q_row.begin()); + std::copy_n(key.begin() + + static_cast(source), + kLinearHead, k_row.begin()); + std::copy_n(value.begin() + + static_cast(source), + kLinearHead, + v.begin() + static_cast(target)); + auto status = kimi_l2_normalize(q_row, 1.0e-6F); + if (!status.ok()) { + failures[head] = std::move(status); + return; + } + status = kimi_l2_normalize(k_row, 1.0e-6F); + if (!status.ok()) { + failures[head] = std::move(status); + return; + } + for (auto& element : q_row) element *= query_scale; + status = kimi_kda_log_decay( + std::span(decay).subspan(target, kLinearHead), + std::span(forget).subspan( + source, kLinearHead), + std::span(*dt_bias.value).subspan( + head * kLinearHead, kLinearHead), + (*a_log.value)[head], -5.0F); + if (!status.ok()) { + failures[head] = std::move(status); + return; + } + head_beta[row] = beta[ + static_cast(row) * kHeads + head]; + } + std::vector raw(q.size()); + auto state = sequence.recurrent(layer).subspan( + head * kLinearHead * kLinearHead, + static_cast(kLinearHead) * kLinearHead); + auto status = kimi_kda_chunk( + raw, state, q, k, v, decay, head_beta, rows, + kLinearHead, kLinearHead); + if (!status.ok()) { + failures[head] = std::move(status); + return; + } + for (std::uint32_t row = 0U; row < rows; ++row) { + const auto source = static_cast(row) * + kLinearHead; + const auto target = static_cast(row) * + kLinearWidth + + head * kLinearHead; + auto raw_row = std::span(raw).subspan( + source, kLinearHead); + round_bf16(raw_row); + status = kimi_kda_output_norm( + std::span(heads_out).subspan( + target, kLinearHead), + raw_row, + std::span(gate).subspan( + target, kLinearHead), + *o_norm.value, 1.0e-5F); + if (!status.ok()) { + failures[head] = std::move(status); + return; + } + round_bf16(std::span(heads_out).subspan( + target, kLinearHead)); + } + }); + if (!replayed.ok()) return replayed; + for (auto& failure : failures) { + if (!failure.ok()) return failure; + } + } else { + for (std::uint32_t row = 0U; row < rows; ++row) { + const auto row_begin = static_cast(row) * + kLinearWidth; + for (std::uint32_t head = 0U; head < kHeads; ++head) { + const auto begin = row_begin + + static_cast(head) * kLinearHead; + auto q = std::span(query).subspan(begin, kLinearHead); + auto k = std::span(key).subspan(begin, kLinearHead); + result = kimi_l2_normalize(q, 1.0e-6F); + if (!result.ok()) return result; + result = kimi_l2_normalize(k, 1.0e-6F); + if (!result.ok()) return result; + for (auto& element : q) element *= query_scale; + std::vector decay(kLinearHead); + result = kimi_kda_log_decay( + decay, + std::span(forget).subspan( + begin, kLinearHead), + std::span(*dt_bias.value).subspan( + static_cast(head) * kLinearHead, + kLinearHead), + (*a_log.value)[head], -5.0F); + if (!result.ok()) return result; + for (auto& element : decay) element = std::exp(element); + std::vector raw(kLinearHead); + auto state = sequence.recurrent(layer).subspan( + static_cast(head) * kLinearHead * + kLinearHead, + static_cast(kLinearHead) * kLinearHead); + result = kimi_kda_step( + raw, state, q, k, + std::span(value).subspan( + begin, kLinearHead), + decay, + beta[static_cast(row) * kHeads + head], + kLinearHead, kLinearHead); + if (!result.ok()) return result; + round_bf16(raw); + result = kimi_kda_output_norm( + std::span(heads_out).subspan( + begin, kLinearHead), + raw, + std::span(gate).subspan( + begin, kLinearHead), + *o_norm.value, 1.0e-5F); + if (!result.ok()) return result; + round_bf16(std::span(heads_out).subspan( + begin, kLinearHead)); + } + } + } + return linear(attention + "o_proj", heads_out, rows, kLinearWidth, + output, layer); + } + + [[nodiscard]] ValidationResult attention_mla( + std::span output, std::span input, + std::uint32_t layer, std::uint32_t position, + const std::string& attention, Glm53SequenceState& sequence) { + ValidationResult result; + std::vector q_rank(kQueryRank), query(kMlaWidth), latent(kKvRank); + const std::array first_bases{ + attention + "q_a_proj", attention + "kv_a_proj_with_mqa"}; + const std::array first{ + {{first_bases[0], kQueryRank, kHidden, input, 1U, q_rank, true}, + {first_bases[1], kKvRank, kHidden, input, 1U, latent, true}}}; + result = linear_batch(first, layer); + if (!result.ok()) return result; + result = norm(q_rank, q_rank, attention + "q_a_layernorm.weight"); + if (!result.ok()) return result; + result = norm(latent, latent, attention + "kv_a_layernorm.weight"); + if (!result.ok()) return result; + auto& cache = sequence.mla(layer); + if (cache.rows() != position) { + return {{"GLM-5.3 physical MLA position is not contiguous"}}; + } + result = cache.append(latent); + if (!result.ok()) return result; + const auto history = position + 1U; + std::vector expanded( + static_cast(history) * kHeads * 2U * kMlaHead); + const std::array second_bases{ + attention + "q_b_proj", attention + "kv_b_proj"}; + const auto latent_storage = cache.materialize(); + const auto latent_history = std::span(latent_storage); + const std::array second{ + {{second_bases[0], kMlaWidth, kQueryRank, q_rank, 1U, query, true}, + {second_bases[1], kHeads * 2U * kMlaHead, kKvRank, latent_history, + history, expanded, true}}}; + result = linear_batch(second, layer); + if (!result.ok()) return result; + std::vector attended(kMlaWidth, 0.0F); + const auto score_scale = 1.0F / std::sqrt(static_cast(kMlaHead)); + std::vector scores(history); + for (std::uint32_t head = 0U; head < kHeads; ++head) { + const auto* q = query.data() + static_cast(head) * kMlaHead; + float highest = -std::numeric_limits::infinity(); + for (std::uint32_t token = 0U; token < history; ++token) { + const auto* kv = expanded.data() + + (static_cast(token) * kHeads + head) * + (2U * kMlaHead); + float score = 0.0F; + for (std::uint32_t column = 0U; column < kMlaHead; ++column) { + score += q[column] * kv[column]; + } + scores[token] = score * score_scale; + highest = std::max(highest, scores[token]); + } + float total = 0.0F; + for (auto& score : scores) { + score = std::exp(score - highest); + total += score; + } + auto* destination = attended.data() + + static_cast(head) * kMlaHead; + for (std::uint32_t token = 0U; token < history; ++token) { + const auto* values = expanded.data() + + (static_cast(token) * kHeads + head) * + (2U * kMlaHead) + kMlaHead; + const auto coefficient = bf16_round_f32(scores[token] / total); + for (std::uint32_t column = 0U; column < kMlaHead; ++column) { + destination[column] += coefficient * values[column]; + } + } + } + round_bf16(attended); + return linear(attention + "o_proj", attended, 1U, kMlaWidth, + output, layer); + } + + [[nodiscard]] ValidationResult attention_mla_page( + std::span output, std::span input, + std::uint32_t rows, std::uint32_t layer, + const std::string& attention, Glm53SequenceState& sequence) { + ValidationResult result; + std::vector q_rank(static_cast(rows) * kQueryRank); + std::vector query(static_cast(rows) * kMlaWidth); + std::vector latent(static_cast(rows) * kKvRank); + const std::array first_bases{ + attention + "q_a_proj", attention + "kv_a_proj_with_mqa"}; + const std::array first{ + {{first_bases[0], kQueryRank, kHidden, input, rows, q_rank, true}, + {first_bases[1], kKvRank, kHidden, input, rows, latent, true}}}; + result = linear_batch(first, layer); + if (!result.ok()) return result; + result = norm_rows(q_rank, q_rank, rows, kQueryRank, + attention + "q_a_layernorm.weight"); + if (!result.ok()) return result; + result = norm_rows(latent, latent, rows, kKvRank, + attention + "kv_a_layernorm.weight"); + if (!result.ok()) return result; + auto& cache = sequence.mla(layer); + const auto history_begin = cache.rows(); + result = cache.append_rows(latent, rows); + if (!result.ok()) return result; + const auto history_rows = cache.rows(); + const auto latent_history = cache.materialize(); + std::vector expanded( + static_cast(history_rows) * kHeads * 2U * kMlaHead); + const std::array second_bases{ + attention + "q_b_proj", attention + "kv_b_proj"}; + const std::array second{ + {{second_bases[0], kMlaWidth, kQueryRank, q_rank, rows, query, true}, + {second_bases[1], kHeads * 2U * kMlaHead, kKvRank, + latent_history, history_rows, expanded, true}}}; + result = linear_batch(second, layer); + if (!result.ok()) return result; + std::vector attended( + static_cast(rows) * kMlaWidth, 0.0F); + const auto score_scale = 1.0F / std::sqrt(static_cast(kMlaHead)); + for (std::uint32_t row = 0U; row < rows; ++row) { + const auto visible = history_begin + row + 1U; + std::vector scores(visible); + for (std::uint32_t head = 0U; head < kHeads; ++head) { + const auto* q = query.data() + + (static_cast(row) * kHeads + head) * kMlaHead; + float highest = -std::numeric_limits::infinity(); + for (std::uint32_t token = 0U; token < visible; ++token) { + const auto* kv = expanded.data() + + (static_cast(token) * kHeads + head) * + (2U * kMlaHead); + float score = 0.0F; + for (std::uint32_t column = 0U; column < kMlaHead; ++column) { + score += q[column] * kv[column]; + } + scores[token] = score * score_scale; + highest = std::max(highest, scores[token]); + } + float total = 0.0F; + for (auto& score : scores) { + score = std::exp(score - highest); + total += score; + } + auto* destination = attended.data() + + (static_cast(row) * kHeads + head) * kMlaHead; + for (std::uint32_t token = 0U; token < visible; ++token) { + const auto* values = expanded.data() + + (static_cast(token) * kHeads + head) * + (2U * kMlaHead) + kMlaHead; + const auto coefficient = + bf16_round_f32(scores[token] / total); + for (std::uint32_t column = 0U; column < kMlaHead; ++column) { + destination[column] += coefficient * values[column]; + } + } + } + } + round_bf16(attended); + return linear(attention + "o_proj", attended, rows, kMlaWidth, + output, layer); + } + + [[nodiscard]] ValidationResult swiglu_block( + std::span output, std::span input, + const std::string& prefix, std::uint32_t inner, + std::uint32_t layer) { + ValidationResult result; + std::vector gate(inner), up(inner), activated(inner); + const std::array bases{ + prefix + "gate_proj", prefix + "up_proj"}; + const std::array projections{ + {{bases[0], inner, kHidden, input, 1U, gate, true}, + {bases[1], inner, kHidden, input, 1U, up, true}}}; + result = linear_batch(projections, layer); + if (!result.ok()) return result; + for (std::size_t index = 0U; index < inner; ++index) { + const auto g = std::min(gate[index], 10.0F); + const auto u = std::clamp(up[index], -10.0F, 10.0F); + activated[index] = g * sigmoid(g) * u; + } + round_bf16(activated); + return linear(prefix + "down_proj", activated, 1U, inner, + output, layer); + } + + [[nodiscard]] ValidationResult swiglu_block_page( + std::span output, std::span input, + const std::string& prefix, std::uint32_t rows, + std::uint32_t inner, std::uint32_t layer) { + ValidationResult result; + std::vector gate(static_cast(rows) * inner); + std::vector up(gate.size()), activated(gate.size()); + const std::array bases{ + prefix + "gate_proj", prefix + "up_proj"}; + const std::array projections{ + {{bases[0], inner, kHidden, input, rows, gate, true}, + {bases[1], inner, kHidden, input, rows, up, true}}}; + result = linear_batch(projections, layer); + if (!result.ok()) return result; + for (std::size_t index = 0U; index < gate.size(); ++index) { + const auto g = std::min(gate[index], 10.0F); + const auto u = std::clamp(up[index], -10.0F, 10.0F); + activated[index] = g * sigmoid(g) * u; + } + round_bf16(activated); + return linear(prefix + "down_proj", activated, rows, inner, + output, layer); + } + + [[nodiscard]] ValidationResult feedforward( + std::span output, std::span input, + std::uint32_t layer, const std::string& prefix, + std::uint64_t route_request = 0U, + std::uint32_t route_position = 0U, + bool schedule_prefetch = false) { + if (layer != kMtpLayer && !glm53_moe_layer(layer)) { + return swiglu_block(output, input, prefix + "mlp.", 12288U, layer); + } + ValidationResult result; + std::vector logits(288U); + // The reference router explicitly promotes both operands to F32. + result = linear(prefix + "mlp.gate", input, 1U, kHidden, logits, + layer, false); + if (!result.ok()) return result; + auto bias = host_tensor( + prefix + "mlp.gate.e_score_correction_bias", 288U); + if (!bias.ok()) return {std::move(bias.errors)}; + std::array selected{}; + result = kimi_route_topk(selected, logits, *bias.value, 2.5F); + if (!result.ok()) return result; + observe_route(layer, selected, route_request, route_position, + schedule_prefetch); + if (host_moe_active) { + return host_moe(prefix + "mlp.", selected, input, output); + } + return weights->moe(slot_for(layer), prefix + "mlp.", selected, + input, output); + } + + [[nodiscard]] ValidationResult feedforward_page( + std::span output, std::span input, + std::uint32_t rows, std::uint32_t layer, const std::string& prefix, + std::span route_requests = {}, + std::span route_positions = {}, + bool schedule_prefetch = false) { + if ((!route_requests.empty() || !route_positions.empty()) && + (route_requests.size() != rows || route_positions.size() != rows)) { + return {{"GLM-5.3 route-observation page has an invalid shape"}}; + } + if (layer != kMtpLayer && !glm53_moe_layer(layer)) { + return swiglu_block_page(output, input, prefix + "mlp.", rows, + 12288U, layer); + } + ValidationResult result; + std::vector logits(static_cast(rows) * 288U); + result = linear(prefix + "mlp.gate", input, rows, kHidden, logits, + layer, false); + if (!result.ok()) return result; + auto bias = host_tensor( + prefix + "mlp.gate.e_score_correction_bias", 288U); + if (!bias.ok()) return {std::move(bias.errors)}; + std::vector> selected_rows(rows); + for (std::uint32_t row = 0U; row < rows; ++row) { + auto& selected = selected_rows[row]; + result = kimi_route_topk( + selected, + std::span(logits).subspan( + static_cast(row) * 288U, 288U), + *bias.value, 2.5F); + if (!result.ok()) return result; + if (!route_requests.empty()) { + observe_route(layer, selected, route_requests[row], + route_positions[row], schedule_prefetch); + } + } + if (host_moe_active) { + return host_moe_page(prefix + "mlp.", selected_rows, input, + output); + } + for (std::uint32_t row = 0U; row < rows; ++row) { + result = weights->moe( + slot_for(layer), prefix + "mlp.", selected_rows[row], + input.subspan(static_cast(row) * kHidden, kHidden), + output.subspan(static_cast(row) * kHidden, kHidden)); + if (!result.ok()) return result; + } + return result; + } + + [[nodiscard]] ValidationResult initialize_streams( + std::uint32_t token, std::span streams) { + ValidationResult result; + if (streams.size() != static_cast(kMhc) * kHidden) { + result.errors.emplace_back( + "GLM-5.3 token streams have an invalid shape"); + return result; + } + auto embedding = checkpoint->read_f32_row( + "model.language_model.embed_tokens.weight", token); + if (!embedding.ok()) return {std::move(embedding.errors)}; + for (std::uint32_t stream = 0U; stream < kMhc; ++stream) { + std::copy(embedding.value.begin(), embedding.value.end(), + streams.begin() + static_cast( + stream * kHidden)); + } + return result; + } + + [[nodiscard]] ValidationResult forward_layer( + std::span streams, std::uint32_t layer, + std::uint32_t position, Glm53SequenceState& sequence) { + ValidationResult result; + if (streams.size() != static_cast(kMhc) * kHidden || + layer >= kLayers) { + result.errors.emplace_back( + "GLM-5.3 layer command has an invalid shape"); + return result; + } + std::vector collapsed(kHidden), normalized(kHidden), branch(kHidden); + const auto prefix = "model.language_model.layers." + + std::to_string(layer) + "."; + Dsv4MhcMix mix; + result = mhc_pre(collapsed, mix, streams, prefix + "hc_attn"); + if (!result.ok()) return result; + result = norm(normalized, collapsed, prefix + "input_layernorm.weight"); + if (!result.ok()) return result; + const auto attention = prefix + "self_attn."; + result = glm53_kda_layer(layer) + ? attention_kda(branch, normalized, layer, attention, sequence) + : attention_mla(branch, normalized, layer, position, attention, + sequence); + if (!result.ok()) return result; + std::vector transitioned(streams.size()); + result = dsv4_mhc_post_f32(transitioned, branch, streams, mix, kMhc); + if (!result.ok()) return result; + round_bf16(transitioned); + std::copy(transitioned.begin(), transitioned.end(), streams.begin()); + + result = mhc_pre(collapsed, mix, streams, prefix + "hc_ffn"); + if (!result.ok()) return result; + result = norm(normalized, collapsed, + prefix + "post_attention_layernorm.weight"); + if (!result.ok()) return result; + result = feedforward( + branch, normalized, layer, prefix, + route_request_key(&sequence, position), position, true); + if (!result.ok()) return result; + std::fill(transitioned.begin(), transitioned.end(), 0.0F); + result = dsv4_mhc_post_f32(transitioned, branch, streams, mix, kMhc); + if (!result.ok()) return result; + round_bf16(transitioned); + std::copy(transitioned.begin(), transitioned.end(), streams.begin()); + return result; + } + + [[nodiscard]] ValidationResult forward_layer_resident( + std::span streams, std::uint32_t layer, + std::uint32_t position, Glm53SequenceState& sequence, + DeviceSequenceState& device_sequence) { + if (!resident_execution_active || !device_sequence.ready || + streams.size() != static_cast(kMhc) * kHidden) { + return {{"GLM-5.3 resident layer command is not admissible"}}; + } + const auto device = device_for(layer); + auto result = cuda.dsv4_mhc_begin_device( + device, resident_layers[layer].attention, streams); + if (!result.ok()) return result; + const auto prefix = "model.language_model.layers." + + std::to_string(layer) + "."; + const auto attention = prefix + "self_attn."; + if (glm53_kda_layer(layer)) { + CudaGlm53KdaRequest request; + request.state = &device_sequence.kda[layer]; + request.heads = kHeads; + request.head_dim = kLinearHead; + request.convolution_kernel = 4U; + request.mhc_source_destination = true; + result = weights->kda_decode( + slot_for(layer), attention, request, {}); + } else { + CudaGlm53MlaRequest request; + request.state = &device_sequence.mla[layer]; + request.position = position; + request.maximum_context = config.maximum_context_tokens; + request.heads = kHeads; + request.head_dim = kMlaHead; + request.query_rank = kQueryRank; + request.key_value_rank = kKvRank; + result = weights->mla_decode_mhc( + slot_for(layer), attention, request); + } + if (!result.ok()) return result; + result = cuda.dsv4_mhc_transition_next_device( + device, resident_layers[layer].feedforward); + if (!result.ok()) return result; + if (host_moe_active) { + std::vector normalized(kHidden), branch(kHidden); + result = cuda.dsv4_mhc_download_layer_input(device, normalized); + if (!result.ok()) return result; + if (glm53_moe_layer(layer)) { + std::vector logits(288U); + result = weights->router_mhc( + slot_for(layer), prefix + "mlp.gate", logits); + if (!result.ok()) return result; + auto bias = host_tensor( + prefix + "mlp.gate.e_score_correction_bias", 288U); + if (!bias.ok()) return {std::move(bias.errors)}; + std::array selected{}; + result = kimi_route_topk(selected, logits, *bias.value, 2.5F); + if (!result.ok()) return result; + observe_route( + layer, selected, route_request_key(&sequence, position), + position, false); + result = host_moe(prefix + "mlp.", selected, normalized, + branch); + } else { + result = swiglu_block(branch, normalized, prefix + "mlp.", + 12288U, layer); + } + if (!result.ok()) return result; + return cuda.dsv4_mhc_finish(device, branch, streams); + } + if (glm53_moe_layer(layer)) { + std::vector logits(288U); + result = weights->router_mhc( + slot_for(layer), prefix + "mlp.gate", logits); + if (!result.ok()) return result; + auto bias = host_tensor( + prefix + "mlp.gate.e_score_correction_bias", 288U); + if (!bias.ok()) return {std::move(bias.errors)}; + std::array selected{}; + result = kimi_route_topk(selected, logits, *bias.value, 2.5F); + if (!result.ok()) return result; + observe_route( + layer, selected, route_request_key(&sequence, position), + position, true); + result = weights->moe( + slot_for(layer), prefix + "mlp.", selected, {}, {}, true); + } else { + result = weights->swiglu_mhc( + slot_for(layer), prefix + "mlp.", 12288U); + } + if (!result.ok()) return result; + return cuda.dsv4_mhc_finish_device(device, streams); + } + + [[nodiscard]] ValidationResult forward_layer_page( + std::span streams, std::uint32_t rows, std::uint32_t layer, + Glm53SequenceState& sequence) { + ValidationResult result; + const auto stream_columns = static_cast(kMhc) * kHidden; + if (rows == 0U || layer >= kLayers || + streams.size() != static_cast(rows) * stream_columns) { + result.errors.emplace_back( + "GLM-5.3 layer page has an invalid shape"); + return result; + } + const auto hidden_elements = static_cast(rows) * kHidden; + std::vector collapsed(hidden_elements), normalized(hidden_elements), + branch(hidden_elements); + std::vector mixes(rows); + const auto prefix = "model.language_model.layers." + + std::to_string(layer) + "."; + for (std::uint32_t row = 0U; row < rows; ++row) { + result = mhc_pre( + std::span(collapsed).subspan( + static_cast(row) * kHidden, kHidden), + mixes[row], + std::span(streams).subspan( + static_cast(row) * stream_columns, + stream_columns), + prefix + "hc_attn"); + if (!result.ok()) return result; + } + result = norm_rows(normalized, collapsed, rows, kHidden, + prefix + "input_layernorm.weight"); + if (!result.ok()) return result; + const auto attention = prefix + "self_attn."; + result = glm53_kda_layer(layer) + ? attention_kda_page(branch, normalized, rows, layer, attention, + sequence) + : attention_mla_page(branch, normalized, rows, layer, attention, + sequence); + if (!result.ok()) return result; + for (std::uint32_t row = 0U; row < rows; ++row) { + const auto stream_begin = + static_cast(row) * stream_columns; + std::vector transitioned(stream_columns); + result = dsv4_mhc_post_f32( + transitioned, + std::span(branch).subspan( + static_cast(row) * kHidden, kHidden), + std::span(streams).subspan( + stream_begin, stream_columns), + mixes[row], kMhc); + if (!result.ok()) return result; + round_bf16(transitioned); + std::copy(transitioned.begin(), transitioned.end(), + streams.begin() + static_cast(stream_begin)); + } + for (std::uint32_t row = 0U; row < rows; ++row) { + result = mhc_pre( + std::span(collapsed).subspan( + static_cast(row) * kHidden, kHidden), + mixes[row], + std::span(streams).subspan( + static_cast(row) * stream_columns, + stream_columns), + prefix + "hc_ffn"); + if (!result.ok()) return result; + } + result = norm_rows(normalized, collapsed, rows, kHidden, + prefix + "post_attention_layernorm.weight"); + if (!result.ok()) return result; + std::vector route_requests(rows); + std::vector route_positions(rows); + const auto position_base = sequence.token_count(); + for (std::uint32_t row = 0U; row < rows; ++row) { + route_positions[row] = position_base + row; + route_requests[row] = route_request_key( + &sequence, route_positions[row]); + } + result = feedforward_page(branch, normalized, rows, layer, prefix, + route_requests, route_positions); + if (!result.ok()) return result; + for (std::uint32_t row = 0U; row < rows; ++row) { + const auto stream_begin = + static_cast(row) * stream_columns; + std::vector transitioned(stream_columns); + result = dsv4_mhc_post_f32( + transitioned, + std::span(branch).subspan( + static_cast(row) * kHidden, kHidden), + std::span(streams).subspan( + stream_begin, stream_columns), + mixes[row], kMhc); + if (!result.ok()) return result; + round_bf16(transitioned); + std::copy(transitioned.begin(), transitioned.end(), + streams.begin() + static_cast(stream_begin)); + } + return result; + } + + // Independent sequence rows share the layer's resident weights while + // retaining disjoint recurrent/MLA state. This is the decode batch shape: + // unlike prompt pages, rows are not causally related to one another. + [[nodiscard]] ValidationResult forward_layer_sequences( + std::span streams, std::uint32_t layer, + std::span positions, + std::span sequences, + std::span device_sequences) { + ValidationResult result; + const auto rows = static_cast(sequences.size()); + const auto stream_columns = static_cast(kMhc) * kHidden; + if (rows == 0U || positions.size() != rows || + device_sequences.size() != rows || + streams.size() != static_cast(rows) * stream_columns) { + return {{"GLM-5.3 independent layer batch has an invalid shape"}}; + } + if (rows == 1U && device_sequences.front() != nullptr && + resident_execution_active && + (glm53_kda_layer(layer) || resident_mla_enabled()) && + (glm53_moe_layer(layer) || host_moe_active)) { + return forward_layer_resident( + streams, layer, positions.front(), *sequences.front(), + *device_sequences.front()); + } + const auto hidden_elements = static_cast(rows) * kHidden; + std::vector collapsed(hidden_elements), normalized(hidden_elements), + branch(hidden_elements); + std::vector mixes(rows); + const auto prefix = "model.language_model.layers." + + std::to_string(layer) + "."; + for (std::uint32_t row = 0U; row < rows; ++row) { + result = mhc_pre( + std::span(collapsed).subspan( + static_cast(row) * kHidden, kHidden), + mixes[row], + std::span(streams).subspan( + static_cast(row) * stream_columns, + stream_columns), + prefix + "hc_attn"); + if (!result.ok()) return result; + } + result = norm_rows(normalized, collapsed, rows, kHidden, + prefix + "input_layernorm.weight"); + if (!result.ok()) return result; + const auto attention = prefix + "self_attn."; + for (std::uint32_t row = 0U; row < rows; ++row) { + auto destination = std::span(branch).subspan( + static_cast(row) * kHidden, kHidden); + const auto input = std::span(normalized).subspan( + static_cast(row) * kHidden, kHidden); + result = glm53_kda_layer(layer) + ? attention_kda(destination, input, layer, attention, + *sequences[row], + device_sequences[row] == nullptr + ? nullptr + : &device_sequences[row]->kda[layer]) + : attention_mla(destination, input, layer, positions[row], + attention, *sequences[row]); + if (!result.ok()) return result; + } + for (std::uint32_t row = 0U; row < rows; ++row) { + const auto stream_begin = + static_cast(row) * stream_columns; + std::vector transitioned(stream_columns); + result = dsv4_mhc_post_f32( + transitioned, + std::span(branch).subspan( + static_cast(row) * kHidden, kHidden), + std::span(streams).subspan( + stream_begin, stream_columns), + mixes[row], kMhc); + if (!result.ok()) return result; + round_bf16(transitioned); + std::copy(transitioned.begin(), transitioned.end(), + streams.begin() + static_cast(stream_begin)); + } + for (std::uint32_t row = 0U; row < rows; ++row) { + result = mhc_pre( + std::span(collapsed).subspan( + static_cast(row) * kHidden, kHidden), + mixes[row], + std::span(streams).subspan( + static_cast(row) * stream_columns, + stream_columns), + prefix + "hc_ffn"); + if (!result.ok()) return result; + } + result = norm_rows(normalized, collapsed, rows, kHidden, + prefix + "post_attention_layernorm.weight"); + if (!result.ok()) return result; + std::vector route_requests(rows); + for (std::uint32_t row = 0U; row < rows; ++row) { + route_requests[row] = route_request_key(sequences[row], + positions[row]); + } + result = feedforward_page(branch, normalized, rows, layer, prefix, + route_requests, positions, true); + if (!result.ok()) return result; + for (std::uint32_t row = 0U; row < rows; ++row) { + const auto stream_begin = + static_cast(row) * stream_columns; + std::vector transitioned(stream_columns); + result = dsv4_mhc_post_f32( + transitioned, + std::span(branch).subspan( + static_cast(row) * kHidden, kHidden), + std::span(streams).subspan( + stream_begin, stream_columns), + mixes[row], kMhc); + if (!result.ok()) return result; + round_bf16(transitioned); + std::copy(transitioned.begin(), transitioned.end(), + streams.begin() + static_cast(stream_begin)); + } + return result; + } + + [[nodiscard]] ValidationResult collapse_streams_page( + std::span streams, std::uint32_t rows, + std::span collapsed) const { + const auto stream_columns = static_cast(kMhc) * kHidden; + if (rows == 0U || + streams.size() != static_cast(rows) * stream_columns || + collapsed.size() != static_cast(rows) * kHidden) { + return {{"GLM-5.3 residual collapse has an invalid shape"}}; + } + for (std::uint32_t row = 0U; row < rows; ++row) { + const auto stream_base = + static_cast(row) * stream_columns; + const auto hidden_base = static_cast(row) * kHidden; + for (std::size_t column = 0U; column < kHidden; ++column) { + collapsed[hidden_base + column] = 0.25F * + (streams[stream_base + column] + + streams[stream_base + kHidden + column] + + streams[stream_base + 2U * kHidden + column] + + streams[stream_base + 3U * kHidden + column]); + } + } + round_bf16(collapsed); + return {}; + } + + [[nodiscard]] ValidationResult finish_streams( + std::span streams, std::span logits) { + ValidationResult result; + if (streams.size() != static_cast(kMhc) * kHidden || + logits.empty()) { + result.errors.emplace_back( + "GLM-5.3 final text state has an invalid shape"); + return result; + } + std::vector collapsed(kHidden), normalized(kHidden); + result = collapse_streams_page(streams, 1U, collapsed); + if (!result.ok()) return result; + result = norm(normalized, collapsed, "model.language_model.norm.weight"); + if (!result.ok()) return result; + if (lm_head_ranges.size() > 1U && projection_workers != nullptr && + lm_head_ranges.size() == devices.size() && + logits.size() == kVocabulary) { + std::vector shard_results(devices.size()); + auto dispatched = projection_workers->parallel_for_addressed( + devices.size(), [&](std::size_t slot) { + const auto range = lm_head_ranges[slot]; + const Glm53WeightCache::LinearRequest request{ + "lm_head", range.count, kHidden, normalized, 1U, + logits.subspan(static_cast(range.begin), + static_cast(range.count)), + true, kVocabulary, range.begin}; + shard_results[slot] = weights->matmul_batch( + slot, std::span( + &request, 1U)); + }); + if (!dispatched.ok()) return dispatched; + for (auto& shard_result : shard_results) { + append(result.errors, std::move(shard_result.errors)); + } + if (result.ok()) { + tensor_parallel_head_batches.fetch_add( + 1U, std::memory_order_relaxed); + } + return result; + } + return linear("lm_head", normalized, 1U, kHidden, logits, kLayers - 1U); + } + + [[nodiscard]] ValidationResult project_lm_head_page( + std::span normalized, std::uint32_t rows, + std::span logits) { + if (rows == 0U || normalized.size() != + static_cast(rows) * kHidden || + logits.size() != static_cast(rows) * kVocabulary) { + return {{"GLM-5.3 LM-head page has an invalid shape"}}; + } + if (lm_head_ranges.size() > 1U && projection_workers != nullptr && + lm_head_ranges.size() == devices.size()) { + std::vector> shards(devices.size()); + std::vector shard_results(devices.size()); + auto dispatched = projection_workers->parallel_for_addressed( + devices.size(), [&](std::size_t slot) { + const auto range = lm_head_ranges[slot]; + shards[slot].resize(static_cast(rows) * + range.count); + const Glm53WeightCache::LinearRequest request{ + "lm_head", range.count, kHidden, normalized, rows, + shards[slot], true, kVocabulary, range.begin}; + shard_results[slot] = weights->matmul_batch( + slot, std::span( + &request, 1U)); + }); + if (!dispatched.ok()) return dispatched; + for (std::size_t slot = 0U; slot < devices.size(); ++slot) { + if (!shard_results[slot].ok()) return shard_results[slot]; + const auto range = lm_head_ranges[slot]; + for (std::uint32_t row = 0U; row < rows; ++row) { + std::copy_n( + shards[slot].begin() + static_cast( + static_cast(row) * range.count), + range.count, + logits.begin() + static_cast( + static_cast(row) * kVocabulary + + range.begin)); + } + } + tensor_parallel_head_batches.fetch_add(1U, + std::memory_order_relaxed); + return {}; + } + return linear("lm_head", normalized, rows, kHidden, logits, + kLayers - 1U); + } + + [[nodiscard]] ValidationResult finish_streams_page( + std::span streams, std::uint32_t rows, + std::span logits) { + const auto stream_columns = static_cast(kMhc) * kHidden; + if (rows == 0U || + streams.size() != static_cast(rows) * stream_columns || + logits.size() != static_cast(rows) * kVocabulary) { + return {{"GLM-5.3 final sequence batch has an invalid shape"}}; + } + std::vector collapsed(static_cast(rows) * kHidden); + auto result = collapse_streams_page(streams, rows, collapsed); + if (!result.ok()) return result; + std::vector normalized(collapsed.size()); + result = norm_rows(normalized, collapsed, rows, kHidden, + "model.language_model.norm.weight"); + if (!result.ok()) return result; + return project_lm_head_page(normalized, rows, logits); + } + + [[nodiscard]] ValidationResult forward_token_batch( + std::span tokens, + std::span positions, + std::span sequences, + std::span device_sequences, + std::span logits, std::span base_hidden = {}) { + const auto rows = static_cast(tokens.size()); + if (rows == 0U || positions.size() != rows || + sequences.size() != rows || device_sequences.size() != rows || + logits.size() != static_cast(rows) * kVocabulary || + (!base_hidden.empty() && + base_hidden.size() != static_cast(rows) * kHidden)) { + return {{"GLM-5.3 decode batch has an invalid shape"}}; + } + const auto stream_columns = static_cast(kMhc) * kHidden; + std::vector streams(static_cast(rows) * + stream_columns); + for (std::uint32_t row = 0U; row < rows; ++row) { + auto result = initialize_streams( + tokens[row], std::span(streams).subspan( + static_cast(row) * stream_columns, + stream_columns)); + if (!result.ok()) return result; + } + for (std::uint32_t layer = 0U; layer < kLayers; ++layer) { + auto result = forward_layer_sequences( + streams, layer, positions, sequences, device_sequences); + if (!result.ok()) return result; + } + if (!base_hidden.empty()) { + auto collapsed = collapse_streams_page(streams, rows, base_hidden); + if (!collapsed.ok()) return collapsed; + } + auto result = finish_streams_page(streams, rows, logits); + if (result.ok()) { + for (std::uint32_t row = 0U; row < rows; ++row) { + sequences[row]->set_token_count(positions[row] + 1U); + } + } + return result; + } + + [[nodiscard]] ValidationResult forward_mtp( + std::uint32_t next_token, std::span previous_hidden, + std::uint32_t position, Glm53SequenceState& sequence, + std::span logits, std::span feedback_hidden) { + if (previous_hidden.size() != kHidden || logits.size() != kVocabulary || + feedback_hidden.size() != kHidden || + sequence.mla(kMtpLayer).rows() != position) { + return {{"GLM-5.3 MTP command has an invalid sequence shape"}}; + } + const std::string prefix = "model.language_model.layers.45."; + auto embedding = checkpoint->read_f32_row( + "model.language_model.embed_tokens.weight", next_token); + if (!embedding.ok()) return {std::move(embedding.errors)}; + if (position == 0U) { + std::fill(embedding.value.begin(), embedding.value.end(), 0.0F); + } + std::vector normalized_embedding(kHidden); + std::vector normalized_hidden(kHidden); + auto result = norm(normalized_embedding, embedding.value, + prefix + "enorm.weight"); + if (!result.ok()) return result; + result = norm(normalized_hidden, previous_hidden, + prefix + "hnorm.weight"); + if (!result.ok()) return result; + std::vector fused(static_cast(2U) * kHidden); + std::copy(normalized_embedding.begin(), normalized_embedding.end(), + fused.begin()); + std::copy(normalized_hidden.begin(), normalized_hidden.end(), + fused.begin() + kHidden); + std::vector hidden(kHidden); + result = linear(prefix + "eh_proj", fused, 1U, 2U * kHidden, + hidden, kMtpLayer); + if (!result.ok()) return result; + + // The standalone MTP block deliberately disables mHC and follows the + // ordinary BF16 residual contract: residual <- input, attention, + // add+norm, MoE, final add. Its shared head then applies its own norm + // before reusing the target LM head. + std::vector residual = hidden; + std::vector normalized(kHidden), branch(kHidden); + result = norm(normalized, hidden, prefix + "input_layernorm.weight"); + if (!result.ok()) return result; + result = attention_mla(branch, normalized, kMtpLayer, position, + prefix + "self_attn.", sequence); + if (!result.ok()) return result; + for (std::size_t index = 0U; index < kHidden; ++index) { + residual[index] = bf16_round_f32(residual[index] + branch[index]); + } + result = norm(normalized, residual, + prefix + "post_attention_layernorm.weight"); + if (!result.ok()) return result; + result = feedforward(branch, normalized, kMtpLayer, prefix); + if (!result.ok()) return result; + for (std::size_t index = 0U; index < kHidden; ++index) { + feedback_hidden[index] = + bf16_round_f32(residual[index] + branch[index]); + } + result = norm(normalized, feedback_hidden, + prefix + "shared_head.norm.weight"); + if (!result.ok()) return result; + result = project_lm_head_page(normalized, 1U, logits); + return result; + } + + [[nodiscard]] ValidationResult prepare_mtp_prompt( + std::span prompt, + std::span base_hidden, Glm53SequenceState& sequence) { + if (base_hidden.size() != prompt.size() * kHidden) { + return {{"GLM-5.3 MTP prefill hidden-state extent is invalid"}}; + } + auto& cache = sequence.mla(kMtpLayer); + const auto required_rows = prompt.empty() ? 0U : + static_cast(prompt.size() - 1U); + if (cache.rows() > required_rows) { + return {{"GLM-5.3 MTP prefix state is ahead of the prompt"}}; + } + std::vector ignored_logits(kVocabulary); + std::vector feedback(kHidden); + for (std::uint32_t position = cache.rows(); position < required_rows; + ++position) { + auto result = forward_mtp( + prompt[position + 1U], + base_hidden.subspan(static_cast(position) * kHidden, + kHidden), + position, sequence, ignored_logits, feedback); + if (!result.ok()) return result; + } + return {}; + } + + [[nodiscard]] ValidationResult forward_token( + std::uint32_t token, std::uint32_t position, + std::span logits, Glm53SequenceState& sequence) { + ValidationResult result; + std::vector streams(static_cast(kMhc) * kHidden); + result = initialize_streams(token, streams); + if (!result.ok()) return result; + for (std::uint32_t layer = 0U; layer < kLayers; ++layer) { + result = forward_layer(streams, layer, position, sequence); + if (!result.ok()) return result; + if (config.load_progress) { + std::cerr << "\r[glm53] layer " << (layer + 1U) << '/' << kLayers + << std::flush; + } + } + if (config.load_progress) { + std::cerr << '\r' << std::string(32U, ' ') << '\r'; + } + result = finish_streams(streams, logits); + if (result.ok()) sequence.set_token_count(position + 1U); + return result; + } + + [[nodiscard]] ValidationResult forward_prompt( + std::span tokens, std::span logits, + Glm53SequenceState& sequence, + std::vector* base_hidden_rows = nullptr, + bool all_row_logits = false) { + ValidationResult result; + if (tokens.empty() || + logits.size() != (all_row_logits + ? tokens.size() * kVocabulary : kVocabulary)) { + result.errors.emplace_back("GLM-5.3 prefill has an invalid shape"); + return result; + } + const auto position_base = sequence.token_count(); + const auto stream_columns = static_cast(kMhc) * kHidden; + std::vector streams(tokens.size() * stream_columns); + std::vector encode_results(tokens.size()); + const auto encode = [&](std::size_t position) { + encode_results[position] = initialize_streams( + tokens[position], + std::span(streams).subspan( + position * stream_columns, stream_columns)); + }; + if (phase_scheduler_enabled() && kda_workers != nullptr && + tokens.size() >= kda_workers->size()) { + result = kda_workers->parallel_for(tokens.size(), encode); + if (!result.ok()) return result; + parallel_encode_pages.fetch_add(1U, std::memory_order_relaxed); + } else { + for (std::size_t position = 0U; position < tokens.size(); + ++position) { + encode(position); + } + } + for (auto& encoded : encode_results) { + if (!encoded.ok()) return encoded; + } + // Prompt rows are page/layer-major. Recurrent KDA and causal MLA state + // still advance in token order inside each layer, while the active + // layer's routed experts remain reusable in the bounded CUDA cache. + for (std::uint32_t layer = 0U; layer < kLayers; ++layer) { + result = forward_layer_page( + streams, static_cast(tokens.size()), layer, + sequence); + if (!result.ok()) return result; + if (config.load_progress) { + std::cerr << "\r[glm53-prefill] layer " << (layer + 1U) << '/' + << kLayers << " rows " << tokens.size() << std::flush; + } + } + if (config.load_progress) { + std::cerr << '\r' << std::string(48U, ' ') << '\r'; + } + if (base_hidden_rows != nullptr) { + const auto old_size = base_hidden_rows->size(); + base_hidden_rows->resize(old_size + tokens.size() * kHidden); + result = collapse_streams_page( + streams, static_cast(tokens.size()), + std::span(*base_hidden_rows).subspan(old_size)); + if (!result.ok()) return result; + } + result = all_row_logits + ? finish_streams_page( + streams, static_cast(tokens.size()), logits) + : finish_streams( + std::span(streams).last(stream_columns), logits); + if (result.ok()) { + sequence.set_token_count( + position_base + static_cast(tokens.size())); + } + return result; + } + + void complete_request(const std::shared_ptr& request) { + if (request->decoding) { + request->result.metrics.decode_seconds = + now_seconds() - request->decode_started; + if (request->streamed != nullptr) { + request->streamed->finish(request->on_token); + request->result.text = request->streamed->text(); + request->result.stopped = request->result.stopped || + request->streamed->stopped(); + } + const auto cache = weights->stats(); + std::cerr << "[glm53-decode-cache] misses=" + << (cache.misses - request->decode_cache_start.misses) + << " evictions=" + << (cache.evictions - + request->decode_cache_start.evictions) + << " useful_prefetches=" + << (cache.useful_prefetches - + request->decode_cache_start.useful_prefetches) + << '\n'; + } + { + std::scoped_lock lock(request->completion_mutex); + request->done = true; + } + request->completion.notify_all(); + } + + [[nodiscard]] bool prepare_request( + const std::shared_ptr& request) { + auto warm = wait_for_warmup(); + if (!warm.ok()) { + request->result.errors = std::move(warm.errors); + complete_request(request); + return false; + } + auto reset = reset_sequence(request->sequence); + if (!reset.ok()) { + request->result.errors = std::move(reset.errors); + complete_request(request); + return false; + } + request->result.prompt_token_ids = request->prompt; + request->result.metrics.prompt_tokens = request->prompt.size(); + request->logits.resize(kVocabulary); + const auto reused = restore_prefix( + request->prompt, request->sequence, request->logits, + request->base_hidden); + request->result.metrics.reused_prompt_tokens = reused; + request->prefill_cursor = reused; + request->prefill_started = now_seconds(); + request->prepared = true; + return true; + } + + [[nodiscard]] ValidationResult prepare_device_sequence( + Glm53SequenceState& sequence, DeviceSequenceState& device_sequence) { + ValidationResult result; + for (std::uint32_t layer = 0U; layer < kLayers; ++layer) { + const auto attention = "model.language_model.layers." + + std::to_string(layer) + ".self_attn."; + if (!glm53_kda_layer(layer)) { + const auto cache_floats = + static_cast(config.maximum_context_tokens) * + kKvRank; + std::vector packed( + cache_floats + kQueryRank + kKvRank, 0.0F); + const auto latent = sequence.mla(layer).materialize(); + if (latent.size() > cache_floats) { + return {{"GLM-5.3 resident MLA cache exceeds its admitted " + "context"}}; + } + std::copy(latent.begin(), latent.end(), packed.begin()); + auto q_norm = host_tensor( + attention + "q_a_layernorm.weight", kQueryRank); + auto kv_norm = host_tensor( + attention + "kv_a_layernorm.weight", kKvRank); + if (!q_norm.ok() || !kv_norm.ok()) { + append(result.errors, std::move(q_norm.errors)); + append(result.errors, std::move(kv_norm.errors)); + return result; + } + std::copy(q_norm.value->begin(), q_norm.value->end(), + packed.begin() + + static_cast(cache_floats)); + std::copy(kv_norm.value->begin(), kv_norm.value->end(), + packed.begin() + static_cast( + cache_floats + kQueryRank)); + result = cuda.upload_buffer( + device_for(layer), + std::as_bytes(std::span(packed)), + device_sequence.mla[layer]); + if (!result.ok()) return result; + continue; + } + const std::array tap_names{ + attention + "q_conv1d.weight", + attention + "k_conv1d.weight", + attention + "v_conv1d.weight"}; + std::array>, 3U> taps; + for (std::size_t projection = 0U; projection < taps.size(); + ++projection) { + auto loaded = host_tensor( + tap_names[projection], + static_cast(kLinearWidth) * 4U); + if (!loaded.ok()) return {std::move(loaded.errors)}; + taps[projection] = std::move(loaded.value); + } + auto a_log = host_tensor(attention + "A_log", kHeads); + auto dt_bias = host_tensor(attention + "dt_bias", kLinearWidth); + auto o_norm = host_tensor(attention + "o_norm.weight", kLinearHead); + if (!a_log.ok() || !dt_bias.ok() || !o_norm.ok()) { + append(result.errors, std::move(a_log.errors)); + append(result.errors, std::move(dt_bias.errors)); + append(result.errors, std::move(o_norm.errors)); + return result; + } + const auto recurrent = sequence.recurrent(layer); + const auto convolution_elements = + static_cast(3U) * kLinearWidth * 3U; + const auto tap_elements = + static_cast(3U) * kLinearWidth * 4U; + std::vector packed( + recurrent.size() + convolution_elements + tap_elements + + kHeads + kLinearWidth + kLinearHead + kKdaWorkspaceFloats); + auto destination = packed.begin(); + destination = std::copy(recurrent.begin(), recurrent.end(), + destination); + for (std::uint32_t projection = 0U; projection < 3U; + ++projection) { + const auto history = sequence.convolution(layer, projection); + destination = std::copy(history.begin(), history.end(), + destination); + } + for (const auto& tap : taps) { + destination = std::copy(tap->begin(), tap->end(), destination); + } + destination = std::copy(a_log.value->begin(), a_log.value->end(), + destination); + destination = std::copy(dt_bias.value->begin(), dt_bias.value->end(), + destination); + static_cast(std::copy(o_norm.value->begin(), + o_norm.value->end(), destination)); + result = cuda.upload_buffer( + device_for(layer), std::as_bytes(std::span(packed)), + device_sequence.kda[layer]); + if (!result.ok()) return result; + } + device_sequence.ready = true; + return result; + } + + void finish_prefill(const std::shared_ptr& request) { + request->result.metrics.prefill_tokens = + request->prompt.size() - + request->result.metrics.reused_prompt_tokens; + request->result.metrics.prefill_seconds = + now_seconds() - request->prefill_started; + if (mtp_enabled() && request->sampling.temperature == 0.0 && + request->maximum_new_tokens > 1U) { + auto mtp = prepare_mtp_prompt( + request->prompt, request->base_hidden, request->sequence); + if (!mtp.ok()) { + request->result.errors = std::move(mtp.errors); + complete_request(request); + return; + } + request->mtp_ready = true; + } + store_prefix(request->prompt, request->sequence, request->logits, + request->base_hidden); + // The prompt cache remains host/COW F32. Decode state is admitted once + // after that immutable snapshot, then never read back per token. + if (request->maximum_new_tokens > 1U && fused_kda_enabled()) { + auto prepared = prepare_device_sequence( + request->sequence, request->device_sequence); + if (!prepared.ok()) { + request->result.errors = std::move(prepared.errors); + complete_request(request); + return; + } + } + request->counts.assign(kVocabulary, 0U); + request->generator.seed(request->sampling.seed); + request->streamed = + std::make_unique(request->stop); + request->position = static_cast(request->prompt.size()); + request->decode_cache_start = weights->stats(); + request->decode_started = now_seconds(); + request->decoding = true; + if (request->maximum_new_tokens == 0U) { + complete_request(request); + } + } + + void advance_prefill(const std::shared_ptr& request, + std::size_t maximum_rows) { + if (request->done || request->decoding) return; + if (request->prefill_cursor == request->prompt.size()) { + finish_prefill(request); + return; + } + const auto count = std::min( + maximum_rows, request->prompt.size() - request->prefill_cursor); + auto prefill = forward_prompt( + std::span(request->prompt).subspan( + request->prefill_cursor, count), + request->logits, request->sequence, &request->base_hidden); + if (!prefill.ok()) { + request->result.errors = std::move(prefill.errors); + complete_request(request); + return; + } + request->prefill_cursor += count; + if (request->prefill_cursor == request->prompt.size()) { + finish_prefill(request); + } + } + + [[nodiscard]] bool publish_draw( + const std::shared_ptr& request, + const TokenLogprob& drawn, std::uint32_t& forward_token_id) { + if (drawn.token == 154820U || drawn.token == 154827U || + drawn.token == 154829U) { + request->result.stopped = true; + complete_request(request); + return false; + } + request->result.generated_token_ids.push_back(drawn.token); + request->result.logprobs.push_back(drawn); + request->sampled.push_back(drawn.token); + ++request->counts[drawn.token]; + auto piece = tokenizer.decode_token(drawn.token); + if (!piece.ok()) { + request->result.errors = std::move(piece.errors); + complete_request(request); + return false; + } + request->streamed->append(drawn.token, piece.value, + request->on_token); + if (request->streamed->stopped() || + request->streamed->cancelled() || + request->result.generated_token_ids.size() == + request->maximum_new_tokens) { + complete_request(request); + return false; + } + forward_token_id = drawn.token; + return true; + } + + [[nodiscard]] bool sample_request( + const std::shared_ptr& request, + std::uint32_t& forward_token_id) { + auto drawn = sample_logits( + request->logits, request->sampling, + SamplingHistory{request->counts, request->sampled}, + request->generator); + if (!drawn.ok()) { + request->result.errors = std::move(drawn.errors); + complete_request(request); + return false; + } + return publish_draw(request, drawn, forward_token_id); + } + + [[nodiscard]] bool try_mtp_step( + const std::shared_ptr& request, + std::uint32_t first_token) { + if (!request->mtp_ready || request->sampling.temperature != 0.0 || + request->sampling.xtc_probability != 0.0 || + request->sampling.future_entropy_candidates != 0U || + request->base_hidden.size() < kHidden || request->done) { + return false; + } + Glm53SequenceState mtp_after_first = request->sequence; + std::vector draft_logits(kVocabulary), draft_feedback(kHidden); + const auto mtp_position = + static_cast(mtp_after_first.mla(kMtpLayer).rows()); + auto status = forward_mtp( + first_token, + std::span(request->base_hidden).last(kHidden), + mtp_position, mtp_after_first, draft_logits, draft_feedback); + if (!status.ok()) { + request->result.errors = std::move(status.errors); + complete_request(request); + return true; + } + ++mtp_drafts; + auto draft_generator = request->generator; + auto draft = sample_logits( + draft_logits, request->sampling, + SamplingHistory{request->counts, request->sampled}, + draft_generator); + if (!draft.ok()) { + request->result.errors = std::move(draft.errors); + complete_request(request); + return true; + } + + Glm53SequenceState verified = request->sequence; + const std::array candidates{ + first_token, draft.token}; + std::vector verification_logits( + static_cast(2U) * kVocabulary); + std::vector verification_hidden; + status = forward_prompt(candidates, verification_logits, verified, + &verification_hidden, true); + if (!status.ok()) { + request->result.errors = std::move(status.errors); + complete_request(request); + return true; + } + auto target_generator = request->generator; + auto target = sample_logits( + std::span(verification_logits).first(kVocabulary), + request->sampling, + SamplingHistory{request->counts, request->sampled}, + target_generator); + if (!target.ok()) { + request->result.errors = std::move(target.errors); + complete_request(request); + return true; + } + if (target.token == draft.token) { + verified.copy_mla_from(kMtpLayer, mtp_after_first); + request->sequence = std::move(verified); + request->generator = std::move(target_generator); + request->base_hidden.insert( + request->base_hidden.end(), verification_hidden.begin(), + verification_hidden.end()); + std::copy_n(verification_logits.begin() + kVocabulary, + kVocabulary, request->logits.begin()); + request->position += 2U; + request->result.metrics.decode_tokens += 2U; + request->iteration += 2U; + ++mtp_accepted; + std::uint32_t ignored = 0U; + static_cast(publish_draw(request, target, ignored)); + if (!request->done) { + // Keep the draft cache aligned through the accepted token. + // Its next proposal is deliberately discarded; the next + // target sample remains the sole source of published tokens. + std::vector ignored_logits(kVocabulary); + std::vector ignored_feedback(kHidden); + status = forward_mtp( + target.token, + std::span(verification_hidden).subspan( + 0U, kHidden), + mtp_position + 1U, request->sequence, ignored_logits, + ignored_feedback); + if (!status.ok()) { + request->result.errors = std::move(status.errors); + complete_request(request); + } + } + return true; + } + + // A rejected second token must leave the target exactly after the + // first token. COW makes the retry cheap in state memory; execution is + // intentionally repeated rather than trying to extract a mutable + // intermediate snapshot from the two-row verification page. + std::vector first_hidden; + status = forward_prompt( + std::span(&first_token, 1U), request->logits, + request->sequence, &first_hidden); + if (!status.ok()) { + request->result.errors = std::move(status.errors); + complete_request(request); + return true; + } + request->sequence.copy_mla_from(kMtpLayer, mtp_after_first); + request->base_hidden.insert(request->base_hidden.end(), + first_hidden.begin(), first_hidden.end()); + ++request->position; + ++request->result.metrics.decode_tokens; + ++request->iteration; + return true; + } + + void scheduler_loop() { + for (;;) { + { + std::unique_lock lock(scheduler_mutex); + scheduler_ready.wait(lock, [&] { + return scheduler_stopping || !pending_requests.empty() || + !active_requests.empty(); + }); + if (scheduler_stopping && pending_requests.empty() && + active_requests.empty()) { + return; + } + // A fresh queue gets a tiny admission window so requests that + // arrived together become one prefill/decode cohort. Once a + // cohort is active there is no delay: iteration admission is + // immediate. Two milliseconds is below network jitter while + // avoiding a model- or hardware-specific batching timeout. + if (active_requests.empty() && pending_requests.size() == 1U && + !scheduler_stopping) { + static_cast(scheduler_ready.wait_for( + lock, std::chrono::milliseconds(2), [&] { + return scheduler_stopping || + pending_requests.size() > 1U; + })); + } + while (!pending_requests.empty() && + active_requests.size() < scheduler_capacity) { + active_requests.push_back(pending_requests.front()); + pending_requests.pop_front(); + } + } + for (auto& request : active_requests) { + if (!request->prepared && !request->done) { + static_cast(prepare_request(request)); + } + } + std::size_t live = 0U; + std::size_t decoding = 0U; + for (const auto& request : active_requests) { + if (!request->done) ++live; + if (!request->done && request->decoding) ++decoding; + } + if (live > 0U) { + scheduler_iterations.fetch_add(1U, std::memory_order_relaxed); + if (live > 1U) { + scheduler_batched_iterations.fetch_add( + 1U, std::memory_order_relaxed); + } + std::vector> step_requests; + std::vector step_tokens; + std::vector step_positions; + std::vector step_sequences; + std::vector step_device_sequences; + for (auto& request : active_requests) { + if (request->done || !request->decoding) continue; + std::uint32_t token = 0U; + if (sample_request(request, token)) { + step_requests.push_back(request); + step_tokens.push_back(token); + step_positions.push_back(request->position); + step_sequences.push_back(&request->sequence); + step_device_sequences.push_back( + request->device_sequence.ready + ? &request->device_sequence : nullptr); + } + } + if (!step_requests.empty()) { + const bool mtp_handled = step_requests.size() == 1U && + try_mtp_step(step_requests.front(), + step_tokens.front()); + if (!mtp_handled) { + if (step_requests.size() > 1U) { + for (auto& request : step_requests) { + request->mtp_ready = false; + } + } + std::vector step_logits( + step_requests.size() * kVocabulary); + std::vector step_hidden( + step_requests.size() * kHidden); + const bool capture = profiler_capture_enabled() && + !profiler_captured.exchange( + true, std::memory_order_relaxed); + if (capture) { + const auto started = cuda.profiler_start(); + if (!started.ok()) { + std::cerr << "[glm53-profile] " + << started.errors.front() << '\n'; + } + } + auto step = forward_token_batch( + step_tokens, step_positions, step_sequences, + step_device_sequences, step_logits, step_hidden); + if (capture) { + const auto stopped = cuda.profiler_stop(); + if (!stopped.ok()) { + std::cerr << "[glm53-profile] " + << stopped.errors.front() << '\n'; + } + } + if (!step.ok()) { + for (auto& request : step_requests) { + request->result.errors = step.errors; + complete_request(request); + } + } else { + for (std::size_t row = 0U; + row < step_requests.size(); ++row) { + auto& request = step_requests[row]; + std::copy_n( + step_logits.begin() + + static_cast( + row * kVocabulary), + kVocabulary, request->logits.begin()); + request->base_hidden.insert( + request->base_hidden.end(), + step_hidden.begin() + + static_cast(row * kHidden), + step_hidden.begin() + + static_cast( + (row + 1U) * kHidden)); + ++request->position; + ++request->result.metrics.decode_tokens; + ++request->iteration; + } + } + } + } + // Decode has latency priority. A newly admitted prompt gets a + // single-token chunk while decoders are live; with no decode + // work, a page-sized chunk retains the wide prefill route. + const std::size_t prefill_rows = decoding == 0U ? 64U : 1U; + for (auto& request : active_requests) { + if (!request->done && !request->decoding) { + advance_prefill(request, prefill_rows); + } + } + } + active_requests.erase( + std::remove_if(active_requests.begin(), active_requests.end(), + [](const auto& request) { return request->done; }), + active_requests.end()); + } + } + + [[nodiscard]] Glm53GenerationResult schedule( + std::vector prompt, + std::uint32_t maximum_new_tokens, const SamplingOptions& sampling, + std::span stop, + const TokenStreamCallback& on_token) { + auto request = std::make_shared(); + request->prompt = std::move(prompt); + request->maximum_new_tokens = maximum_new_tokens; + request->sampling = sampling; + request->stop.assign(stop.begin(), stop.end()); + request->on_token = on_token; + { + std::scoped_lock lock(scheduler_mutex); + if (scheduler_stopping) { + request->result.errors.emplace_back( + "GLM-5.3 iteration scheduler is stopping"); + return std::move(request->result); + } + pending_requests.push_back(request); + } + scheduler_ready.notify_one(); + std::unique_lock lock(request->completion_mutex); + request->completion.wait(lock, [&] { return request->done; }); + return std::move(request->result); + } +}; + +Glm53Runtime::Glm53Runtime() : impl_(std::make_unique()) {} +Glm53Runtime::~Glm53Runtime() = default; +Glm53Runtime::Glm53Runtime(Glm53Runtime&&) noexcept = default; +Glm53Runtime& Glm53Runtime::operator=(Glm53Runtime&&) noexcept = default; + +ValidationResult Glm53Runtime::initialize( + const std::string& model_directory, const Glm53RuntimeConfig& config) { + ValidationResult result; + if (impl_->ready) { + result.errors.emplace_back("GLM-5.3 runtime is already initialized"); + return result; + } + if (config.maximum_context_tokens == 0U || + config.maximum_context_tokens > kExactSparseContext) { + result.errors.push_back( + "GLM-5.3 text context must be within [1, 2048]; above 2048 the " + "checkpoint's exact k-pool sparse indexer is required"); + return result; + } + impl_->config = config; + impl_->prefix_cache_limit = + prefix_cache_entries(config.maximum_context_tokens); + impl_->scheduler_capacity = std::max( + 1U, std::min(32U, impl_->prefix_cache_limit)); + impl_->devices = resolve_runtime_devices(config.devices); + result = validate_common_runtime_config( + impl_->devices, config.vram_cache_fraction, + config.sampling_temperature, "GLM-5.3"); + if (!result.ok()) return result; + auto device_plan = plan_runtime_devices( + impl_->devices, config.vram_cache_fraction, kDeviceWorkspaceReserve, + kMinimumDeviceBudget, "GLM-5.3"); + if (!device_plan.ok()) return {std::move(device_plan.errors)}; + auto tokenizer = ModelTokenizer::load(model_directory + "/tokenizer.json"); + if (!tokenizer.ok()) return {std::move(tokenizer.errors)}; + // The tokenizer has 154,820 base pieces plus 36 added special tokens. + // The checkpoint pads its embedding and output matrices to 154,880 rows; + // those 24 padding rows are deliberately not tokenizable. + if (tokenizer.value.vocabulary_size() != 154856U) { + result.errors.emplace_back( + "GLM-5.3 tokenizer must expose 154856 usable token ids"); + return result; + } + auto checkpoint = Glm53CheckpointReader::open(model_directory); + if (!checkpoint.ok()) return {std::move(checkpoint.errors)}; + result = impl_->cuda.initialize(impl_->devices); + if (!result.ok()) return result; + for (std::size_t slot = 0U; slot < impl_->devices.size(); ++slot) { + result = impl_->cuda.reserve_weight_arena( + impl_->devices[slot], device_plan.value.weight_capacities[slot]); + if (!result.ok()) return result; + } + impl_->tokenizer = std::move(tokenizer.value); + impl_->checkpoint = std::move(checkpoint.value); + impl_->weight_capacities = device_plan.value.weight_capacities; + if (impl_->devices.size() > 1U && + !cross_gpu_projections_enabled(impl_->devices)) { + // PCIe/PHB systems pay a full activation bridge for every owner + // change. Use capacity-weighted contiguous pipeline stages so a token + // crosses once. Best-rank P2P (NVLink/NVSwitch) keeps the fine-grained + // schedule, which the TP executor can consume without redistributing + // layer ownership when that topology is available. + impl_->device_schedule = contiguous_layer_schedule( + kLayers, impl_->weight_capacities); + } else { + impl_->device_schedule = std::move( + device_plan.value.weighted_schedule); + } + if (impl_->device_schedule.empty()) { + return {{"GLM-5.3 could not derive a topology-aware layer schedule"}}; + } + impl_->resident_execution_active = fused_kda_enabled(); + if (impl_->resident_execution_active) { + for (const auto device : impl_->devices) { + if (!impl_->cuda.validate_dsv4_mhc_device(device).ok()) { + impl_->resident_execution_active = false; + break; + } + } + } + if (impl_->resident_execution_active) { + // mHC weights use the same arena as cached linears. Reserve their + // exact order of magnitude per discovered layer owner before the + // cache fills the arena; no device-count or VRAM-size assumption is + // embedded here. + constexpr std::uint64_t per_layer_mhc_reserve = 4ULL << 20U; + std::vector resident_reserve(impl_->devices.size()); + for (std::uint32_t layer = 0U; layer < kLayers; ++layer) { + resident_reserve[impl_->slot_for(layer)] += + per_layer_mhc_reserve; + } + for (std::size_t slot = 0U; slot < impl_->devices.size(); ++slot) { + if (impl_->weight_capacities[slot] <= resident_reserve[slot] + + kMinimumDeviceBudget) { + impl_->resident_execution_active = false; + break; + } + } + if (impl_->resident_execution_active) { + for (std::size_t slot = 0U; slot < impl_->devices.size(); ++slot) { + impl_->weight_capacities[slot] -= resident_reserve[slot]; + } + } + } + impl_->weights = std::make_unique( + *impl_->checkpoint, impl_->cuda, impl_->devices, + impl_->weight_capacities); + const std::string expert_prefix = + "model.language_model.layers.3.mlp.experts.0."; + const auto expert_bytes = + impl_->checkpoint->cuda_linear_storage_bytes( + expert_prefix + "gate_proj") + + impl_->checkpoint->cuda_linear_storage_bytes( + expert_prefix + "up_proj") + + impl_->checkpoint->cuda_linear_storage_bytes( + expert_prefix + "down_proj"); + const auto cache_bytes = std::accumulate( + impl_->weight_capacities.begin(), impl_->weight_capacities.end(), + std::uint64_t{0U}); + const auto routed_bytes = std::accumulate( + impl_->checkpoint->manifest().tensors.begin(), + impl_->checkpoint->manifest().tensors.end(), std::uint64_t{0U}, + [](std::uint64_t total, const Glm53ManifestTensor& tensor) { + return tensor.role == Glm53TensorRole::RoutedExpert + ? total + tensor.source_bytes : total; + }); + const auto& hardware = host_hardware_profile(); + const auto host_width = std::min( + hardware.worker_threads(0.5), hardware.usable_cpu_ids.size()); + const auto model_parallel_width = static_cast( + impl_->checkpoint->config().experts_per_token) * 2U; + const auto override = host_moe_override(); + const bool host_instruction_support = +#if STRATA_GLM53_HOST_AVX2 + __builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma"); +#else + false; +#endif + const bool host_moe_admitted = override > 0 || + (override < 0 && host_instruction_support && + host_width >= model_parallel_width && + cache_bytes != 0U && routed_bytes > 2U * cache_bytes); + if (host_moe_admitted && host_width != 0U) { + std::vector cpus(hardware.usable_cpu_ids.begin(), + hardware.usable_cpu_ids.begin() + + static_cast(host_width)); + impl_->host_moe_workers = std::make_unique( + std::move(cpus), std::chrono::milliseconds(1)); + impl_->host_moe_active = + impl_->host_moe_workers->size() == host_width; + } + if (expert_bytes != 0U) { + const auto cache_experts = static_cast( + cache_bytes / expert_bytes); + const auto host_window = static_cast( + host_hardware_profile().worker_threads(0.25)) * + impl_->checkpoint->config().experts_per_token; + impl_->prefetch_queue_limit = std::max( + 1U, std::min(cache_experts, host_window)); + impl_->prefetch_prediction_limit = std::min( + impl_->checkpoint->config().experts_per_token, + impl_->prefetch_queue_limit); + impl_->prefetch_minimum_confidence = + 1.0 - static_cast( + impl_->checkpoint->config().experts_per_token) / + static_cast( + impl_->checkpoint->config().routed_experts); + } + if (impl_->devices.size() > 1U) { + auto worker_cpus = projection_worker_cpus(impl_->devices); + if (worker_cpus.size() == impl_->devices.size()) { + impl_->projection_workers = std::make_unique( + std::move(worker_cpus)); + } + } + impl_->full_tensor_parallel_active = + full_tensor_parallel_enabled() && impl_->devices.size() == 2U && + impl_->projection_workers != nullptr && + cross_gpu_projections_enabled(impl_->devices); + if ((tensor_parallel_head_enabled() || + impl_->full_tensor_parallel_active) && + impl_->projection_workers != nullptr) { + impl_->lm_head_ranges = weighted_row_ranges( + kVocabulary, impl_->weight_capacities, 128U); + } + if (impl_->devices.size() > 1U || config.verbose) { + std::uint32_t hops = 0U; + for (std::uint32_t layer = 1U; layer < kLayers; ++layer) { + if (impl_->slot_for(layer) != impl_->slot_for(layer - 1U)) ++hops; + } + std::cerr << "[glm53-topology] mode=" + << (impl_->full_tensor_parallel_active + ? "high-speed-peer-tp2" + : (cross_gpu_projections_enabled(impl_->devices) + ? "high-speed-peer" + : "contiguous-pipeline")) + << " activation_hops=" << hops << " layers=" << kLayers + << '\n'; + std::cerr << "[glm53-resident] mode=" + << (impl_->resident_execution_active + ? "fused-layer" + : "host-boundary-fallback") + << '\n'; + std::cerr << "[glm53-expert-tier] mode=" + << (impl_->host_moe_active ? "host-fp8" : "cuda-lru") + << " workers=" + << (impl_->host_moe_workers == nullptr + ? 0U : impl_->host_moe_workers->size()) + << " routed_gib=" + << static_cast(routed_bytes) / + static_cast(1ULL << 30U) + << " cuda_cache_gib=" + << static_cast(cache_bytes) / + static_cast(1ULL << 30U) + << '\n'; + } + if (replay_ssm_enabled() || phase_scheduler_enabled()) { + auto worker_cpus = compute_worker_cpus(); + if (!worker_cpus.empty()) { + impl_->kda_workers = std::make_unique( + std::move(worker_cpus), std::chrono::milliseconds(1)); + } + } + impl_->ready = true; + try { + // Keep API/server startup lazy-fast while warming independent device + // spines in the background. The first generation joins this work; an + // idle server usually reaches full residency before its first request. + impl_->warmup_thread = std::thread([state = impl_.get()] { + state->warmup_result = state->warmup(); + }); + impl_->prefetch_threads.reserve(impl_->devices.size()); + for (std::size_t worker = 0U; worker < impl_->devices.size(); ++worker) { + impl_->prefetch_threads.emplace_back([state = impl_.get()] { + state->prefetch_loop(); + }); + } + impl_->scheduler_thread = std::thread([state = impl_.get()] { + state->scheduler_loop(); + }); + } catch (const std::system_error& error) { + impl_->ready = false; + result.errors.push_back( + "GLM-5.3 could not start background spine warmup: " + + std::string(error.what())); + } + return result; +} + +Glm53GenerationResult Glm53Runtime::generate_chat_stream( + std::span messages, std::uint32_t maximum_new_tokens, + const SamplingOptions& sampling, std::span stop, + const TokenStreamCallback& on_token) { + Glm53GenerationResult result; + if (!impl_->ready) { + result.errors.emplace_back("GLM-5.3 runtime is not initialized"); + return result; + } + std::string error; + if (!validate_sampling_options(sampling, error)) { + result.errors.push_back(std::move(error)); + return result; + } + if (!validate_chat_messages(messages, error)) { + result.errors.push_back(std::move(error)); + return result; + } + for (const auto& message : messages) { + for (const auto& part : message.parts) { + if (part.kind != ChatContentKind::Text) { + result.errors.emplace_back( + "GLM-5.3 vision is not implemented; this runtime supports text-only messages"); + return result; + } + } + } + auto encoded = impl_->tokenizer.encode( + render_glm53_chat_prompt(messages, "max", true)); + if (!encoded.ok()) { + result.errors = std::move(encoded.errors); + return result; + } + if (encoded.value.empty() || encoded.value.size() + maximum_new_tokens > + impl_->config.maximum_context_tokens) { + result.errors.emplace_back( + "GLM-5.3 prompt and requested generation exceed the admitted text context"); + return result; + } + const auto mtp_drafts_before = + impl_->mtp_drafts.load(std::memory_order_relaxed); + const auto mtp_accepted_before = + impl_->mtp_accepted.load(std::memory_order_relaxed); + const auto prefetch_requests_before = + impl_->prefetch_requests.load(std::memory_order_relaxed); + const auto cache_before = impl_->weights->stats(); + result = impl_->schedule(std::move(encoded.value), maximum_new_tokens, + sampling, stop, on_token); + const auto request_mtp_drafts = + impl_->mtp_drafts.load(std::memory_order_relaxed) - mtp_drafts_before; + const auto request_mtp_accepted = + impl_->mtp_accepted.load(std::memory_order_relaxed) - + mtp_accepted_before; + if (request_mtp_drafts != 0U) { + std::cerr << "[glm53-mtp] drafts=" << request_mtp_drafts + << " accepted=" << request_mtp_accepted + << " acceptance=" + << (100.0 * static_cast(request_mtp_accepted) / + static_cast(request_mtp_drafts)) + << "%\n"; + } + const auto request_prefetches = + impl_->prefetch_requests.load(std::memory_order_relaxed) - + prefetch_requests_before; + if (request_prefetches != 0U) { + const auto cache_after = impl_->weights->stats(); + std::cerr << "[glm53-residency] predictions=" << request_prefetches + << " completed=" + << impl_->prefetch_completed.load(std::memory_order_relaxed) + << " dropped=" + << impl_->prefetch_dropped.load(std::memory_order_relaxed) + << " errors=" + << impl_->prefetch_errors.load(std::memory_order_relaxed) + << " useful=" + << (cache_after.useful_prefetches - + cache_before.useful_prefetches) + << " demand_misses=" + << (cache_after.misses - cache_before.misses) + << " evictions=" + << (cache_after.evictions - cache_before.evictions) + << '\n'; + } + if (impl_->config.verbose) { + std::cerr << "[glm53-projection] parallel_batches=" + << impl_->parallel_projection_batches.load( + std::memory_order_relaxed) + << " parallel_requests=" + << impl_->parallel_projection_requests.load( + std::memory_order_relaxed) + << " tensor_parallel_head_batches=" + << impl_->tensor_parallel_head_batches.load( + std::memory_order_relaxed) + << " parallel_encode_pages=" + << impl_->parallel_encode_pages.load( + std::memory_order_relaxed) + << " prefix_cache_hits=" + << impl_->prefix_cache_hits.load(std::memory_order_relaxed) + << " prefix_cache_tokens=" + << impl_->prefix_cache_tokens.load(std::memory_order_relaxed) + << " scheduler_iterations=" + << impl_->scheduler_iterations.load(std::memory_order_relaxed) + << " scheduler_batched_iterations=" + << impl_->scheduler_batched_iterations.load( + std::memory_order_relaxed) + << " mtp_drafts=" + << impl_->mtp_drafts.load(std::memory_order_relaxed) + << " mtp_accepted=" + << impl_->mtp_accepted.load(std::memory_order_relaxed) + << " prefetch_requests=" + << impl_->prefetch_requests.load(std::memory_order_relaxed) + << " prefetch_completed=" + << impl_->prefetch_completed.load(std::memory_order_relaxed) + << " prefetch_dropped=" + << impl_->prefetch_dropped.load(std::memory_order_relaxed) + << " prefetch_errors=" + << impl_->prefetch_errors.load(std::memory_order_relaxed) + << " prefetch_queue_limit=" + << impl_->prefetch_queue_limit + << " host_moe_calls=" + << impl_->host_moe_calls.load(std::memory_order_relaxed) + << " host_moe_ms=" + << static_cast(impl_->host_moe_nanoseconds.load( + std::memory_order_relaxed)) / 1.0e6 + << '\n'; + const auto cache = impl_->weights->stats(); + std::cerr << "[glm53-cache] hits=" << cache.hits + << " misses=" << cache.misses + << " evictions=" << cache.evictions + << " prefetches=" << cache.prefetches + << " useful_prefetches=" << cache.useful_prefetches + << " failed_prefetches=" << cache.failed_prefetches + << '\n'; + } + return result; +} + +} // namespace strata diff --git a/src/models/glm53/glm53_sequence.cpp b/src/models/glm53/glm53_sequence.cpp new file mode 100644 index 0000000..84a9c06 --- /dev/null +++ b/src/models/glm53/glm53_sequence.cpp @@ -0,0 +1,182 @@ +#include "strata/models/glm53/glm53_sequence.hpp" + +#include "strata/models/glm53/glm53_manifest.hpp" + +#include +#include +#include + +namespace strata { + +Glm53PagedRows::Glm53PagedRows(std::uint32_t columns, + std::uint32_t page_rows) { + static_cast(reset(columns, page_rows)); +} + +ValidationResult Glm53PagedRows::reset(std::uint32_t columns, + std::uint32_t page_rows) { + if (columns == 0U || page_rows == 0U || + static_cast(columns) * page_rows > + std::numeric_limits::max()) { + return {{"GLM-5.3 physical page geometry is invalid"}}; + } + columns_ = columns; + page_rows_ = page_rows; + rows_ = 0U; + pages_.clear(); + return {}; +} + +ValidationResult Glm53PagedRows::ensure_append_page() { + if (columns_ == 0U || page_rows_ == 0U) { + return {{"GLM-5.3 physical page table is not initialized"}}; + } + const auto page_index = rows_ / page_rows_; + try { + if (page_index == pages_.size()) { + pages_.push_back(std::make_shared( + static_cast(columns_) * page_rows_)); + } else if (!pages_[page_index].unique()) { + pages_[page_index] = std::make_shared(*pages_[page_index]); + } + } catch (const std::bad_alloc&) { + return {{"GLM-5.3 could not allocate a physical MLA page"}}; + } + return {}; +} + +ValidationResult Glm53PagedRows::append(std::span row_values) { + if (row_values.size() != columns_) { + return {{"GLM-5.3 physical MLA row has an invalid width"}}; + } + auto result = ensure_append_page(); + if (!result.ok()) return result; + const auto page_index = rows_ / page_rows_; + const auto page_row = rows_ % page_rows_; + std::copy(row_values.begin(), row_values.end(), + pages_[page_index]->values.begin() + + static_cast(page_row * columns_)); + ++rows_; + return {}; +} + +ValidationResult Glm53PagedRows::append_rows(std::span values, + std::uint32_t row_count) { + if (values.size() != static_cast(row_count) * columns_) { + return {{"GLM-5.3 physical MLA page append has an invalid shape"}}; + } + for (std::uint32_t row_index = 0U; row_index < row_count; ++row_index) { + auto result = append(values.subspan( + static_cast(row_index) * columns_, columns_)); + if (!result.ok()) return result; + } + return {}; +} + +ValidationResult Glm53PagedRows::truncate(std::uint32_t rows) { + if (rows > rows_) return {{"GLM-5.3 physical MLA truncate grows state"}}; + rows_ = rows; + const auto keep = rows == 0U ? 0U : (rows + page_rows_ - 1U) / page_rows_; + pages_.resize(keep); + return {}; +} + +std::span Glm53PagedRows::row(std::uint32_t index) const noexcept { + if (index >= rows_ || columns_ == 0U || page_rows_ == 0U) return {}; + const auto page_index = index / page_rows_; + const auto page_row = index % page_rows_; + return std::span(pages_[page_index]->values) + .subspan(static_cast(page_row) * columns_, columns_); +} + +std::vector Glm53PagedRows::materialize() const { + std::vector result(static_cast(rows_) * columns_); + for (std::uint32_t index = 0U; index < rows_; ++index) { + const auto source = row(index); + std::copy(source.begin(), source.end(), + result.begin() + static_cast( + static_cast(index) * columns_)); + } + return result; +} + +std::uint64_t Glm53PagedRows::private_bytes() const noexcept { + std::uint64_t result = 0U; + for (const auto& page : pages_) { + if (page.unique()) result += page->values.size() * sizeof(float); + } + return result; +} + +ValidationResult Glm53SequenceState::reset( + std::uint32_t maximum_context_tokens, std::uint32_t mla_page_rows) { + if (maximum_context_tokens == 0U || mla_page_rows == 0U) { + return {{"GLM-5.3 sequence geometry is invalid"}}; + } + maximum_context_tokens_ = maximum_context_tokens; + mla_page_rows_ = std::min(maximum_context_tokens, mla_page_rows); + token_count_ = 0U; + recurrent_.fill({}); + for (auto& layer : convolution_) layer.fill({}); + for (std::uint32_t layer = 0U; layer < kGlm53LayerCount; ++layer) { + auto result = mla_[layer].reset(kGlm53MlaRank, mla_page_rows_); + if (!result.ok()) return result; + } + return {}; +} + +std::span Glm53SequenceState::writable(Buffer& buffer, + std::size_t elements) { + if (buffer == nullptr) { + buffer = std::make_shared>(elements, 0.0F); + } else if (!buffer.unique()) { + buffer = std::make_shared>(*buffer); + } + return *buffer; +} + +std::span Glm53SequenceState::recurrent(std::uint32_t layer) { + if (layer >= kGlm53LayerCount || !glm53_kda_layer(layer)) return {}; + return writable(recurrent_[layer], + static_cast(kGlm53KdaHeads) * kGlm53KdaHeadWidth * + kGlm53KdaHeadWidth); +} + +std::span Glm53SequenceState::convolution( + std::uint32_t layer, std::uint32_t projection) { + if (layer >= kGlm53LayerCount || projection >= 3U || + !glm53_kda_layer(layer)) return {}; + return writable(convolution_[layer][projection], + static_cast(kGlm53KdaWidth) * 3U); +} + +Glm53PagedRows& Glm53SequenceState::mla(std::uint32_t layer) { + return mla_.at(layer); +} + +const Glm53PagedRows& Glm53SequenceState::mla(std::uint32_t layer) const { + return mla_.at(layer); +} + +void Glm53SequenceState::copy_mla_from( + std::uint32_t layer, const Glm53SequenceState& source) { + mla_.at(layer) = source.mla_.at(layer); +} + +std::uint64_t Glm53SequenceState::private_bytes() const noexcept { + std::uint64_t result = 0U; + for (const auto& buffer : recurrent_) { + if (buffer != nullptr && buffer.unique()) + result += buffer->size() * sizeof(float); + } + for (const auto& layer : convolution_) { + for (const auto& buffer : layer) { + if (buffer != nullptr && buffer.unique()) + result += buffer->size() * sizeof(float); + } + } + for (const auto& pages : mla_) result += pages.private_bytes(); + return result; +} + +} // namespace strata diff --git a/src/platform/cuda_backend_stub.cpp b/src/platform/cuda_backend_stub.cpp index 525096d..4345142 100644 --- a/src/platform/cuda_backend_stub.cpp +++ b/src/platform/cuda_backend_stub.cpp @@ -25,12 +25,22 @@ const char* cuda_matmul_route_name(CudaMatmulRoute route) noexcept { case CudaMatmulRoute::PackedOffsetInt: return "packed_offset_int"; case CudaMatmulRoute::Nvfp4Group16: return "nvfp4_group16"; case CudaMatmulRoute::Fp8TensorPage: return "fp8_tensor_page"; + case CudaMatmulRoute::Fp8F32TensorPage: + return "fp8_f32_tensor_page"; case CudaMatmulRoute::Fp8E4m3Block128: return "fp8_e4m3_block128"; + case CudaMatmulRoute::Fp8E4m3Block128F32: + return "fp8_e4m3_block128_f32"; case CudaMatmulRoute::Fp4E2m1Group32: return "fp4_e2m1_group32"; case CudaMatmulRoute::Fp8RegisterFed: return "fp8_register_fed"; + case CudaMatmulRoute::Fp8F32RegisterFed: + return "fp8_f32_register_fed"; case CudaMatmulRoute::Fp4RegisterFed: return "fp4_register_fed"; case CudaMatmulRoute::GemmaMarlin: return "gemma_marlin"; case CudaMatmulRoute::MoePlainBf16: return "moe_plain_bf16"; + case CudaMatmulRoute::MoeFp8E4m3Block128F32: + return "moe_fp8_e4m3_block128_f32"; + case CudaMatmulRoute::MoeFp8F32RegisterFed: + return "moe_fp8_f32_register_fed"; case CudaMatmulRoute::MoeNvfp4Group16: return "moe_nvfp4_group16"; case CudaMatmulRoute::MoeFp4E2m1Group32: return "moe_fp4_e2m1_group32"; case CudaMatmulRoute::MoePackedInt4: return "moe_packed_int4"; @@ -96,6 +106,10 @@ std::vector CudaBackend::available_devices() { return {}; } ParseResult CudaBackend::device_memory(int) { return {{}, {"CUDA support was not compiled into this build"}}; } +int CudaBackend::device_numa_node(int) noexcept { return -1; } +bool CudaBackend::high_speed_peer_access_supported(int, int) noexcept { + return false; +} std::uint64_t CudaBackend::weight_storage_bytes( std::uint64_t weight_bytes, std::uint64_t scale_bytes) noexcept { @@ -152,6 +166,11 @@ ValidationResult CudaBackend::allocate_buffer(int, std::uint64_t, CudaBuffer&) { return cuda_unavailable(); } +ValidationResult CudaBackend::glm53_kda_decode( + const CudaGlm53KdaRequest&, std::span) { + return cuda_unavailable(); +} + ValidationResult CudaBackend::upload_gemma4_kv( const CudaBuffer&, std::span, std::span, std::uint32_t, std::uint32_t, @@ -187,6 +206,11 @@ ValidationResult CudaBackend::matmul(const CudaWeight&, std::span, return {{"CUDA support was not compiled into this build"}}; } +ValidationResult CudaBackend::matmul_batch( + std::span) { + return {{"CUDA support was not compiled into this build"}}; +} + ValidationResult CudaBackend::matmul_softcap( const CudaWeight&, std::span, float, std::span) { return {{"CUDA support was not compiled into this build"}}; @@ -216,6 +240,14 @@ bool CudaBackend::dsv4_fp8_tensor_page_supported(int) const noexcept { return false; } +bool CudaBackend::fp8_f32_tensor_page_supported(int) const noexcept { + return false; +} + +bool CudaBackend::fp8_f32_register_fed_supported(int) const noexcept { + return false; +} + ValidationResult CudaBackend::validate_dsv4_mhc_device(int) const { return {{"DeepSeek device mHC requires a CUDA-enabled build"}}; } @@ -308,6 +340,11 @@ ValidationResult CudaBackend::dsv4_mhc_finish_device( return {{"DeepSeek device mHC requires a CUDA-enabled build"}}; } +ValidationResult CudaBackend::dsv4_mhc_download_layer_input( + int, std::span) { + return {{"DeepSeek device mHC requires a CUDA-enabled build"}}; +} + ValidationResult CudaBackend::dsv4_mhc_device_view( int, CudaDsv4MhcDeviceView&) { return {{"DeepSeek device mHC requires a CUDA-enabled build"}}; @@ -441,7 +478,7 @@ ValidationResult CudaBackend::collect_deepseek_moe_rows( ValidationResult CudaBackend::enqueue_moe( int, std::span, std::uint32_t, - std::span, const CudaMoeExpert*) { + std::span, const CudaMoeExpert*, float) { return {{"CUDA support was not compiled into this build"}}; } @@ -473,7 +510,7 @@ ValidationResult CudaBackend::collect_moe( ValidationResult CudaBackend::matmul_impl( const CudaWeight&, std::span, std::uint32_t, std::uint32_t, std::uint64_t, std::span, float, bool, - CudaMatmulProfile*, bool) { + CudaMatmulProfile*, bool, const std::byte*, std::byte*, bool) { return {{"CUDA support was not compiled into this build"}}; } diff --git a/src/platform/hardware_profile.cpp b/src/platform/hardware_profile.cpp index c99e6e6..4cfd989 100644 --- a/src/platform/hardware_profile.cpp +++ b/src/platform/hardware_profile.cpp @@ -38,21 +38,31 @@ namespace { // The process's CPU affinity mask, not the machine's CPU count: a cgroup or a // taskset makes those differ, and threads beyond the mask only contend. -[[nodiscard]] std::size_t read_usable_cpus() noexcept { +[[nodiscard]] std::vector read_usable_cpu_ids() noexcept { cpu_set_t set; CPU_ZERO(&set); if (sched_getaffinity(0, sizeof(set), &set) == 0) { - const int count = CPU_COUNT(&set); - if (count > 0) return static_cast(count); + std::vector cpus; + cpus.reserve(static_cast(CPU_COUNT(&set))); + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + if (CPU_ISSET(cpu, &set)) cpus.push_back(cpu); + } + if (!cpus.empty()) return cpus; } const auto concurrency = std::thread::hardware_concurrency(); - return concurrency == 0U ? 1U : static_cast(concurrency); + const auto count = concurrency == 0U ? 1U : concurrency; + std::vector cpus(count); + for (std::size_t index = 0U; index < cpus.size(); ++index) { + cpus[index] = static_cast(index); + } + return cpus; } [[nodiscard]] HardwareProfile probe() { HardwareProfile profile; profile.host_memory_bytes = read_host_memory_bytes(); - profile.usable_cpus = read_usable_cpus(); + profile.usable_cpu_ids = read_usable_cpu_ids(); + profile.usable_cpus = profile.usable_cpu_ids.size(); profile.numa = NumaTopology::detect(); std::size_t smallest = 0U; for (const auto& cpus : profile.numa.node_cpus) { diff --git a/tests/test_cuda_backend.cpp b/tests/test_cuda_backend.cpp index 5b04a31..4985153 100644 --- a/tests/test_cuda_backend.cpp +++ b/tests/test_cuda_backend.cpp @@ -5,6 +5,7 @@ #include "strata/models/deepseek/deepseek_kv_cache.hpp" #include "strata/models/deepseek/deepseek_host_expert.hpp" #include "strata/models/deepseek/deepseek_ops.hpp" +#include "strata/models/kimi_k3/kimi_k3_ops.hpp" #include "strata/models/deepseek/deepseek_attention_kv.hpp" #include "strata/platform/numerics.hpp" @@ -154,6 +155,39 @@ strata::CudaWeight upload_fp8( return result; } +strata::CudaWeight upload_glm_fp8( + strata::CudaBackend& backend, int device, std::uint64_t rows, + std::uint64_t columns, std::uint8_t seed) { + strata::CudaWeightDescriptor descriptor; + descriptor.encoding = strata::CudaWeightEncoding::Fp8E4m3Block128F32; + descriptor.dtype = strata::SafetensorsDtype::F8E4M3; + descriptor.rows = rows; + descriptor.columns = columns; + descriptor.packed_columns = columns; + descriptor.scale_columns = (columns + 127U) / 128U; + descriptor.group_size = 128U; + constexpr std::array encodings{ + 0x00U, 0x28U, 0xA8U, 0x30U, 0xB0U, 0x38U, 0xB8U, 0x20U}; + std::vector weights(static_cast(rows * columns)); + for (std::size_t index = 0U; index < weights.size(); ++index) { + weights[index] = static_cast( + encodings[(index * 5U + seed) % encodings.size()]); + } + const auto scale_rows = (rows + 127U) / 128U; + std::vector scales( + static_cast(scale_rows * descriptor.scale_columns) * + sizeof(float)); + for (std::size_t index = 0U; index < scales.size() / sizeof(float); ++index) { + const float scale = 0.125F + + static_cast((index + seed) % 3U) * 0.03125F; + std::memcpy(scales.data() + index * sizeof(float), &scale, + sizeof(scale)); + } + strata::CudaWeight result; + REQUIRE(backend.upload(device, descriptor, weights, scales, result).ok()); + return result; +} + strata::CudaWeight upload_int4( strata::CudaBackend& backend, int device, std::uint64_t rows, std::uint64_t columns, std::uint8_t seed) { @@ -1915,6 +1949,322 @@ TEST_CASE("SM86 DeepSeek FP8 page projections match at the BF16 boundary") { } } +TEST_CASE("GLM-5.3 F32-scaled FP8 uses continuous dynamic activation scales") { + const auto devices = strata::CudaBackend::available_devices(); + if (!strata::CudaBackend::compiled() || devices.empty()) return; + strata::CudaBackend backend; + const std::array selected{devices.front()}; + REQUIRE(backend.initialize(selected).ok()); + + strata::CudaWeightDescriptor descriptor; + descriptor.encoding = + strata::CudaWeightEncoding::Fp8E4m3Block128F32; + descriptor.dtype = strata::SafetensorsDtype::F8E4M3; + descriptor.rows = 1U; + descriptor.columns = 128U; + descriptor.packed_columns = 128U; + descriptor.scale_columns = 1U; + descriptor.group_size = 128U; + std::array weights{}; + weights.front() = std::byte{0x38U}; // E4M3 1.0 + std::array scales{}; + const float unit_scale = 1.0F; + std::memcpy(scales.data(), &unit_scale, sizeof(unit_scale)); + strata::CudaWeight weight; + REQUIRE(backend.upload(devices.front(), descriptor, weights, scales, + weight).ok()); + + std::array input{}; + // A power-of-two activation scale would round this to an E4M3 grid point. + // Continuous max/448 scaling maps the maximum to 448 exactly and therefore + // reconstructs 300 exactly before accumulation. + input.front() = 300.0F; + std::array output{}; + REQUIRE(backend.matmul(weight, input, 1U, output).ok()); + REQUIRE_NEAR(output.front(), 300.0F, 1.0e-4F); + REQUIRE(strata::cuda_matmul_route_census() + .counts[static_cast( + strata::CudaMatmulRoute::Fp8E4m3Block128F32)] > 0U); +} + +TEST_CASE("CUDA matmul batches preserve exact outputs with one completion") { + const auto devices = strata::CudaBackend::available_devices(); + if (!strata::CudaBackend::compiled() || devices.empty()) return; + strata::CudaBackend backend; + const std::array selected{devices.front()}; + REQUIRE(backend.initialize(selected).ok()); + + auto first_weight = upload_glm_fp8(backend, devices.front(), 128U, 128U, 3U); + auto second_weight = upload_glm_fp8(backend, devices.front(), 128U, 128U, 7U); + std::array first_input{}; + std::array second_input{}; + for (std::size_t index = 0U; index < first_input.size(); ++index) { + first_input[index] = static_cast(index % 11U) * 0.03125F - 0.125F; + second_input[index] = static_cast(index % 7U) * -0.046875F + 0.25F; + } + std::array first_reference{}, second_reference{}; + std::array first_batched{}, second_batched{}; + const auto before = backend.stats().synchronization_calls; + REQUIRE(backend.matmul(first_weight, first_input, 1U, first_reference, + true).ok()); + REQUIRE(backend.matmul(second_weight, second_input, 1U, second_reference, + true).ok()); + const auto after_sequential = backend.stats().synchronization_calls; + const std::array batch{ + {{&first_weight, first_input, 1U, first_batched, true, false}, + {&second_weight, second_input, 1U, second_batched, true, false}}}; + REQUIRE(backend.matmul_batch(batch).ok()); + const auto after_batch = backend.stats().synchronization_calls; + + REQUIRE(first_batched == first_reference); + REQUIRE(second_batched == second_reference); + REQUIRE(after_sequential - before == 2U); + REQUIRE(after_batch - after_sequential == 1U); +} + +TEST_CASE("GLM-5.3 F32-scaled FP8 page tensor route preserves its BF16 contract") { + const auto devices = strata::CudaBackend::available_devices(); + if (!strata::CudaBackend::compiled() || devices.empty()) return; + const int device = devices.front(); + strata::CudaBackend backend; + const std::array selected{device}; + REQUIRE(backend.initialize(selected).ok()); + if (!backend.fp8_f32_tensor_page_supported(device)) return; + + constexpr std::uint32_t batch = 17U; + constexpr std::uint32_t columns = 256U; + constexpr std::uint32_t rows = 128U; + constexpr std::uint32_t scale_columns = columns / 128U; + strata::CudaWeightDescriptor descriptor; + descriptor.encoding = + strata::CudaWeightEncoding::Fp8E4m3Block128F32; + descriptor.dtype = strata::SafetensorsDtype::F8E4M3; + descriptor.rows = rows; + descriptor.columns = columns; + descriptor.packed_columns = columns; + descriptor.scale_columns = scale_columns; + descriptor.group_size = 128U; + + constexpr std::array codes{ + 0x00U, 0x20U, 0xA0U, 0x28U, 0xA8U, + 0x30U, 0xB0U, 0x38U, 0xB8U, 0x40U}; + std::vector weight_bytes( + static_cast(rows) * columns); + for (std::size_t index = 0U; index < weight_bytes.size(); ++index) { + weight_bytes[index] = + static_cast(codes[(index * 7U + 3U) % codes.size()]); + } + std::array weight_scale_values{0.137F, 0.219F}; + std::vector weight_scales( + scale_columns * sizeof(float)); + std::memcpy(weight_scales.data(), weight_scale_values.data(), + weight_scales.size()); + strata::CudaWeight weight; + REQUIRE(backend.upload(device, descriptor, weight_bytes, weight_scales, + weight).ok()); + strata::CudaWeight prepacked_weight; + REQUIRE(backend.upload(device, descriptor, weight_bytes, weight_scales, + prepacked_weight).ok()); + REQUIRE(backend.prepack_fragment(device, prepacked_weight).ok()); + + std::vector input(static_cast(batch) * columns); + std::array activation_scales{}; + for (std::uint32_t batch_row = 0U; batch_row < batch; ++batch_row) { + for (std::uint32_t group = 0U; group < scale_columns; ++group) { + const float scale = 0.173F + 0.019F * batch_row + 0.031F * group; + activation_scales[batch_row * scale_columns + group] = scale; + const auto base = static_cast(batch_row) * columns + + group * 128U; + // Force max/448 to recover the deliberately non-power-of-two + // scale. Every other value is also exactly E4M3*scale, so the + // independent oracle knows the compact activation codes. + input[base] = 448.0F * scale; + for (std::uint32_t column = 1U; column < 128U; ++column) { + const auto code = codes[(batch_row * 11U + group * 5U + + column * 3U) % codes.size()]; + input[base + column] = decode_e4m3(code) * scale; + } + } + } + + std::vector incumbent(static_cast(batch) * rows); + std::vector tensor(incumbent.size()); + std::vector prepacked_tensor(incumbent.size()); + strata::reset_cuda_matmul_route_census(); + REQUIRE(backend.matmul(weight, input, batch, incumbent, true, nullptr, + false).ok()); + REQUIRE(backend.matmul(weight, input, batch, tensor, true, nullptr, + true).ok()); + REQUIRE(backend.matmul(prepacked_weight, input, batch, prepacked_tensor, + true, nullptr, true).ok()); + REQUIRE(strata::cuda_matmul_route_census() + .counts[static_cast( + strata::CudaMatmulRoute::Fp8F32TensorPage)] == 2U); + for (std::size_t index = 0U; index < tensor.size(); ++index) { + REQUIRE(std::bit_cast(prepacked_tensor[index]) == + std::bit_cast(tensor[index])); + } + + double incumbent_squared = 0.0; + double tensor_squared = 0.0; + double incumbent_maximum = 0.0; + double tensor_maximum = 0.0; + for (std::uint32_t batch_row = 0U; batch_row < batch; ++batch_row) { + for (std::uint32_t output_row = 0U; output_row < rows; ++output_row) { + double oracle = 0.0; + for (std::uint32_t column = 0U; column < columns; ++column) { + const auto group = column / 128U; + const float activation_code = column % 128U == 0U + ? 448.0F + : decode_e4m3(codes[(batch_row * 11U + group * 5U + + (column % 128U) * 3U) % codes.size()]); + const auto weight_code = std::to_integer( + weight_bytes[static_cast(output_row) * + columns + column]); + oracle += static_cast(activation_code) * + activation_scales[batch_row * scale_columns + group] * + decode_e4m3(weight_code) * + weight_scale_values[group]; + } + const auto expected = round_bf16(static_cast(oracle)); + const auto index = static_cast(batch_row) * rows + + output_row; + const double incumbent_error = + std::fabs(static_cast(incumbent[index]) - expected); + const double tensor_error = + std::fabs(static_cast(tensor[index]) - expected); + incumbent_squared += incumbent_error * incumbent_error; + tensor_squared += tensor_error * tensor_error; + incumbent_maximum = std::max(incumbent_maximum, incumbent_error); + tensor_maximum = std::max(tensor_maximum, tensor_error); + } + } + // Reassociation is permitted, but the tensor route must not weaken the + // carried BF16 result relative to an independent FP64 block-scale oracle. + REQUIRE(tensor_maximum <= incumbent_maximum); + REQUIRE(tensor_squared <= incumbent_squared); +} + +TEST_CASE("GLM-5.3 F32-scaled FP8 MoE batch matches scalar expert execution") { + const auto devices = strata::CudaBackend::available_devices(); + if (!strata::CudaBackend::compiled() || devices.empty()) return; + constexpr std::uint64_t hidden_columns = 128U; + constexpr std::uint64_t intermediate_columns = 128U; + const int device = devices.front(); + strata::CudaBackend backend; + const std::array selected{device}; + REQUIRE(backend.initialize(selected, true).ok()); + + auto routed_gate = upload_glm_fp8( + backend, device, intermediate_columns, hidden_columns, 1U); + auto routed_up = upload_glm_fp8( + backend, device, intermediate_columns, hidden_columns, 3U); + auto routed_down = upload_glm_fp8( + backend, device, hidden_columns, intermediate_columns, 5U); + auto shared_gate = upload_glm_fp8( + backend, device, intermediate_columns, hidden_columns, 7U); + auto shared_up = upload_glm_fp8( + backend, device, intermediate_columns, hidden_columns, 2U); + auto shared_down = upload_glm_fp8( + backend, device, hidden_columns, intermediate_columns, 4U); + std::array hidden{}; + for (std::size_t index = 0U; index < hidden.size(); ++index) { + hidden[index] = static_cast(static_cast(index % 19U) - 9) / + 16.0F; + } + const auto expected_routed = reference_expert( + backend, routed_gate, routed_up, routed_down, hidden, + intermediate_columns, 1.0F, true); + const auto expected_shared = reference_expert( + backend, shared_gate, shared_up, shared_down, hidden, + intermediate_columns, 1.0F, false); + + REQUIRE(backend.prepack_fragment(device, routed_gate).ok()); + REQUIRE(backend.prepack_fragment(device, routed_up).ok()); + REQUIRE(backend.prepack_fragment(device, routed_down).ok()); + REQUIRE(backend.prepack_fragment(device, shared_gate).ok()); + REQUIRE(backend.prepack_fragment(device, shared_up).ok()); + REQUIRE(backend.prepack_fragment(device, shared_down).ok()); + + const std::array routed{{ + {&routed_gate, &routed_up, &routed_down, 1.0F}, + }}; + const strata::CudaMoeExpert shared{ + &shared_gate, &shared_up, &shared_down, 1.0F}; + std::array routed_output{}; + std::array shared_output{}; + strata::reset_cuda_matmul_route_census(); + REQUIRE(backend.enqueue_moe(device, hidden, 1U, routed, &shared, 10.0F).ok()); + REQUIRE(backend.collect_moe(device, routed_output, shared_output).ok()); + for (std::size_t index = 0U; index < hidden_columns; ++index) { + REQUIRE_NEAR(routed_output[index], expected_routed[index], 1.0e-4F); + REQUIRE_NEAR(shared_output[index], expected_shared[index], 1.0e-4F); + } + const auto census = strata::cuda_matmul_route_census(); + REQUIRE(census.counts[static_cast( + strata::CudaMatmulRoute::MoeFp8F32RegisterFed)] == 1U); +} + +TEST_CASE("GLM deferred prefetch overlaps a leased MoE command exactly") { + const auto devices = strata::CudaBackend::available_devices(); + if (!strata::CudaBackend::compiled() || devices.empty()) return; + constexpr std::uint64_t columns = 128U; + const int device = devices.front(); + strata::CudaBackend backend; + const std::array selected_device{device}; + REQUIRE(backend.initialize(selected_device, true).ok()); + + auto gate = upload_glm_fp8(backend, device, columns, columns, 1U); + auto up = upload_glm_fp8(backend, device, columns, columns, 3U); + auto down = upload_glm_fp8(backend, device, columns, columns, 5U); + auto shared_gate = upload_glm_fp8(backend, device, columns, columns, 7U); + auto shared_up = upload_glm_fp8(backend, device, columns, columns, 9U); + auto shared_down = upload_glm_fp8(backend, device, columns, columns, 11U); + REQUIRE(backend.prepack_fragment(device, gate).ok()); + REQUIRE(backend.prepack_fragment(device, up).ok()); + REQUIRE(backend.prepack_fragment(device, down).ok()); + REQUIRE(backend.prepack_fragment(device, shared_gate).ok()); + REQUIRE(backend.prepack_fragment(device, shared_up).ok()); + REQUIRE(backend.prepack_fragment(device, shared_down).ok()); + std::array hidden{}; + for (std::size_t index = 0U; index < hidden.size(); ++index) { + hidden[index] = static_cast(index % 13U) / 32.0F; + } + const std::array routed{{ + {&gate, &up, &down, 1.0F}, + }}; + const strata::CudaMoeExpert shared{ + &shared_gate, &shared_up, &shared_down, 1.0F}; + REQUIRE(backend.enqueue_moe( + device, hidden, 1U, routed, &shared, 10.0F).ok()); + + strata::CudaWeightDescriptor descriptor; + descriptor.encoding = strata::CudaWeightEncoding::Fp8E4m3Block128F32; + descriptor.dtype = strata::SafetensorsDtype::F8E4M3; + descriptor.rows = columns; + descriptor.columns = columns; + descriptor.packed_columns = columns; + descriptor.scale_columns = 1U; + descriptor.group_size = 128U; + std::vector payload(columns * columns, std::byte{0x38}); + std::vector scale_values(1U, 1.0F); + strata::CudaWeight prefetched; + REQUIRE(backend.upload( + device, descriptor, payload, std::as_bytes(std::span(scale_values)), + prefetched, + strata::CudaBackend::UploadCompletion::DeferredConcurrent).ok()); + REQUIRE(backend.synchronize_uploads(device).ok()); + std::array routed_output{}; + std::array shared_output{}; + REQUIRE(backend.collect_moe( + device, routed_output, shared_output).ok()); + + std::array projected{}; + REQUIRE(backend.matmul(prefetched, hidden, 1U, projected, true).ok()); + REQUIRE(std::all_of(projected.begin(), projected.end(), + [](float value) { return std::isfinite(value); })); +} + TEST_CASE("native CUDA backend batches reusable target-shape GLM INT4 MoE commands") { const auto devices = strata::CudaBackend::available_devices(); if (!strata::CudaBackend::compiled() || devices.empty()) return; @@ -3422,6 +3772,70 @@ TEST_CASE("MIX-2 register-fed matmul matches the scalar route it replaces") { strata::set_register_fed_matmul(true); } +TEST_CASE("GLM-5.3 continuous-scale register-fed W8A8 matches scalar execution") { + const auto devices = strata::CudaBackend::available_devices(); + if (!strata::CudaBackend::compiled() || devices.empty()) return; + strata::CudaBackend backend; + const int device = devices.front(); + REQUIRE(backend.initialize(std::vector{device}, true).ok()); + constexpr std::array palette{ + 1.0F, -1.0F, 0.5F, 2.0F, -0.5F, + -2.0F, 0.03125F, -0.0625F, 0.0F}; + + const auto compare = [&](std::uint64_t output_rows, + std::uint64_t columns, + std::uint32_t batch, + std::uint8_t seed) { + const auto control = + upload_glm_fp8(backend, device, output_rows, columns, seed); + const auto candidate = + upload_glm_fp8(backend, device, output_rows, columns, seed); + std::vector activation( + static_cast(columns) * batch); + for (std::size_t index = 0U; index < activation.size(); ++index) { + activation[index] = palette[(index * 5U + seed) % palette.size()]; + } + std::vector expected( + static_cast(output_rows) * batch); + std::vector measured(expected.size()); + + strata::set_register_fed_matmul(false); + REQUIRE(backend.matmul(control, activation, batch, expected, true).ok()); + strata::set_register_fed_matmul(true); + REQUIRE(backend.prepack_fragment(device, candidate).ok()); + strata::reset_cuda_matmul_route_census(); + REQUIRE(backend.matmul(candidate, activation, batch, measured, true).ok()); + REQUIRE(strata::cuda_matmul_route_census() + .counts[static_cast( + strata::CudaMatmulRoute::Fp8F32RegisterFed)] > 0U); + + double worst = 0.0; + for (std::size_t index = 0U; index < expected.size(); ++index) { + const double denominator = std::max( + 1.0, std::fabs(static_cast(expected[index]))); + worst = std::max( + worst, + std::fabs(static_cast(measured[index]) - + static_cast(expected[index])) / + denominator); + } + if (!(worst < 2.0e-3)) { + std::fprintf(stderr, + "GLM register-fed mismatch: N %llu K %llu M %u " + "worst relative residual %g\n", + static_cast(output_rows), + static_cast(columns), batch, + worst); + } + REQUIRE(worst < 2.0e-3); + }; + + compare(128U, 128U, 1U, 0x13U); + compare(128U, 256U, 5U, 0x29U); + compare(256U, 512U, 16U, 0x47U); + strata::set_register_fed_matmul(true); +} + TEST_CASE("Gemma 4 shaped MXFP4 register-fed matmul matches identical scalar uploads") { const auto devices = strata::CudaBackend::available_devices(); if (!strata::CudaBackend::compiled() || devices.empty()) return; @@ -3728,3 +4142,163 @@ TEST_CASE("a partially prepacked MXFP4 MoE batch is refused, not half-served") { const auto refused = backend.enqueue_moe(device, hidden, 1U, routed, nullptr); REQUIRE(!refused.ok()); } + +TEST_CASE("native CUDA GLM-5.3 KDA keeps exact recurrent state resident") { + const auto devices = strata::CudaBackend::available_devices(); + if (!strata::CudaBackend::compiled() || devices.empty()) return; + constexpr std::uint32_t heads = 2U; + constexpr std::uint32_t dim = 8U; + constexpr std::uint32_t kernel = 4U; + constexpr std::uint32_t width = heads * dim; + const auto recurrent_floats = static_cast(heads) * dim * dim; + const auto convolution_floats = 3U * width * (kernel - 1U); + const auto tap_floats = 3U * width * kernel; + std::vector packed(recurrent_floats + convolution_floats + + tap_floats + heads + width + dim); + for (std::size_t index = 0U; index < recurrent_floats; ++index) { + packed[index] = static_cast(static_cast(index % 17U) - 8) / + 512.0F; + } + auto* convolution = packed.data() + recurrent_floats; + auto* taps = convolution + convolution_floats; + for (std::size_t index = 0U; index < convolution_floats; ++index) { + convolution[index] = static_cast(static_cast(index % 7U) - 3) / + 128.0F; + } + for (std::size_t index = 0U; index < tap_floats; ++index) { + taps[index] = static_cast(static_cast(index % 11U) - 5) / + 32.0F; + } + auto* a_log = taps + tap_floats; + auto* dt_bias = a_log + heads; + auto* norm = dt_bias + width; + for (std::uint32_t head = 0U; head < heads; ++head) a_log[head] = -0.4F; + for (std::uint32_t index = 0U; index < width; ++index) { + dt_bias[index] = static_cast(index % 5U) / 32.0F; + } + for (std::uint32_t index = 0U; index < dim; ++index) norm[index] = 0.75F; + + auto expected_state = packed; + std::vector query(width), key(width), value(width), forget(width), + gate(width), beta(heads); + for (std::uint32_t index = 0U; index < width; ++index) { + query[index] = static_cast(static_cast(index) - 7) / 32.0F; + key[index] = static_cast(static_cast(index % 9U) - 4) / 24.0F; + value[index] = static_cast(static_cast(index % 13U) - 6) / 20.0F; + forget[index] = static_cast(static_cast(index % 7U) - 3) / 16.0F; + gate[index] = static_cast(static_cast(index % 6U) - 2) / 8.0F; + } + beta[0] = round_bf16(0.45F); + beta[1] = round_bf16(0.65F); + + auto expected_q = query; + auto expected_k = key; + auto expected_v = value; + auto* expected_convolution = expected_state.data() + recurrent_floats; + for (std::uint32_t projection = 0U; projection < 3U; ++projection) { + auto& values = projection == 0U ? expected_q + : projection == 1U ? expected_k : expected_v; + auto history = std::span(expected_convolution + + static_cast(projection) * width * (kernel - 1U), + width * (kernel - 1U)); + std::vector convolved(width); + REQUIRE(strata::kimi_short_conv_step( + convolved, values, + std::span(taps + + static_cast(projection) * width * kernel, + width * kernel), history, kernel).ok()); + for (auto& element : convolved) element = round_bf16(element); + values = std::move(convolved); + } + std::vector expected(width); + for (std::uint32_t head = 0U; head < heads; ++head) { + const auto begin = static_cast(head) * dim; + auto q = std::span(expected_q).subspan(begin, dim); + auto k = std::span(expected_k).subspan(begin, dim); + REQUIRE(strata::kimi_l2_normalize(q, 1.0e-6F).ok()); + REQUIRE(strata::kimi_l2_normalize(k, 1.0e-6F).ok()); + for (auto& element : q) element /= std::sqrt(static_cast(dim)); + std::vector decay(dim); + REQUIRE(strata::kimi_kda_log_decay( + decay, std::span(forget).subspan(begin, dim), + std::span(dt_bias, width).subspan(begin, dim), + a_log[head], -5.0F).ok()); + for (auto& element : decay) element = std::exp(element); + std::vector raw(dim); + REQUIRE(strata::kimi_kda_step( + raw, + std::span(expected_state).subspan( + static_cast(head) * dim * dim, dim * dim), + q, k, std::span(expected_v).subspan(begin, dim), + decay, beta[head], dim, dim).ok()); + for (auto& element : raw) element = round_bf16(element); + REQUIRE(strata::kimi_kda_output_norm( + std::span(expected).subspan(begin, dim), raw, + std::span(gate).subspan(begin, dim), + std::span(norm, dim), 1.0e-5F).ok()); + for (auto& element : std::span(expected).subspan(begin, dim)) { + element = round_bf16(element); + } + } + + strata::CudaBackend backend; + const std::array selected{devices.front()}; + REQUIRE(backend.initialize(selected, true).ok()); + strata::CudaBuffer state; + REQUIRE(backend.upload_buffer(devices.front(), std::as_bytes( + std::span(packed)), state).ok()); + std::vector actual(width); + strata::CudaGlm53KdaRequest request; + request.state = &state; + request.query = query; + request.key = key; + request.value = value; + request.forget = forget; + request.beta = beta; + request.gate = gate; + request.heads = heads; + request.head_dim = dim; + request.convolution_kernel = kernel; + REQUIRE(backend.glm53_kda_decode(request, actual).ok()); + for (std::size_t index = 0U; index < actual.size(); ++index) { + REQUIRE_NEAR(actual[index], expected[index], 2.0e-3F); + } + std::vector measured_state(recurrent_floats + convolution_floats); + REQUIRE(backend.download_buffer( + state, 0U, std::as_writable_bytes(std::span(measured_state))).ok()); + for (std::size_t index = 0U; index < measured_state.size(); ++index) { + REQUIRE_NEAR(measured_state[index], expected_state[index], 2.0e-3F); + } + + // The production checkpoint's KDA o_proj is BF16. Prove the chained + // command consumes the resident KDA row directly and publishes the same + // rounded projection as the standalone boundary. + strata::CudaWeightDescriptor projection_descriptor; + projection_descriptor.encoding = strata::CudaWeightEncoding::Plain; + projection_descriptor.dtype = strata::SafetensorsDtype::Bf16; + projection_descriptor.rows = dim; + projection_descriptor.columns = width; + std::vector projection_bytes( + static_cast(dim) * width * 2U); + for (std::uint32_t row = 0U; row < dim; ++row) { + for (std::uint32_t column = 0U; column < width; ++column) { + const auto encoded = bf16(column == row ? 1.0F : 0.0F); + std::copy(encoded.begin(), encoded.end(), + projection_bytes.begin() + static_cast( + (static_cast(row) * width + column) * 2U)); + } + } + strata::CudaWeight projection; + REQUIRE(backend.upload(devices.front(), projection_descriptor, + projection_bytes, {}, projection).ok()); + strata::CudaBuffer chained_state; + REQUIRE(backend.upload_buffer(devices.front(), std::as_bytes( + std::span(packed)), chained_state).ok()); + request.state = &chained_state; + request.output_projection = &projection; + std::vector projected(dim); + REQUIRE(backend.glm53_kda_decode(request, projected).ok()); + for (std::uint32_t index = 0U; index < dim; ++index) { + REQUIRE_NEAR(projected[index], expected[index], 2.0e-3F); + } +} diff --git a/tests/test_glm53_manifest.cpp b/tests/test_glm53_manifest.cpp new file mode 100644 index 0000000..61babbd --- /dev/null +++ b/tests/test_glm53_manifest.cpp @@ -0,0 +1,138 @@ +#include "test.hpp" + +#include "strata/engine/model_executor.hpp" +#include "strata/models/common/tokenizer.hpp" +#include "strata/models/glm53/glm53_manifest.hpp" +#include "strata/models/glm53/glm53_runtime.hpp" +#include "strata/models/glm53/glm53_sequence.hpp" + +#include +#include +#include +#include + +TEST_CASE("GLM-5.3 schedule partitions KDA, sparse MLA, dense, and MoE layers") { + std::uint32_t kda = 0U; + std::uint32_t sparse = 0U; + std::uint32_t dense = 0U; + std::uint32_t moe = 0U; + for (std::uint32_t layer = 0U; layer < 45U; ++layer) { + if (strata::glm53_kda_layer(layer)) ++kda; else ++sparse; + if (strata::glm53_moe_layer(layer)) ++moe; else ++dense; + } + REQUIRE(kda == 34U); + REQUIRE(sparse == 11U); + REQUIRE(dense == 3U); + REQUIRE(moe == 42U); + REQUIRE(strata::glm53_full_attention_layer(3U)); + REQUIRE(strata::glm53_kda_layer(44U)); +} + +TEST_CASE("GLM-5.3 tensor classifier separates text, MTP, and vision") { + std::int32_t layer = -1; + std::int32_t expert = -1; + const auto classify = [&](std::string_view name) { + return strata::classify_glm53_tensor(name, layer, expert); + }; + REQUIRE(classify("model.language_model.layers.0.self_attn.A_log") == + strata::Glm53TensorRole::KdaAttention); + REQUIRE(layer == 0); + REQUIRE(classify("model.language_model.layers.3.self_attn.q_a_proj.weight") == + strata::Glm53TensorRole::SparseAttention); + REQUIRE(classify("model.language_model.layers.3.self_attn.indexer.wk.weight") == + strata::Glm53TensorRole::AttentionIndexer); + REQUIRE(classify("model.language_model.layers.2.mlp.gate_proj.weight") == + strata::Glm53TensorRole::DenseMlp); + REQUIRE(classify("model.language_model.layers.10.mlp.experts.287.up_proj.weight") == + strata::Glm53TensorRole::RoutedExpert); + REQUIRE(expert == 287); + REQUIRE(classify("model.language_model.layers.10.mlp.shared_experts.down_proj.weight") == + strata::Glm53TensorRole::SharedExpert); + REQUIRE(classify("model.language_model.layers.45.shared_head.norm.weight") == + strata::Glm53TensorRole::Mtp); + REQUIRE(classify("model.visual.blocks.0.attn.qkv.weight") == + strata::Glm53TensorRole::Vision); +} + +TEST_CASE("GLM-5.3 projection assignment is deterministic and capacity weighted") { + constexpr std::array keys{ + "q", "k", "v", "f", "b", "g", "up", "down"}; + constexpr std::array costs{ + 10U, 10U, 10U, 10U, 10U, 10U, 10U, 10U}; + constexpr std::array equal{20U, 20U}; + const auto first = strata::glm53_projection_slots(keys, costs, equal, 1U); + const auto second = strata::glm53_projection_slots(keys, costs, equal, 1U); + REQUIRE(first == second); + REQUIRE(std::count(first.begin(), first.end(), 0U) == 4); + REQUIRE(std::count(first.begin(), first.end(), 1U) == 4); + + constexpr std::array weighted{30U, 10U}; + const auto asymmetric = + strata::glm53_projection_slots(keys, costs, weighted, 0U); + REQUIRE(std::count(asymmetric.begin(), asymmetric.end(), 0U) == 6); + REQUIRE(std::count(asymmetric.begin(), asymmetric.end(), 1U) == 2); + constexpr std::array invalid{30U, 0U}; + REQUIRE(strata::glm53_projection_slots(keys, costs, invalid, 0U).empty()); +} + +TEST_CASE("GLM-5.3 registers as a text-only chat and server model") { + const auto* model = strata::find_model(strata::RuntimeModel::Glm53); + REQUIRE(model != nullptr); + REQUIRE(std::string_view(model->cli_name) == "glm53"); + REQUIRE(model->placement == strata::PlacementModel::Glm53); + const auto executor = model->make(); + REQUIRE(executor != nullptr); + REQUIRE(!executor->accepts_images()); +} + +TEST_CASE("GLM-5.3 physical sequence pages fork copy-on-write") { + strata::Glm53PagedRows rows(4U, 2U); + constexpr std::array first{1, 2, 3, 4}; + constexpr std::array second{5, 6, 7, 8}; + constexpr std::array third{9, 10, 11, 12}; + REQUIRE(rows.append(first).ok()); + REQUIRE(rows.append(second).ok()); + auto fork = rows; + REQUIRE(rows.private_bytes() == 0U); + REQUIRE(fork.private_bytes() == 0U); + REQUIRE(fork.append(third).ok()); + REQUIRE(rows.rows() == 2U); + REQUIRE(fork.rows() == 3U); + REQUIRE(fork.row(0U)[0] == 1.0F); + REQUIRE(fork.row(2U)[3] == 12.0F); + REQUIRE(fork.private_bytes() != 0U); +} + +TEST_CASE("GLM-5.3 recurrent state forks lazily and exactly") { + strata::Glm53SequenceState state; + REQUIRE(state.reset(128U, 16U).ok()); + auto recurrent = state.recurrent(0U); + REQUIRE(!recurrent.empty()); + recurrent[17U] = 3.25F; + auto fork = state; + auto forked = fork.recurrent(0U); + REQUIRE(forked[17U] == 3.25F); + forked[17U] = -2.0F; + REQUIRE(state.recurrent(0U)[17U] == 3.25F); + REQUIRE(fork.recurrent(0U)[17U] == -2.0F); +} + +TEST_CASE("GLM-5.3 chat rendering matches its text-only Jinja contract") { + REQUIRE(strata::render_glm53_user_prompt("x") == + "[gMASK]<|system|>Reasoning Effort: Max" + "<|user|>x<|assistant|>"); + const std::array messages{ + strata::ChatMessage{strata::ChatRole::User, "x"}, + strata::ChatMessage{strata::ChatRole::Assistant, " answer "}, + strata::ChatMessage{strata::ChatRole::User, "y"}, + }; + REQUIRE(strata::render_glm53_chat_prompt(messages, "max", true) == + "[gMASK]<|system|>Reasoning Effort: Max" + "<|user|>x<|assistant|>answer" + "<|user|>y<|assistant|>"); + REQUIRE(strata::render_glm53_chat_prompt( + std::span(messages).first(1U), + "unsupported", true) == + "[gMASK]<|system|>Reasoning Effort: Max" + "<|user|>x<|assistant|>"); +} diff --git a/tests/test_placement.cpp b/tests/test_placement.cpp index 0a7fcc3..4a29f01 100644 --- a/tests/test_placement.cpp +++ b/tests/test_placement.cpp @@ -438,6 +438,22 @@ TEST_CASE("the storage tier is admitted wherever the checkpoint lives") { REQUIRE(solve_placement(inventory, unknown, request).ok()); } +TEST_CASE("an explicitly streamed component is accounted as storage resident") { + auto inventory = make_dense_inventory(1U, 1U); + inventory.prescriptive = false; + auto& item = inventory.items.front(); + item.preferred_tier = strata::PlacementTier::Storage; + item.host_bytes = item.source_bytes; + item.device_bytes = 0U; + const auto planned = solve_placement( + inventory, make_hardware({8U}), make_request(1U)); + REQUIRE(planned.ok()); + REQUIRE(planned.value.host_resident_bytes == 0U); + REQUIRE(planned.value.storage_resident_bytes == kGigabyte); + REQUIRE(planned.value.decode_storage_read_bytes == kGigabyte); + REQUIRE(planned.value.io_dependent); +} + TEST_CASE("the backing block device of a real path resolves") { const auto storage = strata::resolve_backing_storage(STRATA_SOURCE_DIR); REQUIRE(storage.resolved); diff --git a/tests/test_runtime.cpp b/tests/test_runtime.cpp index ee389dc..766dfb6 100644 --- a/tests/test_runtime.cpp +++ b/tests/test_runtime.cpp @@ -127,14 +127,15 @@ TEST_CASE("runtime session cannot generate before initialization") { } // Free test 1 (Phase 2 survey, brief 04 task 2): the only two conformance -// properties assertable across all six RuntimeModel values without a real +// properties assertable across all seven RuntimeModel values without a real // checkpoint. A failed initialize() must never leave impl_ holding a runtime // -- the facade falls through to std::monostate, so generation afterward // must report the same "not initialized" error every model gets from a // session that was never touched, not silently produce output or crash. -TEST_CASE("initialization failure leaves generation disabled, all six models") { - constexpr std::array models{ - strata::RuntimeModel::Glm52, strata::RuntimeModel::DeepSeekV4, +TEST_CASE("initialization failure leaves generation disabled, all seven models") { + constexpr std::array models{ + strata::RuntimeModel::Glm52, strata::RuntimeModel::Glm53, + strata::RuntimeModel::DeepSeekV4, strata::RuntimeModel::Gemma4, strata::RuntimeModel::Laguna, strata::RuntimeModel::Inkling, strata::RuntimeModel::KimiK3, }; @@ -197,14 +198,14 @@ TEST_CASE("generation metrics can express reuse fields as not applicable") { // Without -Wl,--whole-archive on strata_models, every ModelRegistrar is // dropped by the linker -- nothing references those objects by name -- and -// find_model returns null for all six models. Every binary then rejects every +// find_model returns null for all seven models. Every binary then rejects every // --model-type. Nothing else in the suite catches that: the initialization // and rejection tests below pass just as happily against an empty registry, // because "unhandled runtime model: 0" is still a failure and an // uninitialized session still says "not initialized". TEST_CASE("every model is registered, which is what --whole-archive buys") { - REQUIRE(strata::registered_models().size() == 6U); - for (const auto* cli_name : {"glm", "deepseek", "gemma4", "kimi-k3", + REQUIRE(strata::registered_models().size() == 7U); + for (const auto* cli_name : {"glm", "glm53", "deepseek", "gemma4", "kimi-k3", "laguna", "inkling"}) { const auto* found = strata::find_model_by_cli_name(cli_name); REQUIRE(found != nullptr);