cuda: accelerate GLM prefill on DGX Spark

Implement GLM 5.3 Q2 and Q4 inference across Metal and CUDA, including tensor parallelism, RDMA, batching, MTP, snapshots, SSD streaming, and GLM 5.2 shared-path improvements.

Add quantization and validation tools, official quality fixtures, model downloads, regression coverage, release QA, and user documentation.

.gitattributes

@@ -0,0 +1,3 @@
+# Preserve hosted-model continuation bytes exactly; these fixtures may contain
+# model-generated trailing whitespace that is meaningful to tokenizer scoring.
+gguf-tools/quality-testing/data/*/continuations/*.txt -diff

.gitignore

@@ -1,14 +1,42 @@
 /ds4
 /ds4-server
 /ds4-bench
+/ds4-eval
+/ds4-agent
 /ds4_native
+/ds4_agent_test
 /ds4_server_test
 /ds4_test
 /ds4flash.gguf
 /TODO.md
 /gguf/
+/core
+/core.*
+/tests/test_engine_mgpu_placement
+/tests/test_engine_mgpu_runtime
+/tests/test_engine_mgpu_refusal
+/tests/test_engine_correctness
+/tests/test_gpu_args
+/tests/test_gpu_lookup_cache_strict
+/tests/test_gpu_model_cache
+/tests/test_gpu_xdev
+/tests/test_cuda_mixed_batch
+/tests/test_cuda_session_batch
+/tests/cuda_long_context_smoke
+/tests/test_layer_pack
+/tests/test_metal_session_batch
+/tests/test_mxfp4_cuda
+/tests/test_mxfp4_dot
+/tests/test_mxfp4_metal
+/tests/test_q4k_dot
+/tests/test_sampling
+/tests/test_glm53_kda
 *.o
 *.dSYM/
+__pycache__/
+*.pyc
+/gguf-tools/libds4quants.dylib
+/gguf-tools/libds4quants.so
 /misc/
 .*.swp
 .DS_Store

AGENT.md

@@ -7,32 +7,26 @@ Objective-C only where Metal requires it and Metal kernels under `metal/`.
 ## Goals
 
 - Keep the production path as whole-model Metal graph inference.
-- Keep model loading mmap-backed; do not eagerly copy the full GGUF.
+- Always make sure that the SSD streaming, CUDA, distributed inference, Metal default inference are not affected by fixes to other parts of the code.
+- Keep model loading mmap-backed for the Metal default case; do not eagerly copy the full GGUF. Keep the model loading for SSD streaming of routed experts explicit: allocated buffers, fast reads from disk, always try to hide loading of missing routed experts by loading them while performing the inference of the shared expert and routed experts already in RAM. Always try to hide loading of layers for prefill in SSD streaming mode using the inference time of the current layer as the next one is loaded.
 - Keep the CPU backend CPU-only and use it only as reference/debug code.
-- Preserve correctness before speed. Do not keep a faster path with unexplained
-  attention, KV cache, or logits drift.
-- Make long local agent sessions practical through live KV reuse and disk KV
-  checkpoints.
+- Preserve correctness before speed. Do not keep a faster path with unexplained attention, KV cache, or logits drift.
+- Make long local agent sessions practical through live KV reuse and disk KV checkpoints.
 
 ## Quality Rules
 
-- Comment important inference code where the model mechanics, cache lifetime,
-  memory policy, or API orchestration are not obvious from the local code.
+- Keep the implementation small, sharp, easy to understand. Try to write elegant code in a state of grace. Don't settle for the first thing that comes to mind, try to find the most minimal and better working design. Don't introduce slop: very fragile code that just patches specific cases, dead code, useless code and code ways more complicated of how it should be.
+- Comment important inference code where the model mechanics, cache lifetime, memory policy, or API orchestration are not obvious from the local code.
 - Prefer comments beside the implementation over separate design documents.
-- Keep comments instructive and compact: explain why a shape, ordering, cache
-  boundary, or memory choice exists.
+- Keep comments instructive and compact: explain why a shape, ordering, cache boundary, or memory choice exists.
 - Keep public APIs narrow. CLI/server code should not know tensor internals.
-- Do not add permanent semantic variants behind flags. Diagnostic switches are
-  fine when they validate the one release path.
+- Do not add permanent semantic variants behind flags. Diagnostic switches are fine when they validate the one release path.
 - Do not introduce C++.
 
 ## Safety
 
-- Avoid large CPU inference runs on macOS; the CPU path has previously exposed
-  kernel VM failures with very large mappings.
-- Do not run multiple huge model processes concurrently. The instance lock is
-  intentional.
-- Prefer short Metal smoke tests for build verification.
+- Avoid large CPU inference runs on macOS; the CPU path has previously exposed kernel VM failures with very large mappings.
+- Do not run multiple huge model processes concurrently. The instance lock is intentional.
 
 ## Layout
 
@@ -46,8 +40,17 @@ Objective-C only where Metal requires it and Metal kernels under `metal/`.
 - `tests/`: unit and live integration tests.
 - `misc/`: ignored notes, experiments, and old planning material.
 
+This list is not complete, check the files for more info.
+
 ## Testing
 
 Use `make` for build validation. Use `make test` for unit/regression tests when a
 model and Metal are available. Use live server tests only when intentionally
 testing the API surface.
+
+At every major change where one of the following could be affected, make sure to:
+
+1. Test the normal Metal path and that speed is still at the level it was.
+2. Test the SSD streaming path.
+3. Test the distributed inference if it could be affected, but ask the user before doing so.
+4. Check if CUDA could be broken after the change, and ask the user to give you access to the CUDA machine to actually test if everything is still fine.

MODEL_CARD.md

@@ -0,0 +1,243 @@
+# DeepSeek v4 model card synopsis
+
+This document extracts the most important information from the official
+DeepSeek-V4-Flash Hugging Face model card, with emphasis on facts that matter
+for local inference, DS4 development, and benchmark interpretation.
+
+Source: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash
+
+## Model Family
+
+DeepSeek-V4 is a preview model family with two Mixture-of-Experts language
+models:
+
+| Model | Total parameters | Active parameters | Context length |
+|---|---:|---:|---:|
+| DeepSeek-V4-Flash | 284B | 13B | 1M tokens |
+| DeepSeek-V4-Pro | 1.6T | 49B | 1M tokens |
+
+Flash is the smaller and more efficient model. The model card says Flash-Max can
+approach Pro reasoning performance when given a larger thinking budget, while
+remaining behind Pro on pure knowledge and the most complex agentic tasks.
+
+## Architecture
+
+DeepSeek-V4 uses long-context compressed attention. The model card calls the
+hybrid design Compressed Sparse Attention (CSA) plus Heavily Compressed
+Attention (HCA). In DS4 terms, each layer keeps a raw sliding-window KV cache
+for the latest 128 tokens. This is the high-resolution local context.
+
+After that raw window, the model uses layer-dependent compressed KV rows:
+
+| 0-based layer indexes | DS4 ratio | Extra state | Meaning |
+|---|---:|---|---|
+| 0, 1 | none | none | Raw 128-token sliding window only |
+| even layers from 2 onward | 4 | compressed KV + indexer KV | One compressed row per 4 tokens, with an indexer selecting visible compressed rows |
+| odd layers from 3 onward | 128 | compressed KV | One compressed row per 128 tokens |
+
+So, after the first two layers, the model alternates ratio-4 and ratio-128
+compressed attention. A token in a compressed layer attends over both the raw
+latest-128-token window and the older compressed history. The compression here
+is time-axis compression: several token positions are pooled into one KV row.
+The attention rows still use the model attention/value dimensions, so raw and
+compressed rows can be consumed by the same mixed-attention computation.
+
+Ratio-4 layers are the selective compressed-attention layers. They maintain a
+second compressed stream for the indexer, and when the compressed history is
+larger than the configured top-k, DS4 scores the compressed rows and selects up
+to 512 of them for attention. Ratio-128 layers are the heavily compressed path:
+they do not have the indexer stream and use the available ratio-128 compressed
+rows directly.
+
+DS4 validates these details from the GGUF metadata. The relevant fixed
+implementation constants are:
+
+- Layers: 43
+- Raw sliding-window attention: 128 tokens
+- Indexer heads: 64
+- Indexer head dimension: 128
+- Indexer top-k: 512
+
+This is the practical reason the model can expose a 1M-token context without a
+standard full KV cache for every token in every layer. The model card reports
+that, at 1M tokens, DeepSeek-V4-Pro needs much less single-token inference
+compute and KV cache than DeepSeek-V3.2.
+
+The family also uses:
+
+- Manifold-Constrained Hyper-Connections (mHC), intended to improve signal
+  propagation stability across layers.
+- The Muon optimizer, used for faster convergence and training stability.
+- A post-training pipeline with domain expert cultivation followed by unified
+  consolidation via on-policy distillation.
+
+## Precision And Weights
+
+Official download entries include:
+
+| Model | Precision |
+|---|---|
+| DeepSeek-V4-Flash-Base | FP8 Mixed |
+| DeepSeek-V4-Flash | FP4 + FP8 Mixed |
+| DeepSeek-V4-Pro-Base | FP8 Mixed |
+| DeepSeek-V4-Pro | FP4 + FP8 Mixed |
+
+For the instruct models, the model card describes FP4 + FP8 Mixed as using FP4
+for MoE expert parameters and FP8 for most other parameters.
+
+## Reasoning Modes
+
+The instruct models support three reasoning-effort modes:
+
+| Mode | Intended behavior | Output shape |
+|---|---|---|
+| Non-think | Fast, intuitive replies | `</think>` summary |
+| High | Deliberate reasoning for harder tasks | `<think>... </think>` summary |
+| Max | Largest reasoning budget | Special system prompt plus thinking and summary |
+
+The model card recommends using at least a 384K-token context window for Think
+Max.
+
+## Important Flash Benchmarks
+
+### DeepSeek-V4-Flash Across Reasoning Modes
+
+| Benchmark | Non-Think | High | Max |
+|---|---:|---:|---:|
+| GPQA Diamond Pass@1 | 71.2 | 87.4 | 88.1 |
+| MMLU-Pro EM | 83.0 | 86.4 | 86.2 |
+| SimpleQA-Verified Pass@1 | 23.1 | 28.9 | 34.1 |
+| Chinese-SimpleQA Pass@1 | 71.5 | 73.2 | 78.9 |
+| HLE Pass@1 | 8.1 | 29.4 | 34.8 |
+| LiveCodeBench Pass@1 | 55.2 | 88.4 | 91.6 |
+| HMMT 2026 Feb Pass@1 | 40.8 | 91.9 | 94.8 |
+| IMOAnswerBench Pass@1 | 41.9 | 85.1 | 88.4 |
+| SWE Verified Resolved | 73.7 | 78.6 | 79.0 |
+| Terminal Bench 2.0 Acc | 49.1 | 56.6 | 56.9 |
+| MCPAtlas Pass@1 | 64.0 | 67.4 | 69.0 |
+| Toolathlon Pass@1 | 40.7 | 43.5 | 47.8 |
+
+### DeepSeek-V4-Flash-Base
+
+The base-model table reports these Flash-Base scores:
+
+| Benchmark | Shots | Score |
+|---|---:|---:|
+| SuperGPQA EM | 5-shot | 46.5 |
+| MMLU EM | 5-shot | 88.7 |
+| MMLU-Pro EM | 5-shot | 68.3 |
+| Simple-QA verified EM | 25-shot | 30.1 |
+| HumanEval Pass@1 | 0-shot | 69.5 |
+| GSM8K EM | 8-shot | 90.8 |
+| LongBench-V2 EM | 1-shot | 44.7 |
+
+The model card reports SuperGPQA for the base model table, not in the instruct
+reasoning-mode comparison table.
+
+## Chat Template And Encoding
+
+The release does not use a Jinja chat template as the source of truth. The
+official prompt renderer is the Python code in
+`encoding/encoding_dsv4.py`, with examples and tests in the same `encoding`
+directory:
+
+- https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/raw/main/encoding/encoding_dsv4.py
+- https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/raw/main/encoding/test_encoding_dsv4.py
+
+The important special tokens are:
+
+| Purpose | Token |
+|---|---|
+| Beginning of sequence | `<｜begin▁of▁sentence｜>` |
+| End of assistant turn | `<｜end▁of▁sentence｜>` |
+| User turn prefix | `<｜User｜>` |
+| Assistant turn prefix | `<｜Assistant｜>` |
+| Latest reminder prefix | `<｜latest_reminder｜>` |
+| Thinking start | `<think>` |
+| Thinking end / non-thinking marker | `</think>` |
+| DSML tool markup marker | `｜DSML｜` |
+
+The renderer accepts `system`, `user`, `assistant`, `tool`,
+`latest_reminder`, and `developer` roles. The `developer` role is described in
+the Python comments as an internal search-agent role, not as a normal public
+chat role.
+
+Normal chat mode starts with the BOS token, then system text if present, then
+alternating user and assistant markers. In non-thinking chat mode, a new
+assistant generation is opened with:
+
+```text
+<｜Assistant｜></think>
+```
+
+That immediate `</think>` tells the model to skip hidden reasoning and produce
+the visible answer. In thinking mode, a new assistant generation is opened with:
+
+```text
+<｜Assistant｜><think>
+```
+
+Completed assistant thinking turns are rendered as reasoning content inside
+`<think>...</think>`, followed by the visible answer and the EOS token.
+
+By default, the Python renderer drops earlier assistant reasoning content before
+the last user message. If tools are present on any message, it disables that
+reasoning drop and keeps the full reasoning/tool context. `reasoning_effort=max`
+also prepends a special high-effort instruction prefix before the first rendered
+message in thinking mode.
+
+Tool definitions are passed in OpenAI-compatible function schema form, but the
+model is instructed to emit DSML. A tool call is rendered as a DSML
+`tool_calls` block containing one or more `invoke` entries, each with named
+parameters. Parameters carry a `string="true"` flag for raw strings and
+`string="false"` for JSON values such as numbers, booleans, arrays, or objects.
+
+DeepSeek-V4 does not render standalone `tool` role messages. The Python
+preprocessor converts tool results into user content blocks and renders each
+result as:
+
+```text
+<tool_result>...</tool_result>
+```
+
+Tool-result bodies are rendered as raw text. Literal `<`, `>`, and `&` from
+file contents or shell output are preserved; only the exact closing sentinel
+`</tool_result>` is escaped so the wrapper cannot be terminated by data.
+
+When there are multiple tool results, the renderer sorts them to match the
+order of the preceding assistant tool calls.
+
+The same script also defines special task tokens for internal quick tasks such
+as title generation, search-query generation, action selection, authority
+classification, domain classification, and URL-read decisions. Those are
+separate from normal chat/tool rendering.
+
+## Local Running Notes
+
+The model card lists vLLM and SGLang examples for OpenAI-compatible serving.
+For local deployment, it recommends:
+
+- `temperature = 1.0`
+- `top_p = 1.0`
+- At least 384K context for Think Max
+
+These are deployment recommendations from the model card, not necessarily the
+same settings used for deterministic benchmarking. DS4 keeps `top_p=1.0` but
+adds a local `min_p=0.05` default to avoid sampling tokens whose probability is
+far below the best token.
+
+## Licensing
+
+The repository and model weights are licensed under the MIT License.
+
+## Citation
+
+The model card cites:
+
+```bibtex
+@misc{deepseekai2026deepseekv4,
+      title={DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence},
+      author={DeepSeek-AI},
+      year={2026},
+}
+```

Makefile

@@ -6,57 +6,145 @@ NATIVE_CPU_FLAG ?= -mcpu=native
 else
 NATIVE_CPU_FLAG ?= -march=native
 endif
+SAMPLING_TEST := tests/test_sampling
+GLM53_KDA_TEST := tests/test_glm53_kda
+GLM53_KDA_ROCM_TEST := tests/test_glm53_kda_rocm
 
-CFLAGS ?= -O3 -ffast-math $(NATIVE_CPU_FLAG) -Wall -Wextra -std=c99
-OBJCFLAGS ?= -O3 -ffast-math $(NATIVE_CPU_FLAG) -Wall -Wextra -fobjc-arc
+DEBUG_FLAGS ?= -g
+CFLAGS ?= -O3 -ffast-math $(DEBUG_FLAGS) $(NATIVE_CPU_FLAG) -Wall -Wextra -std=c99
+OBJCFLAGS ?= -O3 -ffast-math $(DEBUG_FLAGS) $(NATIVE_CPU_FLAG) -Wall -Wextra -fobjc-arc
+QUALITY_CFLAGS ?= -O3 $(DEBUG_FLAGS) $(NATIVE_CPU_FLAG) -Wall -Wextra -std=c11
 
 LDLIBS ?= -lm -pthread
 METAL_SRCS := $(wildcard metal/*.metal)
+ROCM_SRCS := $(wildcard rocm/*.cuh)
+DS4_TEST_MODEL ?= ds4flash.gguf
+DS4_TEST_MTP ?= gguf/DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.gguf
+DS4_DSPARK_MODEL ?= $(DS4_TEST_MODEL)
+DS4_DSPARK_SUPPORT ?= gguf/DeepSeek-V4-Flash-DSpark-support-0731.gguf
 
 ifeq ($(UNAME_S),Darwin)
 METAL_LDLIBS := $(LDLIBS) -framework Foundation -framework Metal
-CORE_OBJS = ds4.o ds4_metal.o
-CPU_CORE_OBJS = ds4_cpu.o
+CORE_OBJS = ds4.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_metal.o ds4_layer_pack.o
+CPU_CORE_OBJS = ds4_cpu.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
 else
 CFLAGS += -D_GNU_SOURCE -fno-finite-math-only
-CUDA_HOME ?= /usr/local/cuda
+CUDA_HOME ?= $(shell if [ -x /usr/local/cuda/bin/nvcc ]; then \
+	printf '%s' /usr/local/cuda; \
+	elif command -v nvcc >/dev/null 2>&1; then \
+	dirname "$$(dirname "$$(command -v nvcc)")"; \
+	else \
+	printf '%s' /usr/local/cuda; \
+	fi)
 NVCC ?= $(CUDA_HOME)/bin/nvcc
 CUDA_ARCH ?=
 ifneq ($(strip $(CUDA_ARCH)),)
+ifneq ($(filter sm_120 sm_120a,$(strip $(CUDA_ARCH))),)
+NVCC_ARCH_FLAGS := -gencode arch=compute_120a,code=sm_120a -DDS4_CUDA_HAVE_MXF4=1
+else ifneq ($(filter sm_121 sm_121a,$(strip $(CUDA_ARCH))),)
+NVCC_ARCH_FLAGS := -gencode arch=compute_121a,code=sm_121a -DDS4_CUDA_HAVE_MXF4=1
+else
 NVCC_ARCH_FLAGS := -arch=$(CUDA_ARCH)
 endif
-NVCCFLAGS ?= -O3 --use_fast_math $(NVCC_ARCH_FLAGS) -Xcompiler $(NATIVE_CPU_FLAG) -Xcompiler -pthread
+endif
+NVCCFLAGS ?= -O3 -g -lineinfo --use_fast_math $(NVCC_ARCH_FLAGS) -Xcompiler $(NATIVE_CPU_FLAG) -Xcompiler -pthread
+# Vendored llama.cpp mmq prefill tier (cuda/mmq/, see cuda/mmq/VENDOR.md).
+MMQ_INCLUDES := -Icuda/mmq
+MMQ_OBJS := cuda/mmq/ds4_ggml_stubs.o cuda/mmq/ds4_mmq.o cuda/mmq/ds4_mmq_d2r.o cuda/mmq/quantize.o cuda/mmq/mmid.o cuda/mmq/mmvq.o cuda/mmq/ds4_repack.o
+CORE_OBJS = ds4.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_cuda.o ds4_layer_pack.o $(MMQ_OBJS)
+CPU_CORE_OBJS = ds4_cpu.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
 CUDA_LDLIBS ?= -lm -Xcompiler -pthread -L$(CUDA_HOME)/targets/sbsa-linux/lib -L$(CUDA_HOME)/lib64 -lcudart -lcublas
-CORE_OBJS = ds4.o ds4_cuda.o
-CPU_CORE_OBJS = ds4_cpu.o
+HIPCC ?= $(shell command -v hipcc 2>/dev/null || echo /opt/rocm/bin/hipcc)
+ROCM_ARCH ?= gfx1151
+ROCM_HOST_CFLAGS ?= -fPIC
+ROCM_CFLAGS ?= -O3 -ffast-math -g -fno-finite-math-only -pthread -D__HIP_PLATFORM_AMD__ -Wno-unused-command-line-argument --offload-arch=$(ROCM_ARCH)
+ROCM_LDLIBS ?= -lm -pthread -lhipblas -lhipblaslt
+DS4_LINK ?= $(NVCC) $(NVCCFLAGS)
+DS4_LINK_LIBS ?= $(CUDA_LDLIBS)
 METAL_LDLIBS := $(LDLIBS)
 endif
 
-.PHONY: all help clean test cpu cuda cuda-spark cuda-generic cuda-regression
+.PHONY: all help clean test test-rocm test-glm53-kda-rocm test-metal-session-batch test-mxfp4-cuda test-mxfp4-rocm test-cuda-session-batch test-cuda-mixed-batch dspark-acceptance dspark-verify-depth mtp-verify-depth cpu cuda cuda-spark cuda-generic cuda-regression strix-halo rocm
 
 ifeq ($(UNAME_S),Darwin)
-all: ds4 ds4-server ds4-bench
+.PHONY: metal-decode-schedule-bench metal-prefill-variant-bench check-mxfp4-half-lut
+
+all: ds4 ds4-server ds4-bench ds4-eval ds4-agent
 
 help:
 	@echo "DS4 build targets:"
-	@echo "  make              Build Metal ./ds4, ./ds4-server, and ./ds4-bench"
-	@echo "  make cpu          Build CPU-only ./ds4, ./ds4-server, and ./ds4-bench"
+	@echo "  make              Build Metal ./ds4, ./ds4-server, ./ds4-bench, ./ds4-eval, and ./ds4-agent"
+	@echo "  make cpu          Build CPU-only ./ds4, ./ds4-server, ./ds4-bench, ./ds4-eval, and ./ds4-agent"
 	@echo "  make test         Build and run tests"
+	@echo "  make metal-decode-schedule-bench  Build the balanced Metal decode schedule benchmark"
+	@echo "  make metal-prefill-variant-bench  Build the balanced Metal prefill variant benchmark"
+	@echo "  make check-mxfp4-half-lut  Verify the checked-in MXFP4 half LUT matches the generator"
+	@echo "  make test-mxfp4-metal  Check the MXFP4 half LUT, then run Metal MXFP4 exactness tests"
+	@echo "  make dspark-verify-depth  Run DSpark speculative verification smoke if support GGUF is present"
+	@echo "  make mtp-verify-depth  Run legacy MTP speculative verification smoke if MTP GGUF is present"
 	@echo "  make clean        Remove build outputs"
 
-ds4: ds4_cli.o linenoise.o $(CORE_OBJS)
-	$(CC) $(CFLAGS) -o $@ ds4_cli.o linenoise.o $(CORE_OBJS) $(METAL_LDLIBS)
+ds4: ds4_cli.o ds4_help.o linenoise.o ds4_gpu_args.o $(CORE_OBJS)
+	$(CC) $(CFLAGS) -o $@ ds4_cli.o ds4_help.o linenoise.o ds4_gpu_args.o $(CORE_OBJS) $(METAL_LDLIBS)
 
-ds4-server: ds4_server.o rax.o $(CORE_OBJS)
-	$(CC) $(CFLAGS) -o $@ ds4_server.o rax.o $(CORE_OBJS) $(METAL_LDLIBS)
+ds4-server: ds4_server.o ds4_help.o ds4_kvstore.o rax.o ds4_gpu_args.o $(CORE_OBJS)
+	$(CC) $(CFLAGS) -o $@ ds4_server.o ds4_help.o ds4_kvstore.o rax.o ds4_gpu_args.o $(CORE_OBJS) $(METAL_LDLIBS)
 
-ds4-bench: ds4_bench.o $(CORE_OBJS)
-	$(CC) $(CFLAGS) -o $@ ds4_bench.o $(CORE_OBJS) $(METAL_LDLIBS)
+ds4-bench: ds4_bench.o ds4_help.o ds4_gpu_args.o $(CORE_OBJS)
+	$(CC) $(CFLAGS) -o $@ ds4_bench.o ds4_help.o ds4_gpu_args.o $(CORE_OBJS) $(METAL_LDLIBS)
 
-cpu: ds4_cli_cpu.o ds4_server_cpu.o ds4_bench_cpu.o linenoise.o rax.o $(CPU_CORE_OBJS)
-	$(CC) $(CFLAGS) -o ds4 ds4_cli_cpu.o linenoise.o $(CPU_CORE_OBJS) $(LDLIBS)
-	$(CC) $(CFLAGS) -o ds4-server ds4_server_cpu.o rax.o $(CPU_CORE_OBJS) $(LDLIBS)
-	$(CC) $(CFLAGS) -o ds4-bench ds4_bench_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
+ds4-eval: ds4_eval.o ds4_help.o $(CORE_OBJS)
+	$(CC) $(CFLAGS) -o $@ ds4_eval.o ds4_help.o $(CORE_OBJS) $(METAL_LDLIBS)
+
+ds4-agent: ds4_agent.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o ds4_gpu_args.o $(CORE_OBJS)
+	$(CC) $(CFLAGS) -o $@ ds4_agent.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o ds4_gpu_args.o $(CORE_OBJS) $(METAL_LDLIBS)
+
+gguf-tools/quality-testing/score_official: gguf-tools/quality-testing/score_official.c ds4.h ds4_distributed.h ds4_tp.h $(CORE_OBJS) rax.o ds4_gpu_args.o
+	$(CC) $(QUALITY_CFLAGS) -I. -o $@ gguf-tools/quality-testing/score_official.c $(CORE_OBJS) rax.o ds4_gpu_args.o $(METAL_LDLIBS)
+
+tests/test_metal_session_batch.o: tests/test_metal_session_batch.c ds4.h
+	$(CC) $(CFLAGS) -I. -c -o $@ tests/test_metal_session_batch.c
+
+tests/test_metal_session_batch: tests/test_metal_session_batch.o $(CORE_OBJS)
+	$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
+
+test-metal-session-batch: tests/test_metal_session_batch
+	DS4_TEST_MODEL="$(DS4_TEST_MODEL)" ./tests/test_metal_session_batch
+
+speed-bench/metal_decode_schedule_bench.o: speed-bench/metal_decode_schedule_bench.c ds4.h
+	$(CC) $(CFLAGS) -I. -c -o $@ $<
+
+speed-bench/metal_decode_schedule_bench: speed-bench/metal_decode_schedule_bench.o $(CORE_OBJS)
+	$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
+
+metal-decode-schedule-bench: speed-bench/metal_decode_schedule_bench
+
+speed-bench/metal_prefill_variant_bench.o: speed-bench/metal_prefill_variant_bench.c ds4.h
+	$(CC) $(CFLAGS) -I. -c -o $@ $<
+
+speed-bench/metal_prefill_variant_bench: speed-bench/metal_prefill_variant_bench.o $(CORE_OBJS)
+	$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
+
+metal-prefill-variant-bench: speed-bench/metal_prefill_variant_bench
+
+tests/test_mxfp4_metal.o: tests/test_mxfp4_metal.c ds4_gpu.h
+	$(CC) $(CFLAGS) -I. -c -o $@ $<
+
+tests/test_mxfp4_metal: tests/test_mxfp4_metal.o ds4_metal.o
+	$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
+
+check-mxfp4-half-lut:
+	python3 metal/generate_mxfp4_half_lut.py --check
+
+test-mxfp4-metal: check-mxfp4-half-lut tests/test_mxfp4_metal
+	./tests/test_mxfp4_metal
+
+cpu: ds4_cli_cpu.o ds4_server_cpu.o ds4_bench_cpu.o ds4_eval_cpu.o ds4_agent_cpu.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o rax.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS)
+	$(CC) $(CFLAGS) -o ds4 ds4_cli_cpu.o ds4_help.o linenoise.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
+	$(CC) $(CFLAGS) -o ds4-server ds4_server_cpu.o ds4_help.o ds4_kvstore.o rax.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
+	$(CC) $(CFLAGS) -o ds4-bench ds4_bench_cpu.o ds4_help.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
+	$(CC) $(CFLAGS) -o ds4-eval ds4_eval_cpu.o ds4_help.o $(CPU_CORE_OBJS) $(LDLIBS)
+	$(CC) $(CFLAGS) -o ds4-agent ds4_agent_cpu.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
 
 cuda-regression:
 	@echo "cuda-regression requires a CUDA build"
@@ -68,15 +156,21 @@ help:
 	@echo "  make cuda-spark          Build CUDA for DGX Spark / GB10"
 	@echo "  make cuda-generic        Build CUDA for a generic local CUDA GPU"
 	@echo "  make cuda CUDA_ARCH=sm_N Build CUDA with an explicit nvcc -arch value"
-	@echo "  make cpu                 Build CPU-only ./ds4, ./ds4-server, and ./ds4-bench"
+	@echo "  make strix-halo          Build ROCm for Strix Halo / gfx1151"
+	@echo "  make rocm                Alias for make strix-halo"
+	@echo "  make test-mxfp4-rocm     Build and run the synthetic ROCm MXFP4 MoE test"
+	@echo "  make test-rocm           Core regression suite on ROCm-only hosts"
+	@echo "  make cpu                 Build CPU-only ./ds4, ./ds4-server, ./ds4-bench, ./ds4-eval, and ./ds4-agent"
 	@echo "  make test                Build and run tests"
+	@echo "  make dspark-verify-depth Run DSpark speculative verification smoke if support GGUF is present"
+	@echo "  make mtp-verify-depth    Run legacy MTP speculative verification smoke if MTP GGUF is present"
 	@echo "  make clean               Remove build outputs"
 
 cuda-spark:
-	$(MAKE) ds4 ds4-server ds4-bench CUDA_ARCH=
+	$(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent CUDA_ARCH=sm_121
 
 cuda-generic:
-	$(MAKE) ds4 ds4-server ds4-bench CUDA_ARCH=native
+	$(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent CUDA_ARCH=native
 
 cuda:
 	@if [ -z "$(strip $(CUDA_ARCH))" ]; then \
@@ -84,41 +178,123 @@ cuda:
 		echo "       or use make cuda-spark / make cuda-generic"; \
 		exit 2; \
 	fi
-	$(MAKE) ds4 ds4-server ds4-bench CUDA_ARCH="$(CUDA_ARCH)"
+	$(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent CUDA_ARCH="$(CUDA_ARCH)"
 
-ds4: ds4_cli.o linenoise.o $(CORE_OBJS)
-	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+strix-halo:
+	$(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent \
+		CORE_OBJS="ds4.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_rocm.o ds4_rocm_compat.o ds4_rocm_unavailable.o ds4_layer_pack.o" \
+		CFLAGS="$(CFLAGS) $(ROCM_HOST_CFLAGS) -DDS4_ROCM_BUILD" \
+		DS4_LINK="$(HIPCC) $(ROCM_CFLAGS)" \
+		DS4_LINK_LIBS="$(ROCM_LDLIBS)"
 
-ds4-server: ds4_server.o rax.o $(CORE_OBJS)
-	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+rocm: strix-halo
 
-ds4-bench: ds4_bench.o $(CORE_OBJS)
-	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+# Core regression suite for ROCm-only hosts: the CUDA-specific binaries
+# (tests/test_sampling, the CUDA session/mixed-batch oracles) are not part
+# of this target; run them through `make test` / `make cuda-regression` on
+# CUDA hosts.  Everything else mirrors `make test`.
+test-rocm:
+	$(MAKE) -B ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test \
+		tests/test_layer_pack tests/test_engine_mgpu_placement tests/test_gpu_args \
+		ds4 ds4-server ds4-bench ds4-agent \
+		CORE_OBJS="ds4.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_rocm.o ds4_rocm_compat.o ds4_rocm_unavailable.o ds4_layer_pack.o" \
+		CFLAGS="$(CFLAGS) $(ROCM_HOST_CFLAGS) -DDS4_ROCM_BUILD" \
+		DS4_LINK="$(HIPCC) $(ROCM_CFLAGS)" \
+		DS4_LINK_LIBS="$(ROCM_LDLIBS)"
+	./ds4-eval --self-test-extractors
+	./ds4_agent_test
+	./ds4_test
+	./tests/test_layer_pack
+	./tests/test_engine_mgpu_placement
+	./tests/test_gpu_args
+	./tests/test_gpu_args_cli.sh
 
-cpu: ds4_cli_cpu.o ds4_server_cpu.o ds4_bench_cpu.o linenoise.o rax.o $(CPU_CORE_OBJS)
-	$(CC) $(CFLAGS) -o ds4 ds4_cli_cpu.o linenoise.o $(CPU_CORE_OBJS) $(LDLIBS)
-	$(CC) $(CFLAGS) -o ds4-server ds4_server_cpu.o rax.o $(CPU_CORE_OBJS) $(LDLIBS)
-	$(CC) $(CFLAGS) -o ds4-bench ds4_bench_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
+ds4: ds4_cli.o ds4_help.o linenoise.o ds4_gpu_args.o $(CORE_OBJS)
+	$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
+
+ds4-server: ds4_server.o ds4_help.o ds4_kvstore.o rax.o ds4_gpu_args.o $(CORE_OBJS)
+	$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
+
+ds4-bench: ds4_bench.o ds4_help.o ds4_gpu_args.o $(CORE_OBJS)
+	$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
+
+ds4-eval: ds4_eval.o ds4_help.o $(CORE_OBJS)
+	$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
+
+ds4-agent: ds4_agent.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o ds4_gpu_args.o $(CORE_OBJS)
+	$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
+
+gguf-tools/quality-testing/score_official.o: gguf-tools/quality-testing/score_official.c ds4.h
+	$(CC) $(filter-out -ffast-math,$(QUALITY_CFLAGS)) $(ROCM_HOST_CFLAGS) -I. -c -o $@ $<
+
+gguf-tools/quality-testing/score_official: gguf-tools/quality-testing/score_official.o $(CORE_OBJS) rax.o ds4_gpu_args.o
+	$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
+
+cpu: ds4_cli_cpu.o ds4_server_cpu.o ds4_bench_cpu.o ds4_eval_cpu.o ds4_agent_cpu.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o rax.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS)
+	$(CC) $(CFLAGS) -o ds4 ds4_cli_cpu.o ds4_help.o linenoise.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
+	$(CC) $(CFLAGS) -o ds4-server ds4_server_cpu.o ds4_help.o ds4_kvstore.o rax.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
+	$(CC) $(CFLAGS) -o ds4-bench ds4_bench_cpu.o ds4_help.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
+	$(CC) $(CFLAGS) -o ds4-eval ds4_eval_cpu.o ds4_help.o $(CPU_CORE_OBJS) $(LDLIBS)
+	$(CC) $(CFLAGS) -o ds4-agent ds4_agent_cpu.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o ds4_gpu_args_cpu.o $(CPU_CORE_OBJS) $(LDLIBS)
 
 cuda-regression: tests/cuda_long_context_smoke
 	./tests/cuda_long_context_smoke
+
+tests/test_mxfp4_cuda: tests/test_mxfp4_cuda.cu $(MMQ_OBJS)
+	$(NVCC) $(NVCCFLAGS) -std=c++17 $(MMQ_INCLUDES) -o $@ $^ $(CUDA_LDLIBS)
+
+test-mxfp4-cuda: tests/test_mxfp4_cuda
+	./tests/test_mxfp4_cuda
 endif
 
-ds4.o: ds4.c ds4.h ds4_gpu.h
+ds4.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_gpu.h
 	$(CC) $(CFLAGS) -c -o $@ ds4.c
 
-ds4_cli.o: ds4_cli.c ds4.h linenoise.h
+ds4_image.o: ds4_image.c ds4_image.h third_party/iris/jpeg.h third_party/iris/png.h
+	$(CC) $(CFLAGS) -c -o $@ ds4_image.c
+
+ds4_ssd.o: ds4_ssd.c ds4_ssd.h
+	$(CC) $(CFLAGS) -c -o $@ ds4_ssd.c
+
+ds4_cli.o: ds4_cli.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h linenoise.h
 	$(CC) $(CFLAGS) -c -o $@ ds4_cli.c
 
-ds4_server.o: ds4_server.c ds4.h rax.h
+ds4_distributed.o: ds4_distributed.c ds4_distributed.h ds4.h ds4_ssd.h
+	$(CC) $(CFLAGS) -c -o $@ ds4_distributed.c
+
+ds4_tp.o: ds4_tp.c ds4_tp.h ds4.h ds4_ssd.h
+	$(CC) $(CFLAGS) -c -o $@ ds4_tp.c
+
+ds4_help.o: ds4_help.c ds4_help.h
+	$(CC) $(CFLAGS) -c -o $@ ds4_help.c
+
+ds4_gpu_args.o: ds4_gpu_args.c ds4_gpu_args.h ds4_gpu_mgpu.h
+	$(CC) $(CFLAGS) -c -o $@ ds4_gpu_args.c
+
+ds4_server.o: ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h rax.h
 	$(CC) $(CFLAGS) -c -o $@ ds4_server.c
 
-ds4_bench.o: ds4_bench.c ds4.h
+ds4_bench.o: ds4_bench.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h
 	$(CC) $(CFLAGS) -c -o $@ ds4_bench.c
 
-ds4_test.o: tests/ds4_test.c ds4_server.c ds4.h rax.h
+ds4_eval.o: ds4_eval.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h
+	$(CC) $(CFLAGS) -c -o $@ ds4_eval.c
+
+ds4_agent.o: ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h
+	$(CC) $(CFLAGS) -c -o $@ ds4_agent.c
+
+ds4_web.o: ds4_web.c ds4_web.h
+	$(CC) $(CFLAGS) -c -o $@ ds4_web.c
+
+ds4_kvstore.o: ds4_kvstore.c ds4_kvstore.h ds4.h ds4_ssd.h
+	$(CC) $(CFLAGS) -c -o $@ ds4_kvstore.c
+
+ds4_test.o: tests/ds4_test.c ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h rax.h
 	$(CC) $(CFLAGS) -Wno-unused-function -c -o $@ tests/ds4_test.c
 
+ds4_agent_test.o: tests/ds4_agent_test.c ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h
+	$(CC) $(CFLAGS) -Wno-unused-function -c -o $@ tests/ds4_agent_test.c
+
 tests/cuda_long_context_smoke.o: tests/cuda_long_context_smoke.c ds4_gpu.h
 	$(CC) $(CFLAGS) -I. -c -o $@ tests/cuda_long_context_smoke.c
 
@@ -128,36 +304,272 @@ rax.o: rax.c rax.h rax_malloc.h
 linenoise.o: linenoise.c linenoise.h
 	$(CC) $(CFLAGS) -c -o $@ linenoise.c
 
-ds4_cpu.o: ds4.c ds4.h ds4_gpu.h
-	$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4.c
+ds4_cpu.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_gpu.h
+	$(CC) $(CFLAGS) -Wno-unused-function -DDS4_NO_GPU -c -o $@ ds4.c
 
-ds4_cli_cpu.o: ds4_cli.c ds4.h linenoise.h
+ds4_cli_cpu.o: ds4_cli.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h linenoise.h
 	$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_cli.c
 
-ds4_server_cpu.o: ds4_server.c ds4.h rax.h
+ds4_gpu_args_cpu.o: ds4_gpu_args.c ds4_gpu_args.h ds4_gpu_mgpu.h
+	$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_gpu_args.c
+
+ds4_server_cpu.o: ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h rax.h
 	$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_server.c
 
-ds4_bench_cpu.o: ds4_bench.c ds4.h
+ds4_bench_cpu.o: ds4_bench.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h
 	$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_bench.c
 
+ds4_eval_cpu.o: ds4_eval.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h
+	$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_eval.c
+
+ds4_agent_cpu.o: ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h
+	$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_agent.c
+
 ds4_metal.o: ds4_metal.m ds4_gpu.h $(METAL_SRCS)
 	$(CC) $(OBJCFLAGS) -c -o $@ ds4_metal.m
 
-ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_iq2_tables_cuda.inc
-	$(NVCC) $(NVCCFLAGS) -c -o $@ ds4_cuda.cu
+tests/test_glm53_kda.o: tests/test_glm53_kda.c ds4_gpu.h
+	$(CC) $(CFLAGS) -I. -c -o $@ tests/test_glm53_kda.c
 
-tests/cuda_long_context_smoke: tests/cuda_long_context_smoke.o ds4_cuda.o
-	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+tests/test_glm53_vision_engine.o: tests/test_glm53_vision_engine.c ds4.h ds4_image.h
+	$(CC) $(CFLAGS) -I. -c -o $@ tests/test_glm53_vision_engine.c
 
-ds4_test: ds4_test.o rax.o $(CORE_OBJS)
+tests/test_glm53_vision_engine: tests/test_glm53_vision_engine.o $(CORE_OBJS)
 ifeq ($(UNAME_S),Darwin)
-	$(CC) $(CFLAGS) -o $@ ds4_test.o rax.o $(CORE_OBJS) $(METAL_LDLIBS)
+	$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
 else
-	$(NVCC) $(NVCCFLAGS) -o $@ ds4_test.o rax.o $(CORE_OBJS) $(CUDA_LDLIBS)
+	$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
 endif
 
-test: ds4_test
+tests/test_glm53_vision_prompt.o: tests/test_glm53_vision_prompt.c ds4.h
+	$(CC) $(CFLAGS) -I. -c -o $@ tests/test_glm53_vision_prompt.c
+
+tests/test_glm53_vision_prompt: tests/test_glm53_vision_prompt.o $(CORE_OBJS)
+ifeq ($(UNAME_S),Darwin)
+	$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
+else
+	$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
+endif
+
+ifeq ($(UNAME_S),Darwin)
+$(GLM53_KDA_TEST): tests/test_glm53_kda.o ds4_metal.o
+	$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
+else
+$(GLM53_KDA_TEST): tests/test_glm53_kda.o ds4_cuda.o $(MMQ_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+endif
+
+.PHONY: test-glm53-kda
+test-glm53-kda: $(GLM53_KDA_TEST)
+	./$(GLM53_KDA_TEST)
+
+tests/test_glm53_kda_rocm.o: tests/test_glm53_kda.c ds4_gpu.h
+	$(CC) $(filter-out -ffast-math,$(CFLAGS)) $(ROCM_HOST_CFLAGS) -DDS4_ROCM_BUILD -I. -c -o $@ $<
+
+$(GLM53_KDA_ROCM_TEST): tests/test_glm53_kda_rocm.o ds4_rocm.o
+	$(HIPCC) $(ROCM_CFLAGS) -o $@ $^ $(ROCM_LDLIBS)
+
+test-glm53-kda-rocm: $(GLM53_KDA_ROCM_TEST)
+	./$(GLM53_KDA_ROCM_TEST)
+
+ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_gpu_mgpu.h ds4_glm53_vision_gpu.cuh ds4_iq2_tables_cuda.inc cuda/mmq/ds4_mmq.h
+	$(NVCC) $(NVCCFLAGS) -c -o $@ ds4_cuda.cu
+
+# Vendored mmq pieces (see cuda/mmq/VENDOR.md).  ds4_mmq.cu transitively
+# pulls in mmq.cuh which has heavy template instantiation -- each piece
+# compiles in its own TU and links in.
+cuda/mmq/ds4_ggml_stubs.o: cuda/mmq/ds4_ggml_stubs.cu cuda/mmq/ds4_ggml_stubs.h cuda/mmq/common.cuh
+	$(NVCC) $(NVCCFLAGS) -std=c++17 $(MMQ_INCLUDES) -c -o $@ $<
+
+cuda/mmq/ds4_mmq.o: cuda/mmq/ds4_mmq.cu cuda/mmq/ds4_mmq.h cuda/mmq/ds4_mmq_d2r.cuh cuda/mmq/mmq.cuh cuda/mmq/common.cuh cuda/mmq/ds4_ggml_stubs.h cuda/mmq/quantize.cuh cuda/mmq/mmid.cuh cuda/mmq/vecdotq.cuh cuda/mmq/mma.cuh
+	$(NVCC) $(NVCCFLAGS) -std=c++17 $(MMQ_INCLUDES) -c -o $@ $<
+
+cuda/mmq/ds4_mmq_d2r.o: cuda/mmq/ds4_mmq_d2r.cu cuda/mmq/ds4_mmq_d2r.cuh cuda/mmq/mmq.cuh cuda/mmq/common.cuh cuda/mmq/ds4_ggml_stubs.h cuda/mmq/vecdotq.cuh cuda/mmq/mma.cuh
+	$(NVCC) $(NVCCFLAGS) -std=c++17 $(MMQ_INCLUDES) -c -o $@ $<
+
+cuda/mmq/quantize.o: cuda/mmq/quantize.cu cuda/mmq/quantize.cuh cuda/mmq/common.cuh cuda/mmq/ds4_ggml_stubs.h cuda/mmq/mmq.cuh
+	$(NVCC) $(NVCCFLAGS) -std=c++17 $(MMQ_INCLUDES) -c -o $@ $<
+
+cuda/mmq/mmid.o: cuda/mmq/mmid.cu cuda/mmq/mmid.cuh cuda/mmq/common.cuh cuda/mmq/ds4_ggml_stubs.h
+	$(NVCC) $(NVCCFLAGS) -std=c++17 $(MMQ_INCLUDES) -c -o $@ $<
+
+cuda/mmq/mmvq.o: cuda/mmq/mmvq.cu cuda/mmq/mmvq.cuh cuda/mmq/common.cuh cuda/mmq/ds4_ggml_stubs.h cuda/mmq/quantize.cuh cuda/mmq/vecdotq.cuh cuda/mmq/unary.cuh
+	$(NVCC) $(NVCCFLAGS) -std=c++17 $(MMQ_INCLUDES) -c -o $@ $<
+
+cuda/mmq/ds4_repack.o: cuda/mmq/ds4_repack.cu cuda/mmq/ds4_repack.h
+	$(NVCC) $(NVCCFLAGS) -std=c++17 -c -o $@ $<
+
+ds4_rocm.o: ds4_rocm.cu ds4_gpu.h ds4_glm53_vision_gpu.cuh ds4_iq2_tables_cuda.inc $(ROCM_SRCS)
+	$(HIPCC) $(ROCM_CFLAGS) -c -o $@ ds4_rocm.cu
+
+tests/test_mxfp4_rocm.o: tests/test_mxfp4_rocm.c ds4_gpu.h
+	$(CC) $(filter-out -ffast-math,$(CFLAGS)) $(ROCM_HOST_CFLAGS) -DDS4_ROCM_BUILD -I. -c -o $@ $<
+
+tests/test_mxfp4_rocm: tests/test_mxfp4_rocm.o ds4_rocm.o
+	$(HIPCC) $(ROCM_CFLAGS) -o $@ $^ $(ROCM_LDLIBS)
+
+tests/bench_mxfp4_rocm.o: tests/bench_mxfp4_rocm.c ds4_gpu.h
+	$(CC) $(filter-out -ffast-math,$(CFLAGS)) $(ROCM_HOST_CFLAGS) -DDS4_ROCM_BUILD -I. -c -o $@ $<
+
+tests/bench_mxfp4_rocm: tests/bench_mxfp4_rocm.o ds4_rocm.o
+	$(HIPCC) $(ROCM_CFLAGS) -o $@ $^ $(ROCM_LDLIBS)
+
+test-mxfp4-rocm: tests/test_mxfp4_rocm
+	./tests/test_mxfp4_rocm
+
+ds4_rocm_compat.o: ds4_rocm_compat.cu ds4_gpu.h ds4_gpu_mgpu.h ds4_gpu_args.h
+	$(HIPCC) $(ROCM_CFLAGS) -c -o $@ ds4_rocm_compat.cu
+
+ds4_rocm_unavailable.o: ds4_rocm_unavailable.cu
+	$(HIPCC) $(ROCM_CFLAGS) -c -o $@ ds4_rocm_unavailable.cu
+
+tests/cuda_long_context_smoke: tests/cuda_long_context_smoke.o ds4_cuda.o $(MMQ_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+
+tests/test_layer_pack.o: tests/test_layer_pack.c ds4_layer_pack.h
+	$(CC) $(CFLAGS) -I. -c -o $@ $<
+
+tests/test_layer_pack: tests/test_layer_pack.o ds4_layer_pack.o
+	$(CC) $(CFLAGS) -o $@ $^ $(LDLIBS)
+
+tests/test_gpu_args.o: tests/test_gpu_args.c ds4_gpu_args.h ds4_gpu_mgpu.h
+	$(CC) $(CFLAGS) -I. -DDS4_NO_GPU -c -o $@ $<
+
+tests/test_gpu_args: tests/test_gpu_args.o ds4_gpu_args_cpu.o
+	$(CC) $(CFLAGS) -o $@ $^ $(LDLIBS)
+
+ds4_cpu_test_hooks.o: ds4.c ds4.h ds4_image.h ds4_gpu.h ds4_gpu_mgpu.h ds4_layer_pack.h
+	$(CC) $(CFLAGS) -Wno-unused-function -DDS4_NO_GPU -DDS4_TEST_HOOKS -c -o $@ ds4.c
+
+tests/test_engine_mgpu_placement.o: tests/test_engine_mgpu_placement.c ds4.h ds4_gpu_mgpu.h ds4_layer_pack.h
+	$(CC) $(CFLAGS) -I. -c -o $@ $<
+
+tests/test_engine_mgpu_placement: tests/test_engine_mgpu_placement.o ds4_cpu_test_hooks.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
+	$(CC) $(CFLAGS) -o $@ $^ $(LDLIBS)
+
+tests/test_sampling.o: tests/test_sampling.c ds4.h
+	$(CC) $(CFLAGS) -fno-finite-math-only -DDS4_TEST_HOOKS -I. -c -o $@ $<
+
+tests/test_sampling: tests/test_sampling.o ds4_cpu_test_hooks.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
+	$(CC) $(CFLAGS) -o $@ $^ $(LDLIBS)
+
+ifneq ($(UNAME_S),Darwin)
+tests/test_gpu_xdev.o: tests/test_gpu_xdev.c ds4_gpu.h ds4_gpu_mgpu.h
+	$(CC) $(CFLAGS) -I. -I$(CUDA_HOME)/include -c -o $@ $<
+
+tests/test_gpu_xdev: tests/test_gpu_xdev.o ds4_cuda.o $(MMQ_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+
+tests/test_gpu_model_cache.o: tests/test_gpu_model_cache.c ds4_gpu.h
+	$(CC) $(CFLAGS) -I. -I$(CUDA_HOME)/include -c -o $@ $<
+
+tests/test_gpu_model_cache: tests/test_gpu_model_cache.o ds4_cuda.o $(MMQ_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+
+tests/test_gpu_lookup_cache_strict.o: tests/test_gpu_lookup_cache_strict.c ds4_gpu.h ds4_gpu_mgpu.h
+	$(CC) $(CFLAGS) -I. -I$(CUDA_HOME)/include -c -o $@ $<
+
+tests/test_gpu_lookup_cache_strict: tests/test_gpu_lookup_cache_strict.o ds4_cuda.o $(MMQ_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+
+ds4_cuda_test_hooks.o: ds4.c ds4.h ds4_gpu.h ds4_gpu_mgpu.h ds4_layer_pack.h
+	$(CC) $(CFLAGS) -Wno-unused-function -DDS4_TEST_HOOKS -I$(CUDA_HOME)/include -c -o $@ ds4.c
+
+tests/test_engine_mgpu_refusal.o: tests/test_engine_mgpu_refusal.c ds4.h ds4_gpu_mgpu.h
+	$(CC) $(CFLAGS) -I. -I$(CUDA_HOME)/include -c -o $@ $<
+
+tests/test_engine_mgpu_refusal: tests/test_engine_mgpu_refusal.o ds4_gpu_args.o ds4_kvstore.o rax.o $(CORE_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+
+tests/test_engine_mgpu_runtime.o: tests/test_engine_mgpu_runtime.c ds4.h ds4_gpu_mgpu.h
+	$(CC) $(CFLAGS) -DDS4_TEST_HOOKS -I. -I$(CUDA_HOME)/include -c -o $@ $<
+
+tests/test_engine_mgpu_runtime: tests/test_engine_mgpu_runtime.o ds4_cuda_test_hooks.o ds4_gpu_args.o ds4_kvstore.o rax.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_cuda.o ds4_layer_pack.o $(MMQ_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+
+tests/test_engine_correctness.o: tests/test_engine_correctness.c ds4.h ds4_gpu_mgpu.h
+	$(CC) $(CFLAGS) -I. -I$(CUDA_HOME)/include -c -o $@ $<
+
+tests/test_engine_correctness: tests/test_engine_correctness.o ds4_gpu_args.o ds4_kvstore.o rax.o $(CORE_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+
+tests/test_cuda_session_batch.o: tests/test_cuda_session_batch.c ds4.h ds4_gpu_args.h ds4_gpu_mgpu.h
+	$(CC) $(CFLAGS) -I. -I$(CUDA_HOME)/include -c -o $@ $<
+
+tests/test_cuda_session_batch: tests/test_cuda_session_batch.o ds4_gpu_args.o ds4_kvstore.o rax.o $(CORE_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+
+test-cuda-session-batch: tests/test_cuda_session_batch
+	DS4_TEST_MODEL="$(DS4_TEST_MODEL)" ./tests/test_cuda_session_batch
+
+tests/test_cuda_mixed_batch.o: tests/test_cuda_mixed_batch.c ds4.h ds4_gpu_args.h ds4_gpu_mgpu.h
+	$(CC) $(CFLAGS) -DDS4_TEST_HOOKS -I. -I$(CUDA_HOME)/include -c -o $@ $<
+
+tests/test_cuda_mixed_batch: tests/test_cuda_mixed_batch.o ds4_cuda_test_hooks.o ds4_gpu_args.o ds4_kvstore.o rax.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_cuda.o ds4_layer_pack.o $(MMQ_OBJS)
+	$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)
+
+test-cuda-mixed-batch: tests/test_cuda_mixed_batch
+	DS4_TEST_MODEL="$(DS4_TEST_MODEL)" ./tests/test_cuda_mixed_batch
+endif
+
+ds4_test: ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS)
+ifeq ($(UNAME_S),Darwin)
+	$(CC) $(CFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(METAL_LDLIBS)
+else
+	$(DS4_LINK) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(DS4_LINK_LIBS)
+endif
+
+ds4_agent_test: ds4_agent_test.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS)
+ifeq ($(UNAME_S),Darwin)
+	$(CC) $(CFLAGS) -o $@ ds4_agent_test.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS) $(METAL_LDLIBS)
+else
+	$(DS4_LINK) -o $@ ds4_agent_test.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS) $(DS4_LINK_LIBS)
+endif
+
+test: ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test \
+	tests/test_layer_pack tests/test_engine_mgpu_placement tests/test_gpu_args \
+	$(SAMPLING_TEST) ds4 ds4-server ds4-bench ds4-agent
+	./ds4-eval --self-test-extractors
+	./ds4_agent_test
 	./ds4_test
+	./tests/test_layer_pack
+	./tests/test_engine_mgpu_placement
+	./tests/test_gpu_args
+	./tests/test_gpu_args_cli.sh
+	./tests/test_sampling
+
+dspark-acceptance: ds4
+	DS4_DSPARK_MODEL="$(DS4_DSPARK_MODEL)" \
+	DS4_DSPARK_SUPPORT="$(DS4_DSPARK_SUPPORT)" \
+	sh tests/dspark_acceptance_fixture.sh
+
+dspark-verify-depth: ds4_test
+	@if [ ! -f "$(DS4_TEST_MODEL)" ]; then \
+		echo "dspark-verify-depth: skipped, missing model $(DS4_TEST_MODEL)"; \
+	elif [ ! -f "$(DS4_DSPARK_SUPPORT)" ]; then \
+		echo "dspark-verify-depth: skipped, missing DSpark support $(DS4_DSPARK_SUPPORT)"; \
+		echo "dspark-verify-depth: run ./download_model.sh ds4f-dspark or set DS4_DSPARK_SUPPORT=FILE"; \
+	else \
+		DS4_TEST_MODEL="$(DS4_TEST_MODEL)" DS4_TEST_DSPARK="$(DS4_DSPARK_SUPPORT)" ./ds4_test --dspark-verify-depth; \
+	fi
+
+mtp-verify-depth: ds4_test
+	@if [ ! -f "$(DS4_TEST_MODEL)" ]; then \
+		echo "mtp-verify-depth: skipped, missing model $(DS4_TEST_MODEL)"; \
+	elif [ ! -f "$(DS4_TEST_MTP)" ]; then \
+		echo "mtp-verify-depth: skipped, missing MTP support $(DS4_TEST_MTP)"; \
+		echo "mtp-verify-depth: set DS4_TEST_MTP=FILE to a legacy support GGUF"; \
+	else \
+		DS4_TEST_MODEL="$(DS4_TEST_MODEL)" DS4_TEST_MTP="$(DS4_TEST_MTP)" ./ds4_test --mtp-verify-depth; \
+	fi
+
+q4k-dot-test: tests/test_q4k_dot.c
+	$(CC) -O2 -Wall -Wextra -std=c99 -o tests/test_q4k_dot tests/test_q4k_dot.c -lm -pthread
+	./tests/test_q4k_dot
+
+mxfp4-dot-test: tests/test_mxfp4_dot.c
+	$(CC) -O2 -Wall -Wextra -std=c99 -o tests/test_mxfp4_dot tests/test_mxfp4_dot.c -lm
+	./tests/test_mxfp4_dot
 
 clean:
-	rm -f ds4 ds4-server ds4-bench ds4_cpu ds4_native ds4_server_test ds4_test *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o
+	rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official gguf-tools/quality-testing/score_official.o speed-bench/metal_decode_schedule_bench speed-bench/metal_prefill_variant_bench speed-bench/*.o tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_mxfp4_rocm tests/test_mxfp4_cuda tests/test_metal_session_batch tests/test_glm53_kda tests/test_glm53_kda_rocm tests/test_glm53_vision_engine tests/test_glm53_vision_prompt tests/test_gpu_xdev tests/test_gpu_model_cache tests/test_gpu_lookup_cache_strict tests/test_engine_mgpu_refusal tests/test_engine_mgpu_runtime tests/test_engine_correctness tests/test_sampling tests/test_cuda_session_batch tests/test_cuda_mixed_batch tests/*.o *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o

STRIXHALO.md

@@ -0,0 +1,146 @@
+# DS4 on Strix Halo
+
+This is the minimal setup for DS4 ROCm inference on a
+Strix Halo machine with 128 GB RAM and Radeon 8060S (`gfx1151`).
+
+## 1. Install ROCm
+
+On Ubuntu 26.04 LTS, install the ROCm compiler/runtime and libraries used by the Strix Halo backend:
+
+```sh
+sudo apt-get update
+sudo apt-get install -y \
+  hipcc rocminfo rocm-smi \
+  libamdhip64-dev \
+  libhipblas-dev libhipblaslt-dev \
+  librocblas-dev \
+  librocwmma-dev \
+  libhipcub-dev
+```
+
+The backend uses rocWMMA. On this Ubuntu 26.04 setup, `librocwmma-dev`
+installs the top-level rocWMMA headers but misses `rocwmma/internal/`.
+No Ubuntu package currently provides those internal headers. Install a complete
+matching rocWMMA header tree:
+
+```sh
+git clone --depth 1 --branch rocm-7.1.0 https://github.com/ROCm/rocWMMA.git /tmp/rocWMMA-rocm-7.1.0
+sudo mkdir -p /usr/local/include
+sudo cp -a /tmp/rocWMMA-rocm-7.1.0/library/include/rocwmma /usr/local/include/
+```
+
+If ROCm is installed under `/usr` but tooling expects `/opt/rocm`, add these
+compatibility links:
+
+```sh
+sudo mkdir -p /opt/rocm/bin
+sudo ln -sf /usr/bin/hipcc /opt/rocm/bin/hipcc
+sudo ln -sfn /usr/lib/x86_64-linux-gnu /opt/rocm/lib
+sudo ln -sfn /usr/include /opt/rocm/include
+```
+
+## 2. Enable ROCm access
+
+The user running DS4 must be able to open `/dev/kfd` and the DRM render node:
+
+```sh
+sudo usermod -aG render,video "$USER"
+```
+
+Log out and back in, or reboot. Verify:
+
+```sh
+rocminfo | grep -A80 'Name:                    gfx1151'
+```
+
+If DS4 says `no ROCm-capable device is detected`, check that `rocminfo` can open
+`/dev/kfd` and that `groups` includes `render`.
+
+## 3. Increase GPU-visible memory
+
+A 128 GB Strix Halo system may initially expose only about 62 GB of GPU-visible
+memory. DS4 needs the larger GTT aperture for the 80.76 GiB model plus runtime
+buffers.
+
+Use these kernel parameters:
+
+```text
+amd_iommu=off amdgpu.gttsize=126976 ttm.pages_limit=32505856 ttm.page_pool_size=32505856
+```
+
+On Ubuntu with GRUB:
+
+```sh
+sudo cp /etc/default/grub /etc/default/grub.bak
+sudoedit /etc/default/grub
+```
+
+Set:
+
+```text
+GRUB_CMDLINE_LINUX_DEFAULT="quiet splash amd_iommu=off amdgpu.gttsize=126976 ttm.pages_limit=32505856 ttm.page_pool_size=32505856"
+```
+
+Then:
+
+```sh
+sudo update-grub
+sudo reboot
+```
+
+After reboot, verify:
+
+```sh
+cat /proc/cmdline
+sudo dmesg | grep -Ei 'GTT|gttsize|TTM|VRAM'
+rocminfo | grep -A80 'Name:                    gfx1151'
+```
+
+Expected signs:
+
+```text
+amdgpu:  126976M of GTT memory ready
+rocminfo gfx1151 pool: 130023424 KB
+```
+
+## 4. Build DS4
+
+Use the normal Strix Halo target. It builds the standard binary names:
+
+```sh
+make strix-halo -j"$(nproc)"
+```
+
+`make rocm` is an alias for `make strix-halo`.
+
+After changes to the ROCm MXFP4 or routed-MoE kernels, run the standalone
+CPU-oracle regression:
+
+```sh
+make test-mxfp4-rocm
+```
+
+This test does not require a full model GGUF. It covers resident decode and
+batched routed-MoE execution at 1, 3, 32, 128, and 512 tokens.
+
+## 5. Use the right GGUF
+
+Use the standard IQ2XXS/Q2K/Q8 imatrix GGUF:
+
+```text
+DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf
+```
+
+Avoid the mixed IQ2/IQ4 or IQ2/Q4 GGUFs on this machine for now. They put much
+more memory pressure on the ROCm path and can trigger system OOM instead of a
+clean DS4 failure.
+
+## 6. Run DS4
+
+Run it normally:
+
+```sh
+./ds4 -m gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf
+```
+
+The ROCm build uses the Strix Halo backend automatically.

cuda/mmq/VENDOR.md

@@ -0,0 +1,149 @@
+# Vendored llama.cpp mmq kernels
+
+This directory contains source files copied verbatim from
+[llama.cpp's `ggml-cuda` backend](https://github.com/ggml-org/llama.cpp/tree/master/ggml/src/ggml-cuda),
+plus a thin ds4-side adapter (`ds4_ggml_stubs.{h,cu}` and `ds4_mmq.{h,cu}`)
+that lets the templated CUDA kernels compile and link without the full ggml
+runtime.
+
+## Why these files are vendored, not submoduled
+
+ds4 is a flat, self-contained C/CUDA codebase with no third-party build
+dependencies. A submodule would force consumers to clone, configure, and
+build all of llama.cpp just to use a few quantized-matmul kernels. Vendoring
+keeps ds4 self-contained at the cost of a periodic re-sync.
+
+## Upstream pin
+
+| Field         | Value                                                                                                              |
+|---------------|--------------------------------------------------------------------------------------------------------------------|
+| Source        | https://github.com/ggml-org/llama.cpp                                                                              |
+| Commit        | `5c0e9468378eba6bf3cc1989ff5d62fbbe4d9e3a`                                                                         |
+| Commit date   | 2026-05-14                                                                                                         |
+| Commit title  | `ggml-hexagon: cpy: add contiguous fast-path in reshape copy (#23076)`                                             |
+| License       | MIT (`LICENSE` at the repository root, copyright "2023-2026 The ggml authors"). Compatible with ds4's MIT license. |
+
+## File inventory
+
+| File                  | Origin in llama.cpp                          | Status                                                                   | Lines |
+|-----------------------|----------------------------------------------|--------------------------------------------------------------------------|-------|
+| `mmq.cuh`             | `ggml/src/ggml-cuda/mmq.cuh`                 | verbatim                                                                 |  4176 |
+| `mma.cuh`             | `ggml/src/ggml-cuda/mma.cuh`                 | verbatim                                                                 |  1456 |
+| `vecdotq.cuh`         | `ggml/src/ggml-cuda/vecdotq.cuh`             | verbatim                                                                 |  1317 |
+| `quantize.cuh`        | `ggml/src/ggml-cuda/quantize.cuh`            | verbatim                                                                 |    41 |
+| `quantize.cu`         | `ggml/src/ggml-cuda/quantize.cu`             | verbatim                                                                 |   443 |
+| `mmid.cuh`            | `ggml/src/ggml-cuda/mmid.cuh`                | verbatim                                                                 |     5 |
+| `mmid.cu`             | `ggml/src/ggml-cuda/mmid.cu`                 | verbatim                                                                 |   164 |
+| `mmvq.cuh`            | `ggml/src/ggml-cuda/mmvq.cuh`                | patched (Step 6): `mul_mat_vec_q_switch_type` proto exposed; ggml-tensor entries gated on `DS4_MMVQ_INCLUDE_GGML_ENTRIES` | ~36 |
+| `mmvq.cu`             | `ggml/src/ggml-cuda/mmvq.cu`                 | patched: `mul_mat_vec_q_switch_type` promoted from `static`; `ggml_cuda_mul_mat_vec_q` + `ggml_cuda_op_mul_mat_vec_q` gated on `DS4_MMVQ_INCLUDE_GGML_ENTRIES` | 1163 |
+| `unary.cuh`           | `ggml/src/ggml-cuda/unary.cuh`               | verbatim (needed by `mmvq.cu` for inline GLU epilogues)                  |   114 |
+| `common.cuh`          | `ggml/src/ggml-cuda/common.cuh`              | verbatim                                                                 |  1489 |
+| `ggml-common.h`       | `ggml/src/ggml-common.h`                     | verbatim                                                                 |  1900 |
+| `vendors/cuda.h`      | `ggml/src/ggml-cuda/vendors/cuda.h`          | verbatim                                                                 |    28 |
+| `ggml.h`              | (new)                                        | redirect to `ds4_ggml_stubs.h`                                           |     5 |
+| `ggml-impl.h`         | (new)                                        | redirect to `ds4_ggml_stubs.h`                                           |     5 |
+| `ggml-cuda.h`         | (new)                                        | redirect to `ds4_ggml_stubs.h`                                           |     5 |
+| `ds4_ggml_stubs.h`    | (new)                                        | shim: ggml_type enum, macros, info struct, type_size lookups             | ~280 |
+| `ds4_ggml_stubs.cu`   | (new)                                        | shim impls: `ggml_cuda_info`, naive pool, `ggml_backend_cuda_context::*` | ~110 |
+| `ds4_mmq.h`           | (new)                                        | public C ABI for ds4 to call                                             |  ~70 |
+| `ds4_mmq.cu`          | (new)                                        | host wrappers; Phase 0 instantiates `mul_mat_q_case<Q8_0>` only          | ~120 |
+
+**Total vendored:** ~11,000 lines of CUDA. **Total shim/adapter:** ~600 lines.
+
+## What llama.cpp's `mmq.cu` does that we don't vendor
+
+The upstream `ggml/src/ggml-cuda/mmq.cu` (372 lines) is the ggml-backend
+dispatch entry. It talks to `ggml_tensor` / `ggml_backend_buffer_get_usage`,
+does its own activation Q8_1 quantization, and wires into the ggml op graph.
+We **don't vendor it.** Instead `ds4_mmq.cu` provides ds4-style entries
+that bypass the ggml graph and call the per-type `mul_mat_q_case<T>` directly.
+
+## Symbol-resolution table
+
+Symbols the vendored files reference, and how they resolve in this directory:
+
+| Symbol category                                | Resolution                                                                                       |
+|------------------------------------------------|--------------------------------------------------------------------------------------------------|
+| `GGML_ASSERT`, `GGML_ABORT`, `GGML_UNUSED`     | Macros in `ds4_ggml_stubs.h`                                                                     |
+| `GGML_PAD`, `GGML_UNUSED_VARS`                 | Macros in `ds4_ggml_stubs.h`                                                                     |
+| `GGML_CUDA_CC_*` constants                     | Defined in vendored `common.cuh`                                                                 |
+| `GGML_TYPE_*` enum values                      | Defined in `ds4_ggml_stubs.h::enum ggml_type`                                                    |
+| `GGML_CUDA_MAX_DEVICES`, `GGML_CUDA_NAME`      | Macros in `ds4_ggml_stubs.h`                                                                     |
+| `GGML_MAX_DIMS`, `GGML_MAX_SRC`                | Macros in `ds4_ggml_stubs.h`                                                                     |
+| `block_q8_0`, `block_q2_K`, `block_iq2_xxs`    | Defined in vendored `ggml-common.h` (gated by `GGML_COMMON_IMPL_CUDA`)                           |
+| `ggml_half`, `ggml_half2`                      | `uint16_t` typedefs in `ds4_ggml_stubs.h`                                                        |
+| `ggml_type_size`, `ggml_blck_size`             | Inline implementations in `ds4_ggml_stubs.h`                                                     |
+| `ggml_cuda_info()`                             | Singleton in `ds4_ggml_stubs.cu`, populated via `cudaGetDeviceProperties`                        |
+| `ggml_cuda_get_device`, `ggml_cuda_set_device` | Thin wrappers in `ds4_ggml_stubs.cu`                                                             |
+| `ggml_time_us`                                 | `std::chrono::steady_clock` in `ds4_ggml_stubs.cu`                                               |
+| `ggml_cuda_pool`, `ggml_cuda_pool_alloc`       | Defined in vendored `common.cuh`. Concrete `ds4_naive_pool` subclass in `ds4_ggml_stubs.cu`      |
+| `ggml_backend_cuda_context`                    | Defined in vendored `common.cuh`. `new_pool_for_device` and `~dtor` provided in `ds4_ggml_stubs.cu` |
+| `ggml_tensor`                                  | Forward declaration in `ds4_ggml_stubs.h`. Never dereferenced - only held as pointer.            |
+| `ggml_glu_op`                                  | Enum stub in `ds4_ggml_stubs.h` (fusion args present in common.cuh but unused by ds4)            |
+| `CUDA_CHECK`, `CUBLAS_CHECK`                   | Macros in `ds4_ggml_stubs.h`                                                                     |
+| `ggml_cuda_launch_mm_ids_helper`               | Defined in vendored `mmid.cu`                                                                    |
+| `ggml_cuda_should_use_mmq`                     | Defined in vendored `mmq.cu` - **we re-implement in `ds4_mmq.cu`** since we don't vendor mmq.cu  |
+| `ggml_cuda_mul_mat_q*`                         | Upstream's host entry. **Replaced** by `ds4_mmq_q8_0_dense` and family.                          |
+
+## Things we deliberately do NOT support
+
+- **CUDA graphs.** ds4 manages its own streams; `USE_CUDA_GRAPH` is undefined.
+- **HIP / MUSA / AMD backends.** vendored code's HIP/MUSA `#ifdef` branches are dead in our build but kept for upstream diff-cleanliness.
+- **The full `ggml_tensor` type.** No tensor introspection - shapes and strides come in via raw arguments to `ds4_mmq_*`.
+- **`ggml_op` graph evaluation.** We call kernels directly.
+
+## Re-syncing with upstream
+
+When upstream lands a bugfix or perf improvement we want, the procedure is:
+
+```sh
+cd /tmp/llama-research && git fetch && git checkout <NEW_COMMIT>
+cd /Users/ent/code/ds4-mmq-lift/cuda/mmq
+cp /tmp/llama-research/ggml/src/ggml-cuda/{mmq.cuh,mma.cuh,vecdotq.cuh,quantize.cuh,quantize.cu,mmid.cuh,mmid.cu,common.cuh} .
+cp /tmp/llama-research/ggml/src/ggml-common.h .
+cp /tmp/llama-research/ggml/src/ggml-cuda/vendors/cuda.h vendors/
+# Verify the shim still covers all referenced symbols:
+grep -hoE 'GGML_[A-Z_]+|ggml_[a-z_]+' *.cuh *.cu *.h | sort -u > /tmp/symbols.new
+diff /tmp/symbols.last /tmp/symbols.new  # check for newly-introduced names
+# Update this VENDOR.md's commit pin and run `make cuda CUDA_ARCH=sm_120`.
+```
+
+If the new symbols are minor (e.g., new `GGML_CUDA_CC_*` constants), they
+likely come from `common.cuh` which we vendor and don't need any shim
+changes. If they're new `ggml_*` host functions (rare, but possible if
+upstream adds a new helper), extend `ds4_ggml_stubs.h`.
+
+## Testing matrix
+
+| Test                                          | Status                       | Phase     |
+|-----------------------------------------------|------------------------------|-----------|
+| `nvcc -c cuda/mmq/ds4_mmq.cu` builds cleanly  | **passes** (sm_120, nvcc 13) | Phase 0   |
+| Q8_0 dense parity vs CPU reference            | **passes** (4 shapes)        | Phase 2   |
+| Q2_K dense parity                             | **passes** (4 shapes)        | Phase 3   |
+| IQ2_XXS dense parity                          | **passes** (4 shapes)        | Phase 3   |
+| MoE `_id` Q8_0 parity                         | **passes** (3 shapes)        | Phase 4   |
+| MoE `_id` Q2_K parity                         | **passes** (3 shapes)        | Phase 4   |
+| MoE `_id` IQ2_XXS parity                      | **passes** (4 shapes)        | Phase 4   |
+| `make ds4-bench` with mmq integration         | **builds and runs**          | Phase 5/6 |
+| Frontier sweep, ctx 2k-16k, V4 Flash IQ2XXS   | **see results below**        | Phase 7   |
+
+## Validated performance
+
+PRO 6000 Blackwell (sm_120), CUDA 13.0, V4 Flash IQ2_XXS GGUF (86.7 GB),
+default `DS4_CUDA_MMQ_MOE_MIN_TOKENS=2` (legacy decode preserved):
+
+| ctx    | baseline pf t/s | mmq pf t/s | speedup    | baseline gen t/s | mmq gen t/s | gen ratio |
+|--------|-----------------|------------|------------|------------------|-------------|-----------|
+|  2048  | 373.21          | 1033.42    | **2.77x**  | 40.48            | 39.33       | 0.972x    |
+|  4096  | 366.25          | 1041.39    | **2.84x**  | 39.64            | 38.64       | 0.975x    |
+|  6144  | 364.81          | 1025.24    | **2.81x**  | 39.50            | 38.55       | 0.976x    |
+|  8192  | 364.01          | 1026.71    | **2.82x**  | 38.81            | 37.88       | 0.976x    |
+| 10240  | 361.75          | 1019.53    | **2.82x**  | 38.45            | 37.57       | 0.977x    |
+| 12288  | 360.52          | 1013.17    | **2.81x**  | 38.31            | 37.22       | 0.972x    |
+| 14336  | 359.15          | 1004.45    | **2.80x**  | 38.09            | 37.18       | 0.976x    |
+| 16384  | 357.99          | 1001.29    | **2.80x**  | 38.71            | 37.86       | 0.978x    |
+
+Sustained ~2.80x prefill speedup across the swept context range; gen
+within 2.5% of baseline (run-to-run variance).  See `local/docs/`
+(auto-round companion repo) for the full Phase 0-7 execution log,
+parity-test output, and detailed plan.

cuda/mmq/ds4_ggml_stubs.cu

@@ -0,0 +1,160 @@
+// SPDX-License-Identifier: MIT
+// Implementations of the ggml-API stubs declared in ds4_ggml_stubs.h plus
+// the bodies of ggml_backend_cuda_context / ggml_cuda_info /
+// new_pool_for_device that the vendored common.cuh declares without
+// defining.
+//
+// Phase 0: pool is plain cudaMallocAsync / cudaFreeAsync. Phase 4 swaps
+// in ds4's existing cuda_tmp_alloc slab allocator.
+
+#include "common.cuh"   // pulls in ds4_ggml_stubs.h via redirect headers
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+
+#include <chrono>
+#include <cstdio>
+#include <cstdlib>
+#include <memory>
+#include <mutex>
+
+// ----------------------------------------------------------------------------
+// Device info singleton.
+//
+// Common.cuh declares `const ggml_cuda_device_info & ggml_cuda_info();` -
+// we provide the body. The struct layout (common.cuh:1091) is:
+//   { int device_count; cuda_device_info devices[GGML_CUDA_MAX_DEVICES];
+//     std::array<float, GGML_CUDA_MAX_DEVICES> default_tensor_split; }
+// where cuda_device_info has { cc, nsm, smpb, smpbo, integrated, vmm,
+// vmm_granularity, total_vram, warp_size, supports_cooperative_launch }.
+// ----------------------------------------------------------------------------
+
+const ggml_cuda_device_info & ggml_cuda_info() {
+    static ggml_cuda_device_info info;
+    static std::once_flag once;
+    std::call_once(once, []{
+        int count = 0;
+        cudaError_t err = cudaGetDeviceCount(&count);
+        if (err != cudaSuccess) {
+            fprintf(stderr, "ggml_cuda_info: cudaGetDeviceCount failed: %s\n", cudaGetErrorString(err));
+            count = 0;
+        }
+        if (count > GGML_CUDA_MAX_DEVICES) count = GGML_CUDA_MAX_DEVICES;
+        info.device_count = count;
+        for (int i = 0; i < count; i++) {
+            cudaDeviceProp p;
+            CUDA_CHECK(cudaGetDeviceProperties(&p, i));
+            info.devices[i].cc                          = p.major * 100 + p.minor * 10;
+            info.devices[i].nsm                         = p.multiProcessorCount;
+            info.devices[i].smpb                        = p.sharedMemPerBlock;
+            info.devices[i].smpbo                       = p.sharedMemPerBlockOptin;
+            info.devices[i].integrated                  = p.integrated != 0;
+            info.devices[i].vmm                         = false;
+            info.devices[i].vmm_granularity             = 0;
+            info.devices[i].total_vram                  = p.totalGlobalMem;
+            info.devices[i].warp_size                   = p.warpSize;
+            info.devices[i].supports_cooperative_launch = p.cooperativeLaunch != 0;
+        }
+    });
+    return info;
+}
+
+// ggml_cuda_get_device / ggml_cuda_set_device are declared (not defined) in
+// common.cuh. We provide thin wrappers.
+
+int ggml_cuda_get_device() {
+    int dev = 0;
+    cudaGetDevice(&dev);
+    return dev;
+}
+
+void ggml_cuda_set_device(int device) {
+    int cur = -1;
+    cudaGetDevice(&cur);
+    if (cur != device) {
+        CUDA_CHECK(cudaSetDevice(device));
+    }
+}
+
+int64_t ggml_time_us() {
+    using clock = std::chrono::steady_clock;
+    static const auto t0 = clock::now();
+    return std::chrono::duration_cast<std::chrono::microseconds>(clock::now() - t0).count();
+}
+
+// ----------------------------------------------------------------------------
+// ggml_cuda_error: invoked by the CUDA_CHECK / CUBLAS_CHECK macros defined
+// in common.cuh on the error path. Marked [[noreturn]] in the declaration
+// (common.cuh:155) - abort() satisfies that contract.
+// ----------------------------------------------------------------------------
+
+[[noreturn]] void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg) {
+    fprintf(stderr, "CUDA error: %s\n  call: %s\n  in: %s at %s:%d\n", msg, stmt, func, file, line);
+    fflush(stderr);
+    abort();
+}
+
+// ----------------------------------------------------------------------------
+// Concrete pool wrapping cudaMallocAsync / cudaFreeAsync.
+// ----------------------------------------------------------------------------
+
+namespace {
+
+/* Thread-local stream that ds4_naive_pool uses for cudaMallocAsync /
+ * cudaFreeAsync.  Defaults to cudaStreamPerThread (preserves prior
+ * behaviour).  Step 8 / CUDA Graphs sets this to the capture stream
+ * before allocating, so the alloc node lives on the captured stream
+ * and capture is not invalidated.  Set via ds4_pool_set_stream() from
+ * ds4_mmq.cu wrappers; an explicit stream=0 means the legacy default
+ * stream so pool ops stay ordered with legacy-stream kernels. */
+static thread_local cudaStream_t t_ds4_pool_stream = cudaStreamPerThread;
+
+} // anonymous namespace
+
+extern "C" void ds4_pool_set_stream(cudaStream_t stream) {
+    t_ds4_pool_stream = stream;
+}
+
+extern "C" cudaStream_t ds4_pool_get_stream(void) {
+    return t_ds4_pool_stream;
+}
+
+namespace {
+
+struct ds4_naive_pool : public ggml_cuda_pool {
+    int device;
+
+    explicit ds4_naive_pool(int device) : device(device) {}
+
+    void * alloc(size_t size, size_t * actual_size) override {
+        ggml_cuda_set_device(device);
+        void * ptr = nullptr;
+        CUDA_CHECK(cudaMallocAsync(&ptr, size, t_ds4_pool_stream));
+        if (actual_size) *actual_size = size;
+        return ptr;
+    }
+
+    void free(void * ptr, size_t /*size*/) override {
+        if (!ptr) return;
+        ggml_cuda_set_device(device);
+        CUDA_CHECK(cudaFreeAsync(ptr, t_ds4_pool_stream));
+    }
+};
+
+} // anonymous namespace
+
+std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int device, int /*stream_no*/) {
+    return std::unique_ptr<ggml_cuda_pool>(new ds4_naive_pool(device));
+}
+
+ggml_backend_cuda_context::~ggml_backend_cuda_context() {
+    if (copy_event) {
+        cudaEventDestroy(copy_event);
+        copy_event = nullptr;
+    }
+    // streams[][], cublas_handles[], and pools[][] are owned-by-value
+    // (cudaStream_t and cublasHandle_t are opaque handles - destroying the
+    // context "should" tear them down, but in our shim ds4 manages streams
+    // externally and we leave them alone. The pools auto-destruct via
+    // unique_ptr.).
+}

cuda/mmq/ds4_ggml_stubs.h

@@ -0,0 +1,303 @@
+// SPDX-License-Identifier: MIT
+// ds4_ggml_stubs.h - minimal ggml-API stubs for ds4's vendored mmq kernels.
+//
+// The mmq.cuh / mma.cuh / vecdotq.cuh / quantize.cuh / mmid.cuh / common.cuh
+// files in this directory are vendored verbatim from llama.cpp's ggml-cuda
+// backend (MIT, copyright 2023-2026 The ggml authors). They transitively
+// #include "ggml.h", "ggml-impl.h", "ggml-cuda.h" - in ds4 those names
+// resolve to thin redirect headers in this directory which all #include this
+// stubs file.
+//
+// This file declares the minimum surface of the ggml API that the vendored
+// CUDA code references, EXCLUDING what's already provided by common.cuh
+// itself (compute-capability constants, MMA flags, ggml_cuda_device_info,
+// ggml_cuda_pool, ggml_cuda_pool_alloc, ggml_backend_cuda_context, the
+// CUDA_CHECK / CUBLAS_CHECK macros, ggml_cuda_get_device, ggml_cuda_set_device,
+// ggml_cuda_info). Those names live in common.cuh and we let it own them.
+//
+// Things this header DOES provide:
+//   * GGML_ASSERT / GGML_ABORT / GGML_UNUSED / GGML_UNUSED_VARS / GGML_PAD
+//   * GGML_MAX_DIMS / GGML_MAX_SRC / GGML_CUDA_NAME / GGML_CUDA_MAX_DEVICES /
+//     GGML_CUDA_MAX_STREAMS / GGML_LOG_DEBUG
+//   * enum ggml_type (all 21 mmq type codes - we only USE a subset for V4
+//     Flash but the switch in mmq.cu's downstream replacement must compile)
+//   * enum ggml_glu_op (just for the unused mm_fusion_args fields)
+//   * struct ggml_tensor (complete enough for common.cuh's
+//     ggml_cuda_concurrent_event::is_valid() to compile - we never call it)
+//   * int64_t ggml_nbytes(const ggml_tensor *) (stub - never called)
+//   * int64_t ggml_time_us() (used by USE_CUDA_GRAPH paths we disable)
+//   * inline ggml_type_size() / ggml_blck_size() lookups
+//
+// Things ggml-common.h (vendored) owns:
+//   * ggml_half / ggml_half2 typedefs
+//   * GGML_EXTENSION macro
+//   * block_q*, block_iq* struct definitions
+
+#pragma once
+
+#include <cassert>
+#include <cstdarg>
+#include <cstddef>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+
+// ----------------------------------------------------------------------------
+// Macros
+// ----------------------------------------------------------------------------
+
+#ifndef GGML_ASSERT
+#define GGML_ASSERT(cond) \
+    do { if (!(cond)) { \
+        fprintf(stderr, "GGML_ASSERT(%s) failed at %s:%d\n", #cond, __FILE__, __LINE__); \
+        abort(); \
+    } } while (0)
+#endif
+
+#ifndef GGML_ABORT
+#define GGML_ABORT(fmt, ...) \
+    do { \
+        fprintf(stderr, "GGML_ABORT: " fmt " at %s:%d\n", ##__VA_ARGS__, __FILE__, __LINE__); \
+        abort(); \
+    } while (0)
+#endif
+
+#ifndef GGML_UNUSED
+#define GGML_UNUSED(x) ((void)(x))
+#endif
+
+// Variadic GGML_UNUSED_VARS: drop up to 12 unused names without warnings.
+#ifndef GGML_UNUSED_VARS
+#define GGML_UNUSED_VARS_1(_1)                                             GGML_UNUSED(_1)
+#define GGML_UNUSED_VARS_2(_1,_2)                                          GGML_UNUSED(_1); GGML_UNUSED(_2)
+#define GGML_UNUSED_VARS_3(_1,_2,_3)                                       GGML_UNUSED_VARS_2(_1,_2); GGML_UNUSED(_3)
+#define GGML_UNUSED_VARS_4(_1,_2,_3,_4)                                    GGML_UNUSED_VARS_3(_1,_2,_3); GGML_UNUSED(_4)
+#define GGML_UNUSED_VARS_5(_1,_2,_3,_4,_5)                                 GGML_UNUSED_VARS_4(_1,_2,_3,_4); GGML_UNUSED(_5)
+#define GGML_UNUSED_VARS_6(_1,_2,_3,_4,_5,_6)                              GGML_UNUSED_VARS_5(_1,_2,_3,_4,_5); GGML_UNUSED(_6)
+#define GGML_UNUSED_VARS_7(_1,_2,_3,_4,_5,_6,_7)                           GGML_UNUSED_VARS_6(_1,_2,_3,_4,_5,_6); GGML_UNUSED(_7)
+#define GGML_UNUSED_VARS_8(_1,_2,_3,_4,_5,_6,_7,_8)                        GGML_UNUSED_VARS_7(_1,_2,_3,_4,_5,_6,_7); GGML_UNUSED(_8)
+#define GGML_UNUSED_VARS_9(_1,_2,_3,_4,_5,_6,_7,_8,_9)                     GGML_UNUSED_VARS_8(_1,_2,_3,_4,_5,_6,_7,_8); GGML_UNUSED(_9)
+#define GGML_UNUSED_VARS_10(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10)                GGML_UNUSED_VARS_9(_1,_2,_3,_4,_5,_6,_7,_8,_9); GGML_UNUSED(_10)
+#define GGML_UNUSED_VARS_11(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10,_11)            GGML_UNUSED_VARS_10(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10); GGML_UNUSED(_11)
+#define GGML_UNUSED_VARS_12(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10,_11,_12)        GGML_UNUSED_VARS_11(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10,_11); GGML_UNUSED(_12)
+#define GGML_UNUSED_VARS_PICK(_1,_2,_3,_4,_5,_6,_7,_8,_9,_10,_11,_12,NAME,...) NAME
+#define GGML_UNUSED_VARS(...) \
+    GGML_UNUSED_VARS_PICK(__VA_ARGS__, \
+        GGML_UNUSED_VARS_12, GGML_UNUSED_VARS_11, GGML_UNUSED_VARS_10, \
+        GGML_UNUSED_VARS_9, GGML_UNUSED_VARS_8, GGML_UNUSED_VARS_7, \
+        GGML_UNUSED_VARS_6, GGML_UNUSED_VARS_5, GGML_UNUSED_VARS_4, \
+        GGML_UNUSED_VARS_3, GGML_UNUSED_VARS_2, GGML_UNUSED_VARS_1)(__VA_ARGS__)
+#endif
+
+#ifndef GGML_PAD
+#define GGML_PAD(x, n) (((x) + (n) - 1) / (n) * (n))
+#endif
+
+#ifndef GGML_MAX_DIMS
+#define GGML_MAX_DIMS 4
+#endif
+
+#ifndef GGML_MAX_SRC
+#define GGML_MAX_SRC  10
+#endif
+
+#ifndef GGML_CUDA_NAME
+#define GGML_CUDA_NAME "DS4_CUDA"
+#endif
+
+#ifndef GGML_CUDA_MAX_DEVICES
+#define GGML_CUDA_MAX_DEVICES 16
+#endif
+
+#ifndef GGML_CUDA_MAX_STREAMS
+#define GGML_CUDA_MAX_STREAMS 8
+#endif
+
+#ifndef GGML_LOG_DEBUG
+#define GGML_LOG_DEBUG(...) ((void)0)
+#endif
+
+// Cuda-graphs are explicitly disabled - ds4 manages its own streams.
+#undef GGML_CUDA_USE_GRAPHS
+#undef GGML_HIP_GRAPHS
+#undef GGML_MUSA_GRAPHS
+
+// GGML_EXTENSION: ggml-common.h provides the canonical definition. We leave
+// it undefined here so the vendored header's `#define GGML_EXTENSION
+// __extension__` wins.
+
+// ----------------------------------------------------------------------------
+// Quantization type enum.
+//
+// Order matches llama.cpp's enum ggml_type. Values are pinned because the
+// mmq switch uses them as case labels.
+// ----------------------------------------------------------------------------
+
+enum ggml_type {
+    GGML_TYPE_F32     = 0,
+    GGML_TYPE_F16     = 1,
+    GGML_TYPE_Q4_0    = 2,
+    GGML_TYPE_Q4_1    = 3,
+    // GGML_TYPE_Q4_2 / Q4_3 deprecated
+    GGML_TYPE_Q5_0    = 6,
+    GGML_TYPE_Q5_1    = 7,
+    GGML_TYPE_Q8_0    = 8,
+    GGML_TYPE_Q8_1    = 9,
+    GGML_TYPE_Q2_K    = 10,
+    GGML_TYPE_Q3_K    = 11,
+    GGML_TYPE_Q4_K    = 12,
+    GGML_TYPE_Q5_K    = 13,
+    GGML_TYPE_Q6_K    = 14,
+    GGML_TYPE_Q8_K    = 15,
+    GGML_TYPE_IQ2_XXS = 16,
+    GGML_TYPE_IQ2_XS  = 17,
+    GGML_TYPE_IQ3_XXS = 18,
+    GGML_TYPE_IQ1_S   = 19,
+    GGML_TYPE_IQ4_NL  = 20,
+    GGML_TYPE_IQ3_S   = 21,
+    GGML_TYPE_IQ2_S   = 22,
+    GGML_TYPE_IQ4_XS  = 23,
+    GGML_TYPE_I8      = 24,
+    GGML_TYPE_I16     = 25,
+    GGML_TYPE_I32     = 26,
+    GGML_TYPE_I64     = 27,
+    GGML_TYPE_F64     = 28,
+    GGML_TYPE_IQ1_M   = 29,
+    GGML_TYPE_BF16    = 30,
+    GGML_TYPE_MXFP4   = 39,
+    GGML_TYPE_NVFP4   = 40,
+    GGML_TYPE_Q1_0    = 41,
+    GGML_TYPE_COUNT   = 42,
+};
+
+enum ggml_glu_op {
+    GGML_GLU_OP_REGLU,
+    GGML_GLU_OP_GEGLU,
+    GGML_GLU_OP_SWIGLU,
+    GGML_GLU_OP_SWIGLU_OAI, // referenced by mmvq.cu's fused-GLU epilogue
+    GGML_GLU_OP_COUNT,
+};
+
+// ----------------------------------------------------------------------------
+// ggml_tensor: complete enough for common.cuh's
+// ggml_cuda_concurrent_event::is_valid() to compile cleanly. We NEVER
+// instantiate or dereference one of these - the concurrent path is
+// disabled.
+//
+// Field set matches the upstream order/types so cudaGraph node_properties
+// (which holds a `ggml_tensor node` by value inside `#ifdef USE_CUDA_GRAPH`)
+// also compiles. Sizes are conservative for storage; ds4 never copies into
+// these.
+// ----------------------------------------------------------------------------
+
+struct ggml_tensor {
+    enum ggml_type type;
+    int32_t op;                              // enum ggml_op (opaque to us)
+    int32_t flags;
+    int64_t ne[GGML_MAX_DIMS];               // shape
+    size_t  nb[GGML_MAX_DIMS];               // stride bytes
+    int32_t op_params[16];                   // GGML_MAX_OP_PARAMS / sizeof(int32_t)
+    struct ggml_tensor * src[GGML_MAX_SRC];
+    struct ggml_tensor * view_src;
+    size_t  view_offs;
+    void *  data;
+    char    name[64];                        // GGML_MAX_NAME
+    void *  extra;
+    char    padding[8];
+};
+
+// ggml_nbytes: byte size of tensor data. We never call this; provide a
+// stub so common.cuh's is_valid() compiles. If anything does call it the
+// returned 0 will surface as an immediate logic error.
+static inline int64_t ggml_nbytes(const struct ggml_tensor * /*t*/) { return 0; }
+
+// Microsecond timer (used only inside USE_CUDA_GRAPH paths we disable).
+int64_t ggml_time_us();
+
+// ----------------------------------------------------------------------------
+// Inline size traits.
+//
+// Lookup tables hand-aligned with llama.cpp's ggml_type_traits in
+// ggml/src/ggml.c. Only types referenced by the vendored switch are needed.
+// ----------------------------------------------------------------------------
+
+inline size_t ggml_type_size(enum ggml_type t) {
+    switch (t) {
+        case GGML_TYPE_F32:     return 4;
+        case GGML_TYPE_F16:     return 2;
+        case GGML_TYPE_BF16:    return 2;
+        case GGML_TYPE_I8:      return 1;
+        case GGML_TYPE_I16:     return 2;
+        case GGML_TYPE_I32:     return 4;
+        case GGML_TYPE_I64:     return 8;
+        case GGML_TYPE_F64:     return 8;
+        case GGML_TYPE_Q4_0:    return 18;
+        case GGML_TYPE_Q4_1:    return 20;
+        case GGML_TYPE_Q5_0:    return 22;
+        case GGML_TYPE_Q5_1:    return 24;
+        case GGML_TYPE_Q8_0:    return 34;
+        case GGML_TYPE_Q8_1:    return 36;
+        case GGML_TYPE_Q2_K:    return 84;
+        case GGML_TYPE_Q3_K:    return 110;
+        case GGML_TYPE_Q4_K:    return 144;
+        case GGML_TYPE_Q5_K:    return 176;
+        case GGML_TYPE_Q6_K:    return 210;
+        case GGML_TYPE_Q8_K:    return 292;
+        case GGML_TYPE_IQ2_XXS: return 66;
+        case GGML_TYPE_IQ2_XS:  return 74;
+        case GGML_TYPE_IQ2_S:   return 82;
+        case GGML_TYPE_IQ3_XXS: return 98;
+        case GGML_TYPE_IQ3_S:   return 110;
+        case GGML_TYPE_IQ1_S:   return 50;
+        case GGML_TYPE_IQ1_M:   return 56;
+        case GGML_TYPE_IQ4_NL:  return 18;
+        case GGML_TYPE_IQ4_XS:  return 136;
+        case GGML_TYPE_MXFP4:   return 17;
+        case GGML_TYPE_NVFP4:   return 18;
+        case GGML_TYPE_Q1_0:    return 36;
+        default:                return 0;
+    }
+}
+
+inline int64_t ggml_blck_size(enum ggml_type t) {
+    switch (t) {
+        case GGML_TYPE_F32:
+        case GGML_TYPE_F16:
+        case GGML_TYPE_BF16:
+        case GGML_TYPE_I8:
+        case GGML_TYPE_I16:
+        case GGML_TYPE_I32:
+        case GGML_TYPE_I64:
+        case GGML_TYPE_F64:
+            return 1;
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q5_1:
+        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_Q8_1:
+        case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_MXFP4:
+            return 32;
+        case GGML_TYPE_NVFP4:
+            return 64;
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+        case GGML_TYPE_Q8_K:
+        case GGML_TYPE_IQ2_XXS:
+        case GGML_TYPE_IQ2_XS:
+        case GGML_TYPE_IQ2_S:
+        case GGML_TYPE_IQ3_XXS:
+        case GGML_TYPE_IQ3_S:
+        case GGML_TYPE_IQ1_S:
+        case GGML_TYPE_IQ1_M:
+        case GGML_TYPE_IQ4_XS:
+        case GGML_TYPE_Q1_0:
+            return 256;
+        default:
+            return 1;
+    }
+}

cuda/mmq/ds4_mmq.h

@@ -0,0 +1,890 @@
+// SPDX-License-Identifier: MIT
+// ds4_mmq.h - public C ABI for ds4's quantized matmul kernels.
+//
+// All functions are extern "C" so ds4.c / ds4_cuda.cu can call them
+// without C++ compilation. Functions return 0 on success and non-zero on
+// failure (with stderr error message). Device pointers are caller-owned.
+//
+// Phase 0: skeleton only. Q8_0 dense entry compiles and instantiates
+// mul_mat_q_case<Q8_0> but is not yet wired into ds4_cuda.cu.
+// Phase 1: Q8_1 activation quantizer wrapper added.
+// Phase 2: Q8_0 dense entry verified against cublas+dequant baseline.
+// Phase 3: Q2_K + IQ2_XXS dense entries.
+// Phase 4: MoE _id variants of all three.
+
+#pragma once
+
+#include <cuda_runtime.h>
+#include <stddef.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// One-time init. Sets the current CUDA device and triggers lazy population
+// of the device-info singleton. Safe to call repeatedly.
+//
+//   device: CUDA device ordinal (0 for the primary GPU).
+// Returns 0 on success.
+int ds4_mmq_init(int device);
+void ds4_mmq_set_aligned_q81_scratch(void *ptr, size_t bytes);
+
+// Query whether ds4_mmq is willing to handle a given matmul. Returns
+//   1 if mmq is faster than dequant+cublas for this shape on this device,
+//   0 otherwise (caller should fall back to its existing dequant+cublas path).
+//
+// Wraps ggml_cuda_should_use_mmq. type_x uses ds4 quant codes which match
+// ggml's enum:
+//   8  = Q8_0
+//   10 = Q2_K
+//   16 = IQ2_XXS
+//
+//   ne11:      batch dimension (number of activation columns / tokens).
+//   n_experts: 0 for dense matmul, >0 for MoE (e.g. 256 for V4 Flash).
+int ds4_mmq_should_use(int type_x, int64_t ne11, int64_t n_experts);
+
+// Dense matmul entry points. Per-type wrappers that all share the same
+// underlying mul_mat_q template, parameterised by the weight quant type.
+//
+// All three variants compute:
+//
+//   out[col, row] = sum_k W[row, k] * X[k, col]      0 <= row < M, 0 <= col < N
+//
+// Layouts (matching ggml + llama.cpp mmq conventions, all on device):
+//   W:       [M rows, K cols], row-major, packed in the type-specific block
+//            format. K must be a multiple of 256.
+//   X_f32:   [N rows, K cols] F32 row-major (logical [K, N] with K
+//            innermost - i.e. for each "column" col of the logical [K, N]
+//            matrix, K contiguous floats live at X[col*K .. col*K + K]).
+//   out_f32: caller-allocated, M*N floats. mmq writes in column-major:
+//            out[col*M + row]. Callers expecting row-major must transpose.
+//
+// Returns 0 on success, non-zero on validation or launch failure.
+
+int ds4_mmq_q8_0_dense(
+    const void  * W_q8_0,
+    const float * X_f32,
+    float       * out_f32,
+    int           M,
+    int           N,
+    int           K,
+    cudaStream_t  stream);
+
+// v0.5 flat-pool p5a: same contract as ds4_mmq_q8_0_dense but the
+// activation arrives PRE-QUANTIZED in the block_q8_1_mmq D4 layout
+// (producer-emitted, op-for-op equal to quantize_mmq_q8_1<D4>; the fused
+// own out_a kernel is the producer).  Y_q8_mmq must hold N*(K/128)
+// 144-byte blocks (ib = kseg*N + row) plus mmq_x_max tail slack, which
+// this entry zeroes.  Requires K such that GGML_PAD(K, MATRIX_ROW_PADDING)
+// == K.  Returns 0 on success; callers fall back to the quantizing entry
+// on any non-zero.
+int ds4_mmq_q8_0_dense_preq(
+    const void  * W_q8_0,
+    const void  * Y_q8_mmq,
+    size_t        y_bytes,
+    float       * out_f32,
+    int           M,
+    int           N,
+    int           K,
+    cudaStream_t  stream);
+
+// p5a verify instrument: reference quantize (dense_impl parameters) into a
+// caller buffer, for byte-diffing producer emits.
+int ds4_mmq_q8_0_quantize_ref(
+    const float * X,
+    void        * y,
+    size_t        y_bytes,
+    int           N,
+    int           K,
+    cudaStream_t  stream);
+
+// Dense Q8_0 D2R on the kind-5 aligned artifact (weight server
+// --repack-q8-aligned).  Same in/out contract as ds4_mmq_q8_0_dense but W is
+// the ALIGNED artifact base ([half dq[nblk]][pad64][int8 qs]), and the shape
+// must satisfy M % 128 == 0 && K % 1024 == 0.  Callers gate on n_tok scale
+// and K <= 4096 (o_proj's K=8192 measured faster on mmq).
+int ds4_mmq_q8_0_dense_d2r(
+    const void  * W_aligned,
+    const float * X_f32,
+    float       * out_f32,
+    int           M,
+    int           N,
+    int           K,
+    cudaStream_t  stream);
+
+// flat-pool p5c: D2R dense over a producer-quantized token-major Y
+// (block_q8_1_mmq D4, ib = kseg*N + row, requires GGML_PAD(K) == K).
+// y_bytes must cover the payload plus max(mmq_x_max, 128) slack blocks;
+// the slack is zeroed here each call (S1.1a).  Returns 0 on success;
+// callers fall back to the quantizing entry on any non-zero.
+int ds4_mmq_q8_0_dense_d2r_preq(
+    const void  * W_aligned,
+    const void  * Y_q8_mmq,
+    size_t        y_bytes,
+    float       * out_f32,
+    int           M,
+    int           N,
+    int           K,
+    cudaStream_t  stream);
+
+int ds4_mmq_q2_K_dense(
+    const void  * W_q2_K,
+    const float * X_f32,
+    float       * out_f32,
+    int           M,
+    int           N,
+    int           K,
+    cudaStream_t  stream);
+
+int ds4_mmq_iq2_xxs_dense(
+    const void  * W_iq2_xxs,
+    const float * X_f32,
+    float       * out_f32,
+    int           M,
+    int           N,
+    int           K,
+    cudaStream_t  stream);
+
+int ds4_mmq_q4_K_dense(
+    const void  * W_q4_K,
+    const float * X_f32,
+    float       * out_f32,
+    int           M,
+    int           N,
+    int           K,
+    cudaStream_t  stream);
+
+int ds4_mmq_mxfp4_dense(
+    const void  * W_mxfp4,
+    const float * X_f32,
+    float       * out_f32,
+    int           M,
+    int           N,
+    int           K,
+    cudaStream_t  stream);
+
+// MoE matmul entry points. For each (token, slot-within-token's-top-k) pair
+// the kernel computes:
+//
+//   out[col, row] = sum_k W[ids[token, slot], row, k] * X[token, k]
+//
+// where col = token * n_expert_used + slot, row in [0, M).  The caller is
+// responsible for any downstream sum-weighted-by-router-weights reduction
+// across the n_expert_used dimension (Phase 5 wires this into ds4's
+// existing moe_sum_kernel).
+//
+// Layouts:
+//   W:       device pointer, [n_experts, M rows, K cols] in the
+//            type-specific block format.  Per-expert slab is M*K/blck
+//            blocks stored contiguously; experts are stacked.
+//   X_f32:   device pointer, [n_tokens, K] F32 row-major (K innermost).
+//   ids:     device pointer, [n_tokens, n_expert_used] int32_t row-major.
+//            ids[t*n_expert_used + s] is the expert id for token t's
+//            s-th routing slot.  Values must be in [0, n_experts).
+//   out_f32: caller-allocated, M * n_tokens * n_expert_used floats.
+//            Column-major: out[col*M + row].
+//
+// K must be a multiple of 256.  n_expert_used must be one of the values
+// the vendored mm_ids_helper template specialises on: 2, 4, 6, 8, 16, 32
+// (or any other value, which falls back to the generic path).  For V4
+// Flash, n_expert_used = 6.
+//
+// Returns 0 on success, non-zero on validation or launch failure.
+
+int ds4_mmq_q8_0_moe(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_q2_K_moe(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_iq2_xxs_moe(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+// ds4 (P4 Inc3): same contract as ds4_mmq_q2_K_moe but W_soa is the aligned
+// row-pair-SoA artifact (weight server --repack-q2k-aligned, layout in
+// ds4_mmq_q2_k_aligned_bytes' comment) instead of the raw block stream.  The
+// tile loader reads the SoA sections directly -- bit-identical output to the
+// raw path, no derepack scratch.
+// CONTRACT DIFFERENCE (P3): unlike the raw entries, the output is NOT
+// nonfinite-sanitized; the routed-MoE consumers (moe_mmq_swiglu / moe_sum
+// with guard_nonfinite=1) sanitize at read, so the standalone whole-buffer
+// pass is skipped.  New callers must sanitize at consumption.
+int ds4_mmq_q2_K_moe_soa(
+    const void    * W_soa,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_q4_K_moe(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_mxfp4_moe(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+// Paired MoE entries. Compute gate AND up over the same activation in a
+// single call so the Q8_1 quantize of X (and the mm_ids_helper bookkeeping)
+// happens once instead of twice. Both weights must be the same quant type
+// and the same shape (M, K, n_experts); out_a / out_b have the same layout
+// as a single ds4_mmq_<type>_moe call. Saves one launch of
+// quantize_mmq_q8_1_cuda and one ggml_cuda_launch_mm_ids_helper per MoE
+// block. See ds4_mmq.cu / routed_moe_launch for the wiring.
+//
+// Returns 0 on success; on error neither output is guaranteed valid.
+
+int ds4_mmq_iq2_xxs_moe_pair(
+    const void    * W_a,
+    const void    * W_b,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_a,
+    float         * out_b,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+// ds4 (P4 Inc3): same contract as ds4_mmq_iq2_xxs_moe_pair but over the
+// aligned-SoA artifacts (weight server --repack-iq2-aligned); see
+// ds4_mmq_q2_K_moe_soa.
+int ds4_mmq_iq2_xxs_moe_pair_soa(
+    const void    * Wa_soa,
+    const void    * Wb_soa,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_a,
+    float         * out_b,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+/* v0.5 inc-9 (F7): fused target-prefill pipeline over the aligned-SoA
+ * IQ2_XXS gate/up and Q2_K down artifacts.  Builds the expert-major
+ * assignment map once, runs the paired gate/up MMQs, computes clamp +
+ * SwiGLU + router weighting in mid_f32, then gathers and quantizes those
+ * rows for the Q2_K down MMQ through the same ids_dst/expert_bounds.  No
+ * second mm_ids_helper; gate/up/mid/down keep the pair-major layout. */
+int ds4_mmq_iq2_xxs_q2_K_moe_fused_soa(
+    const void    * W_gate_soa,
+    const void    * W_up_soa,
+    const void    * W_down_soa,
+    const float   * X_f32,
+    const int32_t * ids,
+    const float   * router_weights,
+    float         * gate_f32,
+    float         * up_f32,
+    float         * mid_f32,
+    float         * down_f32,
+    int             expert_mid_dim,
+    int             expert_in_dim,
+    int             out_dim,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    float           clamp,
+    cudaStream_t    stream);
+
+/* Aligned-artifact production fast path: gate/up stay in registers, weighted
+ * SwiGLU is quantized directly into down_q8_scratch, and only the pair-major
+ * down output is materialized.  Caller-owned scratch keeps this hot path free
+ * of stream-ordered allocations; all three ranges must be distinct, remain
+ * live until return, and be sized to their LOGICAL segments (never an owning
+ * arena's capacity). */
+/* flat-pool p5c: input_q8_ext, when non-NULL, is a producer-emitted
+ * TOKEN-COMPACT block_q8_1_mmq buffer of X (ib = kseg*n_tokens + row,
+ * n_tokens rows) and the internal activation quantize is skipped; the
+ * gate/up D2R kernel reads it through the p5b ids_src indirection.
+ * Ignored (classic gathered quantize) when DS4_MMQ_NO_YIND is set. */
+int ds4_mmq_iq2_xxs_q2_K_moe_fused_direct_soa(
+    const void    * W_gate_soa,
+    const void    * W_up_soa,
+    const void    * W_down_soa,
+    const float   * X_f32,
+    const int32_t * ids,
+    const float   * router_weights,
+    void          * input_q8_scratch,
+    size_t          input_q8_scratch_bytes,
+    void          * down_q8_scratch,
+    size_t          down_q8_scratch_bytes,
+    void          * work_scratch,
+    size_t          work_scratch_bytes,
+    const void    * input_q8_ext,
+    size_t          input_q8_ext_bytes,
+    float         * down_f32,
+    int             expert_mid_dim,
+    int             expert_in_dim,
+    int             out_dim,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    float           clamp,
+    cudaStream_t    stream);
+
+int ds4_mmq_iq2_xxs_q2_K_moe_fused_direct_scratch_sizes(
+    int     expert_mid_dim,
+    int     expert_in_dim,
+    int     n_tokens,
+    int     n_experts,
+    int     n_expert_used,
+    size_t *input_q8_bytes,
+    size_t *down_q8_bytes,
+    size_t *work_bytes);
+
+int ds4_mmq_q4_K_moe_pair(
+    const void    * W_a,
+    const void    * W_b,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_a,
+    float         * out_b,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_mxfp4_moe_pair(
+    const void    * W_a,
+    const void    * W_b,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_a,
+    float         * out_b,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+// MoE vector matmul entries (Step 6). Same signature and semantics as the
+// ds4_mmq_<type>_moe entries above, but route through llama.cpp's mmvq
+// kernels instead of mmq. mmvq is structurally optimised for small batch
+// counts (single-token decode, short prefill), where mmq's tile-based
+// approach wastes work on empty columns.
+//
+// Constraints:
+//   - n_tokens * something must fit under mmvq's per-arch batch cap
+//     (MMVQ_MAX_BATCH_SIZE = 8 on Blackwell). Specifically, ncols_dst as
+//     computed by the wrapper must be <= 8. The wrapper rejects with -1
+//     if the request is too large.
+//   - K must be a multiple of 256 (same as the mmq path).
+//
+// Unlike the mmq path, mmvq consumes a CANONICAL block_q8_1 buffer (not
+// the interleaved block_q8_1_mmq the mmq path uses). The wrapper builds
+// the canonical buffer internally; callers cannot reuse a Q8_1 buffer
+// previously built for the mmq path.
+//
+// Returns 0 on success, non-zero on validation or launch failure.
+
+int ds4_mmq_q8_0_moe_vec(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_q2_K_moe_vec(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_iq2_xxs_moe_vec(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_q4_K_moe_vec(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_mxfp4_moe_vec(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+// Aligned-SoA IQ2_XXS decode matvec (megakernel program M1-Inc1).
+//
+// block_iq2_xxs is 66 bytes, so the raw expert stream is only 2-byte aligned
+// and every 32-bit code word costs two 16-bit loads.  W_aligned is a repacked
+// copy of the SAME bytes with the block scales split out and the code stream
+// 64B-aligned:
+//
+//   [ __half dq[nblk] ][ pad to 64B ][ uint2 qs[nblk * 8] ]
+//
+// where nblk = n_experts * M * (K / 256) and block linear order matches the
+// raw tensor byte order (expert-major, then row, then block).  The weight
+// server builds this layout (--repack-iq2-aligned); use
+// ds4_mmq_iq2_xxs_aligned_bytes to size or validate an artifact.
+//
+// Semantics and output layout are identical to ds4_mmq_iq2_xxs_moe_vec at
+// n_tokens == 1 (the only supported width; other widths return non-zero so
+// the caller can fall back).
+uint64_t ds4_mmq_iq2_xxs_aligned_bytes(int M, int K, int n_experts);
+
+// M1-Inc2b: exact inverse of the weight-server repack.  Fills raw_out
+// (nblk * 66 bytes, raw block_iq2_xxs byte stream) from an aligned artifact,
+// device->device on `stream`.  Lets the batched/mmq raw-layout consumers run
+// from a device scratch instead of the client mmap when the raw spans were
+// excluded from the upload.
+int ds4_mmq_iq2_xxs_aligned_derepack(
+    const void    * W_aligned,
+    void          * raw_out,
+    int             M,
+    int             K,
+    int             n_experts,
+    cudaStream_t    stream);
+
+// M2 moe-down: aligned row-pair-SoA Q2_K routed-expert decode matvec.  Twin
+// of mul_mat_vec_q_moe<GGML_TYPE_Q2_K, 2> at the down-leg call shape
+// (n_expert_used == 1: each (token, slot) assignment is its own "token") with
+// bit-identical outputs; only the weight layout changes.  W_aligned is the
+// weight-server --repack-q2k-aligned artifact (DERIVED_Q2_K_ALIGNED_MOE,
+// REPLACES the raw range; byte-neutral):
+//
+//   npair = n_experts * (M/2) * (K/256)
+//   [ uint2 dm2[npair] ][ pad 64B ][ int4 sc4[npair*2] ][ pad 64B ]
+//   [ uint2 qs2[npair*16] ]
+//
+// keyed to the kernel's rows_per_block == 2 (see ds4_mmq.cu for the exact
+// field packing).  Use ds4_mmq_q2_k_aligned_bytes to size or validate an
+// artifact.  Unsupported shapes return non-zero so the caller can fall back.
+uint64_t ds4_mmq_q2_k_aligned_bytes(int M, int K, int n_experts);
+
+int ds4_mmq_q2_K_aligned_moe_vec(
+    const void    * W_aligned,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+// Exact inverse of the weight-server repack: fills raw_out (nblk * 84 bytes,
+// raw block_q2_K stream) from an aligned artifact, device->device on
+// `stream`, for the batched/mmq raw-layout consumers (same role as
+// ds4_mmq_iq2_xxs_aligned_derepack).
+int ds4_mmq_q2_K_aligned_derepack(
+    const void    * W_aligned,
+    void          * raw_out,
+    int             M,
+    int             K,
+    int             n_experts,
+    cudaStream_t    stream);
+
+// M1-Inc3: aligned-SoA Q8_0 dense decode matvec.  Artifact layout
+// [__half dq[nblk]][pad to 64B][int8 qs[nblk*32]], nblk = M * (K/32), block
+// order equal to the raw tensor byte order (weight server
+// --repack-q8-aligned; raw spans stay served — dense artifacts duplicate,
+// they do not replace).  n_tokens == 1 and K % 1024 == 0 only; other shapes
+// return non-zero so the caller can fall back to ds4_mmq_q8_0_dense_vec.
+uint64_t ds4_mmq_q8_0_aligned_bytes(int M, int K);
+
+int ds4_mmq_q8_0_aligned_dense_vec(
+    const void  * W_aligned,
+    const float * X_f32,
+    float       * out_f32,
+    int           M,
+    int           N,
+    int           K,
+    cudaStream_t  stream);
+
+// Decode-only pair: quantize the shared activation once and compute both
+// aligned projections in one row-concatenated launch.
+int ds4_mmq_q8_0_aligned_dense_vec_pair(
+    const void  * W0_aligned,
+    const void  * W1_aligned,
+    const float * X_f32,
+    float       * out0_f32,
+    float       * out1_f32,
+    int           M0,
+    int           M1,
+    int           K,
+    cudaStream_t  stream);
+
+int ds4_mmq_iq2_xxs_aligned_moe_vec(
+    const void    * W_aligned,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+// M1-Inc2 variants over the same aligned artifacts (n_tokens == 1 only).
+//
+// _pair_vec: one activation quantize + one launch computes both raw gate and
+// up outputs (nonfinite zeroed in-kernel; no sanitize pass needed).  The
+// caller still runs its clamp-aware SwiGLU.
+//
+// _gate_up_mid_vec: additionally folds the clamp/SwiGLU/router-weight
+// epilogue (identical semantics to ds4_mmq_moe_gate_up_mid_q8_1_qwarp32) and
+// writes mid[slot * M + row] directly.
+int ds4_mmq_iq2_xxs_aligned_moe_pair_vec(
+    const void    * W_gate_aligned,
+    const void    * W_up_aligned,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * gate_out,
+    float         * up_out,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(
+    const void    * W_gate_aligned,
+    const void    * W_up_aligned,
+    const float   * X_f32,
+    const int32_t * ids,
+    const float   * weights,
+    float         * mid_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    float           clamp,
+    cudaStream_t    stream);
+
+// Fused down+sum vector entries for routed MoE with top_k=6. These preserve
+// the canonical Q8_1 activation quantization used by the regular _moe_vec
+// path, but avoid materializing [token, slot, out_dim] down results and avoid
+// the separate slot-sum kernel. The caller must already have baked router
+// weights into X_f32, so the helper computes:
+//
+//   out[token, row] = sum_slot W[ids[token, slot], row, :] @ X[token, slot, :]
+//
+// Layouts:
+//   W:       [n_experts, M rows, K cols] in Q2_K or Q4_K blocks
+//   X_f32:   [n_tokens * 6, K] F32 row-major
+//   ids:     [n_tokens, 6] int32 row-major
+//   out_f32: [n_tokens, M] F32 row-major
+//
+// Constraints:
+//   - n_expert_used must be 6
+//   - K must be a multiple of 256
+//
+// Returns 0 on success, non-zero on validation or launch failure.
+
+int ds4_mmq_q2_K_moe_down_sum6_vec(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_q4_K_moe_down_sum6_vec(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_mxfp4_moe_down_sum6_vec(
+    const void    * W,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+// Fused gate+up+SwiGLU vector entries for routed MoE with top_k=6. These use
+// canonical Q8_1 activation quantization and compute weighted mid directly:
+//
+//   mid[token, slot, row] =
+//       silu(clamp_gate(W_gate[expert,row] @ X[token]))
+//     * clamp_up(W_up[expert,row] @ X[token])
+//     * router_weight[token, slot]
+//
+// Clamp semantics match ds4_cuda.cu: gate is capped above, up is clamped to
+// [-clamp, clamp], and no clamp is applied when clamp <= 1e-6.
+//
+// Layouts:
+//   W_gate/W_up: [n_experts, M rows, K cols] in IQ2_XXS or Q4_K blocks
+//   X_f32:       [n_tokens, K] F32 row-major
+//   ids:         [n_tokens, 6] int32 row-major
+//   weights:     [n_tokens, 6] F32 row-major
+//   mid_f32:     [n_tokens * 6, M] F32 row-major
+//
+// Constraints:
+//   - n_expert_used must be 6
+//   - K must be a multiple of 256
+//
+// Returns 0 on success, non-zero on validation or launch failure.
+
+int ds4_mmq_iq2_xxs_moe_gate_up_mid_vec(
+    const void    * W_gate,
+    const void    * W_up,
+    const float   * X_f32,
+    const int32_t * ids,
+    const float   * weights,
+    float         * mid_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    float           clamp,
+    cudaStream_t    stream);
+
+int ds4_mmq_q4_K_moe_gate_up_mid_vec(
+    const void    * W_gate,
+    const void    * W_up,
+    const float   * X_f32,
+    const int32_t * ids,
+    const float   * weights,
+    float         * mid_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    float           clamp,
+    cudaStream_t    stream);
+
+int ds4_mmq_mxfp4_moe_gate_up_mid_vec(
+    const void    * W_gate,
+    const void    * W_up,
+    const float   * X_f32,
+    const int32_t * ids,
+    const float   * weights,
+    float         * mid_f32,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    float           clamp,
+    cudaStream_t    stream);
+
+// Pair-fused MoE vector matmul entries (Step 6). Computes
+//
+//   out[col, row] = (W_a[ids, row, :] @ X[token, :])
+//                 * silu(W_b[ids, row, :] @ X[token, :])
+//
+// in a SINGLE mmvq launch via mmvq's built-in fusion (fusion.gate = W_b,
+// fusion.glu_op = GGML_GLU_OP_SWIGLU). The kernel applies silu to the
+// fusion.gate matmul and multiplies into the main matmul: pass the
+// SwiGLU "up" weights as W_a and the SwiGLU "gate" weights as W_b to
+// match ds4's expected silu(gate)*up semantics. The DeepSeek V4 clamp
+// and router-weight multiplication are NOT applied by the kernel - the
+// caller is expected to apply them as a small post-process (or to skip
+// clamp if clamp==0).
+//
+// Constraints:
+//   - n_tokens = 1 ONLY. mmvq supports fusion only at ncols_dst = 1.
+//   - K must be a multiple of 256.
+//
+// Returns 0 on success, non-zero on validation or launch failure.
+
+int ds4_mmq_iq2_xxs_moe_pair_vec(
+    const void    * W_a,
+    const void    * W_b,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_silu,
+    int             M,
+    int             K,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_q4_K_moe_pair_vec(
+    const void    * W_a,
+    const void    * W_b,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_silu,
+    int             M,
+    int             K,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+// Raw pair MoE vector entries. They quantize X to canonical Q8_1 once, then
+// run the same mmvq matvec kernel twice to produce the unfused gate and up
+// outputs. This preserves the caller's clamp-aware SwiGLU epilogue while
+// avoiding the duplicate Q8_1 quantize/scratch setup of two separate
+// ds4_mmq_*_moe_vec calls.
+//
+// Layout and constraints match ds4_mmq_*_moe_vec:
+//   out_[token * n_expert_used + slot, row]
+//   K must be a multiple of 256.
+
+int ds4_mmq_iq2_xxs_moe_pair_raw_vec(
+    const void    * W_a,
+    const void    * W_b,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_a,
+    float         * out_b,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+int ds4_mmq_q4_K_moe_pair_raw_vec(
+    const void    * W_a,
+    const void    * W_b,
+    const float   * X_f32,
+    const int32_t * ids,
+    float         * out_a,
+    float         * out_b,
+    int             M,
+    int             K,
+    int             n_tokens,
+    int             n_experts,
+    int             n_expert_used,
+    cudaStream_t    stream);
+
+// Dense vector matmul entry (Step 6). Same shape semantics as
+// ds4_mmq_q8_0_dense but routed through mmvq for batch counts that
+// favour the vec path (n_tokens <= 8 on Blackwell).
+//
+// Returns 0 on success, non-zero on validation or launch failure.
+
+int ds4_mmq_q8_0_dense_vec(
+    const void  * W_q8_0,
+    const float * X_f32,
+    float       * out_f32,
+    int           M,
+    int           N,
+    int           K,
+    cudaStream_t  stream);
+
+int ds4_mmq_q4_K_dense_pair_vec(
+    const void  * W0_q4_K,
+    const void  * W1_q4_K,
+    const float * X_f32,
+    float       * out0_f32,
+    float       * out1_f32,
+    int           M,
+    int           K,
+    cudaStream_t  stream);
+
+// Set the thread-local stream that the internal cuda pool uses for
+// cudaMallocAsync / cudaFreeAsync.  Defaults to cudaStreamPerThread.
+// Step 8 (CUDA Graphs) calls this with the capture stream so pool
+// allocations land on the captured stream and don't invalidate capture.
+// Pass NULL to reset to cudaStreamPerThread.
+void ds4_pool_set_stream(cudaStream_t stream);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif

cuda/mmq/ds4_mmq_d2r.cuh

@@ -0,0 +1,91 @@
+// SPDX-License-Identifier: MIT
+// Internal launcher for the gated D2R Q2_K MoE down-GEMM path.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+#include <stddef.h>
+#include <stdint.h>
+
+bool ds4_mmq_q2_K_moe_d2r_available(int cc);
+bool ds4_mmq_iq2_xxs_moe_d2r_available(int cc);
+
+size_t ds4_mmq_q2_K_moe_d2r_scratch_bytes(int64_t ncols_max, int n_experts);
+size_t ds4_mmq_iq2_xxs_moe_d2r_pair_scratch_bytes(int64_t ncols_max, int n_experts);
+size_t ds4_mmq_iq2_xxs_moe_d2r_fused_scratch_bytes(
+    int64_t ncols_max, int n_experts);
+
+int ds4_mmq_q2_K_moe_d2r_launch(
+    const void    * W_soa,
+    int64_t         soa_blocks,
+    const void    * q8,
+    const int32_t * ids_dst,
+    const int32_t * expert_bounds,
+    float         * out,
+    int             M,
+    int             K,
+    int64_t         ne_get_rows,
+    int             n_experts,
+    void          * worklist_scratch,
+    size_t          worklist_scratch_bytes,
+    cudaStream_t    stream);
+
+bool ds4_mmq_q8_0_dense_d2r_available(int cc);
+
+// Dense Q8_0 D2R on the kind-5 aligned artifact (--repack-q8-aligned).
+// q8 = block_q8_1_mmq D4 activation buffer ([k128][col], stride N cols,
+// over-allocated >= 128 blocks past N*K/128 for the guarded last col tile).
+int ds4_mmq_q8_0_dense_d2r_launch(
+    const void   * W_aligned,
+    const void   * q8,
+    float        * out,
+    int            M,
+    int            N,
+    int            K,
+    cudaStream_t   stream);
+
+int ds4_mmq_iq2_xxs_moe_d2r_pair_launch(
+    const void    * gate_soa,
+    const void    * up_soa,
+    int64_t         soa_blocks,
+    const void    * q8,
+    const int32_t * ids_dst,
+    const int32_t * expert_bounds,
+    float         * out_gate,
+    float         * out_up,
+    int             M,
+    int             K,
+    int64_t         ne_get_rows,
+    int             n_experts,
+    void          * worklist_scratch,
+    size_t          worklist_scratch_bytes,
+    cudaStream_t    stream);
+
+// Complete target-prefill gate/up path: both IQ2_XXS projections share one
+// activation tile, then sanitize + clamp + SwiGLU + routing weight are folded
+// directly into the expert-major Q8_1 D2S6 input consumed by Q2_K down.
+//
+// flat-pool p5b: with ids_src != NULL, input_q8 is TOKEN-COMPACT (n_tokens
+// rows quantized once, no gather) and each assignment column stages its
+// token's row through ids_src[col]; with ids_src == NULL, input_q8 is the
+// classic slot-gathered buffer (ne_get_rows rows) and n_tokens is ignored.
+int ds4_mmq_iq2_xxs_moe_d2r_fused_launch(
+    const void    * gate_soa,
+    const void    * up_soa,
+    int64_t         soa_blocks,
+    const void    * input_q8,
+    const int32_t * ids_src,
+    int             n_tokens,
+    const int32_t * ids_dst,
+    const int32_t * expert_bounds,
+    const float   * router_weights,
+    void          * down_q8,
+    int             M,
+    int             K,
+    int64_t         ne_get_rows,
+    int             n_experts,
+    float           clamp,
+    void          * worklist_scratch,
+    size_t          worklist_scratch_bytes,
+    cudaStream_t    stream);

cuda/mmq/ds4_repack.h

@@ -0,0 +1,126 @@
+/* ds4_repack: the aligned-SoA derived-weight layout library shared by the
+ * weight server (tools/ds4_weight_server.cu) and the engine's in-process
+ * self-load build (ds4_cuda.cu).  The GGUF tensor catalog walk, the repack
+ * candidate predicates, the S5 parallel repack driver, the FNV-1a artifact
+ * hash, and the three aligned repack builds (IQ2_XXS gate/up, Q2_K down,
+ * Q8_0 dense) live here so both producers emit bit-identical artifacts
+ * (gate: --repack-hash / DS4_WS_REPACK_HASH FNV lines).  Layout contracts
+ * are shared with the consuming kernels in cuda/mmq/ds4_mmq.h. */
+#ifndef DS4_REPACK_H
+#define DS4_REPACK_H
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+struct ds4_repack_file {
+    int fd = -1;
+    int direct_fd = -1;
+    const uint8_t *data = nullptr;
+    uint64_t size = 0;
+    uint64_t direct_align = 1;
+};
+
+struct ds4_repack_span {
+    uint64_t off = 0;
+    uint64_t end = 0;
+};
+
+struct ds4_repack_tensor {
+    std::string name;
+    uint32_t type = 0;
+    uint32_t ndim = 0;
+    uint64_t dims[8] = {};
+    uint64_t elements = 0;
+    uint64_t off = 0;
+    uint64_t bytes = 0;
+};
+
+/* Aligned-artifact kinds.  Numbering is shared with the weight-server
+ * manifest format and the engine's cuda_derived_kind (ds4_cuda.cu); the
+ * layout comments live on the builders in ds4_repack.cu. */
+enum {
+    DS4_REPACK_Q8_0_F16_COLMAJOR = 2,
+    DS4_REPACK_IQ2_XXS_ALIGNED_MOE = 4,
+    DS4_REPACK_Q8_0_ALIGNED_DENSE = 5,
+    DS4_REPACK_Q2_K_ALIGNED_MOE = 6,
+};
+
+struct ds4_repack_artifact {
+    const ds4_repack_tensor *t = nullptr; /* borrowed from the caller's catalog */
+    uint32_t kind = 0;                    /* DS4_REPACK_* */
+    uint64_t bytes = 0;                   /* artifact byte size */
+    uint64_t in_dim = 0;                  /* dims[0] */
+    uint64_t out_dim = 0;                 /* dims[1] */
+    uint32_t group_count = 0;             /* dims[2] for MoE kinds, 1 for dense,
+                                             0 for f16 colmajor (the engine's
+                                             lazy-cache lookup passes 0) */
+    void *dev = nullptr;                  /* device artifact buffer */
+};
+
+/* Optional device-allocation hooks (weight server VMM backend).  alloc must
+ * fill art->dev with at least art->bytes of device memory and return true;
+ * free releases an allocation made by alloc.  Leave both NULL for plain
+ * cudaMalloc/cudaFree. */
+typedef bool (*ds4_repack_alloc_fn)(void *ctx, ds4_repack_artifact *art);
+typedef void (*ds4_repack_free_fn)(void *ctx, ds4_repack_artifact *art);
+
+struct ds4_repack_build_args {
+    const char *log_prefix = "ds4";  /* stderr line prefix ("ds4_weight_server" | "ds4") */
+    const char *model_id = "base";   /* log label */
+    const char *path = nullptr;      /* GGUF path (each builder re-maps it) */
+    const std::vector<ds4_repack_tensor> *records = nullptr;
+    int device = 0;
+    uint64_t copy_chunk_bytes = 0;   /* staged read chunk (weight server --copy-chunk-mb) */
+    ds4_repack_alloc_fn alloc_fn = nullptr;
+    ds4_repack_free_fn free_fn = nullptr;
+    void *alloc_ctx = nullptr;
+};
+
+bool ds4_repack_map_file(const char *log_prefix, const char *path, ds4_repack_file &m);
+void ds4_repack_unmap_file(ds4_repack_file &m);
+
+/* Parse the GGUF header + tensor table into whole-tensor spans and records
+ * with absolute file offsets.  Either out pointer may be NULL. */
+bool ds4_repack_collect_catalog(const char *log_prefix,
+                                const ds4_repack_file &m,
+                                std::vector<ds4_repack_span> *spans,
+                                std::vector<ds4_repack_tensor> *records);
+
+bool ds4_repack_pread_full(int fd, void *buf, uint64_t bytes, uint64_t offset);
+/* Read [file_off, file_off+bytes) into stage, preferring the O_DIRECT fd;
+ * *payload points at the requested bytes inside stage. */
+bool ds4_repack_read_stage(const ds4_repack_file &m, void *stage, uint64_t stage_bytes,
+                           uint64_t file_off, uint64_t bytes, const char **payload);
+
+/* Candidate predicates for the three aligned repacks. */
+bool ds4_repack_iq2_candidate(const ds4_repack_tensor &t);
+bool ds4_repack_q2k_candidate(const ds4_repack_tensor &t);
+bool ds4_repack_q8_candidate(const ds4_repack_tensor &t);
+bool ds4_repack_q8_f16_candidate(const ds4_repack_tensor &t);
+
+/* S5 driver configuration.  Unset (0 / false) defers to the env knobs
+ * DS4_WS_REPACK_THREADS and DS4_WS_REPACK_HASH. */
+void ds4_repack_set_threads(int n);
+void ds4_repack_set_hash(bool on);
+
+uint64_t ds4_repack_fnv1a(const unsigned char *p, uint64_t n, uint64_t h);
+
+/* Build every candidate's aligned artifact.  On success appends the built
+ * artifacts to out, adds their byte total to *repacked_bytes_out (may be
+ * NULL), and returns true; on any failure releases every artifact this call
+ * allocated and returns false without touching out. */
+bool ds4_repack_build_q8_aligned(const ds4_repack_build_args &a,
+                                 std::vector<ds4_repack_artifact> &out,
+                                 uint64_t *repacked_bytes_out);
+bool ds4_repack_build_iq2_aligned(const ds4_repack_build_args &a,
+                                  std::vector<ds4_repack_artifact> &out,
+                                  uint64_t *repacked_bytes_out);
+bool ds4_repack_build_q2k_aligned(const ds4_repack_build_args &a,
+                                  std::vector<ds4_repack_artifact> &out,
+                                  uint64_t *repacked_bytes_out);
+bool ds4_repack_build_q8_f16(const ds4_repack_build_args &a,
+                             std::vector<ds4_repack_artifact> &out,
+                             uint64_t *repacked_bytes_out);
+
+#endif /* DS4_REPACK_H */

cuda/mmq/ggml-cuda.h

@@ -0,0 +1,4 @@
+// SPDX-License-Identifier: MIT
+// Thin redirect - see ggml.h.
+#pragma once
+#include "ds4_ggml_stubs.h"

cuda/mmq/ggml-impl.h

@@ -0,0 +1,4 @@
+// SPDX-License-Identifier: MIT
+// Thin redirect - see ggml.h.
+#pragma once
+#include "ds4_ggml_stubs.h"

cuda/mmq/ggml.h

@@ -0,0 +1,5 @@
+// SPDX-License-Identifier: MIT
+// Thin redirect so the vendored common.cuh / vecdotq.cuh / mmid.cu / etc.
+// can write #include "ggml.h" unchanged.
+#pragma once
+#include "ds4_ggml_stubs.h"

cuda/mmq/mmid.cu

@@ -0,0 +1,339 @@
+#include "common.cuh"
+#include "mmid.cuh"
+
+// To reduce shared memory use, store "it" and "iex_used" with 22/10 bits each.
+struct mm_ids_helper_store {
+    uint32_t data;
+
+    __device__ mm_ids_helper_store(const uint32_t it, const uint32_t iex_used) {
+        data = (it & 0x003FFFFF) | (iex_used << 22);
+    }
+
+    __device__ uint32_t it() const {
+        return data & 0x003FFFFF;
+    }
+
+    __device__ uint32_t iex_used() const {
+        return data >> 22;
+    }
+};
+static_assert(sizeof(mm_ids_helper_store) == 4, "unexpected size for mm_ids_helper_store");
+
+// Helper function for mul_mat_id, converts ids to a more convenient format.
+// ids_src1 describes how to permute the flattened column indices of src1 in order to get a compact src1 tensor sorted by expert.
+// ids_dst describes the same mapping but for the dst tensor.
+// The upper and lower bounds for the ith expert in the compact src1 tensor are stored in expert_bounds[i:i+1].
+template <int n_expert_used_template>
+__launch_bounds__(ggml_cuda_get_physical_warp_size(), 1)
+static __global__ void mm_ids_helper(
+        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
+        const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
+    const int expert = blockIdx.x;
+
+    extern __shared__ char data_mm_ids_helper[];
+    mm_ids_helper_store * store = (mm_ids_helper_store *) data_mm_ids_helper;
+
+    int nex_prev   = 0; // Number of columns for experts with a lower index.
+    int it_compact = 0; // Running index for the compact slice of this expert.
+
+    if constexpr (n_expert_used_template == 0) {
+        // Generic implementation:
+        for (int it = 0; it < n_tokens; ++it) {
+            int iex_used = -1; // The index at which the expert is used, if any.
+            for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
+                const int expert_used = ids[it*si1 + iex];
+                nex_prev += expert_used < expert;
+                if (expert_used == expert) {
+                    iex_used = iex;
+                }
+            }
+
+            if (iex_used != -1) {
+                store[it_compact] = mm_ids_helper_store(it, iex_used);
+            }
+
+            if (warp_reduce_any<warp_size>(iex_used != -1)) {
+                it_compact++;
+            }
+        }
+    } else {
+        // Implementation optimized for specific numbers of experts used:
+        static_assert(n_expert_used == 6 || warp_size % n_expert_used == 0, "bad n_expert_used");
+        const int neu_padded = n_expert_used == 6 ? 8 : n_expert_used; // Padded to next higher power of 2.
+        for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) {
+            const int it = it0 + threadIdx.x / neu_padded;
+
+            const int iex = threadIdx.x % neu_padded; // The index at which the expert is used, if any.
+            const int expert_used = (neu_padded == n_expert_used || iex < n_expert_used) && it < n_tokens ?
+                ids[it*si1 + iex] : INT_MAX;
+            const int iex_used = expert_used == expert ? iex : -1;
+            nex_prev += expert_used < expert;
+
+            // Whether the threads at this token position have used the expert:
+            const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);
+
+            // Do a scan over threads at lower token positions in warp to get the correct index for writing data:
+            int it_compact_add_lower = 0;
+#pragma unroll
+            for (int offset = neu_padded; offset < warp_size; offset += neu_padded) {
+                const int tmp = __shfl_up_sync(0xFFFFFFFF, it_compact_add_self, offset, warp_size);
+                if (threadIdx.x >= static_cast<unsigned int>(offset)) {
+                    it_compact_add_lower += tmp;
+                }
+            }
+
+            if (iex_used != -1) {
+                store[it_compact + it_compact_add_lower] = mm_ids_helper_store(it, iex_used);
+            }
+
+            // The thread with the highest index in the warp always has the sum over the whole warp, use it to increment all threads:
+            it_compact += __shfl_sync(0xFFFFFFFF, it_compact_add_lower + it_compact_add_self, warp_size - 1, warp_size);
+        }
+    }
+    nex_prev = warp_reduce_sum<warp_size>(nex_prev);
+
+    // Post-Volta independent thread scheduling: the store[] writes above are done by
+    // some warp lanes and read below by other lanes. Without an explicit warp barrier
+    // those shared-memory writes are not guaranteed visible to the cross-lane reads,
+    // so a lane could read a stale/uninitialized store[] slot -> wrong compacted
+    // expert ids -> nondeterministic MoE routing. On GB10 (sm_121) this realized as
+    // the cont-multiseq non-determinism + BOS-spam (compute-sanitizer racecheck:
+    // "RAW hazard at __shared__" between the store writes and these reads).
+    __syncwarp();
+
+    for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) {
+        const mm_ids_helper_store store_it = store[itc];
+        const int it       = store_it.it();
+        const int iex_used = store_it.iex_used();
+        ids_src1[nex_prev + itc] = it*sis1          + iex_used % nchannels_y;
+        ids_dst [nex_prev + itc] = it*n_expert_used + iex_used;
+    }
+
+    if (threadIdx.x != 0) {
+        return;
+    }
+
+    expert_bounds[expert] = nex_prev;
+
+    if (expert < static_cast<int>(gridDim.x) - 1) {
+        return;
+    }
+
+    expert_bounds[gridDim.x] = nex_prev + it_compact;
+}
+
+// ds4 local (P5): large-n variant of mm_ids_helper with NO shared-memory
+// staging.  The smem kernel above stages the compact (it, iex_used) list in
+// n_tokens*4 B of dynamic shared memory, capping n_tokens at smpbo/4 (~25k on
+// GB10); the routed-MoE down matmul passes n_tokens = assignment rows
+// (6x the forward width), so 8192-row prefill chunks hit 48384 "tokens" and
+// the whole MoE block used to fall back to the pre-mmq expert-tile kernels.
+// This variant runs the SAME per-expert scan twice: pass 0 only counts
+// (nex_prev, bucket size), pass 1 re-scans and writes ids_src1/ids_dst
+// directly at their final offsets.  Bit-identical to the smem kernel by
+// construction: identical (it, iex_used) tuples in identical token order and
+// identical output expressions — the (lossless) 22/10-bit store round-trip is
+// simply removed.  The rescan is cheap: the ids array at the target shape is
+// ~190 KB and L2-resident across all expert blocks.  Like the smem kernel,
+// behavior is undefined if one token lists the same expert in multiple slots
+// (the router's top-k is without replacement, so this cannot occur).
+template <int n_expert_used_template>
+__launch_bounds__(ggml_cuda_get_physical_warp_size(), 1)
+static __global__ void mm_ids_helper_global(
+        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
+        const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
+    const int expert = blockIdx.x;
+
+    int nex_prev         = 0; // Number of columns for experts with a lower index.
+    int it_compact_count = 0; // Bucket size for this expert (pass-0 result).
+
+#pragma unroll 1
+    for (int pass = 0; pass < 2; ++pass) {
+        int it_compact = 0; // Running index for the compact slice of this expert.
+
+        if constexpr (n_expert_used_template == 0) {
+            // Generic implementation:
+            for (int it = 0; it < n_tokens; ++it) {
+                int iex_used = -1; // The index at which the expert is used, if any.
+                for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
+                    const int expert_used = ids[it*si1 + iex];
+                    if (pass == 0) {
+                        nex_prev += expert_used < expert;
+                    }
+                    if (expert_used == expert) {
+                        iex_used = iex;
+                    }
+                }
+
+                if (pass == 1 && iex_used != -1) {
+                    ids_src1[nex_prev + it_compact] = it*sis1          + iex_used % nchannels_y;
+                    ids_dst [nex_prev + it_compact] = it*n_expert_used + iex_used;
+                }
+
+                if (warp_reduce_any<warp_size>(iex_used != -1)) {
+                    it_compact++;
+                }
+            }
+        } else {
+            // Implementation optimized for specific numbers of experts used:
+            static_assert(n_expert_used == 6 || warp_size % n_expert_used == 0, "bad n_expert_used");
+            const int neu_padded = n_expert_used == 6 ? 8 : n_expert_used; // Padded to next higher power of 2.
+            for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) {
+                const int it = it0 + threadIdx.x / neu_padded;
+
+                const int iex = threadIdx.x % neu_padded; // The index at which the expert is used, if any.
+                const int expert_used = (neu_padded == n_expert_used || iex < n_expert_used) && it < n_tokens ?
+                    ids[it*si1 + iex] : INT_MAX;
+                const int iex_used = expert_used == expert ? iex : -1;
+                if (pass == 0) {
+                    nex_prev += expert_used < expert;
+                }
+
+                // Whether the threads at this token position have used the expert:
+                const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);
+
+                // Do a scan over threads at lower token positions in warp to get the correct index for writing data:
+                int it_compact_add_lower = 0;
+#pragma unroll
+                for (int offset = neu_padded; offset < warp_size; offset += neu_padded) {
+                    const int tmp = __shfl_up_sync(0xFFFFFFFF, it_compact_add_self, offset, warp_size);
+                    if (threadIdx.x >= static_cast<unsigned int>(offset)) {
+                        it_compact_add_lower += tmp;
+                    }
+                }
+
+                if (pass == 1 && iex_used != -1) {
+                    const int itc = it_compact + it_compact_add_lower;
+                    ids_src1[nex_prev + itc] = it*sis1          + iex_used % nchannels_y;
+                    ids_dst [nex_prev + itc] = it*n_expert_used + iex_used;
+                }
+
+                // The thread with the highest index in the warp always has the sum over the whole warp, use it to increment all threads:
+                it_compact += __shfl_sync(0xFFFFFFFF, it_compact_add_lower + it_compact_add_self, warp_size - 1, warp_size);
+            }
+        }
+
+        if (pass == 0) {
+            it_compact_count = it_compact;
+            nex_prev = warp_reduce_sum<warp_size>(nex_prev);
+        }
+    }
+
+    if (threadIdx.x != 0) {
+        return;
+    }
+
+    expert_bounds[expert] = nex_prev;
+
+    if (expert < static_cast<int>(gridDim.x) - 1) {
+        return;
+    }
+
+    expert_bounds[gridDim.x] = nex_prev + it_compact_count;
+}
+
+// ds4 local: kill switch for the large-n global path (DS4_MMID_LARGE=0).
+// With the switch off the ds4_mmq.cu callers refuse past-cap shapes exactly
+// as before (whole-MoE fallback to the expert-tile kernels).
+bool ds4_mmid_large_enabled(void) {
+    static int cached = -1;
+    if (cached < 0) {
+        const char * env = getenv("DS4_MMID_LARGE");
+        cached = !(env && env[0] == '0');
+    }
+    return cached != 0;
+}
+
+template <int n_expert_used_template>
+static void launch_mm_ids_helper(
+        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
+        const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, cudaStream_t stream) {
+    GGML_ASSERT(n_tokens          < (1 << 22) && "too few bits in mm_ids_helper_store");
+    GGML_ASSERT(n_expert_used_var < (1 << 10) && "too few bits in mm_ids_helper_store");
+
+    const int id = ggml_cuda_get_device();
+    const int warp_size = ggml_cuda_info().devices[id].warp_size;
+    const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
+
+    const dim3 num_blocks(n_experts, 1, 1);
+    const dim3 block_size(warp_size, 1, 1);
+    const size_t nbytes_shared = n_tokens*sizeof(mm_ids_helper_store);
+
+    // ds4 local (P5): past the smem cap, take the two-pass global variant
+    // (bit-identical outputs, see mm_ids_helper_global).  One-shot stderr
+    // line = path proof for the gate harnesses.
+    if (nbytes_shared > smpbo) {
+        static bool logged = false;
+        if (!logged) {
+            logged = true;
+            fprintf(stderr, "ds4: mm_ids_helper large-n global path engaged (P5, n_tokens=%d > cap %zu)\n",
+                    n_tokens, smpbo / sizeof(mm_ids_helper_store));
+        }
+        mm_ids_helper_global<n_expert_used_template><<<num_blocks, block_size, 0, stream>>>
+            (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1);
+        return;
+    }
+
+    CUDA_SET_SHARED_MEMORY_LIMIT(mm_ids_helper<n_expert_used_template>, smpbo);
+    mm_ids_helper<n_expert_used_template><<<num_blocks, block_size, nbytes_shared, stream>>>
+        (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1);
+}
+
+// ds4 local: kill switch for the case-1 fast path below (DS4_MMID_CASE1=0).
+static bool ds4_mmid_case1_enabled() {
+    static int cached = -1;
+    if (cached < 0) {
+        const char * env = getenv("DS4_MMID_CASE1");
+        cached = !(env && env[0] == '0');
+    }
+    return cached != 0;
+}
+
+void ggml_cuda_launch_mm_ids_helper(
+        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
+        const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, cudaStream_t stream) {
+    switch (n_expert_used) {
+        case  1:
+            // ds4 local: the routed-MoE down matmul reinterprets (token, slot)
+            // assignment rows as single-expert "tokens" (ds4_mmq.cu
+            // ds4_mmq_moe_impl, n_expert_used=1).  Without this case it fell to
+            // the generic <0> template: one warp per expert scanning all
+            // assignment rows with a single active lane -> 22.5 ms/launch at
+            // W4096 prefill (2.90 s of a 12k admission).  The optimized template
+            // at neu_padded=1 covers 32 rows/iteration and emits bit-identical
+            // id maps (proto_mm_ids.cu: parity on uniform/skewed/2%-invalid/
+            // decode shapes, 20.4x at the W4096 shape).  DS4_MMID_CASE1=0
+            // reverts to the generic path.
+            if (ds4_mmid_case1_enabled()) {
+                launch_mm_ids_helper< 1>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            } else {
+                launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            }
+            break;
+        case  2:
+            launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            break;
+        case  4:
+            launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            break;
+        case  6:
+            launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            break;
+        case  8:
+            launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            break;
+        case 16:
+            launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            break;
+        case 32:
+            launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            break;
+        default:
+            launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            break;
+    }
+}

cuda/mmq/mmid.cuh

@@ -0,0 +1,9 @@
+#pragma once
+
+void ggml_cuda_launch_mm_ids_helper(
+        const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds,
+        int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, cudaStream_t stream);
+
+// ds4 local (P5): whether the large-n global-memory mm_ids path is enabled
+// (default on; DS4_MMID_LARGE=0 reverts callers to the past-cap refusal).
+bool ds4_mmid_large_enabled(void);

cuda/mmq/mmvq.cuh

@@ -0,0 +1,34 @@
+#pragma once
+#include "common.cuh"
+
+#define MMVQ_MAX_BATCH_SIZE 8 // Max. batch size for which to use MMVQ kernels.
+
+// Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID,
+// based on the quantization type and GPU architecture (compute capability).
+int get_mmvq_mmid_max_batch(ggml_type type, int cc);
+
+// ds4: direct switch over the per-type templated mmvq dispatch.  Was
+// `static` upstream; we promote it so ds4_mmq.cu can call it without
+// constructing ggml_tensor structs.  Args match the original signature
+// exactly.
+void mul_mat_vec_q_switch_type(
+        const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
+        const int ncols_x, const int nrows_x, const int ncols_dst,
+        const int stride_row_x, const int stride_col_y, const int stride_col_dst,
+        const int nchannels_x, const int nchannels_y, const int nchannels_dst,
+        const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
+        const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
+        const int ids_stride, cudaStream_t stream);
+
+#ifdef DS4_MMVQ_INCLUDE_GGML_ENTRIES
+// Original upstream ggml-graph entries.  Disabled by default in the ds4
+// vendor; re-enable for parity comparisons against llama.cpp.
+void ggml_cuda_mul_mat_vec_q(ggml_backend_cuda_context & ctx,
+    const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion = nullptr);
+
+void ggml_cuda_op_mul_mat_vec_q(
+    ggml_backend_cuda_context & ctx,
+    const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i,
+    const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols,
+    const int64_t src1_padded_row_size, cudaStream_t stream);
+#endif

cuda/mmq/quantize.cu

@@ -0,0 +1,443 @@
+#include "quantize.cuh"
+#include <cstdint>
+
+__launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1)
+static __global__ void quantize_q8_1(
+        const float * __restrict__ x, void * __restrict__ vy,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const uint32_t ne1, const uint3 ne2) {
+    const int64_t i0 = (int64_t)blockDim.x*blockIdx.x + threadIdx.x;
+
+    if (i0 >= ne0) {
+        return;
+    }
+
+    const int64_t i3 = fastdiv(blockIdx.z, ne2);
+    const int64_t i2 = blockIdx.z - i3*ne2.z;
+    const int64_t i1 = blockIdx.y;
+
+    const int64_t & i00 = i0;
+    const int64_t & i01 = i1;
+    const int64_t & i02 = i2;
+    const int64_t & i03 = i3;
+
+    const int64_t i_cont = ((i3*ne2.z + i2) * ne1 + i1) * ne0 + i0;
+
+    block_q8_1 * y = (block_q8_1 *) vy;
+
+    const int64_t ib  = i_cont / QK8_1; // block index
+    const int64_t iqs = i_cont % QK8_1; // quant index
+
+    const float xi = i0 < ne00 ? x[i03*s03 + i02*s02 + i01*s01 + i00] : 0.0f;
+    float amax = fabsf(xi);
+    float sum = xi;
+
+    amax = warp_reduce_max<QK8_1>(amax);
+    sum  = warp_reduce_sum<QK8_1>(sum);
+
+    const float  d = amax / 127.0f;
+    const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);
+
+    y[ib].qs[iqs] = q;
+
+    if (iqs > 0) {
+        return;
+    }
+
+    y[ib].ds = make_half2(d, sum);
+}
+
+__device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) {
+    if (!(amax > 0.0f)) {
+        return 0;
+    }
+
+    // FP4 E2M1: max exponent (unbiased) is 2.
+    constexpr int FP4_E2M1_EMAX = 2;
+
+    const float e = log2f(amax);
+
+    // "even" -> round-to-nearest integer, ties-to-even
+    const int e_int = __float2int_rn(e);
+
+    const int shared_exp = e_int - FP4_E2M1_EMAX;
+
+    int biased = shared_exp + 127;
+
+    biased = max(biased, 0);
+    biased = min(biased, 254);
+
+    return static_cast<uint8_t>(biased);
+}
+
+
+static __global__ void quantize_mmq_nvfp4(
+        const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const int64_t ne1, const int64_t ne2) {
+#if defined(BLACKWELL_MMA_AVAILABLE)
+
+    const int64_t i0_base = ((int64_t) blockDim.x * blockIdx.y + threadIdx.x) * QK_NVFP4_SUB;
+    if (i0_base >= ne0) {
+        return;
+    }
+
+    const int64_t i1 = blockIdx.x;
+    const int64_t i2 = blockIdx.z % ne2;
+    const int64_t i3 = blockIdx.z / ne2;
+    const int64_t i01 = ids ? ids[i1] : i1;
+    const int64_t k_block = i0_base / QK_K;
+    const int64_t blocks_per_col = (ne0 + QK_K - 1) / QK_K;
+    if (k_block >= blocks_per_col) {
+        return;
+    }
+
+    const int64_t ib = blockIdx.z * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x;
+    block_fp4_mmq * y = (block_fp4_mmq *) vy;
+    block_fp4_mmq * yb = y + ib;
+
+    const int sub = (i0_base % QK_K) / QK_NVFP4_SUB;
+
+    float vals_raw[QK_NVFP4_SUB];
+    float amax_raw = 0.0f;
+    const int64_t base_idx = i3 * s03 + i2 * s02 + i01 * s01;
+#pragma unroll
+    for (int k = 0; k < QK_NVFP4_SUB; k++) {
+        const int64_t i00 = i0_base + k;
+        if (i00 < ne00) {
+            const float v = x[base_idx + i00];
+            vals_raw[k] = v;
+            amax_raw = fmaxf(amax_raw, fabsf(v));
+        } else {
+            vals_raw[k] = 0.0f;
+        }
+    }
+
+    static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2};
+    const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_raw / 6.0f);
+
+    float best_err = FLT_MAX;
+    uint8_t fp8_code = 0;
+    float subblock_scale = 0.0f;
+
+#pragma unroll // Check +/- 2 to find best code to reduce NVFP4 activation loss. Negligible overhead on Blackwell.
+    for (int i = 0; i < 5; i++) {
+        const int test_code = first_fp8_code + test_offsets[i];
+        if (test_code < 0 || test_code > 0x7e) {
+            continue;
+        }
+        const uint8_t code = (uint8_t) test_code;
+        const float test_scale = ggml_cuda_ue4m3_to_fp32(code);
+        const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f;
+        float cur_err = 0.0f;
+#pragma unroll
+        for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+            const float v = vals_raw[k];
+            const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale);
+            const float err_diff = fabsf(v) - fabsf(kvalues_mxfp4[q & 0x7]) * test_scale;
+            cur_err = fmaf(err_diff, err_diff, cur_err);
+        }
+
+        if (cur_err < best_err) {
+            best_err = cur_err;
+            fp8_code = test_code;
+            subblock_scale = test_scale;
+        }
+    }
+
+    const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+    uint32_t q0 = 0;
+    uint32_t q1 = 0;
+#pragma unroll // this is faster than the previous __nv_fp4x4_e2m1
+    for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) {
+        q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k +  0], inv_scale) << (8 * k);
+        q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k +  8], inv_scale) << (8 * k + 4);
+        q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k +  4], inv_scale) << (8 * k);
+        q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 12], inv_scale) << (8 * k + 4);
+    }
+
+    uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
+    yqs[2 * sub + 0] = q0;
+    yqs[2 * sub + 1] = q1;
+    reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
+#else
+    NO_DEVICE_CODE; // This is for Blackwell NVFP4 activations only.
+#endif // defined(BLACKWELL_MMA_AVAILABLE)
+
+}
+
+// quantize values in the format mxfp4 is stored which is interleaved nibbles
+// i.e. a block a0-a31 is represented as a0a16,a1a17 ...a15a31
+static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
+                                          const int32_t * __restrict__ ids,
+                                          void * __restrict__ vy,
+                                          const int64_t ne00,
+                                          const int64_t s01,
+                                          const int64_t s02,
+                                          const int64_t s03,
+                                          const int64_t ne0,
+                                          const int     ne1,
+                                          const int     ne2) {
+    constexpr int vals_per_scale = 32;
+    constexpr int vals_per_warp  = 2 * vals_per_scale;  // Each warp processes 2 blocks of 32 = 64 values
+
+    const int warp_id = threadIdx.y;
+    const int lane_id_32 = threadIdx.x;
+
+    const int nwarps = blockDim.y;
+
+    const int64_t warp_start_offset = (blockIdx.y * nwarps + warp_id) * vals_per_warp;
+
+    if (warp_start_offset >= ne0) {
+        return;
+    }
+
+    const int64_t i1 = blockIdx.x;
+    const int64_t i2 = blockIdx.z % ne2;
+    const int64_t i3 = blockIdx.z / ne2;
+
+    const int64_t i01 = ids ? ids[i1] : i1;
+    const int64_t i02 = i2;
+    const int64_t i03 = i3;
+
+    block_fp4_mmq * y = (block_fp4_mmq *) vy;
+
+    const int64_t block_fp4_mmq_size = 8 * QK_MXFP4;  // 256 values
+    const int64_t ib0                = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size));
+    const int64_t ib = ib0 + (warp_start_offset / block_fp4_mmq_size) * ne1 + blockIdx.x;
+    const int64_t quad_idx_in_block  = (warp_start_offset % block_fp4_mmq_size) / vals_per_warp;
+
+    const int group_id = lane_id_32 / 4;
+    const int lane_in_group = lane_id_32 % 4;
+    const int base = group_id * 2;
+    char2 * yqs2 = (char2 *) y[ib].qs;
+
+    const int64_t base_pos = i03 * s03 + i02 * s02 + i01 * s01;
+
+    uint8_t scales[2];
+
+#pragma unroll
+    for (int b = 0; b < 2; ++b) {
+        const int64_t i0 = warp_start_offset + b * vals_per_scale + lane_id_32;
+        const float xi = (i0 < ne00) ? x[base_pos + i0] : 0.0f;
+
+        float amax = fabsf(xi);
+#pragma unroll
+        for (int mask = 16; mask > 0; mask >>= 1) {
+            amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, WARP_SIZE));
+        }
+
+        const uint8_t e = compute_e8m0_scale(amax);
+        scales[b] = e;
+        const float inv_s = (amax == 0.0f) ? 0.0f : __frcp_rn(ggml_cuda_e8m0_to_fp32(e));
+
+#if CUDART_VERSION >= 12080
+        const float scaled_val = xi * inv_s;
+
+        const float val0 = __shfl_sync(0xFFFFFFFF, scaled_val, base, WARP_SIZE);
+        const float val1 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 16, WARP_SIZE);
+        const float val2 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 1, WARP_SIZE);
+        const float val3 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 17, WARP_SIZE);
+
+        if (lane_in_group == 0) {
+            __nv_fp4x4_e2m1 fp4_packed(make_float4(val0, val1, val2, val3));
+
+            yqs2[quad_idx_in_block * 16 + b * 8 + group_id] = *(char2 *) &fp4_packed;
+        }
+#else
+        // Fallback: manual FP4 conversion using LUT
+        const uint8_t q_val = ggml_cuda_float_to_fp4_e2m1(xi, inv_s);
+
+        const uint8_t q_lo_0 = __shfl_sync(0xFFFFFFFF, q_val, base,      WARP_SIZE);
+        const uint8_t q_lo_1 = __shfl_sync(0xFFFFFFFF, q_val, base + 1,  WARP_SIZE);
+        const uint8_t q_hi_0 = __shfl_sync(0xFFFFFFFF, q_val, base + 16, WARP_SIZE);
+        const uint8_t q_hi_1 = __shfl_sync(0xFFFFFFFF, q_val, base + 17, WARP_SIZE);
+
+        if (lane_in_group == 0) {
+            char2 q;
+            q.x = (q_hi_0 << 4) | q_lo_0;
+            q.y = (q_hi_1 << 4) | q_lo_1;
+            yqs2[quad_idx_in_block * 16 + b * 8 + group_id] = q;
+        }
+#endif // CUDART_VERSION >= 12080
+    }
+
+    if (lane_id_32 == 0) {
+        // Store 2 scales packed into 1 uint32
+        y[ib].d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
+    }
+}
+
+template <mmq_q8_1_ds_layout ds_layout>
+static __global__ void quantize_mmq_q8_1(
+        const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const int ne1, const int ne2) {
+
+    constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32;
+    constexpr int vals_per_sum   = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32;
+
+    const int64_t i0 = ((int64_t)blockDim.x*blockIdx.y + threadIdx.x)*4;
+
+    if (i0 >= ne0) {
+        return;
+    }
+
+    const int64_t i1 = blockIdx.x;
+    const int64_t i2 = blockIdx.z % ne2;
+    const int64_t i3 = blockIdx.z / ne2;
+
+    const int64_t i00 = i0;
+    const int64_t i01 = ids ? ids[i1] : i1;
+    const int64_t i02 = i2;
+    const int64_t i03 = i3;
+
+    const float4 * x4 = (const float4 *) x;
+
+    block_q8_1_mmq * y = (block_q8_1_mmq *) vy;
+
+    const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel
+    const int64_t ib  = ib0 + (i0 / (4*QK8_1))*ne1 + blockIdx.x;                    // block index in channel
+    const int64_t iqs = i0 % (4*QK8_1);                                             // quant index in block
+
+    // Load 4 floats per thread and calculate max. abs. value between them:
+    const float4 xi = i0 < ne00 ? x4[(i03*s03 + i02*s02 + i01*s01 + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+    float amax = fabsf(xi.x);
+    amax = fmaxf(amax, fabsf(xi.y));
+    amax = fmaxf(amax, fabsf(xi.z));
+    amax = fmaxf(amax, fabsf(xi.w));
+
+    // Exchange max. abs. value between vals_per_scale/4 threads.
+#pragma unroll
+    for (int offset = vals_per_scale/8; offset > 0; offset >>= 1) {
+        amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE));
+    }
+
+    float sum;
+    if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
+        sum = xi.x + xi.y + xi.z + xi.w;
+
+        // Calculate sums across vals_per_sum/4 threads.
+#pragma unroll
+        for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
+            sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE);
+        }
+    }
+
+    const float d_inv = 127.0f / amax;
+    char4 q;
+    q.x = roundf(xi.x*d_inv);
+    q.y = roundf(xi.y*d_inv);
+    q.z = roundf(xi.z*d_inv);
+    q.w = roundf(xi.w*d_inv);
+
+    // Write back 4 int8 values as a single 32 bit value for better memory bandwidth:
+    char4 * yqs4 = (char4 *) y[ib].qs;
+    yqs4[iqs/4] = q;
+
+    if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) {
+        if (iqs % 16 != 0 || iqs >= 96) {
+            return;
+        }
+
+        y[ib].d2s6[2 + iqs/16] = sum;
+
+        if (iqs % 64 != 0) {
+            return;
+        }
+
+        const float d = 1.0f / d_inv;
+
+        y[ib].d2s6[iqs/64] = d;
+
+        return;
+    }
+
+    if (iqs % 32 != 0) {
+        return;
+    }
+
+    const float d = 1.0f / d_inv;
+
+    if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) {
+        y[ib].ds4[iqs/32] = make_half2(d, sum);
+    } else {
+        y[ib].d4[iqs/32]  = d;
+    }
+}
+
+void quantize_row_q8_1_cuda(
+        const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
+    GGML_ASSERT(!ids);
+    GGML_ASSERT(ne0 % QK8_1 == 0);
+
+    const uint3 ne2_fastdiv = init_fastdiv_values(ne2);
+
+    const int64_t block_num_x = (ne0 + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE;
+    const dim3 num_blocks(block_num_x, ne1, ne2*ne3);
+    const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1);
+    quantize_q8_1<<<num_blocks, block_size, 0, stream>>>(x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv);
+    GGML_UNUSED(type_src0);
+}
+
+void quantize_mmq_q8_1_cuda(
+        const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
+    GGML_ASSERT(ne00 % 4 == 0);
+    GGML_ASSERT(ne0 % (4*QK8_1) == 0);
+
+    // ne1 tends to assume the highest values, therefore use it as the "x" dimension of the CUDA grid:
+    const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
+    const dim3 num_blocks(ne1, block_num_y, ne2*ne3);
+    const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+    switch (mmq_get_q8_1_ds_layout(type_src0)) {
+        case MMQ_Q8_1_DS_LAYOUT_D4:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4>
+                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+            break;
+        case MMQ_Q8_1_DS_LAYOUT_DS4:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4>
+                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+            break;
+        case MMQ_Q8_1_DS_LAYOUT_D2S6:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6>
+                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+            break;
+    }
+}
+
+void quantize_mmq_fp4_cuda(
+        const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
+    GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4 || type_src0 == GGML_TYPE_NVFP4);
+    GGML_ASSERT(ne0 > 0);
+
+    if (type_src0 == GGML_TYPE_NVFP4) {
+        GGML_ASSERT(ne00 % QK_NVFP4 == 0);
+        constexpr int nvfp4_block_size = 128;
+        const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size);
+        const dim3 block_size(nvfp4_block_size, 1, 1);
+        const dim3 num_blocks(ne1, block_num_y, ne2 * ne3);
+        quantize_mmq_nvfp4<<<num_blocks, block_size, 0, stream>>>(
+            x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+    } else {
+        GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0);
+
+        constexpr int nwarps = 8;
+        constexpr int vals_per_warp  = 2 * QK_MXFP4;
+        constexpr int vals_per_block = nwarps * vals_per_warp;
+
+        const int64_t block_num_y = (ne0 + vals_per_block - 1) / vals_per_block;
+        const dim3    num_blocks(ne1, block_num_y, ne2 * ne3);
+        const dim3    block_size(WARP_SIZE, nwarps, 1);
+
+        quantize_mmq_mxfp4<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+    }
+}

cuda/mmq/quantize.cuh

@@ -0,0 +1,41 @@
+#pragma once
+
+#include "common.cuh"
+#include "mmq.cuh"
+
+#include <cstdint>
+
+#define CUDA_QUANTIZE_BLOCK_SIZE     256
+#define CUDA_QUANTIZE_BLOCK_SIZE_MMQ 128
+
+static_assert(MATRIX_ROW_PADDING %    CUDA_QUANTIZE_BLOCK_SIZE      == 0, "Risk of out-of-bounds access.");
+static_assert(MATRIX_ROW_PADDING % (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ) == 0, "Risk of out-of-bounds access.");
+
+typedef void (*quantize_cuda_t)(
+        const float * x, const int32_t * ids, void * vy,
+        ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
+        int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
+
+void quantize_row_q8_1_cuda(
+        const float * x, const int32_t * ids, void * vy,
+        ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
+        int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
+
+void quantize_mmq_q8_1_cuda(
+        const float * x, const int32_t * ids, void * vy,
+        ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
+        int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
+
+void quantize_mmq_fp4_cuda(const float *   x,
+                             const int32_t * ids,
+                             void *          vy,
+                             ggml_type       type_src0,
+                             int64_t         ne00,
+                             int64_t         s01,
+                             int64_t         s02,
+                             int64_t         s03,
+                             int64_t         ne0,
+                             int64_t         ne1,
+                             int64_t         ne2,
+                             int64_t         ne3,
+                             cudaStream_t    stream);

cuda/mmq/test/iq2_host_tables.h

@@ -0,0 +1,105 @@
+// SPDX-License-Identifier: MIT
+// iq2_host_tables.h - host-accessible copies of the IQ2 lookup tables.
+// Regenerate from cuda/mmq/ggml-common.h whenever upstream changes them:
+//
+//   awk '/GGML_TABLE_BEGIN.*iq2xxs_grid, 256/,/GGML_TABLE_END/' \
+//       cuda/mmq/ggml-common.h | grep -v GGML_TABLE > iq2xxs_grid.lines
+//
+// In the vendored ggml-common.h the IQ2 tables are gated on
+// GGML_COMMON_IMPL_CUDA and emitted as static const __device__, so they
+// are not host-accessible.  Reproducing them here as plain const arrays
+// gives the CPU parity reference something to look up against.
+
+#pragma once
+#include <cstdint>
+
+namespace iq2_host {
+
+// Bit positions for sign extraction (1, 2, 4, ..., 128).
+static const uint8_t kmask_iq2xs[8] = {1, 2, 4, 8, 16, 32, 64, 128};
+
+// 128-entry sign lookup. Index is a 7-bit pattern of which positions
+// are negative; value is the 8-bit pattern with the 8th bit set to
+// maintain even parity (matches ggml-common.h:503).
+static const uint8_t ksigns_iq2xs[128] = {
+      0, 129, 130,   3, 132,   5,   6, 135, 136,   9,  10, 139,  12, 141, 142,  15,
+    144,  17,  18, 147,  20, 149, 150,  23,  24, 153, 154,  27, 156,  29,  30, 159,
+    160,  33,  34, 163,  36, 165, 166,  39,  40, 169, 170,  43, 172,  45,  46, 175,
+     48, 177, 178,  51, 180,  53,  54, 183, 184,  57,  58, 187,  60, 189, 190,  63,
+    192,  65,  66, 195,  68, 197, 198,  71,  72, 201, 202,  75, 204,  77,  78, 207,
+     80, 209, 210,  83, 212,  85,  86, 215, 216,  89,  90, 219,  92, 221, 222,  95,
+     96, 225, 226,  99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111,
+    240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255,
+};
+
+// 256-entry IQ2_XXS grid: each uint64 holds 8 uint8 quantized values.
+// Indexed by the 8-bit aux8[l] from the encoded super-block.
+// (matches ggml-common.h:550).
+static const uint64_t iq2xxs_grid[256] = {
+    0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08,
+    0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808,
+    0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819,
+    0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819,
+    0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b,
+    0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808,
+    0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08,
+    0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b,
+    0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819,
+    0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08,
+    0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808,
+    0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08,
+    0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808,
+    0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808,
+    0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919,
+    0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819,
+    0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08,
+    0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908,
+    0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819,
+    0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808,
+    0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808,
+    0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908,
+    0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808,
+    0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08,
+    0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819,
+    0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819,
+    0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819,
+    0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908,
+    0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19,
+    0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819,
+    0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b,
+    0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808,
+    0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908,
+    0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08,
+    0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08,
+    0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908,
+    0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819,
+    0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808,
+    0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808,
+    0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19,
+    0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819,
+    0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919,
+    0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b,
+    0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08,
+    0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808,
+    0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908,
+    0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b,
+    0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819,
+    0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08,
+    0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08,
+    0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808,
+    0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b,
+    0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b,
+    0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908,
+    0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819,
+    0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808,
+    0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908,
+    0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b,
+    0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808,
+    0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b,
+    0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b,
+    0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808,
+    0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19,
+    0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908,
+};
+
+} // namespace iq2_host

cuda/mmq/test/proto_gemm_dense_q8_d2r.cu

@@ -0,0 +1,773 @@
+// SPDX-License-Identifier: MIT
+// proto_gemm_dense_q8_d2r.cu - dense Q8_0 decode-to-registers prototype.
+//
+// Targets the dense-q8 census shapes (cmdq8shape15 2026-07-09, N = chunk
+// tokens 4096) where production mul_mat_q<8,128,0> runs far under the
+// cuBLASLt-int8 ceiling at the same shapes:
+//   sexp gate/up  M= 2048 K=4096   mmq 65 TF  vs Lt-int8 156
+//   sexp down     M= 4096 K=2048   mmq 24 TF  vs Lt-int8 111
+//   attn q_b      M=32768 K=1024   mmq 56 TF  vs Lt-int8  79-85
+//   o_proj        M= 4096 K=8192   mmq 96 TF  vs Lt-int8 145
+//
+// Weight layout: byte-neutral SoA repack of raw Q8_0 (34 B/block of 32):
+//   [half d[M*nb]] [int8 qs[M*K]]   (d row-major [M][nb], qs row-major [M][K])
+// Total bytes == raw bytes.  NOTE a "nicer" fully-tiled contiguous-stage
+// layout ([row_tile][k64][h][row][16B]) measured 27% SLOWER at identical
+// instruction counts: every col-tile CTA of a row group reads the same 8 KiB
+// block concurrently and the contiguous span lives on a few L2 slices (slice
+// camping).  Row-major keeps each stage's 128 rows K bytes apart -> spread
+// across slices; its 2-way ldgsts store conflicts are the cheaper poison.
+// This is the REPLACING-artifact candidate (IQ2/Q2K precedent); raw consumers
+// fall back to the host mmap.
+//
+// Kernel: m128n64 CTA (8 warps), q8 activations = production block_q8_1_mmq
+// D4 layout (f32 scale per 32) double-buffered at k128; weights double-
+// buffered at k64 (128 rows x 64 B + 16 B row pad for ldmatrix bank rotation);
+// mma.m16n8k32.s8 with per-k32 float fold acc += (s32)C * d_w(row) * d_a(col).
+// Q8_0 has no mins: no bias terms, no fix staging (d4 read straight from the
+// q8 block header).  Grid is 1-D with triton-style grouped supertiling
+// (group_m row-tiles hot in L2 while col tiles stream) - group_m is a runtime
+// arg, swept by the harness.
+//
+// Build on the GB10/Blackwell box after `make -j6 cuda-spark`:
+//   ar rcs libds4mmq.a cuda/mmq/ds4_ggml_stubs.o cuda/mmq/ds4_mmq.o \
+//          cuda/mmq/quantize.o cuda/mmq/mmid.o cuda/mmq/mmvq.o
+//   /usr/local/cuda/bin/nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 \
+//        -Icuda/mmq cuda/mmq/test/proto_gemm_dense_q8_d2r.cu libds4mmq.a \
+//        -lcudart -lcublas -lcuda -o proto_gemm_dense_q8_d2r
+// argv: [reps=20] [warmup=5] [N=4096]
+
+#include <type_traits>
+
+#include "ds4_mmq.h"
+#include "quantize.cuh"
+
+extern "C" int ds4_cuda_q8_fold_take_q81(const void *src, uint64_t in_dim, void *out) {
+    (void)src;
+    (void)in_dim;
+    (void)out;
+    return 0;
+}
+
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <limits>
+#include <random>
+#include <vector>
+
+#define CHECK(call) do { \
+    cudaError_t err__ = (call); \
+    if (err__ != cudaSuccess) { \
+        fprintf(stderr, "CUDA error %s at %s:%d: %s\n", #call, __FILE__, __LINE__, \
+                cudaGetErrorString(err__)); \
+        exit(1); \
+    } \
+} while (0)
+
+namespace {
+
+constexpr int kMTile   = 128;
+#ifndef PROTO_Q8_NTILE
+#define PROTO_Q8_NTILE 128
+#endif
+constexpr int kNTile   = PROTO_Q8_NTILE;
+// 16-warp CTA, one per SM: 8 row-warps x 2 col-warps.  Same 16 warps/SM
+// residency as 2x8, but per-thread acc halves to 8 fragments (32 regs) --
+// the 2x8-warp NT128 shape spilled ~76 B and the spill reloads were the top
+// stall (long_sb 3.02, local_ld 2.65M).
+constexpr int kRowWarps = 8;
+constexpr int kColWarps = 2;
+constexpr int kWarps   = kRowWarps * kColWarps;
+constexpr int kThreads = 32 * kWarps;
+// 3-stage rings: issues happen at the TOP of each iter for stage i+2 (a
+// buffer no in-flight reader touches), so the post-compute __syncthreads of
+// the 2-stage schedule disappears -- one barrier per k64.  69 KiB total needs
+// the dynamic-smem opt-in (1 CTA/SM is already forced by the 512-thread
+// register budget, so no occupancy cost).
+constexpr int kStages  = 3;   // q8 qs ring (k64 halves) + header ring (k128)
+constexpr int kWStages = 3;   // weight ring, k64 granularity
+constexpr int kNFrag   = kNTile / 8;
+constexpr int kWRowPad = 80;  // 64 B payload + 16 B pad: 20-int ldmatrix stride rotates banks
+
+// q8 staging is SPLIT so NT=128 fits under 48 KiB static smem: the 16 B
+// header (d4 scales) double-buffered at k128 cadence, the 128 B qs payload
+// staged as k64 halves (64 B rows, same padded geometry as the weight ring).
+constexpr int kQ8QsChunks  = kNTile * 4;  // cols x 4 16B-chunks per k64 half
+constexpr int kQ8QsTrips   = (kQ8QsChunks + kThreads - 1) / kThreads;
+constexpr int kWQChunks    = kMTile * 4;  // 128 rows x 4 16B-chunks (64 B payload)
+constexpr int kWQTrips     = kWQChunks / kThreads;
+
+static_assert(kQ8QsTrips == 1, "unexpected q8 qs issue trip count");
+static_assert(kWQTrips == 1, "unexpected weight issue trip count");
+static_assert(kNTile == 128, "16-warp decomposition assumes NT=128");
+static_assert(kNTile <= kThreads, "q8 header prefetch assumes one trip");
+constexpr int kNFragPerWarp = kNFrag / kColWarps;
+
+// Loop invariants live in plain (warp-uniform) registers: everything here
+// derives arithmetically from blockIdx + kernel args, so ptxas can promote
+// the math to the uniform register file.  (The Q2K proto's shared invariant
+// table exists for its memory-sourced expert-worklist lookups - dense has
+// none, and routing these through volatile smem forced them into regular
+// registers, feeding the 128-reg wall.)
+struct DenseQ8Params {
+    const char *wd_row0;       // d plane at cta_row0 (row stride nb*2 = K/16)
+    const char *wq_row0;       // qs plane at cta_row0 (row stride K)
+    const char *q8_tile;       // q8 blocks at col_lo
+    uint32_t q8_k128_stride_bytes;
+    uint32_t wq_row_stride;    // K bytes
+    int k64_iters;
+    int k128_iters;
+};
+
+uint16_t sane_half(std::mt19937 &rng) {
+    const uint16_t mant = (uint16_t)(rng() & 0x03FFu);
+    const uint16_t sign = (uint16_t)((rng() & 1u) << 15);
+    return (uint16_t)(sign | (14u << 10) | mant);
+}
+
+float *device_random_floats(std::mt19937 &rng, size_t n) {
+    std::vector<float> h(n);
+    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
+    for (size_t i = 0; i < n; ++i) {
+        h[i] = dist(rng);
+    }
+    float *d = nullptr;
+    CHECK(cudaMalloc(&d, n * sizeof(float)));
+    CHECK(cudaMemcpy(d, h.data(), n * sizeof(float), cudaMemcpyHostToDevice));
+    return d;
+}
+
+template <typename F>
+float best_time_ms(F &&fn, int warmup, int reps) {
+    cudaEvent_t a, b;
+    CHECK(cudaEventCreate(&a));
+    CHECK(cudaEventCreate(&b));
+    for (int i = 0; i < warmup; ++i) {
+        fn();
+    }
+    CHECK(cudaDeviceSynchronize());
+    float best = std::numeric_limits<float>::infinity();
+    for (int r = 0; r < reps; ++r) {
+        CHECK(cudaEventRecord(a));
+        fn();
+        CHECK(cudaEventRecord(b));
+        CHECK(cudaEventSynchronize(b));
+        float ms = 0.0f;
+        CHECK(cudaEventElapsedTime(&ms, a, b));
+        best = std::min(best, ms);
+    }
+    CHECK(cudaEventDestroy(a));
+    CHECK(cudaEventDestroy(b));
+    return best;
+}
+
+__device__ __forceinline__ void zero_16B(void *dst) {
+    *reinterpret_cast<int4 *>(dst) = make_int4(0, 0, 0, 0);
+}
+
+__device__ __forceinline__ void cp_async_16B(void *dst, const void *src, bool pred) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+    if (pred) {
+        const unsigned smem = static_cast<unsigned>(__cvta_generic_to_shared(dst));
+        asm volatile("cp.async.cg.shared.global [%0], [%1], 16;"
+                     :: "r"(smem), "l"(src));
+    } else {
+        zero_16B(dst);
+    }
+#else
+    if (pred) {
+        *reinterpret_cast<int4 *>(dst) = *reinterpret_cast<const int4 *>(src);
+    } else {
+        zero_16B(dst);
+    }
+#endif
+}
+
+__device__ __forceinline__ void cp_async_4B(void *dst, const void *src, bool pred) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+    if (pred) {
+        const unsigned smem = static_cast<unsigned>(__cvta_generic_to_shared(dst));
+        asm volatile("cp.async.ca.shared.global [%0], [%1], 4;"
+                     :: "r"(smem), "l"(src));
+    } else {
+        *reinterpret_cast<uint32_t *>(dst) = 0u;
+    }
+#else
+    if (pred) {
+        *reinterpret_cast<uint32_t *>(dst) = *reinterpret_cast<const uint32_t *>(src);
+    } else {
+        *reinterpret_cast<uint32_t *>(dst) = 0u;
+    }
+#endif
+}
+
+__device__ __forceinline__ void cp_async_commit() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+    asm volatile("cp.async.commit_group;");
+#endif
+}
+
+template <int KeepGroups>
+__device__ __forceinline__ void cp_async_wait_group() {
+    static_assert(KeepGroups >= 0 && KeepGroups <= 7, "bad cp.async wait_group depth");
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+    asm volatile("cp.async.wait_group %0;" :: "n"(KeepGroups));
+#endif
+}
+
+// CTA is launched as dim3(32, kWarps): the ggml mma primitives
+// (load_ldmatrix, tile get_i/get_j) use threadIdx.x as the LANE.
+__device__ __forceinline__ int d2r_lane() {
+#if defined(__CUDA_ARCH__)
+    return (int)threadIdx.x;
+#else
+    return 0;
+#endif
+}
+
+__device__ __forceinline__ int d2r_warp() {
+#if defined(__CUDA_ARCH__)
+    return (int)threadIdx.y;
+#else
+    return 0;
+#endif
+}
+
+__device__ __forceinline__ int d2r_tid() {
+#if defined(__CUDA_ARCH__)
+    return (int)(threadIdx.y * 32 + threadIdx.x);
+#else
+    return 0;
+#endif
+}
+
+constexpr size_t kSmemQ8HBytes = (size_t)kStages * kNTile * 16;        // headers
+constexpr size_t kSmemQ8QBytes = (size_t)kStages * kNTile * kWRowPad;  // qs k64 halves
+constexpr size_t kSmemWQBytes  = (size_t)kWStages * kMTile * kWRowPad;
+constexpr size_t kSmemWDBytes  = (size_t)kWStages * kMTile * sizeof(uint32_t);
+constexpr size_t kSmemWQOff  = 0;
+constexpr size_t kSmemQ8QOff = kSmemWQOff + kSmemWQBytes;
+constexpr size_t kSmemQ8HOff = kSmemQ8QOff + kSmemQ8QBytes;
+constexpr size_t kSmemWDOff  = kSmemQ8HOff + kSmemQ8HBytes;
+constexpr size_t kSmemTotalBytes = kSmemWDOff + kSmemWDBytes;
+static_assert(kSmemTotalBytes <= 99ull * 1024ull, "dynamic shared memory exceeds sm_121 limit");
+
+// ---- q8 activation prefetch, split staging ----
+// qs k64 halves: 128 cols x 64 B (bytes 16 + (k64&1)*64 .. of each col's
+// 144 B block), padded 80 B rows — same geometry as the weight ring.
+// headers: 128 cols x 16 B (the d4 scale header), k128 cadence.
+
+__device__ __forceinline__ void issue_q8_qs_prefetch(
+        int8_t (&s_q8q)[kStages][kNTile][kWRowPad],
+        const DenseQ8Params &p,
+        int k64_iter) {
+    const bool pred = k64_iter < p.k64_iters;
+    const char *base = p.q8_tile +
+                       (uint64_t)(k64_iter >> 1) * p.q8_k128_stride_bytes +
+                       16u + (uint64_t)(k64_iter & 1) * 64u;
+    const int buf = k64_iter % kStages;
+#pragma unroll
+    for (int trip = 0; trip < kQ8QsTrips; ++trip) {
+        const int chunk = d2r_tid() + trip * kThreads;
+        if (kQ8QsChunks % kThreads != 0 && chunk >= kQ8QsChunks) {
+            break;
+        }
+        const int col = chunk >> 2;
+        const int h = chunk & 3;
+        void *dst = &s_q8q[buf][col][h * 16];
+        const void *src = base + (uint64_t)col * sizeof(block_q8_1_mmq) + (uint64_t)h * 16u;
+        cp_async_16B(dst, src, pred);
+    }
+}
+
+__device__ __forceinline__ void issue_q8_hdr_prefetch(
+        float (&s_q8h)[kStages][kNTile][4],
+        const DenseQ8Params &p,
+        int k128_iter) {
+    if (d2r_tid() >= kNTile) {
+        return;
+    }
+    const bool pred = k128_iter < p.k128_iters;
+    const char *base = p.q8_tile +
+                       (uint64_t)k128_iter * p.q8_k128_stride_bytes;
+    const int col = d2r_tid();
+    void *dst = &s_q8h[k128_iter % kStages][col][0];
+    const void *src = base + (uint64_t)col * sizeof(block_q8_1_mmq);
+    cp_async_16B(dst, src, pred);
+}
+
+// ---- weight prefetch (one k64 stage = 128 rows x 64 B qs + 128 x 4 B d) ----
+
+__device__ __forceinline__ void issue_w_prefetch(
+        int8_t (&s_wq)[kWStages][kMTile][kWRowPad],
+        uint32_t (&s_wd)[kWStages][kMTile],
+        const DenseQ8Params &p,
+        int k64_iter) {
+    const bool pred = k64_iter < p.k64_iters;
+    const int wst = k64_iter % kWStages;
+    // Row-major planes (see layout note at top: contiguous-stage tiling was
+    // 27% slower from L2 slice camping).
+    const char *wq = p.wq_row0 + (uint64_t)k64_iter * 64u;
+    const char *wd = p.wd_row0 + (uint64_t)k64_iter * 4u;
+    const uint32_t wq_row_stride = p.wq_row_stride;
+    const uint32_t wd_row_stride = p.wq_row_stride >> 4;  // nb*2 = K/16
+#pragma unroll
+    for (int trip = 0; trip < kWQTrips; ++trip) {
+        const int chunk = d2r_tid() + trip * kThreads;
+        const int row = chunk >> 2;
+        const int h = chunk & 3;
+        void *dst = &s_wq[wst][row][h * 16];
+        const void *src = wq + (uint64_t)row * wq_row_stride + (uint64_t)h * 16u;
+        cp_async_16B(dst, src, pred);
+    }
+    if (d2r_tid() < kMTile) {
+        const int row = d2r_tid();
+        void *dst = &s_wd[wst][row];
+        const void *src = wd + (uint64_t)row * wd_row_stride;
+        cp_async_4B(dst, src, pred);
+    }
+}
+
+// ---- mainloop ----
+
+template <typename TileA, typename TileB, typename TileC>
+__device__ __forceinline__ void dense_q8_d2r_mainloop(
+        float (&acc)[kNFragPerWarp][4],
+        float (&s_q8h)[kStages][kNTile][4],
+        int8_t (&s_q8q)[kStages][kNTile][kWRowPad],
+        int8_t (&s_wq)[kWStages][kMTile][kWRowPad],
+        uint32_t (&s_wd)[kWStages][kMTile],
+        const DenseQ8Params &p) {
+    static_assert(TileC::ne == 4, "expected m16n8 s32 accumulator fragment");
+    const int k64_iters = p.k64_iters;
+    const int lane = d2r_lane();
+    const int group = lane >> 2;
+    const int wrow = (d2r_warp() >> 1) * 16;          // row-warp
+    const int nf0 = (d2r_warp() & 1) * kNFragPerWarp;  // col-warp half
+    const int c0 = TileC::get_j(0);
+    const int c1 = TileC::get_j(1);
+
+    // Prologue: stages 0 and 1 in flight as two groups.  Loop iters run
+    // wait -> barrier -> issue(i+2) -> compute: with 3 buffers the issue
+    // never aliases the two live stages, so the 2-stage schedule's
+    // post-compute barrier disappears (one barrier per k64).
+    issue_w_prefetch(s_wq, s_wd, p, 0);
+    issue_q8_qs_prefetch(s_q8q, p, 0);
+    issue_q8_hdr_prefetch(s_q8h, p, 0);
+    cp_async_commit();
+    issue_w_prefetch(s_wq, s_wd, p, 1);
+    issue_q8_qs_prefetch(s_q8q, p, 1);
+    issue_q8_hdr_prefetch(s_q8h, p, 1);
+    cp_async_commit();
+
+    for (int i = 0; i < k64_iters; ++i) {
+        // wait -> barrier -> issue: the barrier both publishes stage i's
+        // copies CTA-wide and proves every warp finished reading stage i-1,
+        // whose buffer ((i-1)%3 == (i+2)%3) the issues below overwrite.
+        cp_async_wait_group<1>();
+        __syncthreads();
+        issue_w_prefetch(s_wq, s_wd, p, i + 2);
+        issue_q8_qs_prefetch(s_q8q, p, i + 2);
+        if ((i & 1) == 0) {
+            issue_q8_hdr_prefetch(s_q8h, p, (i >> 1) + 2);
+        }
+        cp_async_commit();
+
+        const int wst = i % kWStages;
+        const int qbuf = i % kStages;          // qs halves ride the k64 cadence
+        const int hbuf = (i >> 1) % kStages;   // headers ride the k128 cadence
+        const uint32_t dw0_bits = s_wd[wst][wrow + group];
+        const uint32_t dw1_bits = s_wd[wst][wrow + group + 8];
+        const float2 dw0 = __half22float2(*reinterpret_cast<const half2 *>(&dw0_bits));
+        const float2 dw1 = __half22float2(*reinterpret_cast<const half2 *>(&dw1_bits));
+
+        // t-phased: each A load overlaps the previous phase's B/MMA/fold tail
+        // (fusing both A loads up front measured -27%: hard A dependency at the
+        // k64 head + burstier MIO).
+#pragma unroll
+        for (int t = 0; t < 2; ++t) {
+            TileA A;
+            ggml_cuda_mma::load_ldmatrix(
+                A, reinterpret_cast<const int *>(&s_wq[wst][wrow][t * 32]),
+                kWRowPad / (int)sizeof(int));
+            const int tq8 = (i & 1) * 2 + t;  // k32 index within the q8 k128 block
+            const float dwr0 = t ? dw0.y : dw0.x;
+            const float dwr1 = t ? dw1.y : dw1.x;
+#pragma unroll
+            for (int nf = 0; nf < kNFragPerWarp; ++nf) {
+                TileB B;
+                TileC C;
+                ggml_cuda_mma::load_ldmatrix(
+                    B, reinterpret_cast<const int *>(&s_q8q[qbuf][(nf0 + nf) * 8][t * 32]),
+                    kWRowPad / (int)sizeof(int));
+                ggml_cuda_mma::mma(C, A, B);
+                const float da0 = s_q8h[hbuf][(nf0 + nf) * 8 + c0][tq8];
+                const float da1 = s_q8h[hbuf][(nf0 + nf) * 8 + c1][tq8];
+                acc[nf][0] = fmaf((float)C.x[0], dwr0 * da0, acc[nf][0]);
+                acc[nf][1] = fmaf((float)C.x[1], dwr0 * da1, acc[nf][1]);
+                acc[nf][2] = fmaf((float)C.x[2], dwr1 * da0, acc[nf][2]);
+                acc[nf][3] = fmaf((float)C.x[3], dwr1 * da1, acc[nf][3]);
+            }
+        }
+
+    }
+}
+
+__global__ __launch_bounds__(kThreads, 1)
+void dense_q8_d2r_kernel(const char * __restrict__ wd_plane,
+                         const char * __restrict__ wq_plane,
+                         const block_q8_1_mmq * __restrict__ q8,
+                         float * __restrict__ outf,
+                         __half * __restrict__ outh,
+                         int M, int N, int K, int group_m) {
+#if defined(TURING_MMA_AVAILABLE)
+    using tile_A = ggml_cuda_mma::tile<16, 8, int>;
+    using tile_B = ggml_cuda_mma::tile<8, 8, int>;
+    using tile_C = ggml_cuda_mma::tile<16, 8, int>;
+
+    extern __shared__ __align__(16) char smem_dyn[];
+    auto &s_wq  = *reinterpret_cast<int8_t (*)[kWStages][kMTile][kWRowPad]>(smem_dyn + kSmemWQOff);
+    auto &s_q8q = *reinterpret_cast<int8_t (*)[kStages][kNTile][kWRowPad]>(smem_dyn + kSmemQ8QOff);
+    auto &s_q8h = *reinterpret_cast<float (*)[kStages][kNTile][4]>(smem_dyn + kSmemQ8HOff);
+    auto &s_wd  = *reinterpret_cast<uint32_t (*)[kWStages][kMTile]>(smem_dyn + kSmemWDOff);
+
+    // Triton-style grouped supertile order: group_m row-tiles stay L2-hot
+    // while the col tiles stream past them.
+    const int num_m = M / kMTile;
+    const int num_n = (N + kNTile - 1) / kNTile;
+    const int width = group_m * num_n;
+    const int g = (int)blockIdx.x / width;
+    const int rem = (int)blockIdx.x - g * width;
+    int gsize = num_m - g * group_m;
+    if (gsize > group_m) {
+        gsize = group_m;
+    }
+    const int pid_m = g * group_m + rem % gsize;
+    const int pid_n = rem / gsize;
+
+    const int cta_row0 = pid_m * kMTile;
+    const int col_lo = pid_n * kNTile;
+
+    DenseQ8Params p;
+    p.wd_row0 = wd_plane + (uint64_t)cta_row0 * (uint64_t)(K >> 4);
+    p.wq_row0 = wq_plane + (uint64_t)cta_row0 * (uint64_t)K;
+    p.q8_tile = reinterpret_cast<const char *>(q8) +
+                (uint64_t)col_lo * sizeof(block_q8_1_mmq);
+    p.q8_k128_stride_bytes = (uint32_t)((uint64_t)N * sizeof(block_q8_1_mmq));
+    p.wq_row_stride = (uint32_t)K;
+    p.k64_iters = K >> 6;
+    p.k128_iters = K >> 7;
+
+    float acc[kNFragPerWarp][tile_C::ne] = {};
+
+    dense_q8_d2r_mainloop<tile_A, tile_B, tile_C>(acc, s_q8h, s_q8q, s_wq, s_wd, p);
+
+    const int out_col_lo = col_lo + (d2r_warp() & 1) * (kNFragPerWarp * 8);
+    const int out_N = N;
+    const int out_M = M;
+    const int out_row0 = cta_row0 + ((d2r_warp() >> 1) << 4);
+    float *outf_base = outf;
+    __half *outh_base = outh;
+#pragma unroll
+    for (int nf = 0; nf < kNFragPerWarp; ++nf) {
+        const int col_frag0 = out_col_lo + nf * 8;
+#pragma unroll
+        for (int l = 0; l < tile_C::ne; ++l) {
+            const int row = out_row0 + tile_C::get_i(l);
+            const int col = col_frag0 + tile_C::get_j(l);
+            if (col < out_N) {
+                const float v = isfinite(acc[nf][l]) ? acc[nf][l] : 0.0f;
+                if (outh_base) {
+                    outh_base[(uint64_t)col * (uint64_t)out_M + (uint64_t)row] = __float2half(v);
+                } else {
+                    outf_base[(uint64_t)col * (uint64_t)out_M + (uint64_t)row] = v;
+                }
+            }
+        }
+    }
+#else
+    (void)wd_plane;
+    (void)wq_plane;
+    (void)q8;
+    (void)outf;
+    (void)outh;
+    (void)M;
+    (void)N;
+    (void)K;
+    (void)group_m;
+#endif
+}
+
+int launch_dense_q8_d2r(const void *wd_plane, const void *wq_plane,
+                        const block_q8_1_mmq *q8, float *outf, __half *outh,
+                        int M, int N, int K, int group_m, cudaStream_t stream) {
+    if (!wd_plane || !wq_plane || !q8 || (!outf && !outh) ||
+        M <= 0 || (M & (kMTile - 1)) != 0 || K % 256 != 0 || N <= 0 || group_m <= 0) {
+        fprintf(stderr, "launch_dense_q8_d2r: bad args M=%d N=%d K=%d group_m=%d\n",
+                M, N, K, group_m);
+        return -1;
+    }
+    const int num_m = M / kMTile;
+    const int num_n = (N + kNTile - 1) / kNTile;
+    static int smem_opted = 0;
+    if (!smem_opted) {
+        cudaFuncSetAttribute(dense_q8_d2r_kernel,
+                             cudaFuncAttributeMaxDynamicSharedMemorySize,
+                             (int)kSmemTotalBytes);
+        smem_opted = 1;
+    }
+    const dim3 grid((unsigned)(num_m * num_n), 1, 1);
+    const dim3 block(32, kWarps, 1);
+    dense_q8_d2r_kernel<<<grid, block, kSmemTotalBytes, stream>>>(
+        (const char *)wd_plane, (const char *)wq_plane, q8, outf, outh, M, N, K, group_m);
+    const cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess) {
+        fprintf(stderr, "dense_q8_d2r_kernel launch failed: %s\n", cudaGetErrorString(err));
+        return -2;
+    }
+    return 0;
+}
+
+void quantize_q8_mmq_dense(const float *x, block_q8_1_mmq *q8, int K, int N,
+                           cudaStream_t stream) {
+    quantize_mmq_q8_1_cuda(
+        x, /*ids=*/nullptr, (void *)q8,
+        GGML_TYPE_Q8_0, /*ne00=*/K, /*s01=*/(int64_t)K, /*s02=*/0, /*s03=*/0,
+        /*ne0=*/K, /*ne1=*/N, /*ne2=*/1, /*ne3=*/1, stream);
+    CHECK(cudaGetLastError());
+}
+
+struct ParityStats {
+    double max_abs = 0.0;
+    double max_rel_rms = 0.0;
+    int max_row = -1;
+    int max_col = -1;
+};
+
+ParityStats check_parity(const float *ref, const float *got, int M, int N) {
+    std::vector<double> ss((size_t)M, 0.0);
+    for (int c = 0; c < N; ++c) {
+        const float *rp = ref + (size_t)c * M;
+        for (int r = 0; r < M; ++r) {
+            ss[(size_t)r] += (double)rp[r] * (double)rp[r];
+        }
+    }
+    ParityStats st;
+    for (int c = 0; c < N; ++c) {
+        const float *rp = ref + (size_t)c * M;
+        const float *gp = got + (size_t)c * M;
+        for (int r = 0; r < M; ++r) {
+            const double abs_err = std::abs((double)gp[r] - (double)rp[r]);
+            const double rms = std::sqrt(ss[(size_t)r] / std::max(1, N));
+            const double rel = abs_err / std::max(1.0e-6, rms);
+            if (abs_err > st.max_abs) {
+                st.max_abs = abs_err;
+                st.max_row = r;
+                st.max_col = c;
+            }
+            st.max_rel_rms = std::max(st.max_rel_rms, rel);
+        }
+    }
+    return st;
+}
+
+void run_shape(std::mt19937 &rng, const char *label, int M, int K, int N,
+               int warmup, int reps, bool do_parity) {
+    const int nb = K / 32;
+    const uint64_t raw_bytes = (uint64_t)M * nb * 34u;
+    const uint64_t d_bytes = (uint64_t)M * nb * 2u;
+
+    // Raw Q8_0 blocks (34 B: half d + 32 int8) for the production reference,
+    // and the byte-neutral SoA planes for the D2R kernel, from the same data.
+    std::vector<uint8_t> h_raw(raw_bytes);
+    std::vector<uint8_t> h_soa(raw_bytes);
+    {
+        // Row-major SoA planes: d [M][nb] halves, qs [M][K] int8.
+        uint8_t *raw = h_raw.data();
+        uint16_t *dp = reinterpret_cast<uint16_t *>(h_soa.data());
+        int8_t *qp = reinterpret_cast<int8_t *>(h_soa.data() + d_bytes);
+        for (int r = 0; r < M; ++r) {
+            for (int b = 0; b < nb; ++b) {
+                uint8_t *blk = raw + ((uint64_t)r * nb + b) * 34u;
+                const uint16_t d = sane_half(rng);
+                memcpy(blk, &d, 2);
+                for (int j = 0; j < 32; ++j) {
+                    blk[2 + j] = (uint8_t)(rng() & 0xFFu);
+                }
+                dp[(uint64_t)r * nb + b] = d;
+                memcpy(qp + (uint64_t)r * K + (uint64_t)b * 32u, blk + 2, 32);
+            }
+        }
+    }
+
+    void *d_raw = nullptr;
+    void *d_soa = nullptr;
+    CHECK(cudaMalloc(&d_raw, raw_bytes));
+    CHECK(cudaMalloc(&d_soa, raw_bytes));
+    CHECK(cudaMemcpy(d_raw, h_raw.data(), raw_bytes, cudaMemcpyHostToDevice));
+    CHECK(cudaMemcpy(d_soa, h_soa.data(), raw_bytes, cudaMemcpyHostToDevice));
+    const char *d_wd = (const char *)d_soa;
+    const char *d_wq = (const char *)d_soa + d_bytes;
+
+    float *d_x = device_random_floats(rng, (size_t)N * (size_t)K);
+
+    // +kNTile block slack so a guarded last col tile can read past col N-1
+    // without faulting (values discarded by the epilogue guard).
+    const size_t q8_blocks = (size_t)N * (size_t)(K / 128) + kNTile;
+    block_q8_1_mmq *d_q8 = nullptr;
+    CHECK(cudaMalloc(&d_q8, q8_blocks * sizeof(block_q8_1_mmq)));
+    CHECK(cudaMemset(d_q8, 0, q8_blocks * sizeof(block_q8_1_mmq)));
+    quantize_q8_mmq_dense(d_x, d_q8, K, N, 0);
+
+    float *d_ref = nullptr;
+    float *d_new = nullptr;
+    __half *d_newh = nullptr;
+    CHECK(cudaMalloc(&d_ref, (size_t)M * N * sizeof(float)));
+    CHECK(cudaMalloc(&d_new, (size_t)M * N * sizeof(float)));
+    CHECK(cudaMalloc(&d_newh, (size_t)M * N * sizeof(__half)));
+
+    if (ds4_mmq_q8_0_dense(d_raw, d_x, d_ref, M, N, K, 0) != 0) {
+        fprintf(stderr, "reference ds4_mmq_q8_0_dense failed M=%d K=%d\n", M, K);
+        exit(1);
+    }
+    if (launch_dense_q8_d2r(d_wd, d_wq, d_q8, d_new, nullptr, M, N, K, 8, 0) != 0) {
+        exit(1);
+    }
+    CHECK(cudaDeviceSynchronize());
+
+    if (do_parity) {
+        std::vector<float> ref((size_t)M * N), got((size_t)M * N);
+        CHECK(cudaMemcpy(ref.data(), d_ref, ref.size() * sizeof(float), cudaMemcpyDeviceToHost));
+        CHECK(cudaMemcpy(got.data(), d_new, got.size() * sizeof(float), cudaMemcpyDeviceToHost));
+        const ParityStats st = check_parity(ref.data(), got.data(), M, N);
+        const bool pass = st.max_rel_rms <= 2.0e-3;
+        printf("  parity %-8s M=%5d K=%4d N=%d: max_abs=%.6g at col=%d row=%d  max_abs/rms(row)=%.6g  %s\n",
+               label, M, K, N, st.max_abs, st.max_col, st.max_row, st.max_rel_rms,
+               pass ? "PASS" : "FAIL");
+        if (!pass && getenv("PROTO_Q8_DEBUG")) {
+            // Error structure: per (row-tile, col-tile) fraction of bad elements,
+            // plus per row%16 / col%8 / warp-row histograms of bad positions.
+            const int ntm = M / 128, ntn = (N + 63) / 64;
+            size_t bad_total = 0;
+            std::vector<size_t> bad_rowmod(16, 0), bad_colmod(8, 0), bad_warp(8, 0);
+            std::vector<size_t> tile_bad((size_t)ntm * ntn, 0);
+            for (int c = 0; c < N; ++c) {
+                const float *rp = ref.data() + (size_t)c * M;
+                const float *gp = got.data() + (size_t)c * M;
+                for (int r = 0; r < M; ++r) {
+                    const double e = std::abs((double)gp[r] - (double)rp[r]);
+                    const double mag = std::max(1.0, std::abs((double)rp[r]));
+                    if (e / mag > 1.0e-2) {
+                        ++bad_total;
+                        ++bad_rowmod[r & 15];
+                        ++bad_colmod[c & 7];
+                        ++bad_warp[(r >> 4) & 7];
+                        ++tile_bad[(size_t)(r / 128) * ntn + (c / 64)];
+                    }
+                }
+            }
+            printf("    bad(rel>1e-2) %zu / %zu\n", bad_total, (size_t)M * N);
+            printf("    by row%%16:"); for (int i = 0; i < 16; ++i) printf(" %zu", bad_rowmod[i]); printf("\n");
+            printf("    by col%%8: "); for (int i = 0; i < 8; ++i) printf(" %zu", bad_colmod[i]); printf("\n");
+            printf("    by warp(row/16%%8):"); for (int i = 0; i < 8; ++i) printf(" %zu", bad_warp[i]); printf("\n");
+            printf("    tile map (rows=row-tile, cols=col-tile, %% bad, first 8x16):\n");
+            for (int tm = 0; tm < std::min(ntm, 8); ++tm) {
+                printf("      ");
+                for (int tn = 0; tn < std::min(ntn, 16); ++tn) {
+                    printf("%3zu", tile_bad[(size_t)tm * ntn + tn] * 100 / (128 * 64));
+                }
+                printf("\n");
+            }
+            // A few sample mismatches with ratios (scale bugs show as clean ratios).
+            int shown = 0;
+            for (int c = 0; c < N && shown < 6; ++c) {
+                for (int r = 0; r < M && shown < 6; ++r) {
+                    const float rv = ref[(size_t)c * M + r], gv = got[(size_t)c * M + r];
+                    if (std::abs(gv - rv) / std::max(1.0f, std::abs(rv)) > 1.0e-2) {
+                        printf("    sample r=%d c=%d ref=%.6g got=%.6g got/ref=%.4f\n",
+                               r, c, rv, gv, rv != 0.0f ? gv / rv : 0.0f);
+                        ++shown;
+                    }
+                }
+            }
+        }
+        if (!pass) {
+            exit(2);
+        }
+    }
+
+    const double flop = 2.0 * (double)M * (double)K * (double)N;
+    const float ref_ms = best_time_ms([&] {
+        ds4_mmq_q8_0_dense(d_raw, d_x, d_ref, M, N, K, 0);
+    }, warmup, reps);
+
+    printf("  %-8s M=%5d K=%4d N=%d  mmq-ref %7.3f ms %5.1f TF\n",
+           label, M, K, N, ref_ms, flop / (ref_ms * 1.0e-3) / 1.0e12);
+
+    float best_kernel_ms = std::numeric_limits<float>::infinity();
+    int best_g = 0;
+    const int groups[] = {1, 4, 8, 16};
+    for (int gi = 0; gi < 4; ++gi) {
+        const int g = groups[gi];
+        const float ms = best_time_ms([&] {
+            launch_dense_q8_d2r(d_wd, d_wq, d_q8, d_new, nullptr, M, N, K, g, 0);
+        }, warmup, reps);
+        printf("    d2r f32-out group_m=%2d  %7.3f ms %5.1f TF\n",
+               g, ms, flop / (ms * 1.0e-3) / 1.0e12);
+        if (ms < best_kernel_ms) {
+            best_kernel_ms = ms;
+            best_g = g;
+        }
+    }
+
+    const float f16_ms = best_time_ms([&] {
+        launch_dense_q8_d2r(d_wd, d_wq, d_q8, nullptr, d_newh, M, N, K, best_g, 0);
+    }, warmup, reps);
+    const float full_ms = best_time_ms([&] {
+        quantize_q8_mmq_dense(d_x, d_q8, K, N, 0);
+        launch_dense_q8_d2r(d_wd, d_wq, d_q8, d_new, nullptr, M, N, K, best_g, 0);
+    }, warmup, reps);
+
+    printf("    d2r BEST group_m=%d: f32-out %7.3f ms %5.1f TF | f16-out %7.3f ms %5.1f TF | +q8 %7.3f ms %5.1f TF | vs mmq %.2fx\n",
+           best_g,
+           best_kernel_ms, flop / (best_kernel_ms * 1.0e-3) / 1.0e12,
+           f16_ms, flop / (f16_ms * 1.0e-3) / 1.0e12,
+           full_ms, flop / (full_ms * 1.0e-3) / 1.0e12,
+           ref_ms / best_kernel_ms);
+
+    CHECK(cudaFree(d_raw));
+    CHECK(cudaFree(d_soa));
+    CHECK(cudaFree(d_x));
+    CHECK(cudaFree(d_q8));
+    CHECK(cudaFree(d_ref));
+    CHECK(cudaFree(d_new));
+    CHECK(cudaFree(d_newh));
+}
+
+} // namespace
+
+int main(int argc, char **argv) {
+    const int reps = (argc > 1) ? atoi(argv[1]) : 20;
+    const int warmup = (argc > 2) ? atoi(argv[2]) : 5;
+    const int N = (argc > 3) ? atoi(argv[3]) : 4096;
+
+    printf("### proto_gemm_dense_q8_d2r reps=%d warmup=%d N=%d ###\n", reps, warmup, N);
+    printf("### CTA m%dn%d, warps=%d, q8 x%d stages, weights x%d stages, dyn smem %zu B ###\n",
+           kMTile, kNTile, kWarps, kStages, kWStages, kSmemTotalBytes);
+    printf("### Lt-int8 ceilings at these shapes: gate/up 156, down 111, q_b 79-85, o_proj 145 TF ###\n");
+
+    std::mt19937 rng(0xD84D8021u + 13u * (unsigned)N);
+
+    run_shape(rng, "gate/up", 2048, 4096, N, warmup, reps, true);
+    run_shape(rng, "down", 4096, 2048, N, warmup, reps, true);
+    run_shape(rng, "q_b", 32768, 1024, N, warmup, reps, false);
+    run_shape(rng, "o_proj", 4096, 8192, N, warmup, reps, false);
+    return 0;
+}

cuda/mmq/test/proto_gemm_lib.cu

@@ -0,0 +1,127 @@
+// SPDX-License-Identifier: MIT
+// proto_gemm_lib.cu - NVIDIA-library dense GEMM ceiling at the exact MoE-leg
+// shapes (cuBLASLt heuristics, up to 16 algos, best-of; TN col-major):
+//   down-shape:   m=4096 n=24576 k=2048  (== one W4096 down launch's total FLOP)
+//   gateup-shape: m=2048 n=24576 k=4096  (== one gate or up GEMM of the pair)
+// Kinds: int8->int32 (COMPUTE_32I) | f16->f16 @ f32-acc (COMPUTE_32F) |
+//        f16->f16 @ f16-acc (COMPUTE_16F).
+//
+// This is "what NVIDIA's best dense kernel does on this silicon at our shape" —
+// the practical ceiling any MoE grouped-GEMM rewrite chases (grouping + quant
+// decode overhead comes on top of it).
+//
+// Build: /usr/local/cuda/bin/nvcc -O3 -std=c++17 -arch=sm_121 \
+//        cuda/mmq/test/proto_gemm_lib.cu -o proto_gemm_lib -lcublasLt
+// argv: [reps=10]
+
+#include <cublasLt.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cstdio>
+#include <cstdlib>
+
+#define CHECK(call) do { cudaError_t e_=(call); if(e_!=cudaSuccess){ \
+    fprintf(stderr,"CUDA %s @%d: %s\n",#call,__LINE__,cudaGetErrorString(e_)); exit(1);} } while(0)
+#define LT(call) do { cublasStatus_t s_=(call); if(s_!=CUBLAS_STATUS_SUCCESS){ \
+    fprintf(stderr,"cublasLt %s @%d: status %d\n",#call,__LINE__,(int)s_); goto done; } } while(0)
+
+namespace {
+
+__global__ void fillk(unsigned *p, size_t n, unsigned seed) {
+    for (size_t i = blockIdx.x * (size_t)blockDim.x + threadIdx.x; i < n;
+         i += (size_t)gridDim.x * blockDim.x)
+        p[i] = ((unsigned)i * 2654435761u + seed) & 0x3BFF3BFFu;  // finite small f16 pairs; valid int8 bytes
+}
+
+void bench_case(cublasLtHandle_t lt, void *dA, void *dB, void *dC, void *ws,
+                size_t wsSize, int m, int n, int k, int kind, int reps,
+                const char *label) {
+    const cudaDataType ab = (kind == 0) ? CUDA_R_8I : CUDA_R_16F;
+    const cudaDataType cd = (kind == 0) ? CUDA_R_32I : CUDA_R_16F;
+    const cublasComputeType_t ct = (kind == 0) ? CUBLAS_COMPUTE_32I
+                                : (kind == 1) ? CUBLAS_COMPUTE_32F : CUBLAS_COMPUTE_16F;
+    const cudaDataType st = (kind == 0) ? CUDA_R_32I
+                          : (kind == 1) ? CUDA_R_32F : CUDA_R_16F;
+    cublasLtMatmulDesc_t op = nullptr;
+    cublasLtMatrixLayout_t La = nullptr, Lb = nullptr, Lc = nullptr;
+    cublasLtMatmulPreference_t pref = nullptr;
+    cublasLtMatmulHeuristicResult_t res[16];
+    int nres = 0;
+    const cublasOperation_t ta = CUBLAS_OP_T, tb = CUBLAS_OP_N;
+    const int32_t ai = 1, bi = 0; const float af = 1.f, bf = 0.f;
+    const __half ah = __float2half(1.f), bh = __float2half(0.f);
+    const void *alpha = (kind == 0) ? (const void *)&ai : (kind == 1) ? (const void *)&af : (const void *)&ah;
+    const void *beta  = (kind == 0) ? (const void *)&bi : (kind == 1) ? (const void *)&bf : (const void *)&bh;
+    float best = 1e30f; int besti = -1;
+
+    LT(cublasLtMatmulDescCreate(&op, ct, st));
+    LT(cublasLtMatmulDescSetAttribute(op, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)));
+    LT(cublasLtMatmulDescSetAttribute(op, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)));
+    LT(cublasLtMatrixLayoutCreate(&La, ab, k, m, k));  // stored kxm; opA=T -> mxk
+    LT(cublasLtMatrixLayoutCreate(&Lb, ab, k, n, k));
+    LT(cublasLtMatrixLayoutCreate(&Lc, cd, m, n, m));
+    LT(cublasLtMatmulPreferenceCreate(&pref));
+    LT(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+                                            &wsSize, sizeof(wsSize)));
+    LT(cublasLtMatmulAlgoGetHeuristic(lt, op, La, Lb, Lc, Lc, pref, 16, res, &nres));
+    if (!nres) { printf("  %-14s m=%5d n=%5d k=%5d: NOALGO\n", label, m, n, k); goto done; }
+
+    for (int i = 0; i < nres; ++i) {
+        if (res[i].state != CUBLAS_STATUS_SUCCESS) continue;
+        if (cublasLtMatmul(lt, op, alpha, dA, La, dB, Lb, beta, dC, Lc, dC, Lc,
+                           &res[i].algo, ws, wsSize, 0) != CUBLAS_STATUS_SUCCESS) continue;
+        CHECK(cudaDeviceSynchronize());
+        cudaEvent_t e0, e1; CHECK(cudaEventCreate(&e0)); CHECK(cudaEventCreate(&e1));
+        CHECK(cudaEventRecord(e0));
+        for (int r = 0; r < reps; ++r)
+            cublasLtMatmul(lt, op, alpha, dA, La, dB, Lb, beta, dC, Lc, dC, Lc,
+                           &res[i].algo, ws, wsSize, 0);
+        CHECK(cudaEventRecord(e1)); CHECK(cudaEventSynchronize(e1));
+        float ms = 0; CHECK(cudaEventElapsedTime(&ms, e0, e1)); ms /= (float)reps;
+        if (ms < best) { best = ms; besti = i; }
+        CHECK(cudaEventDestroy(e0)); CHECK(cudaEventDestroy(e1));
+    }
+    if (besti < 0) { printf("  %-14s m=%5d n=%5d k=%5d: ALL-ALGOS-FAILED (n=%d)\n", label, m, n, k, nres); goto done; }
+    printf("  %-14s m=%5d n=%5d k=%5d: best %7.3f ms  %6.1f TFLOPS  (algo %d of %d)\n",
+           label, m, n, k, best, 2.0 * m * (double)n * k / (best * 1e-3) / 1e12, besti, nres);
+done:
+    if (pref) cublasLtMatmulPreferenceDestroy(pref);
+    if (Lc) cublasLtMatrixLayoutDestroy(Lc);
+    if (Lb) cublasLtMatrixLayoutDestroy(Lb);
+    if (La) cublasLtMatrixLayoutDestroy(La);
+    if (op) cublasLtMatmulDescDestroy(op);
+}
+
+} // namespace
+
+int main(int argc, char **argv) {
+    const int reps = (argc > 1) ? atoi(argv[1]) : 10;
+    printf("### proto_gemm_lib  cublasLt %zu  reps=%d ###\n", cublasLtGetVersion(), reps);
+    cublasLtHandle_t lt; if (cublasLtCreate(&lt) != CUBLAS_STATUS_SUCCESS) { fprintf(stderr, "lt create fail\n"); return 1; }
+    const size_t maxA = (size_t)32768 * 8192 * 2;  // covers kxm all shapes, f16
+    const size_t maxB = (size_t)8192 * 24576 * 2;
+    const size_t maxC = (size_t)32768 * 24576 * 4;
+    const size_t wsSize = 256u << 20;
+    void *dA, *dB, *dC, *ws;
+    CHECK(cudaMalloc(&dA, maxA)); CHECK(cudaMalloc(&dB, maxB));
+    CHECK(cudaMalloc(&dC, maxC)); CHECK(cudaMalloc(&ws, wsSize));
+    fillk<<<512, 256>>>((unsigned *)dA, maxA / 4, 0xA5A5u);
+    fillk<<<512, 256>>>((unsigned *)dB, maxB / 4, 0x5A5Au);
+    CHECK(cudaDeviceSynchronize());
+
+    /* MoE-equiv shapes + the dense-q8 census shapes (cmdq8shape15 2026-07-09):
+     * q_b [32768x1024], o_proj [4096x8192], sexp gate/up [2048x4096],
+     * sexp down [4096x2048], all at N=4096 chunk tokens. */
+    const int shapes[][3] = { {4096, 24576, 2048}, {2048, 24576, 4096},
+                              {32768, 4096, 1024}, {4096, 4096, 8192},
+                              {2048, 4096, 4096},  {4096, 4096, 2048} };
+    const char *kn[] = { "int8.s32acc", "f16.f32acc", "f16.f16acc" };
+    for (const auto &s : shapes) {
+        printf("--- shape m=%d n=%d k=%d ---\n", s[0], s[1], s[2]);
+        for (int kind = 0; kind < 3; ++kind)
+            bench_case(lt, dA, dB, dC, ws, wsSize, s[0], s[1], s[2], kind, reps, kn[kind]);
+    }
+    cublasLtDestroy(lt);
+    CHECK(cudaFree(dA)); CHECK(cudaFree(dB)); CHECK(cudaFree(dC)); CHECK(cudaFree(ws));
+    return 0;
+}

cuda/mmq/test/proto_iq2_aligned.cu

@@ -0,0 +1,262 @@
+// proto_iq2_aligned.cu — M1-MoE Inc1 prototype (megakernel program).
+//
+// HYPOTHESIS: the IQ2_XXS decode matvec runs at ~142 GB/s (vs the ~200
+// engine ceiling) because block_iq2_xxs is 66 bytes, so the code stream is
+// only 2-byte aligned and every 32-bit weight word costs two 16-bit loads
+// (get_int_b2).  Repacking into an aligned SoA layout (d[] halves separate,
+// qs[] 64B-aligned per block) allows full-width aligned loads at identical
+// byte count.
+//
+// A/B at the exact production decode shape (per layer, per token):
+//   gate/up experts: M=2048 rows, K=4096, top-6 of n_experts, n_tokens=1
+//   baseline = ds4_mmq_iq2_xxs_moe_vec (the production vec-tier entry)
+//   aligned  = warp-per-row kernel, same dp4a/grid/sign math, SoA weights
+//
+// Correctness: aligned vs baseline output, tolerance (float accum order
+// differs across blocks; per-pair integer math is identical).
+//
+// Build (box):
+//   nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 -I.. \
+//        proto_iq2_aligned.cu ../ds4_mmq.o -lcudart -lcuda -o proto_iq2_aligned
+
+#include "ds4_mmq.h"
+
+#define GGML_COMMON_DECL_CUDA
+#define GGML_COMMON_IMPL_CUDA
+#include "../ggml-common.h"
+
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <random>
+#include <vector>
+
+#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
+    printf("CUDA ERR %s @%d: %s\n", #x, __LINE__, cudaGetErrorString(e_)); exit(1); } } while (0)
+
+namespace {
+
+constexpr int QKK = 256;              // values per IQ2_XXS block
+constexpr int BLOCK_BYTES = 66;       // sizeof(block_iq2_xxs) = 2 + 32*2
+constexpr int PAIRS_PER_BLOCK = 8;    // 8 (q2,aux32) uint32 pairs, 32 values each
+
+// ---------------------------------------------------------------- host q8_1
+struct q8_1_block { __half2 ds; int8_t qs[32]; };   // mirrors block_q8_1
+
+static void quantize_q8_1_host(const float *x, q8_1_block *out, int k) {
+    for (int b = 0; b < k / 32; b++) {
+        float amax = 0.0f;
+        for (int i = 0; i < 32; i++) amax = fmaxf(amax, fabsf(x[b * 32 + i]));
+        const float d = amax / 127.0f;
+        const float id = d != 0.0f ? 1.0f / d : 0.0f;
+        int sum = 0;
+        for (int i = 0; i < 32; i++) {
+            int q = (int)roundf(x[b * 32 + i] * id);
+            if (q > 127) q = 127;
+            if (q < -127) q = -127;
+            out[b].qs[i] = (int8_t)q;
+            sum += q;
+        }
+        out[b].ds = __floats2half2_rn(d, d * (float)sum);
+    }
+}
+
+// ------------------------------------------------------------ device kernel
+__device__ __forceinline__ uint32_t proto_unpack_ksigns(const uint8_t v) {
+    const uint32_t p = __popc(v) & 1;
+    const uint32_t s = v ^ p << 7;
+    return s * 0x01010101;
+}
+
+__device__ __forceinline__ int proto_dp4a(const int a, const int b, int c) {
+    return __dp4a(a, b, c);
+}
+
+// Aligned SoA weights: for (expert e, row r) with nb blocks per row:
+//   dq[(e*M + r)*nb + b]                      : __half block scale
+//   qs[((e*M + r)*nb + b) * 8 + p]            : uint2 (q2, aux32) pair p
+// qs base is 16B-aligned; each block's 8 uint2 = 64B, 64B-aligned stride.
+__global__ void iq2_aligned_moe_vec_kernel(
+        float             *out,        // [n_slots, M]
+        const uint2       *qs,         // aligned code pairs
+        const __half      *dq,         // block scales
+        const q8_1_block  *x8,         // [K/32] quantized activation (shared by slots)
+        const int32_t     *ids,        // [n_slots] expert ids
+        const float       *slot_w,     // [n_slots] router weights (1.0 in proto)
+        int                M,
+        int                nb)         // blocks per row = K/256
+{
+    const int row  = blockIdx.x;
+    const int slot = blockIdx.y;
+    const int lane = threadIdx.x;      // 32 lanes: lane covers (block b, pair p)
+    const long long rbase = ((long long)ids[slot] * M + row) * nb;
+
+    float acc = 0.0f;
+    // 32 lanes cover 4 blocks x 8 pairs per pass.
+    for (int b0 = 0; b0 < nb; b0 += 4) {
+        const int b = b0 + (lane >> 3);
+        const int p = lane & 7;
+        const uint2 cw   = qs[(rbase + b) * PAIRS_PER_BLOCK + p];   // aligned 8B
+        const uint32_t q2 = cw.x, aux32 = cw.y;
+        const uint8_t *aux8 = (const uint8_t *)&q2;
+
+        int sumi = 0;
+        const int q8i = (b * QKK + p * 32) / 32;   // q8_1 block covering these 32 values
+        const int *u = (const int *)x8[q8i].qs;
+#pragma unroll
+        for (int k0 = 0; k0 < 8; k0 += 2) {
+            const uint2 grid_pos = ((const uint2 *)iq2xxs_grid)[aux8[k0 / 2]];
+            const uint32_t signs = proto_unpack_ksigns(aux32 >> (7 * k0 / 2));
+
+            const int signs0 = __vcmpne4(signs & 0x08040201, 0);
+            const int grid0  = __vsub4(grid_pos.x ^ signs0, signs0);
+            sumi = proto_dp4a(grid0, u[k0 + 0], sumi);
+
+            const int signs1 = __vcmpne4(signs & 0x80402010, 0);
+            const int grid1  = __vsub4(grid_pos.y ^ signs1, signs1);
+            sumi = proto_dp4a(grid1, u[k0 + 1], sumi);
+        }
+        const int ls = aux32 >> 27 | 1;
+        sumi = sumi * ls / 8;
+        const float d = __half2float(dq[rbase + b]) * __low2float(x8[q8i].ds);
+        acc += d * (float)sumi;
+    }
+#pragma unroll
+    for (int off = 16; off > 0; off >>= 1)
+        acc += __shfl_down_sync(0xffffffffu, acc, off);
+    if (lane == 0) out[(long long)slot * M + row] = acc * slot_w[slot];
+}
+
+} // namespace
+
+int main(int argc, char **argv) {
+    const int M = 2048, K = 4096;
+    const int n_experts = 256;
+    const int n_slots = 6;
+    const int n_idsets = 32;
+    const int iters = 2000;
+    const int nb = K / QKK;            // 16 blocks/row
+    const long long row_bytes = (long long)nb * BLOCK_BYTES;   // 1056
+    const double slot_gb = (double)n_slots * M * row_bytes / 1e9;
+
+    std::mt19937 rng(argc > 1 ? (uint32_t)atoi(argv[1]) : 1234u);
+
+    // Random IQ2_XXS blocks (same trick as test_mmq_parity: random bytes are
+    // valid IQ2_XXS content; both kernels see identical bits).
+    std::vector<uint8_t> W((size_t)n_experts * M * row_bytes);
+    for (auto &b : W) b = (uint8_t)(rng() & 0xff);
+    // keep block d halves sane (avoid inf/nan): overwrite each block's d
+    for (long long blk = 0; blk < (long long)n_experts * M * nb; blk++) {
+        uint16_t h = (uint16_t)(0x2c00 | (rng() & 0xff));   // ~0.0625..0.1 range fp16
+        memcpy(&W[blk * BLOCK_BYTES], &h, 2);
+    }
+
+    std::vector<float> X(K);
+    for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
+    std::vector<q8_1_block> X8(K / 32);
+    quantize_q8_1_host(X.data(), X8.data(), K);
+
+    std::vector<int32_t> ids(n_idsets * n_slots);
+    for (int s = 0; s < n_idsets; s++)
+        for (int j = 0; j < n_slots; j++)
+            ids[s * n_slots + j] = (s * n_slots + j) % n_experts;
+
+    // Aligned SoA repack
+    const long long nblocks = (long long)n_experts * M * nb;
+    std::vector<uint2>  QS(nblocks * PAIRS_PER_BLOCK);
+    std::vector<__half> DQ(nblocks);
+    for (long long blk = 0; blk < nblocks; blk++) {
+        const uint8_t *src = &W[blk * BLOCK_BYTES];
+        uint16_t h; memcpy(&h, src, 2);
+        DQ[blk] = __ushort_as_half(h);
+        memcpy(&QS[blk * PAIRS_PER_BLOCK], src + 2, 64);
+    }
+
+    // Device buffers
+    uint8_t *dW; float *dX, *dOutBase, *dOutAl; int32_t *dIds; uint2 *dQS; __half *dDQ;
+    q8_1_block *dX8; float *dSlotW;
+    CK(cudaMalloc(&dW, W.size()));
+    CK(cudaMalloc(&dX, sizeof(float) * K * n_slots));
+    CK(cudaMalloc(&dOutBase, sizeof(float) * n_slots * M));
+    CK(cudaMalloc(&dOutAl, sizeof(float) * n_slots * M));
+    CK(cudaMalloc(&dIds, sizeof(int32_t) * ids.size()));
+    CK(cudaMalloc(&dQS, QS.size() * sizeof(uint2)));
+    CK(cudaMalloc(&dDQ, DQ.size() * sizeof(__half)));
+    CK(cudaMalloc(&dX8, X8.size() * sizeof(q8_1_block)));
+    CK(cudaMalloc(&dSlotW, sizeof(float) * n_slots));
+    CK(cudaMemcpy(dW, W.data(), W.size(), cudaMemcpyHostToDevice));
+    // moe_vec's X layout: [n_tokens*n_expert_used rows? or one row]; the vec
+    // entry treats X as [n_tokens, K] with router weights applied later; we
+    // pass the same activation replicated per slot to keep both paths equal.
+    for (int s = 0; s < n_slots; s++)
+        CK(cudaMemcpy(dX + (size_t)s * K, X.data(), sizeof(float) * K, cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dIds, ids.data(), sizeof(int32_t) * ids.size(), cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dQS, QS.data(), QS.size() * sizeof(uint2), cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dDQ, DQ.data(), DQ.size() * sizeof(__half), cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dX8, X8.data(), X8.size() * sizeof(q8_1_block), cudaMemcpyHostToDevice));
+    std::vector<float> slotw(n_slots, 1.0f);
+    CK(cudaMemcpy(dSlotW, slotw.data(), sizeof(float) * n_slots, cudaMemcpyHostToDevice));
+
+    cudaStream_t stream; CK(cudaStreamCreate(&stream));
+    cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
+
+    // ---- baseline: production vec entry -----------------------------------
+    int rc = ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds, dOutBase, M, K, 1, n_experts, n_slots, stream);
+    if (rc != 0) { printf("baseline moe_vec rc=%d\n", rc); return 1; }
+    CK(cudaStreamSynchronize(stream));
+    CK(cudaEventRecord(e0, stream));
+    for (int i = 0; i < iters; i++)
+        (void)ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds + (i % n_idsets) * n_slots, dOutBase,
+                                      M, K, 1, n_experts, n_slots, stream);
+    CK(cudaEventRecord(e1, stream));
+    CK(cudaStreamSynchronize(stream));
+    float ms_base = 0.0f; CK(cudaEventElapsedTime(&ms_base, e0, e1));
+    ms_base /= iters;
+
+    // ---- aligned prototype -------------------------------------------------
+    dim3 grid(M, n_slots, 1);
+    iq2_aligned_moe_vec_kernel<<<grid, 32, 0, stream>>>(dOutAl, dQS, dDQ, dX8, dIds, dSlotW, M, nb);
+    CK(cudaGetLastError());
+    CK(cudaStreamSynchronize(stream));
+    CK(cudaEventRecord(e0, stream));
+    for (int i = 0; i < iters; i++)
+        iq2_aligned_moe_vec_kernel<<<grid, 32, 0, stream>>>(dOutAl, dQS, dDQ, dX8,
+                                                            dIds + (i % n_idsets) * n_slots,
+                                                            dSlotW, M, nb);
+    CK(cudaEventRecord(e1, stream));
+    CK(cudaStreamSynchronize(stream));
+    float ms_al = 0.0f; CK(cudaEventElapsedTime(&ms_al, e0, e1));
+    ms_al /= iters;
+
+    // ---- correctness (vs baseline, last idset used by both loops) ----------
+    // rerun both once on idset 0 for a clean compare
+    (void)ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds, dOutBase, M, K, 1, n_experts, n_slots, stream);
+    iq2_aligned_moe_vec_kernel<<<grid, 32, 0, stream>>>(dOutAl, dQS, dDQ, dX8, dIds, dSlotW, M, nb);
+    CK(cudaStreamSynchronize(stream));
+    std::vector<float> ob(n_slots * M), oa(n_slots * M);
+    CK(cudaMemcpy(ob.data(), dOutBase, sizeof(float) * ob.size(), cudaMemcpyDeviceToHost));
+    CK(cudaMemcpy(oa.data(), dOutAl, sizeof(float) * oa.size(), cudaMemcpyDeviceToHost));
+    double max_rel = 0.0, max_abs = 0.0; int bad = 0;
+    for (size_t i = 0; i < ob.size(); i++) {
+        const double a = ob[i], b = oa[i];
+        const double ad = fabs(a - b);
+        const double rd = ad / (fabs(a) > 1e-3 ? fabs(a) : 1e-3);
+        if (ad > max_abs) max_abs = ad;
+        if (rd > max_rel) max_rel = rd;
+        if (rd > 2e-2 && ad > 1e-2) bad++;
+    }
+
+    printf("PROTO_IQ2_ALIGNED  M=%d K=%d slots=%d experts=%d iters=%d (weights/iter %.1f MB)\n",
+           M, K, n_slots, n_experts, iters, slot_gb * 1000.0);
+    printf("  baseline moe_vec : %.4f ms  -> %6.1f GB/s\n", ms_base, slot_gb / (ms_base / 1e3));
+    printf("  aligned  proto   : %.4f ms  -> %6.1f GB/s   (%+.1f%%)\n",
+           ms_al, slot_gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0));
+    printf("  parity: max_rel=%.3e max_abs=%.3e bad=%d -> %s\n",
+           max_rel, max_abs, bad, bad == 0 ? "PASS" : "FAIL");
+    return bad == 0 ? 0 : 2;
+}

cuda/mmq/test/proto_kv_reencode_idem.cu

@@ -0,0 +1,858 @@
+// SPDX-License-Identifier: MIT
+// proto_kv_reencode_idem.cu - packed-KV restore re-encode idempotence probe.
+//
+// v0.2.4 ships REFUSED packed disk-KV restores: the restore path reads the
+// F32 committed rows from the file and regenerates the fp8/fp4 mirrors by
+// re-encoding, and the kv_crossmode gate observed the resulting packed-server
+// continuations DRIFTING (fp8-only and fp4-only independently) despite the
+// pow2-scale idempotence argument saying re-encode of committed values is
+// value-exact.  This proto isolates the kernel math from the pipeline:
+//   fp8 : quantize(x) -> C1,codes1,scales1 ; dequant==C1 ;
+//         quantize(C1) -> C2,codes2,scales2 ; C2 ?= C1 (bitwise) ; dequant==C1
+//   fp4 : hadamard_qat(x) -> C1 ; fp4_dequant==C1 ;
+//         noHad_encode(C1) -> C2 ; C2 ?= C1 ; dequant==C1
+// plus a third-iteration fixed-point check and a fast-math boundary scan of
+// scale = exp2f(ceilf(log2f(amax/448 or /6))) at exact power-of-two inputs
+// (the one corner where fast-math approximation could flip ceilf and shift
+// the whole block one exponent).
+// Code/scale mismatches with VALUE identity are tallied separately: a block
+// whose committed max sits exactly on a half-boundary legitimately re-picks
+// scale/2 with all codes shifted one exponent, values bit-identical.
+//
+// The first battery runs verbatim copies of the v0.2.4 (PRE-fix) ds4_cuda.cu
+// bodies and is EXPECTED to fail -- kept as the regression repro.  The _fx
+// battery is the fix as shipped in v0.3 (dsv4_pow2_ceil_scale + signbit) and
+// is the gate (exit status).  Build MUST mirror the engine's fast-math flags:
+//   /usr/local/cuda/bin/nvcc -O3 -std=c++17 --use_fast_math -arch=sm_121 \
+//       cuda/mmq/test/proto_kv_reencode_idem.cu -o proto_kv_reencode_idem
+// argv: [n_rows_fp8=131072] [n_rows_fp4=131072]
+
+#include <cuda_runtime.h>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <cmath>
+#include <cstdint>
+
+#define CHECK(call) do { cudaError_t e_=(call); if(e_!=cudaSuccess){ \
+    fprintf(stderr,"CUDA %s @%d: %s\n",#call,__LINE__,cudaGetErrorString(e_)); exit(1);} } while(0)
+
+// ---------------------------------------------------------------------------
+// Verbatim from ds4_cuda.cu (v0.2.4 tag content).
+// ---------------------------------------------------------------------------
+
+#define DS4_OPP_C_FP8_NOPE_DEV       448u    /* 512 - 64 */
+#define DS4_OPP_C_FP8_BLOCKS_DEV     7u      /* 448 / 64  */
+#define DS4_OPP_C_FP8_ROW_BYTES_DEV  704u    /* 448 codes + 64 FP32 tail */
+
+__device__ static float dsv4_e4m3fn_value_dev(int i) {
+    int exp = (i >> 3) & 15;
+    int mant = i & 7;
+    if (exp == 0) return (float)mant * 0.001953125f;
+    return (1.0f + (float)mant * 0.125f) * exp2f((float)exp - 7.0f);
+}
+
+__device__ static float dsv4_e4m3fn_dequant_dev(float x) {
+    float sign = x < 0.0f ? -1.0f : 1.0f;
+    float ax = fminf(fabsf(x), 448.0f);
+    int lo = 0, hi = 126;
+    while (lo < hi) {
+        int mid = (lo + hi + 1) >> 1;
+        if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
+        else hi = mid - 1;
+    }
+    int best = lo;
+    if (best < 126) {
+        float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
+        float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
+        if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
+    }
+    return sign * dsv4_e4m3fn_value_dev(best);
+}
+
+__device__ static unsigned char dsv4_e4m3fn_encode_dev(float x) {
+    const unsigned int sign_bit = (x < 0.0f) ? 0x80u : 0x00u;
+    const float ax = fminf(fabsf(x), 448.0f);
+    int lo = 0, hi = 126;
+    while (lo < hi) {
+        const int mid = (lo + hi + 1) >> 1;
+        if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
+        else hi = mid - 1;
+    }
+    int best = lo;
+    if (best < 126) {
+        const float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
+        const float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
+        if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
+    }
+    return (unsigned char)(sign_bit | (unsigned int)best);
+}
+
+__device__ float dsv4_e4m3fn_decode_table[128];
+
+static int fp8_decode_table_init(void) {
+    float table[128];
+    for (int i = 0; i < 128; i++) {
+        const int exp  = (i >> 3) & 15;
+        const int mant = i & 7;
+        table[i] = (exp == 0)
+            ? (float)mant * 0.001953125f
+            : (1.0f + (float)mant * 0.125f) * ldexpf(1.0f, exp - 7);
+    }
+    return cudaMemcpyToSymbol(dsv4_e4m3fn_decode_table, table,
+                              sizeof(table), 0, cudaMemcpyHostToDevice) == cudaSuccess;
+}
+
+__device__ static float fp8_kv_read(
+        const unsigned char * __restrict__ codes_base,
+        const float         * __restrict__ scale_base,
+        uint32_t                          c,
+        uint32_t                          d) {
+    const uint64_t row_off = (uint64_t)c * DS4_OPP_C_FP8_ROW_BYTES_DEV;
+    if (d < DS4_OPP_C_FP8_NOPE_DEV) {
+        const unsigned char code = codes_base[row_off + d];
+        const float scale = scale_base[(uint64_t)c * DS4_OPP_C_FP8_BLOCKS_DEV + (d >> 6)];
+        const float mag   = dsv4_e4m3fn_decode_table[code & 0x7fu];
+        return ((code & 0x80u) ? -mag : mag) * scale;
+    }
+    const float *tail = (const float *)(codes_base + row_off + DS4_OPP_C_FP8_NOPE_DEV);
+    return tail[d - DS4_OPP_C_FP8_NOPE_DEV];
+}
+
+__global__ static void fp8_kv_dequant_rows_kernel(
+        float * __restrict__ dst,
+        const unsigned char * __restrict__ codes_base,
+        const float         * __restrict__ scale_base,
+        uint32_t n_rows,
+        uint32_t head_dim) {
+    const uint32_t r = blockIdx.x;
+    if (r >= n_rows) return;
+    float *out = dst + (uint64_t)r * head_dim;
+    for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
+        out[d] = fp8_kv_read(codes_base, scale_base, r, d);
+    }
+}
+
+__global__ static void fp8_kv_quantize_kernel(
+        float *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot,
+        unsigned char * __restrict__ codes_base,
+        float * __restrict__ scale_base) {
+    uint32_t row = blockIdx.x;
+    uint32_t tid = threadIdx.x;
+    uint32_t n_nope = head_dim - n_rot;
+    const uint64_t row_stride = (uint64_t)n_nope + (uint64_t)n_rot * sizeof(float);
+    const uint64_t scale_stride = (uint64_t)(n_nope / 64u);
+    float *xr = x + (uint64_t)row * head_dim;
+    unsigned char *codes_row = codes_base
+        ? codes_base + (uint64_t)row * row_stride
+        : (unsigned char *)0;
+    float *scale_row = scale_base
+        ? scale_base + (uint64_t)row * scale_stride
+        : (float *)0;
+    __shared__ float scratch[64];
+    for (uint32_t off = 0; off < n_nope; off += 64) {
+        float v = 0.0f;
+        if (off + tid < n_nope) v = xr[off + tid];
+        scratch[tid] = off + tid < n_nope ? fabsf(v) : 0.0f;
+        __syncthreads();
+        for (uint32_t stride = 32; stride > 0; stride >>= 1) {
+            if (tid < stride) scratch[tid] = fmaxf(scratch[tid], scratch[tid + stride]);
+            __syncthreads();
+        }
+        float scale = exp2f(ceilf(log2f(fmaxf(scratch[0], 1.0e-4f) / 448.0f)));
+        if (off + tid < n_nope) {
+            float clamp = fminf(448.0f, fmaxf(-448.0f, v / scale));
+            xr[off + tid] = dsv4_e4m3fn_dequant_dev(clamp) * scale;
+            if (codes_row) codes_row[off + tid] = dsv4_e4m3fn_encode_dev(clamp);
+        }
+        if (scale_row && tid == 0) scale_row[off / 64u] = scale;
+        __syncthreads();
+    }
+    if (codes_row && tid < n_rot) {
+        float *tail = (float *)(codes_row + (uint64_t)n_nope);
+        tail[tid] = xr[(uint64_t)n_nope + tid];
+    }
+}
+
+__device__ static float dsv4_e2m1fn_value_dev(int i) {
+    switch (i & 7) {
+    case 0: return 0.0f;
+    case 1: return 0.5f;
+    case 2: return 1.0f;
+    case 3: return 1.5f;
+    case 4: return 2.0f;
+    case 5: return 3.0f;
+    case 6: return 4.0f;
+    default: return 6.0f;
+    }
+}
+
+__device__ static int dsv4_e2m1fn_level_dev(float ax) {
+    ax = fminf(ax, 6.0f);
+    int best = 0;
+    float best_diff = fabsf(ax - dsv4_e2m1fn_value_dev(0));
+    for (int i = 1; i < 8; i++) {
+        float diff = fabsf(ax - dsv4_e2m1fn_value_dev(i));
+        if (diff < best_diff || (diff == best_diff && ((i & 1) == 0) && ((best & 1) != 0))) {
+            best = i;
+            best_diff = diff;
+        }
+    }
+    return best;
+}
+
+__device__ static float indexer_fp4_read(
+        const unsigned char * __restrict__ codes_base,
+        const float         * __restrict__ scale_base,
+        uint32_t row, uint32_t d) {
+    const unsigned char byte = codes_base[(uint64_t)row * 64u + (d >> 1u)];
+    const uint32_t nib = ((uint32_t)byte >> (4u * (d & 1u))) & 0xFu;
+    const float sign = (nib & 8u) ? -1.0f : 1.0f;
+    const float scale = scale_base[(uint64_t)row * 4u + (d >> 5u)];
+    return sign * dsv4_e2m1fn_value_dev((int)(nib & 7u)) * scale;
+}
+
+__global__ static void indexer_fp4_dequant_rows_kernel(
+        float * __restrict__ dst,
+        const unsigned char * __restrict__ codes_base,
+        const float         * __restrict__ scale_base,
+        uint32_t n_rows, uint32_t head_dim) {
+    const uint32_t r = blockIdx.x;
+    if (r >= n_rows) return;
+    float *out = dst + (uint64_t)r * head_dim;
+    for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
+        out[d] = indexer_fp4_read(codes_base, scale_base, r, d);
+    }
+}
+
+__global__ static void indexer_fp4_encode_rows_kernel(
+        float * __restrict__ x, uint32_t n_rows, uint32_t head_dim,
+        unsigned char * __restrict__ codes_base, float * __restrict__ scale_base) {
+    uint32_t row = blockIdx.x;
+    uint32_t tid = threadIdx.x;
+    if (row >= n_rows || head_dim != 128u || tid >= 128u) return;
+    __shared__ float absbuf[128];
+    float *xr = x + (uint64_t)row * head_dim;
+    float v = xr[tid];                       /* committed value -- NO Hadamard */
+    uint32_t fp4_block = tid >> 5u;
+    uint32_t lane = tid & 31u;
+    uint32_t block_base = fp4_block * 32u;
+    absbuf[tid] = fabsf(v);
+    __syncthreads();
+    for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
+        if (lane < stride)
+            absbuf[block_base + lane] = fmaxf(absbuf[block_base + lane],
+                                              absbuf[block_base + lane + stride]);
+        __syncthreads();
+    }
+    float amax = fmaxf(absbuf[block_base], 7.052966104933725e-38f);
+    float scale = exp2f(ceilf(log2f(amax / 6.0f)));
+    float clamp = fminf(6.0f, fmaxf(-6.0f, v / scale));
+    int level = dsv4_e2m1fn_level_dev(fabsf(clamp));
+    xr[tid] = ((clamp < 0.0f ? -1.0f : 1.0f) * dsv4_e2m1fn_value_dev(level)) * scale;
+    uint32_t nib = (clamp < 0.0f ? 8u : 0u) | (uint32_t)level;
+    uint32_t hi  = __shfl_down_sync(0xffffffffu, nib, 1u);
+    if ((tid & 1u) == 0u)  codes_base[(uint64_t)row * 64u + (tid >> 1u)] = (unsigned char)(nib | (hi << 4u));
+    if ((tid & 31u) == 0u) scale_base[(uint64_t)row * 4u + (tid >> 5u)] = scale;
+}
+
+__global__ static void indexer_hadamard_fp4_kernel(float *x, uint32_t n_rows, uint32_t head_dim,
+        unsigned char * __restrict__ codes_base,
+        float         * __restrict__ scale_base) {
+    uint32_t row = blockIdx.x;
+    uint32_t tid = threadIdx.x;
+    if (row >= n_rows || head_dim != 128u || tid >= 128u) return;
+
+    __shared__ float vals[128];
+    __shared__ float absbuf[128];
+    float *xr = x + (uint64_t)row * head_dim;
+    vals[tid] = xr[tid];
+    __syncthreads();
+
+    for (uint32_t stride = 1u; stride < 128u; stride <<= 1u) {
+        if ((tid & stride) == 0u) {
+            uint32_t base = (tid & ~(2u * stride - 1u)) + (tid & (stride - 1u));
+            float a = vals[base];
+            float b = vals[base + stride];
+            vals[base] = a + b;
+            vals[base + stride] = a - b;
+        }
+        __syncthreads();
+    }
+
+    float v = vals[tid] * 0.08838834764831845f;
+    uint32_t fp4_block = tid >> 5u;
+    uint32_t lane = tid & 31u;
+    uint32_t block_base = fp4_block * 32u;
+    absbuf[tid] = fabsf(v);
+    __syncthreads();
+
+    for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
+        if (lane < stride) {
+            absbuf[block_base + lane] = fmaxf(absbuf[block_base + lane],
+                                              absbuf[block_base + lane + stride]);
+        }
+        __syncthreads();
+    }
+
+    float amax = fmaxf(absbuf[block_base], 7.052966104933725e-38f);
+    float scale = exp2f(ceilf(log2f(amax / 6.0f)));
+    float clamp = fminf(6.0f, fmaxf(-6.0f, v / scale));
+    int level = dsv4_e2m1fn_level_dev(fabsf(clamp));
+    xr[tid] = ((clamp < 0.0f ? -1.0f : 1.0f) * dsv4_e2m1fn_value_dev(level)) * scale;
+    if (codes_base) {
+        uint32_t nib = (clamp < 0.0f ? 8u : 0u) | (uint32_t)level;
+        uint32_t hi  = __shfl_down_sync(0xffffffffu, nib, 1u);
+        if ((tid & 1u) == 0u)  codes_base[(uint64_t)row * 64u + (tid >> 1u)] = (unsigned char)(nib | (hi << 4u));
+        if ((tid & 31u) == 0u) scale_base[(uint64_t)row * 4u + (tid >> 5u)] = scale;
+    }
+}
+
+// ---------------------------------------------------------------------------
+// FIXED variants (candidate engine change): exact pow2-ceil scale via frexpf
+// (bit ops, immune to fast-math lg2.approx one-sided error at 2^-n) and
+// signbit-based sign handling (preserves -0.0 across the round trip).
+// ---------------------------------------------------------------------------
+
+__device__ static float dsv4_pow2_ceil_scale(float r) {
+    int e; float m = frexpf(r, &e);          /* r = m * 2^e, m in [0.5, 1) */
+    return ldexpf(1.0f, m == 0.5f ? e - 1 : e);
+}
+
+__device__ static float dsv4_e4m3fn_dequant_fx(float x) {
+    float sign = signbit(x) ? -1.0f : 1.0f;
+    float ax = fminf(fabsf(x), 448.0f);
+    int lo = 0, hi = 126;
+    while (lo < hi) {
+        int mid = (lo + hi + 1) >> 1;
+        if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
+        else hi = mid - 1;
+    }
+    int best = lo;
+    if (best < 126) {
+        float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
+        float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
+        if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
+    }
+    return sign * dsv4_e4m3fn_value_dev(best);
+}
+
+__device__ static unsigned char dsv4_e4m3fn_encode_fx(float x) {
+    const unsigned int sign_bit = signbit(x) ? 0x80u : 0x00u;
+    const float ax = fminf(fabsf(x), 448.0f);
+    int lo = 0, hi = 126;
+    while (lo < hi) {
+        const int mid = (lo + hi + 1) >> 1;
+        if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
+        else hi = mid - 1;
+    }
+    int best = lo;
+    if (best < 126) {
+        const float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
+        const float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
+        if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
+    }
+    return (unsigned char)(sign_bit | (unsigned int)best);
+}
+
+__global__ static void fp8_kv_quantize_kernel_fx(
+        float *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot,
+        unsigned char * __restrict__ codes_base,
+        float * __restrict__ scale_base) {
+    uint32_t row = blockIdx.x;
+    uint32_t tid = threadIdx.x;
+    uint32_t n_nope = head_dim - n_rot;
+    const uint64_t row_stride = (uint64_t)n_nope + (uint64_t)n_rot * sizeof(float);
+    const uint64_t scale_stride = (uint64_t)(n_nope / 64u);
+    float *xr = x + (uint64_t)row * head_dim;
+    unsigned char *codes_row = codes_base
+        ? codes_base + (uint64_t)row * row_stride
+        : (unsigned char *)0;
+    float *scale_row = scale_base
+        ? scale_base + (uint64_t)row * scale_stride
+        : (float *)0;
+    __shared__ float scratch[64];
+    for (uint32_t off = 0; off < n_nope; off += 64) {
+        float v = 0.0f;
+        if (off + tid < n_nope) v = xr[off + tid];
+        scratch[tid] = off + tid < n_nope ? fabsf(v) : 0.0f;
+        __syncthreads();
+        for (uint32_t stride = 32; stride > 0; stride >>= 1) {
+            if (tid < stride) scratch[tid] = fmaxf(scratch[tid], scratch[tid + stride]);
+            __syncthreads();
+        }
+        float scale = dsv4_pow2_ceil_scale(fmaxf(scratch[0], 1.0e-4f) / 448.0f);
+        if (off + tid < n_nope) {
+            float clamp = fminf(448.0f, fmaxf(-448.0f, v / scale));
+            xr[off + tid] = dsv4_e4m3fn_dequant_fx(clamp) * scale;
+            if (codes_row) codes_row[off + tid] = dsv4_e4m3fn_encode_fx(clamp);
+        }
+        if (scale_row && tid == 0) scale_row[off / 64u] = scale;
+        __syncthreads();
+    }
+    if (codes_row && tid < n_rot) {
+        float *tail = (float *)(codes_row + (uint64_t)n_nope);
+        tail[tid] = xr[(uint64_t)n_nope + tid];
+    }
+}
+
+__global__ static void indexer_fp4_encode_rows_kernel_fx(
+        float * __restrict__ x, uint32_t n_rows, uint32_t head_dim,
+        unsigned char * __restrict__ codes_base, float * __restrict__ scale_base) {
+    uint32_t row = blockIdx.x;
+    uint32_t tid = threadIdx.x;
+    if (row >= n_rows || head_dim != 128u || tid >= 128u) return;
+    __shared__ float absbuf[128];
+    float *xr = x + (uint64_t)row * head_dim;
+    float v = xr[tid];
+    uint32_t fp4_block = tid >> 5u;
+    uint32_t lane = tid & 31u;
+    uint32_t block_base = fp4_block * 32u;
+    absbuf[tid] = fabsf(v);
+    __syncthreads();
+    for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
+        if (lane < stride)
+            absbuf[block_base + lane] = fmaxf(absbuf[block_base + lane],
+                                              absbuf[block_base + lane + stride]);
+        __syncthreads();
+    }
+    float amax = fmaxf(absbuf[block_base], 7.052966104933725e-38f);
+    float scale = dsv4_pow2_ceil_scale(amax / 6.0f);
+    float clamp = fminf(6.0f, fmaxf(-6.0f, v / scale));
+    int level = dsv4_e2m1fn_level_dev(fabsf(clamp));
+    xr[tid] = ((signbit(clamp) ? -1.0f : 1.0f) * dsv4_e2m1fn_value_dev(level)) * scale;
+    uint32_t nib = (signbit(clamp) ? 8u : 0u) | (uint32_t)level;
+    uint32_t hi  = __shfl_down_sync(0xffffffffu, nib, 1u);
+    if ((tid & 1u) == 0u)  codes_base[(uint64_t)row * 64u + (tid >> 1u)] = (unsigned char)(nib | (hi << 4u));
+    if ((tid & 31u) == 0u) scale_base[(uint64_t)row * 4u + (tid >> 5u)] = scale;
+}
+
+__global__ static void indexer_hadamard_fp4_kernel_fx(float *x, uint32_t n_rows, uint32_t head_dim,
+        unsigned char * __restrict__ codes_base,
+        float         * __restrict__ scale_base) {
+    uint32_t row = blockIdx.x;
+    uint32_t tid = threadIdx.x;
+    if (row >= n_rows || head_dim != 128u || tid >= 128u) return;
+
+    __shared__ float vals[128];
+    __shared__ float absbuf[128];
+    float *xr = x + (uint64_t)row * head_dim;
+    vals[tid] = xr[tid];
+    __syncthreads();
+
+    for (uint32_t stride = 1u; stride < 128u; stride <<= 1u) {
+        if ((tid & stride) == 0u) {
+            uint32_t base = (tid & ~(2u * stride - 1u)) + (tid & (stride - 1u));
+            float a = vals[base];
+            float b = vals[base + stride];
+            vals[base] = a + b;
+            vals[base + stride] = a - b;
+        }
+        __syncthreads();
+    }
+
+    float v = vals[tid] * 0.08838834764831845f;
+    uint32_t fp4_block = tid >> 5u;
+    uint32_t lane = tid & 31u;
+    uint32_t block_base = fp4_block * 32u;
+    absbuf[tid] = fabsf(v);
+    __syncthreads();
+
+    for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
+        if (lane < stride) {
+            absbuf[block_base + lane] = fmaxf(absbuf[block_base + lane],
+                                              absbuf[block_base + lane + stride]);
+        }
+        __syncthreads();
+    }
+
+    float amax = fmaxf(absbuf[block_base], 7.052966104933725e-38f);
+    float scale = dsv4_pow2_ceil_scale(amax / 6.0f);
+    float clamp = fminf(6.0f, fmaxf(-6.0f, v / scale));
+    int level = dsv4_e2m1fn_level_dev(fabsf(clamp));
+    xr[tid] = ((signbit(clamp) ? -1.0f : 1.0f) * dsv4_e2m1fn_value_dev(level)) * scale;
+    if (codes_base) {
+        uint32_t nib = (signbit(clamp) ? 8u : 0u) | (uint32_t)level;
+        uint32_t hi  = __shfl_down_sync(0xffffffffu, nib, 1u);
+        if ((tid & 1u) == 0u)  codes_base[(uint64_t)row * 64u + (tid >> 1u)] = (unsigned char)(nib | (hi << 4u));
+        if ((tid & 31u) == 0u) scale_base[(uint64_t)row * 4u + (tid >> 5u)] = scale;
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Proto scaffolding (not engine code).
+// ---------------------------------------------------------------------------
+
+// Deterministic per-element fill: xorshift hash of (row,lane) -> [-1,1),
+// scaled per row-band regime.  Bands (of 8) for the fp8 input:
+//   0 uniform | 1 x1e-5 | 2 x1e4 | 3 pow2 magnitude ladder 2^((row%40)-20)
+//   4 block amax forced EXACTLY 448*2^k   (k=(row%40)-20; lane0 of block)
+//   5 block amax forced EXACTLY 224*2^k   (post-encode halving boundary)
+//   6 block amax 433*2^k (e4m3 rounds it UP to 448*2^k -> boundary amax')
+//   7 tiny/zero: x1e-6, every 3rd row all-zero (sub-1e-4-floor blocks)
+// Exact 2^k as a float, immune to any exp2f approximation question.
+__device__ static float exact_pow2(int k) {
+    return __uint_as_float((uint32_t)(127 + k) << 23);
+}
+
+__device__ static float hash_unit(uint32_t row, uint32_t lane, uint32_t salt) {
+    uint32_t h = row * 2654435761u ^ (lane + 1u) * 2246822519u ^ salt * 3266489917u;
+    h ^= h >> 15; h *= 2246822519u; h ^= h >> 13; h *= 3266489917u; h ^= h >> 16;
+    return ((float)(h & 0xFFFFFFu) / 8388608.0f) - 1.0f;   // [-1,1)
+}
+
+__global__ static void fill_fp8_input_kernel(float *x, uint32_t n_rows, uint32_t salt) {
+    uint32_t row = blockIdx.x;
+    uint32_t tid = threadIdx.x;          // 512 threads
+    if (row >= n_rows || tid >= 512u) return;
+    uint32_t band = row / (n_rows / 8u + 1u);
+    float u = hash_unit(row, tid, salt);
+    float v = u;
+    int k = (int)(row % 40u) - 20;
+    switch (band) {
+    case 0: break;
+    case 1: v = u * 1e-5f; break;
+    case 2: v = u * 1e4f; break;
+    case 3: v = u * exact_pow2(k); break;
+    case 4: v = (tid < 448u && (tid & 63u) == 0u) ? 448.0f * exact_pow2(k)
+                                                  : u * 100.0f * exact_pow2(k); break;
+    case 5: v = (tid < 448u && (tid & 63u) == 0u) ? 224.0f * exact_pow2(k)
+                                                  : u * 100.0f * exact_pow2(k); break;
+    case 6: v = (tid < 448u && (tid & 63u) == 0u) ? 433.0f * exact_pow2(k)
+                                                  : u * 100.0f * exact_pow2(k); break;
+    default:
+        v = (row % 3u == 0u) ? 0.0f : u * 1e-6f;
+        break;
+    }
+    x[(uint64_t)row * 512u + tid] = v;
+}
+
+__global__ static void fill_fp4_input_kernel(float *x, uint32_t n_rows, uint32_t salt) {
+    uint32_t row = blockIdx.x;
+    uint32_t tid = threadIdx.x;          // 128 threads
+    if (row >= n_rows || tid >= 128u) return;
+    uint32_t band = row / (n_rows / 6u + 1u);
+    float u = hash_unit(row, tid, salt ^ 0x9e3779b9u);
+    float v = u;
+    int k = (int)(row % 60u) - 30;
+    switch (band) {
+    case 0: break;
+    case 1: v = u * 1e-6f; break;
+    case 2: v = u * 1e5f; break;
+    case 3: v = u * exact_pow2(k); break;
+    case 4: v = (u > 0.0f ? 1.0f : -1.0f) * exact_pow2(k); break;  // pure pow2 lanes
+    default:
+        v = (row % 3u == 0u) ? 0.0f : u * 1e-30f;
+        break;
+    }
+    x[(uint64_t)row * 128u + tid] = v;
+}
+
+struct mismatch_rec { uint32_t idx_hi, idx_lo, a_bits, b_bits; };
+#define MAX_RECS 8
+
+__global__ static void compare_f32_kernel(const float *a, const float *b, uint64_t n,
+                                          unsigned long long *count,
+                                          mismatch_rec *recs, int *nrec) {
+    uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    uint32_t ab = __float_as_uint(a[i]);
+    uint32_t bb = __float_as_uint(b[i]);
+    if (ab != bb) {
+        atomicAdd(count, 1ull);
+        int slot = atomicAdd(nrec, 1);
+        if (slot < MAX_RECS) {
+            recs[slot].idx_hi = (uint32_t)(i >> 32);
+            recs[slot].idx_lo = (uint32_t)(i & 0xFFFFFFFFu);
+            recs[slot].a_bits = ab;
+            recs[slot].b_bits = bb;
+        }
+    }
+}
+
+__global__ static void compare_u8_kernel(const unsigned char *a, const unsigned char *b,
+                                         uint64_t n, unsigned long long *count) {
+    uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    if (a[i] != b[i]) atomicAdd(count, 1ull);
+}
+
+// Fast-math boundary scan: the exact-pow2 inputs the re-encode presents when
+// a block's committed max lands on 448*2^k (fp8) or {6,3,1.5}*2^k (fp4).
+// out[j] = scale computed by the shipped formula; host prints deviations.
+__global__ static void boundary_scan_kernel(float *out) {
+    int idx = threadIdx.x;               // 61 k values: -30..30
+    if (idx >= 61) return;
+    int k = idx - 30;
+    float p2k = exact_pow2(k);
+    // fp8 formula at amax = 448*2^k, 224*2^k, and one-ulp neighbours of 448*2^k
+    float a448  = 448.0f * p2k;
+    float a224  = 224.0f * p2k;
+    float a448u = __uint_as_float(__float_as_uint(a448) + 1u);
+    float a448d = __uint_as_float(__float_as_uint(a448) - 1u);
+    out[idx*8 + 0] = exp2f(ceilf(log2f(fmaxf(a448,  1.0e-4f) / 448.0f)));
+    out[idx*8 + 1] = exp2f(ceilf(log2f(fmaxf(a224,  1.0e-4f) / 448.0f)));
+    out[idx*8 + 2] = exp2f(ceilf(log2f(fmaxf(a448u, 1.0e-4f) / 448.0f)));
+    out[idx*8 + 3] = exp2f(ceilf(log2f(fmaxf(a448d, 1.0e-4f) / 448.0f)));
+    // fp4 formula at amax = 6*2^k, 3*2^k, 1.5*2^k, one-ulp above 6*2^k
+    float a6  = 6.0f * p2k;
+    float a3  = 3.0f * p2k;
+    float a15 = 1.5f * p2k;
+    float a6u = __uint_as_float(__float_as_uint(a6) + 1u);
+    out[idx*8 + 4] = exp2f(ceilf(log2f(fmaxf(a6,  7.052966104933725e-38f) / 6.0f)));
+    out[idx*8 + 5] = exp2f(ceilf(log2f(fmaxf(a3,  7.052966104933725e-38f) / 6.0f)));
+    out[idx*8 + 6] = exp2f(ceilf(log2f(fmaxf(a15, 7.052966104933725e-38f) / 6.0f)));
+    out[idx*8 + 7] = exp2f(ceilf(log2f(fmaxf(a6u, 7.052966104933725e-38f) / 6.0f)));
+}
+
+static unsigned long long report_f32(const char *tag, const float *a, const float *b,
+                                     uint64_t n, uint32_t row_dim) {
+    unsigned long long *d_count; mismatch_rec *d_recs; int *d_nrec;
+    CHECK(cudaMalloc(&d_count, sizeof(*d_count)));
+    CHECK(cudaMalloc(&d_recs, MAX_RECS * sizeof(mismatch_rec)));
+    CHECK(cudaMalloc(&d_nrec, sizeof(int)));
+    CHECK(cudaMemset(d_count, 0, sizeof(*d_count)));
+    CHECK(cudaMemset(d_nrec, 0, sizeof(int)));
+    uint64_t blocks = (n + 255) / 256;
+    compare_f32_kernel<<<(uint32_t)blocks, 256>>>(a, b, n, d_count, d_recs, d_nrec);
+    CHECK(cudaDeviceSynchronize());
+    unsigned long long count; mismatch_rec recs[MAX_RECS]; int nrec;
+    CHECK(cudaMemcpy(&count, d_count, sizeof(count), cudaMemcpyDeviceToHost));
+    CHECK(cudaMemcpy(recs, d_recs, sizeof(recs), cudaMemcpyDeviceToHost));
+    CHECK(cudaMemcpy(&nrec, d_nrec, sizeof(nrec), cudaMemcpyDeviceToHost));
+    printf("  %-28s mismatches=%llu / %llu\n", tag, count, (unsigned long long)n);
+    for (int i = 0; i < nrec && i < MAX_RECS; i++) {
+        uint64_t idx = ((uint64_t)recs[i].idx_hi << 32) | recs[i].idx_lo;
+        float fa, fb;
+        memcpy(&fa, &recs[i].a_bits, 4); memcpy(&fb, &recs[i].b_bits, 4);
+        printf("    [%d] row=%llu lane=%u  a=%08x(%.9g)  b=%08x(%.9g)\n",
+               i, (unsigned long long)(idx / row_dim), (uint32_t)(idx % row_dim),
+               recs[i].a_bits, fa, recs[i].b_bits, fb);
+    }
+    cudaFree(d_count); cudaFree(d_recs); cudaFree(d_nrec);
+    return count;
+}
+
+static unsigned long long report_u8(const char *tag, const unsigned char *a,
+                                    const unsigned char *b, uint64_t n) {
+    unsigned long long *d_count;
+    CHECK(cudaMalloc(&d_count, sizeof(*d_count)));
+    CHECK(cudaMemset(d_count, 0, sizeof(*d_count)));
+    uint64_t blocks = (n + 255) / 256;
+    compare_u8_kernel<<<(uint32_t)blocks, 256>>>(a, b, n, d_count);
+    CHECK(cudaDeviceSynchronize());
+    unsigned long long count;
+    CHECK(cudaMemcpy(&count, d_count, sizeof(count), cudaMemcpyDeviceToHost));
+    printf("  %-28s byte-mismatches=%llu / %llu (level-shift benign iff values match)\n",
+           tag, count, (unsigned long long)n);
+    cudaFree(d_count);
+    return count;
+}
+
+int main(int argc, char **argv) {
+    uint32_t n_fp8 = argc > 1 ? (uint32_t)atoi(argv[1]) : 131072u;
+    uint32_t n_fp4 = argc > 2 ? (uint32_t)atoi(argv[2]) : 131072u;
+    if (!fp8_decode_table_init()) { fprintf(stderr, "decode table init failed\n"); return 1; }
+
+    int fails = 0;      /* FIXED-variant failures: the actual gate */
+    int orig_fail = 0;  /* shipped-kernel failures: expected repro of the bug */
+
+    // ---------------- fp8 comp path ----------------
+    {
+        const uint32_t HD = 512u, NROT = 64u;
+        uint64_t xn = (uint64_t)n_fp8 * HD;
+        uint64_t cb = (uint64_t)n_fp8 * DS4_OPP_C_FP8_ROW_BYTES_DEV;
+        uint64_t sn = (uint64_t)n_fp8 * DS4_OPP_C_FP8_BLOCKS_DEV;
+        float *x, *d1, *x2, *d2, *s1, *s2;
+        unsigned char *c1, *c2;
+        CHECK(cudaMalloc(&x,  xn * 4)); CHECK(cudaMalloc(&d1, xn * 4));
+        CHECK(cudaMalloc(&x2, xn * 4)); CHECK(cudaMalloc(&d2, xn * 4));
+        CHECK(cudaMalloc(&c1, cb)); CHECK(cudaMalloc(&c2, cb));
+        CHECK(cudaMalloc(&s1, sn * 4)); CHECK(cudaMalloc(&s2, sn * 4));
+
+        fill_fp8_input_kernel<<<n_fp8, 512>>>(x, n_fp8, 7u);
+        fp8_kv_quantize_kernel<<<n_fp8, 64>>>(x, n_fp8, HD, NROT, c1, s1);      // -> C1
+        fp8_kv_dequant_rows_kernel<<<n_fp8, 128>>>(d1, c1, s1, n_fp8, HD);      // D1
+        CHECK(cudaMemcpy(x2, x, xn * 4, cudaMemcpyDeviceToDevice));
+        fp8_kv_quantize_kernel<<<n_fp8, 64>>>(x2, n_fp8, HD, NROT, c2, s2);     // -> C2
+        fp8_kv_dequant_rows_kernel<<<n_fp8, 128>>>(d2, c2, s2, n_fp8, HD);      // D2
+        CHECK(cudaDeviceSynchronize());
+
+        printf("FP8 comp path (%u rows, head_dim %u):\n", n_fp8, HD);
+        unsigned long long m_read  = report_f32("read-identity D1 vs C1", d1, x,  xn, HD);
+        unsigned long long m_idem  = report_f32("VALUE idempotence C2 vs C1", x2, x, xn, HD);
+        unsigned long long m_read2 = report_f32("read-identity D2 vs C1", d2, x,  xn, HD);
+        report_u8("codes2 vs codes1", c1, c2, cb);
+        report_f32("scales2 vs scales1", s1, s2, sn, DS4_OPP_C_FP8_BLOCKS_DEV);
+        // third iteration: fixed point
+        fp8_kv_quantize_kernel<<<n_fp8, 64>>>(x2, n_fp8, HD, NROT, c2, s2);
+        CHECK(cudaDeviceSynchronize());
+        unsigned long long m_fix = report_f32("fixed-point C3 vs C2(=x2)", x2, x, xn, HD);
+        int ok = (m_read == 0 && m_idem == 0 && m_read2 == 0 && m_fix == 0);
+        printf("REENC-IDEM FP8: %s\n\n", ok ? "PASS" : "FAIL (repro)");
+        if (!ok) orig_fail++;
+        cudaFree(x); cudaFree(d1); cudaFree(x2); cudaFree(d2);
+        cudaFree(c1); cudaFree(c2); cudaFree(s1); cudaFree(s2);
+    }
+
+    // ---------------- fp4 indexer path ----------------
+    {
+        const uint32_t HD = 128u;
+        uint64_t xn = (uint64_t)n_fp4 * HD;
+        uint64_t cb = (uint64_t)n_fp4 * 64u;
+        uint64_t sn = (uint64_t)n_fp4 * 4u;
+        float *x, *d1, *x2, *d2, *s1, *s2;
+        unsigned char *c1, *c2;
+        CHECK(cudaMalloc(&x,  xn * 4)); CHECK(cudaMalloc(&d1, xn * 4));
+        CHECK(cudaMalloc(&x2, xn * 4)); CHECK(cudaMalloc(&d2, xn * 4));
+        CHECK(cudaMalloc(&c1, cb)); CHECK(cudaMalloc(&c2, cb));
+        CHECK(cudaMalloc(&s1, sn * 4)); CHECK(cudaMalloc(&s2, sn * 4));
+
+        fill_fp4_input_kernel<<<n_fp4, 128>>>(x, n_fp4, 7u);
+        indexer_hadamard_fp4_kernel<<<n_fp4, 128>>>(x, n_fp4, HD, c1, s1);       // QAT -> C1
+        indexer_fp4_dequant_rows_kernel<<<n_fp4, 128>>>(d1, c1, s1, n_fp4, HD);  // D1
+        CHECK(cudaMemcpy(x2, x, xn * 4, cudaMemcpyDeviceToDevice));
+        indexer_fp4_encode_rows_kernel<<<n_fp4, 128>>>(x2, n_fp4, HD, c2, s2);   // load-path re-encode -> C2
+        indexer_fp4_dequant_rows_kernel<<<n_fp4, 128>>>(d2, c2, s2, n_fp4, HD);  // D2
+        CHECK(cudaDeviceSynchronize());
+
+        printf("FP4 indexer path (%u rows, head_dim %u):\n", n_fp4, HD);
+        unsigned long long m_read  = report_f32("read-identity D1 vs C1", d1, x,  xn, HD);
+        unsigned long long m_idem  = report_f32("VALUE idempotence C2 vs C1", x2, x, xn, HD);
+        unsigned long long m_read2 = report_f32("read-identity D2 vs C1", d2, x,  xn, HD);
+        report_u8("codes2 vs codes1", c1, c2, cb);
+        report_f32("scales2 vs scales1", s1, s2, sn, 4u);
+        indexer_fp4_encode_rows_kernel<<<n_fp4, 128>>>(x2, n_fp4, HD, c2, s2);
+        CHECK(cudaDeviceSynchronize());
+        unsigned long long m_fix = report_f32("fixed-point C3 vs C2(=x2)", x2, x, xn, HD);
+        int ok = (m_read == 0 && m_idem == 0 && m_read2 == 0 && m_fix == 0);
+        printf("REENC-IDEM FP4: %s\n\n", ok ? "PASS" : "FAIL (repro)");
+        if (!ok) orig_fail++;
+        cudaFree(x); cudaFree(d1); cudaFree(x2); cudaFree(d2);
+        cudaFree(c1); cudaFree(c2); cudaFree(s1); cudaFree(s2);
+    }
+
+    // ---------------- FIXED fp8 comp path ----------------
+    {
+        const uint32_t HD = 512u, NROT = 64u;
+        uint64_t xn = (uint64_t)n_fp8 * HD;
+        uint64_t cb = (uint64_t)n_fp8 * DS4_OPP_C_FP8_ROW_BYTES_DEV;
+        uint64_t sn = (uint64_t)n_fp8 * DS4_OPP_C_FP8_BLOCKS_DEV;
+        float *x, *d1, *x2, *d2, *s1, *s2;
+        unsigned char *c1, *c2;
+        CHECK(cudaMalloc(&x,  xn * 4)); CHECK(cudaMalloc(&d1, xn * 4));
+        CHECK(cudaMalloc(&x2, xn * 4)); CHECK(cudaMalloc(&d2, xn * 4));
+        CHECK(cudaMalloc(&c1, cb)); CHECK(cudaMalloc(&c2, cb));
+        CHECK(cudaMalloc(&s1, sn * 4)); CHECK(cudaMalloc(&s2, sn * 4));
+
+        fill_fp8_input_kernel<<<n_fp8, 512>>>(x, n_fp8, 7u);
+        fp8_kv_quantize_kernel_fx<<<n_fp8, 64>>>(x, n_fp8, HD, NROT, c1, s1);
+        fp8_kv_dequant_rows_kernel<<<n_fp8, 128>>>(d1, c1, s1, n_fp8, HD);
+        CHECK(cudaMemcpy(x2, x, xn * 4, cudaMemcpyDeviceToDevice));
+        fp8_kv_quantize_kernel_fx<<<n_fp8, 64>>>(x2, n_fp8, HD, NROT, c2, s2);
+        fp8_kv_dequant_rows_kernel<<<n_fp8, 128>>>(d2, c2, s2, n_fp8, HD);
+        CHECK(cudaDeviceSynchronize());
+
+        printf("FIXED FP8 comp path (frexpf scale + signbit; %u rows):\n", n_fp8);
+        unsigned long long m_read  = report_f32("read-identity D1 vs C1", d1, x,  xn, HD);
+        unsigned long long m_idem  = report_f32("VALUE idempotence C2 vs C1", x2, x, xn, HD);
+        unsigned long long m_read2 = report_f32("read-identity D2 vs C1", d2, x,  xn, HD);
+        unsigned long long m_code  = report_u8("codes2 vs codes1", c1, c2, cb);
+        unsigned long long m_scale = report_f32("scales2 vs scales1", s1, s2, sn, DS4_OPP_C_FP8_BLOCKS_DEV);
+        fp8_kv_quantize_kernel_fx<<<n_fp8, 64>>>(x2, n_fp8, HD, NROT, c2, s2);
+        CHECK(cudaDeviceSynchronize());
+        unsigned long long m_fix = report_f32("fixed-point C3 vs C2(=x2)", x2, x, xn, HD);
+        int ok = (m_read == 0 && m_idem == 0 && m_read2 == 0 && m_fix == 0);
+        printf("REENC-IDEM FP8-FIXED: %s%s\n\n", ok ? "PASS" : "FAIL",
+               (ok && (m_code || m_scale)) ? " (codes level-shift, values exact)" : "");
+        if (!ok) fails++;
+        cudaFree(x); cudaFree(d1); cudaFree(x2); cudaFree(d2);
+        cudaFree(c1); cudaFree(c2); cudaFree(s1); cudaFree(s2);
+    }
+
+    // ---------------- FIXED fp4 indexer path ----------------
+    {
+        const uint32_t HD = 128u;
+        uint64_t xn = (uint64_t)n_fp4 * HD;
+        uint64_t cb = (uint64_t)n_fp4 * 64u;
+        uint64_t sn = (uint64_t)n_fp4 * 4u;
+        float *x, *d1, *x2, *d2, *s1, *s2;
+        unsigned char *c1, *c2;
+        CHECK(cudaMalloc(&x,  xn * 4)); CHECK(cudaMalloc(&d1, xn * 4));
+        CHECK(cudaMalloc(&x2, xn * 4)); CHECK(cudaMalloc(&d2, xn * 4));
+        CHECK(cudaMalloc(&c1, cb)); CHECK(cudaMalloc(&c2, cb));
+        CHECK(cudaMalloc(&s1, sn * 4)); CHECK(cudaMalloc(&s2, sn * 4));
+
+        fill_fp4_input_kernel<<<n_fp4, 128>>>(x, n_fp4, 7u);
+        indexer_hadamard_fp4_kernel_fx<<<n_fp4, 128>>>(x, n_fp4, HD, c1, s1);
+        indexer_fp4_dequant_rows_kernel<<<n_fp4, 128>>>(d1, c1, s1, n_fp4, HD);
+        CHECK(cudaMemcpy(x2, x, xn * 4, cudaMemcpyDeviceToDevice));
+        indexer_fp4_encode_rows_kernel_fx<<<n_fp4, 128>>>(x2, n_fp4, HD, c2, s2);
+        indexer_fp4_dequant_rows_kernel<<<n_fp4, 128>>>(d2, c2, s2, n_fp4, HD);
+        CHECK(cudaDeviceSynchronize());
+
+        printf("FIXED FP4 indexer path (frexpf scale + signbit; %u rows):\n", n_fp4);
+        unsigned long long m_read  = report_f32("read-identity D1 vs C1", d1, x,  xn, HD);
+        unsigned long long m_idem  = report_f32("VALUE idempotence C2 vs C1", x2, x, xn, HD);
+        unsigned long long m_read2 = report_f32("read-identity D2 vs C1", d2, x,  xn, HD);
+        unsigned long long m_code  = report_u8("codes2 vs codes1", c1, c2, cb);
+        unsigned long long m_scale = report_f32("scales2 vs scales1", s1, s2, sn, 4u);
+        indexer_fp4_encode_rows_kernel_fx<<<n_fp4, 128>>>(x2, n_fp4, HD, c2, s2);
+        CHECK(cudaDeviceSynchronize());
+        unsigned long long m_fix = report_f32("fixed-point C3 vs C2(=x2)", x2, x, xn, HD);
+        int ok = (m_read == 0 && m_idem == 0 && m_read2 == 0 && m_fix == 0);
+        printf("REENC-IDEM FP4-FIXED: %s%s\n\n", ok ? "PASS" : "FAIL",
+               (ok && (m_code || m_scale)) ? " (codes level-shift, values exact)" : "");
+        if (!ok) fails++;
+        cudaFree(x); cudaFree(d1); cudaFree(x2); cudaFree(d2);
+        cudaFree(c1); cudaFree(c2); cudaFree(s1); cudaFree(s2);
+    }
+
+    // ---------------- fast-math boundary scan ----------------
+    {
+        float *d_out;
+        CHECK(cudaMalloc(&d_out, 61 * 8 * sizeof(float)));
+        boundary_scan_kernel<<<1, 64>>>(d_out);
+        CHECK(cudaDeviceSynchronize());
+        float out[61 * 8];
+        CHECK(cudaMemcpy(out, d_out, sizeof(out), cudaMemcpyDeviceToHost));
+        cudaFree(d_out);
+        int devs = 0, ulps = 0;
+        printf("Fast-math boundary scan (exact-pow2 amax inputs, k=-30..30):\n");
+        for (int idx = 0; idx < 61; idx++) {
+            int k = idx - 30;
+            float p2k  = ldexpf(1.0f, k);
+            float p2k1 = ldexpf(1.0f, k - 1);
+            // expectations from exact real arithmetic; skip cases where the
+            // 1e-4 amax floor binds instead.  `ulp` cases (one-ulp-off amax)
+            // probe approximation sensitivity, not a re-encode reachable
+            // state -- tallied separately as INFO.
+            struct { const char *name; float got, want; int skip, ulp; } cases[8] = {
+                {"fp8 448*2^k ", out[idx*8+0], p2k,  448.0f*p2k < 1.0e-4f, 0},
+                {"fp8 224*2^k ", out[idx*8+1], p2k1, 224.0f*p2k < 1.0e-4f, 0},
+                {"fp8 448u    ", out[idx*8+2], 2.0f*p2k, 448.0f*p2k < 1.0e-4f, 1},
+                {"fp8 448d    ", out[idx*8+3], p2k,  448.0f*p2k < 1.0e-4f, 1},
+                {"fp4 6*2^k   ", out[idx*8+4], p2k,  0, 0},
+                {"fp4 3*2^k   ", out[idx*8+5], p2k1, 0, 0},
+                {"fp4 1.5*2^k ", out[idx*8+6], ldexpf(1.0f, k-2), 0, 0},
+                {"fp4 6u      ", out[idx*8+7], 2.0f*p2k, 0, 1},
+            };
+            for (int c = 0; c < 8; c++) {
+                if (cases[c].skip) continue;
+                if (cases[c].got != cases[c].want) {
+                    printf("  %s k=%3d %s got=%.9g want=%.9g\n",
+                           cases[c].ulp ? "ulp-info " : "DEVIATION",
+                           k, cases[c].name, cases[c].got, cases[c].want);
+                    if (cases[c].ulp) ulps++; else devs++;
+                }
+            }
+        }
+        printf("boundary scan: exact-case deviations=%d%s, ulp-info=%d\n\n", devs,
+               devs ? " (MECHANISM CANDIDATE -- fast-math log2f/ceilf at pow2)" : "", ulps);
+    }
+
+    printf("PROTO-KV-REENC-IDEM: shipped kernels %s, fixed variants %s\n",
+           orig_fail ? "REPRODUCE the bug" : "PASS (bug not reproduced?!)",
+           fails == 0 ? "PASS" : "FAIL");
+    return fails != 0;
+}

cuda/mmq/test/proto_m2_comp.cu

@@ -0,0 +1,537 @@
+/* proto_m2_comp.cu — M2-Inc5 prototype: fused compressor pair+store.
+ * Replaces the per-compressor-event 5-launch chain
+ *   matmul_f16_splitk(kv) + combine + matmul_f16_splitk(sc) + combine
+ * + compressor_store
+ * with ONE cooperative kernel, at the three decode widths
+ *   out 1024 (primary ratio-4), out 512 (primary ratio-2),
+ *   out 256 (indexer ratio-4, head_dim 128); in_dim 4096, ksplit 2/4/8.
+ *
+ * Gate: kv_cur/sc_cur/state_kv/state_score all BIT-IDENTICAL to the ref
+ * chain across widths x ape f16/f32 x pos0 sweep; capture==eager; timing.
+ *
+ * Bit-exactness: phase 1 replicates the ref split-K tile exactly like
+ * M2-Inc3 (proto_m2_router.cu), generalized to seg in {2048,1024,512}:
+ * a lane's A = seg/256 accumulators hold chunks lane+32a (identical values
+ * at identical lanes to ref warps a), each reduced by the verbatim
+ * warp_sum_f32; the ref's cross-warp shfl_down tree over [s0..s7,0-padded]
+ * imposes the parenthesization
+ *   A=2: s0+s1
+ *   A=4: (s0+s2)+(s1+s3)
+ *   A=8: ((s0+s4)+(s2+s6))+((s1+s5)+(s3+s7))
+ * (identical up to x+0.0f==x on exact zeros).  The tail combines partials
+ * in kseg order (verbatim combine kernel) and applies the verbatim store
+ * body (state copy + ape add) to the just-combined values.
+ *
+ * Build (GB10):
+ *   nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 proto_m2_comp.cu -o proto_m2_comp
+ */
+#include <cooperative_groups.h>
+#include <cuda_fp16.h>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <cstdint>
+#include <cmath>
+
+#define CHECK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
+    fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(e_), __FILE__, __LINE__); exit(1); } } while (0)
+
+#define N_IN 4096u
+
+/* ------------------------------------------------------------------ */
+/* Vendored production device functions + kernels (ds4_cuda.cu, verbatim) */
+
+struct ds4_decode_scalars {
+    uint32_t pos0;
+    uint32_t raw_row;
+    uint32_t raw_start;
+    uint32_t n_raw;
+    uint32_t n_comp;
+    uint32_t emit_phase;
+    uint32_t comp_row;
+    uint32_t index_row;
+    uint32_t flags;
+    uint32_t token;
+};
+
+__device__ static float warp_sum_f32(float v) {
+    for (int offset = 16; offset > 0; offset >>= 1) {
+        v += __shfl_down_sync(0xffffffffu, v, offset);
+    }
+    return v;
+}
+
+__device__ static float block_reduce_sum_256(float v) {
+    v = warp_sum_f32(v);
+    __shared__ float s[8];
+    const uint32_t lane = threadIdx.x & 31u;
+    const uint32_t warp = threadIdx.x >> 5u;
+    if (lane == 0u) s[warp] = v;
+    __syncthreads();
+    float r = 0.0f;
+    if (warp == 0u) {
+        r = (lane < 8u) ? s[lane] : 0.0f;
+        r += __shfl_down_sync(0xffffffffu, r, 4);
+        r += __shfl_down_sync(0xffffffffu, r, 2);
+        r += __shfl_down_sync(0xffffffffu, r, 1);
+    }
+    return r;
+}
+
+__global__ static void matmul_f16_splitk_kernel(
+        float *partial,
+        const __half *w,
+        const float *x,
+        uint64_t in_dim,
+        uint64_t out_dim,
+        uint32_t ksplit,
+        int vec_ok) {
+    const uint64_t row = (uint64_t)blockIdx.x;
+    const uint32_t kseg = blockIdx.y;
+    if (row >= out_dim) return;
+    uint64_t seg = (in_dim + (uint64_t)ksplit - 1u) / (uint64_t)ksplit;
+    seg = (seg + 7u) & ~(uint64_t)7u;
+    const uint64_t k0 = (uint64_t)kseg * seg;
+    if (k0 >= in_dim) {
+        if (threadIdx.x == 0u) partial[row * (uint64_t)ksplit + (uint64_t)kseg] = 0.0f;
+        return;
+    }
+    uint64_t k1 = k0 + seg;
+    if (k1 > in_dim) k1 = in_dim;
+    const __half *wr = w + row * in_dim;
+    const uint32_t tid = threadIdx.x;
+    const uint32_t nt = blockDim.x;
+    float sum = 0.0f;
+    if (vec_ok) {
+        const uint64_t vend = k0 + ((k1 - k0) & ~(uint64_t)7u);
+        for (uint64_t i = k0 + (uint64_t)tid * 8u; i < vend; i += (uint64_t)nt * 8u) {
+            const uint4 wv = *reinterpret_cast<const uint4 *>(wr + i);
+            const float4 xa = *reinterpret_cast<const float4 *>(x + i);
+            const float4 xb = *reinterpret_cast<const float4 *>(x + i + 4);
+            const __half2 *h = reinterpret_cast<const __half2 *>(&wv);
+            const float2 w0 = __half22float2(h[0]);
+            const float2 w1 = __half22float2(h[1]);
+            const float2 w2 = __half22float2(h[2]);
+            const float2 w3 = __half22float2(h[3]);
+            sum += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+                 + w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
+        }
+        for (uint64_t i = vend + tid; i < k1; i += nt) sum += __half2float(wr[i]) * x[i];
+    } else {
+        for (uint64_t i = k0 + tid; i < k1; i += nt) sum += __half2float(wr[i]) * x[i];
+    }
+    sum = block_reduce_sum_256(sum);
+    if (threadIdx.x == 0u) partial[row * (uint64_t)ksplit + (uint64_t)kseg] = sum;
+}
+
+__global__ static void matmul_f16_splitk_combine_kernel(
+        float *out,
+        const float *partial,
+        uint64_t out_dim,
+        uint32_t ksplit) {
+    const uint64_t row = (uint64_t)blockIdx.x * (uint64_t)blockDim.x + (uint64_t)threadIdx.x;
+    if (row >= out_dim) return;
+    const float *p = partial + row * (uint64_t)ksplit;
+    float s = 0.0f;
+    for (uint32_t k = 0u; k < ksplit; k++) s += p[k];
+    out[row] = s;
+}
+
+__device__ static float model_scalar_dev(const void *base, uint64_t offset, uint32_t type, uint64_t idx) {
+    const char *p = (const char *)base + offset;
+    if (type == 1u) return __half2float(((const __half *)p)[idx]);
+    return ((const float *)p)[idx];
+}
+
+__global__ static void compressor_store_kernel(
+        const float *kv,
+        const float *sc,
+        float *state_kv,
+        float *state_score,
+        const void *model_map,
+        uint64_t ape_offset,
+        uint32_t ape_type,
+        uint32_t head_dim,
+        uint32_t ratio,
+        uint32_t pos0,
+        uint32_t n_tokens,
+        const struct ds4_decode_scalars * __restrict__ s_override) {
+    if (s_override) {
+        pos0 = s_override->pos0;
+    }
+    uint32_t coff = ratio == 4u ? 2u : 1u;
+    uint32_t width = coff * head_dim;
+    uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    uint64_t n = (uint64_t)n_tokens * width;
+    if (gid >= n) return;
+    uint32_t t = gid / width;
+    uint32_t j = gid - (uint64_t)t * width;
+    uint32_t pos_mod = (pos0 + t) % ratio;
+    uint32_t dst_row = ratio == 4u ? ratio + pos_mod : pos_mod;
+    state_kv[(uint64_t)dst_row * width + j] = kv[(uint64_t)t * width + j];
+    state_score[(uint64_t)dst_row * width + j] =
+        sc[(uint64_t)t * width + j] + model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)pos_mod * width + j);
+}
+
+/* ------------------------------------------------------------------ */
+/* FUSED: pair matmul + combine + store as one cooperative kernel.
+ * Template KS = ksplit in {2,4,8}; SEG = 4096/KS; A = SEG/256 accumulators
+ * per lane. */
+
+template <uint32_t KS>
+__global__ static void comp_pair_store_fused_kernel(
+        float *kv_cur,            /* [width] */
+        float *sc_cur,            /* [width] */
+        float *partials,          /* [2*width*KS] dedicated scratch */
+        float *state_kv,
+        float *state_score,
+        const __half *w_kv,       /* [width][4096] */
+        const __half *w_sc,       /* [width][4096] */
+        const float *x,           /* [4096] */
+        const void *ape_map,
+        uint64_t ape_offset,
+        uint32_t ape_type,
+        uint32_t head_dim,
+        uint32_t ratio,
+        uint32_t width,           /* == coff*head_dim == out_dim */
+        uint32_t pos0,
+        const struct ds4_decode_scalars * __restrict__ s_override) {
+    constexpr uint32_t SEG = N_IN / KS;
+    constexpr uint32_t A = SEG / 256u;
+    const uint32_t lane = threadIdx.x & 31u;
+    const uint32_t warp = threadIdx.x >> 5u;
+
+    /* Phase 1: one warp per (matmul m, row, kseg) tile. */
+    const uint32_t tiles_per_m = width * KS;
+    const uint32_t nwarp = gridDim.x * (blockDim.x >> 5u);
+    for (uint32_t t = blockIdx.x * (blockDim.x >> 5u) + warp;
+         t < 2u * tiles_per_m; t += nwarp) {
+        const uint32_t m = t < tiles_per_m ? 0u : 1u;
+        const uint32_t tt = m ? t - tiles_per_m : t;
+        const uint32_t row = tt / KS;
+        const uint32_t kseg = tt % KS;
+        const uint64_t k0 = (uint64_t)kseg * SEG;
+        const __half *wr = (m ? w_sc : w_kv) + (uint64_t)row * N_IN;
+        float s[A];
+        #pragma unroll
+        for (uint32_t a = 0; a < A; a++) {
+            const uint64_t i = k0 + ((uint64_t)lane + 32u * a) * 8u;
+            const uint4 wv = *reinterpret_cast<const uint4 *>(wr + i);
+            const float4 xa = *reinterpret_cast<const float4 *>(x + i);
+            const float4 xb = *reinterpret_cast<const float4 *>(x + i + 4);
+            const __half2 *h = reinterpret_cast<const __half2 *>(&wv);
+            const float2 w0 = __half22float2(h[0]);
+            const float2 w1 = __half22float2(h[1]);
+            const float2 w2 = __half22float2(h[2]);
+            const float2 w3 = __half22float2(h[3]);
+            float sum = 0.0f;
+            sum += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+                 + w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
+            s[a] = warp_sum_f32(sum);
+        }
+        /* ref block_reduce_sum_256 cross-warp tree over [s0..s(A-1), 0...]:
+         * shfl_down 4,2,1 -> fixed parenthesization (zero-adds are identity). */
+        float partial;
+        if (KS == 8u) {          /* A == 2 */
+            partial = s[0] + s[1 % A];
+        } else if (KS == 4u) {   /* A == 4 */
+            partial = (s[0] + s[2 % A]) + (s[1 % A] + s[3 % A]);
+        } else {                 /* KS == 2, A == 8 */
+            partial = ((s[0] + s[4 % A]) + (s[2 % A] + s[6 % A]))
+                    + ((s[1 % A] + s[5 % A]) + (s[3 % A] + s[7 % A]));
+        }
+        if (lane == 0u) partials[(uint64_t)m * tiles_per_m + (uint64_t)row * KS + kseg] = partial;
+    }
+
+    cooperative_groups::this_grid().sync();
+
+    /* Phase 2: combine (kseg ascending, verbatim) + verbatim store body. */
+    const uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
+    const uint32_t n2 = 2u * width;
+    if (gid >= n2) return;
+    const uint32_t m = gid < width ? 0u : 1u;
+    const uint32_t row = m ? gid - width : gid;
+    const float *p = partials + (uint64_t)m * tiles_per_m + (uint64_t)row * KS;
+    float s = 0.0f;
+    for (uint32_t k = 0u; k < KS; k++) s += p[k];
+    uint32_t p0 = s_override ? s_override->pos0 : pos0;
+    const uint32_t pos_mod = p0 % ratio;
+    const uint32_t dst_row = ratio == 4u ? ratio + pos_mod : pos_mod;
+    if (m == 0u) {
+        kv_cur[row] = s;
+        state_kv[(uint64_t)dst_row * width + row] = s;
+    } else {
+        sc_cur[row] = s;
+        state_score[(uint64_t)dst_row * width + row] =
+            s + model_scalar_dev(ape_map, ape_offset, ape_type, (uint64_t)pos_mod * width + row);
+    }
+}
+
+/* ------------------------------------------------------------------ */
+/* Host harness */
+
+struct cfg_t {
+    const char *name;
+    uint32_t head_dim, ratio, width, ksplit;
+};
+static const cfg_t CFGS[3] = {
+    { "primary-r4 (1024,k2)", 512, 4, 1024, 2 },
+    { "primary-r2 (512,k4)",  512, 2, 512,  4 },
+    { "index-r4   (256,k8)",  128, 4, 256,  8 },
+};
+
+static void fill_rand(float *p, size_t n, unsigned seed, float scale) {
+    unsigned s = seed;
+    for (size_t i = 0; i < n; i++) {
+        s = s * 1664525u + 1013904223u;
+        p[i] = ((float)(s >> 8) / (float)(1u << 24) - 0.5f) * 2.0f * scale;
+    }
+}
+
+struct bufs {
+    __half *w_kv, *w_sc;
+    float *x, *kv_cur, *sc_cur, *state_kv, *state_score, *part_ref, *part_fu;
+    char *ape;             /* f32 or f16 depending on ape_type */
+    ds4_decode_scalars *scalars;
+};
+
+static void *fused_ptr(uint32_t ks) {
+    if (ks == 2u) return (void *)comp_pair_store_fused_kernel<2u>;
+    if (ks == 4u) return (void *)comp_pair_store_fused_kernel<4u>;
+    return (void *)comp_pair_store_fused_kernel<8u>;
+}
+
+static void run_ref(const bufs &b, const cfg_t &c, uint32_t ape_type, uint32_t pos0, int via_scalars) {
+    dim3 sg1(c.width, c.ksplit, 1);
+    matmul_f16_splitk_kernel<<<sg1, 256>>>(b.part_ref, b.w_kv, b.x, N_IN, c.width, c.ksplit, 1);
+    matmul_f16_splitk_combine_kernel<<<(c.width + 255u) / 256u, 256>>>(b.kv_cur, b.part_ref, c.width, c.ksplit);
+    matmul_f16_splitk_kernel<<<sg1, 256>>>(b.part_ref, b.w_sc, b.x, N_IN, c.width, c.ksplit, 1);
+    matmul_f16_splitk_combine_kernel<<<(c.width + 255u) / 256u, 256>>>(b.sc_cur, b.part_ref, c.width, c.ksplit);
+    compressor_store_kernel<<<(c.width + 255u) / 256u, 256>>>(
+            b.kv_cur, b.sc_cur, b.state_kv, b.state_score,
+            b.ape, 0, ape_type, c.head_dim, c.ratio,
+            via_scalars ? 0xdeadbeefu : pos0, 1,
+            via_scalars ? b.scalars : NULL);
+    CHECK(cudaGetLastError());
+}
+
+static int run_fused(const bufs &b, const cfg_t &c, uint32_t ape_type, uint32_t pos0, int via_scalars,
+                     unsigned grid_blocks, cudaStream_t stream) {
+    float *kv_cur = b.kv_cur, *sc_cur = b.sc_cur, *partials = b.part_fu;
+    float *state_kv = b.state_kv, *state_score = b.state_score, *x = b.x;
+    const __half *w_kv = b.w_kv, *w_sc = b.w_sc;
+    const void *ape = b.ape;
+    uint64_t ape_off = 0;
+    uint32_t at = ape_type, hd = c.head_dim, ra = c.ratio, wi = c.width;
+    uint32_t p0 = via_scalars ? 0xdeadbeefu : pos0;
+    const ds4_decode_scalars *sc_ptr = via_scalars ? b.scalars : NULL;
+    void *args[] = {
+        (void *)&kv_cur, (void *)&sc_cur, (void *)&partials,
+        (void *)&state_kv, (void *)&state_score,
+        (void *)&w_kv, (void *)&w_sc, (void *)&x,
+        (void *)&ape, (void *)&ape_off, (void *)&at,
+        (void *)&hd, (void *)&ra, (void *)&wi, (void *)&p0, (void *)&sc_ptr };
+    cudaError_t err = cudaLaunchCooperativeKernel(
+            fused_ptr(c.ksplit), dim3(grid_blocks, 1, 1), dim3(256, 1, 1), args, 0, stream);
+    if (err != cudaSuccess) {
+        fprintf(stderr, "coop launch failed (%s, %u blk): %s\n", c.name, grid_blocks, cudaGetErrorString(err));
+        return 0;
+    }
+    return 1;
+}
+
+int main() {
+    int dev = 0, coop = 0, n_sm = 0;
+    CHECK(cudaGetDevice(&dev));
+    CHECK(cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, dev));
+    CHECK(cudaDeviceGetAttribute(&n_sm, cudaDevAttrMultiProcessorCount, dev));
+    printf("coop=%d n_sm=%d\n", coop, n_sm);
+    if (!coop) return 1;
+    unsigned cap[3];
+    for (int i = 0; i < 3; i++) {
+        int bps = 0;
+        CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&bps, fused_ptr(CFGS[i].ksplit), 256, 0));
+        cap[i] = (unsigned)(bps * n_sm);
+        printf("  %s: blocks/sm=%d cap=%u\n", CFGS[i].name, bps, cap[i]);
+    }
+
+    const uint32_t WMAX = 1024;
+    bufs b = {};
+    CHECK(cudaMalloc(&b.w_kv, (size_t)WMAX * N_IN * sizeof(__half)));
+    CHECK(cudaMalloc(&b.w_sc, (size_t)WMAX * N_IN * sizeof(__half)));
+    CHECK(cudaMalloc(&b.x, N_IN * sizeof(float)));
+    CHECK(cudaMalloc(&b.kv_cur, WMAX * sizeof(float)));
+    CHECK(cudaMalloc(&b.sc_cur, WMAX * sizeof(float)));
+    CHECK(cudaMalloc(&b.state_kv, 8u * WMAX * sizeof(float)));
+    CHECK(cudaMalloc(&b.state_score, 8u * WMAX * sizeof(float)));
+    CHECK(cudaMalloc(&b.part_ref, (size_t)WMAX * 8u * sizeof(float)));
+    CHECK(cudaMalloc(&b.part_fu, 2u * (size_t)WMAX * 8u * sizeof(float)));
+    CHECK(cudaMalloc(&b.ape, 4u * WMAX * sizeof(float)));
+    CHECK(cudaMalloc(&b.scalars, sizeof(ds4_decode_scalars)));
+
+    float *hbuf = (float *)malloc((size_t)WMAX * N_IN * sizeof(float));
+    __half *hh = (__half *)malloc((size_t)WMAX * N_IN * sizeof(__half));
+    float *hape = (float *)malloc(4u * WMAX * sizeof(float));
+    __half *hape16 = (__half *)malloc(4u * WMAX * sizeof(__half));
+
+    const size_t OUTB = WMAX * sizeof(float);
+    const size_t STB = 8u * WMAX * sizeof(float);
+    float *ref_out = (float *)malloc(2 * OUTB + 2 * STB);
+    float *fu_out = (float *)malloc(2 * OUTB + 2 * STB);
+
+    int total_bad = 0, n_cases = 0;
+    const unsigned grids[] = { 48u, 96u, 128u };
+    const uint32_t poss[] = { 0u, 1u, 2u, 3u, 7u, 1023u };
+    for (int ci = 0; ci < 3; ci++) {
+        const cfg_t &c = CFGS[ci];
+        for (unsigned seed = 1; seed <= 3; seed++) {
+            fill_rand(hbuf, (size_t)c.width * N_IN, seed * 7919u + ci, 0.05f);
+            for (size_t i = 0; i < (size_t)c.width * N_IN; i++) hh[i] = __float2half(hbuf[i]);
+            CHECK(cudaMemcpy(b.w_kv, hh, (size_t)c.width * N_IN * sizeof(__half), cudaMemcpyHostToDevice));
+            fill_rand(hbuf, (size_t)c.width * N_IN, seed * 104729u + ci, 0.05f);
+            for (size_t i = 0; i < (size_t)c.width * N_IN; i++) hh[i] = __float2half(hbuf[i]);
+            CHECK(cudaMemcpy(b.w_sc, hh, (size_t)c.width * N_IN * sizeof(__half), cudaMemcpyHostToDevice));
+            fill_rand(hbuf, N_IN, seed * 31u, 1.0f);
+            CHECK(cudaMemcpy(b.x, hbuf, N_IN * sizeof(float), cudaMemcpyHostToDevice));
+            fill_rand(hape, (size_t)c.ratio * c.width, seed * 71u, 0.5f);
+            for (size_t i = 0; i < (size_t)c.ratio * c.width; i++) hape16[i] = __float2half(hape[i]);
+            for (uint32_t ape_type = 0; ape_type <= 1; ape_type++) {
+                if (ape_type == 0) CHECK(cudaMemcpy(b.ape, hape, (size_t)c.ratio * c.width * sizeof(float), cudaMemcpyHostToDevice));
+                else CHECK(cudaMemcpy(b.ape, hape16, (size_t)c.ratio * c.width * sizeof(__half), cudaMemcpyHostToDevice));
+                for (uint32_t pi = 0; pi < 6; pi++) {
+                    const uint32_t pos0 = poss[pi];
+                    const int via_scalars = (int)(pi & 1u);
+                    ds4_decode_scalars sc = {};
+                    sc.pos0 = pos0;
+                    CHECK(cudaMemcpy(b.scalars, &sc, sizeof(sc), cudaMemcpyHostToDevice));
+                    /* ref */
+                    CHECK(cudaMemset(b.kv_cur, 0xCB, OUTB));
+                    CHECK(cudaMemset(b.sc_cur, 0xCB, OUTB));
+                    CHECK(cudaMemset(b.state_kv, 0xCB, STB));
+                    CHECK(cudaMemset(b.state_score, 0xCB, STB));
+                    run_ref(b, c, ape_type, pos0, via_scalars);
+                    CHECK(cudaDeviceSynchronize());
+                    CHECK(cudaMemcpy(ref_out, b.kv_cur, OUTB, cudaMemcpyDeviceToHost));
+                    CHECK(cudaMemcpy(ref_out + WMAX, b.sc_cur, OUTB, cudaMemcpyDeviceToHost));
+                    CHECK(cudaMemcpy(ref_out + 2 * WMAX, b.state_kv, STB, cudaMemcpyDeviceToHost));
+                    CHECK(cudaMemcpy(ref_out + 10 * WMAX, b.state_score, STB, cudaMemcpyDeviceToHost));
+                    for (unsigned g : grids) {
+                        if (g > cap[ci]) continue;
+                        CHECK(cudaMemset(b.kv_cur, 0xCB, OUTB));
+                        CHECK(cudaMemset(b.sc_cur, 0xCB, OUTB));
+                        CHECK(cudaMemset(b.state_kv, 0xCB, STB));
+                        CHECK(cudaMemset(b.state_score, 0xCB, STB));
+                        if (!run_fused(b, c, ape_type, pos0, via_scalars, g, 0)) return 1;
+                        CHECK(cudaDeviceSynchronize());
+                        CHECK(cudaMemcpy(fu_out, b.kv_cur, OUTB, cudaMemcpyDeviceToHost));
+                        CHECK(cudaMemcpy(fu_out + WMAX, b.sc_cur, OUTB, cudaMemcpyDeviceToHost));
+                        CHECK(cudaMemcpy(fu_out + 2 * WMAX, b.state_kv, STB, cudaMemcpyDeviceToHost));
+                        CHECK(cudaMemcpy(fu_out + 10 * WMAX, b.state_score, STB, cudaMemcpyDeviceToHost));
+                        n_cases++;
+                        if (memcmp(ref_out, fu_out, 2 * OUTB + 2 * STB)) {
+                            total_bad++;
+                            printf("MISMATCH %s seed%u ape%u pos%u g%u\n", c.name, seed, ape_type, pos0, g);
+                            for (int i = 0; i < 6; i++)
+                                printf("  kv[%d] ref=%.9g fu=%.9g  sc ref=%.9g fu=%.9g\n", i,
+                                       ref_out[i], fu_out[i], ref_out[WMAX + i], fu_out[WMAX + i]);
+                        }
+                    }
+                }
+            }
+        }
+    }
+    printf("parity: %d/%d cases bit-identical\n", n_cases - total_bad, n_cases);
+    if (total_bad) return 1;
+
+    /* capture/replay bit-identity (one config) */
+    {
+        const cfg_t &c = CFGS[2];
+        const unsigned g = cap[2] < 128u ? cap[2] : 128u;
+        ds4_decode_scalars sc = {};
+        sc.pos0 = 5;
+        CHECK(cudaMemcpy(b.scalars, &sc, sizeof(sc), cudaMemcpyHostToDevice));
+        if (!run_fused(b, c, 1, 0, 1, g, 0)) return 1;
+        CHECK(cudaDeviceSynchronize());
+        CHECK(cudaMemcpy(ref_out, b.kv_cur, OUTB, cudaMemcpyDeviceToHost));
+        CHECK(cudaMemcpy(ref_out + 2 * WMAX, b.state_score, STB, cudaMemcpyDeviceToHost));
+        cudaStream_t cs;
+        CHECK(cudaStreamCreate(&cs));
+        cudaGraph_t graph;
+        cudaGraphExec_t gexec;
+        CHECK(cudaStreamBeginCapture(cs, cudaStreamCaptureModeGlobal));
+        if (!run_fused(b, c, 1, 0, 1, g, cs)) return 1;
+        CHECK(cudaStreamEndCapture(cs, &graph));
+        CHECK(cudaGraphInstantiate(&gexec, graph, NULL, NULL, 0));
+        for (int r = 0; r < 3; r++) {
+            CHECK(cudaMemset(b.kv_cur, 0xCB, OUTB));
+            CHECK(cudaGraphLaunch(gexec, cs));
+            CHECK(cudaStreamSynchronize(cs));
+            CHECK(cudaMemcpy(fu_out, b.kv_cur, OUTB, cudaMemcpyDeviceToHost));
+            CHECK(cudaMemcpy(fu_out + 2 * WMAX, b.state_score, STB, cudaMemcpyDeviceToHost));
+            if (memcmp(ref_out, fu_out, OUTB) || memcmp(ref_out + 2 * WMAX, fu_out + 2 * WMAX, STB)) {
+                printf("capture-replay MISMATCH\n");
+                return 1;
+            }
+        }
+        printf("capture/replay: bit-identical (3 replays)\n");
+        CHECK(cudaGraphExecDestroy(gexec));
+        CHECK(cudaGraphDestroy(graph));
+        CHECK(cudaStreamDestroy(cs));
+    }
+
+    /* timing: rotating 16-layer weight sets (L2-defeated) per config */
+    {
+        const int NL = 16, ITERS = 2000, WARM = 200;
+        cudaEvent_t e0, e1;
+        CHECK(cudaEventCreate(&e0));
+        CHECK(cudaEventCreate(&e1));
+        float ms;
+        for (int ci = 0; ci < 3; ci++) {
+            const cfg_t &c = CFGS[ci];
+            __half *wl[NL][2];
+            for (int l = 0; l < NL; l++)
+                for (int m = 0; m < 2; m++) {
+                    CHECK(cudaMalloc(&wl[l][m], (size_t)c.width * N_IN * sizeof(__half)));
+                    CHECK(cudaMemcpy(wl[l][m], m ? b.w_sc : b.w_kv,
+                                     (size_t)c.width * N_IN * sizeof(__half), cudaMemcpyDeviceToDevice));
+                }
+            bufs bl = b;
+            for (int it = 0; it < WARM; it++) {
+                bl.w_kv = wl[it % NL][0]; bl.w_sc = wl[it % NL][1];
+                run_ref(bl, c, 1, 3, 1);
+            }
+            CHECK(cudaDeviceSynchronize());
+            CHECK(cudaEventRecord(e0));
+            for (int it = 0; it < ITERS; it++) {
+                bl.w_kv = wl[it % NL][0]; bl.w_sc = wl[it % NL][1];
+                run_ref(bl, c, 1, 3, 1);
+            }
+            CHECK(cudaEventRecord(e1));
+            CHECK(cudaDeviceSynchronize());
+            CHECK(cudaEventElapsedTime(&ms, e0, e1));
+            printf("%s ref 5-launch : %8.3f us/event\n", c.name, ms * 1000.0f / ITERS);
+            for (unsigned g : grids) {
+                if (g > cap[ci]) continue;
+                for (int it = 0; it < WARM; it++) {
+                    bl.w_kv = wl[it % NL][0]; bl.w_sc = wl[it % NL][1];
+                    run_fused(bl, c, 1, 3, 1, g, 0);
+                }
+                CHECK(cudaDeviceSynchronize());
+                CHECK(cudaEventRecord(e0));
+                for (int it = 0; it < ITERS; it++) {
+                    bl.w_kv = wl[it % NL][0]; bl.w_sc = wl[it % NL][1];
+                    run_fused(bl, c, 1, 3, 1, g, 0);
+                }
+                CHECK(cudaEventRecord(e1));
+                CHECK(cudaDeviceSynchronize());
+                CHECK(cudaEventElapsedTime(&ms, e0, e1));
+                printf("%s fused g=%3u  : %8.3f us/event\n", c.name, g, ms * 1000.0f / ITERS);
+            }
+            for (int l = 0; l < NL; l++)
+                for (int m = 0; m < 2; m++) CHECK(cudaFree(wl[l][m]));
+        }
+    }
+
+    printf("PROTO_M2_COMP OK\n");
+    return 0;
+}

cuda/mmq/test/proto_m2_q2k.cu

@@ -0,0 +1,741 @@
+// proto_m2_q2k.cu — M2 megakernel program, moe-down Q2K aligned-SoA proto.
+//
+// Target: mul_mat_vec_q_moe<GGML_TYPE_Q2_K, 2> (cuda/mmq/mmvq.cu:608), the
+// production decode down-proj kernel: 43 launches/step at 87.3 us = 3.755
+// ms/step reading raw 84-byte block_q2_K stacks at ~190 GB/s vs the 216-233
+// GB/s device substrate (nsys33d).  This proto vendors the production kernel
+// as the bit-exact reference and races two aligned-SoA layout twins that keep
+// the SAME lane->(block,iqs) mapping and float accumulation order (layout-only
+// change => bit-identical outputs, M1-Inc4 gate recipe):
+//
+//   V1: plain SoA sections  [half2 dm[nblk]] [pad64] [u8 scales[nblk*16]]
+//       [pad64] [u32 qs[nblk*16]], block order = raw tensor byte order.
+//   V2: row-pair SoA keyed to rows_per_block=2 (row pair = 2*blockIdx.x):
+//       [uint2 dm2[npair*nb]] [pad64] [int4 sc4[npair*nb*2]] [pad64]
+//       [uint2 qs2[npair*nb*16]]; per lane-iter one 8B qs load (both rows),
+//       one 16B scales load (both rows' 8B window), one 8B dm load.
+//
+// Parity: full-buffer bitwise compare vs the raw reference across seeds x id
+// patterns (random experts, duplicates, -1 router ids, all -1) x ncols_dst
+// {1, 6, 8}, plus graph capture/replay == eager for both variants.
+//
+// Timing: rotating L-layer weight pool (defeats L2 the way the 43-layer model
+// does), production shape M=4096 K=2048 top-6, 6 distinct experts per launch.
+//
+// build (GB10): nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 \
+//                 -o proto_m2_q2k proto_m2_q2k.cu
+
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <cstdint>
+#include <vector>
+#include <random>
+
+#define CHECK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
+    fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(e_), __FILE__, __LINE__); exit(1); } } while (0)
+
+// ---------------------------------------------------------------------------
+// Vendored ggml pieces (verbatim math).
+// ---------------------------------------------------------------------------
+#define QK_K   256
+#define QI2_K  16      // QK_K / (4*QR2_K)
+#define QR2_K  4
+#define QK8_1  32
+#define QI8_1  8       // QK8_1 / (4*QR8_1)
+
+typedef struct {
+    __half2 ds;        // d (delta), s (d * sum of quants)
+    int8_t  qs[QK8_1];
+} block_q8_1;
+static_assert(sizeof(block_q8_1) == 36, "block_q8_1 size");
+
+typedef struct {
+    uint8_t scales[QK_K/16];
+    uint8_t qs[QK_K/4];
+    __half2 dm;        // d, dmin
+} block_q2_K;
+static_assert(sizeof(block_q2_K) == 84, "block_q2_K size");
+
+static __device__ __forceinline__ int get_int_b4(const void *x, const int &i32) {
+    return ((const int *)x)[i32];   // 4-byte aligned
+}
+
+static __device__ __forceinline__ int dp4a_(const int a, const int b, int c) {
+    return __dp4a(a, b, c);
+}
+
+template<int width = 32>
+static __device__ __forceinline__ float warp_reduce_sum(float x) {
+#pragma unroll
+    for (int offset = width/2; offset > 0; offset >>= 1) {
+        x += __shfl_xor_sync(0xffffffff, x, offset, width);
+    }
+    return x;
+}
+
+// vec_dot_q2_K_q8_1_impl_mmvq, verbatim (vecdotq.cuh:364).
+static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq(
+    const int & v, const int * __restrict__ u, const uint8_t * __restrict__ scales,
+    const __half2 & dm2, const float * __restrict__ d8) {
+
+    float sumf_d = 0.0f;
+    float sumf_m = 0.0f;
+
+#pragma unroll
+    for (int i = 0; i < QR2_K; ++i) {
+        const int sc = scales[2*i];
+
+        const int vi = (v >> (2*i)) & 0x03030303;
+
+        sumf_d += d8[i] * (dp4a_(vi, u[i], 0) * (sc & 0xF));
+
+        int m = sc >> 4;
+        m |= m <<  8;
+        m |= m << 16;
+        sumf_m += d8[i] * dp4a_(m, u[i], 0);
+    }
+
+    const float2 dm2f = __half22float2(dm2);
+
+    return dm2f.x*sumf_d - dm2f.y*sumf_m;
+}
+
+// Same impl but the four scale bytes come from two pre-loaded 32-bit window
+// words (V2's 16B scales load).  sc value chain identical: byte (lo + 2*i) of
+// the 8-byte window == scales[scale_offset + 2*i] of the raw block.
+static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq_w(
+    const int & v, const int * __restrict__ u, const uint32_t w0, const uint32_t w1,
+    const int lo, const __half2 & dm2, const float * __restrict__ d8) {
+
+    float sumf_d = 0.0f;
+    float sumf_m = 0.0f;
+
+#pragma unroll
+    for (int i = 0; i < QR2_K; ++i) {
+        const int bidx = lo + 2*i;                       // 0..7 within the window
+        const uint32_t w = (bidx < 4) ? w0 : w1;
+        const int sc = (int)((w >> ((bidx & 3) * 8)) & 0xFFu);
+
+        const int vi = (v >> (2*i)) & 0x03030303;
+
+        sumf_d += d8[i] * (dp4a_(vi, u[i], 0) * (sc & 0xF));
+
+        int m = sc >> 4;
+        m |= m <<  8;
+        m |= m << 16;
+        sumf_m += d8[i] * dp4a_(m, u[i], 0);
+    }
+
+    const float2 dm2f = __half22float2(dm2);
+
+    return dm2f.x*sumf_d - dm2f.y*sumf_m;
+}
+
+// vec_dot_q2_K_q8_1, verbatim (vecdotq.cuh:814).
+static __device__ __forceinline__ float vec_dot_q2_K_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+    const block_q2_K * bq2_K = (const block_q2_K *) vbq + kbx;
+
+    const int bq8_offset = QR2_K * (iqs / QI8_1);
+    const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2);
+
+    const uint8_t * scales = bq2_K->scales + scale_offset;
+
+    const int v = get_int_b4(bq2_K->qs, iqs);
+    int    u[QR2_K];
+    float d8[QR2_K];
+
+#pragma unroll
+    for (int i = 0; i < QR2_K; ++ i) {
+        u[i]  = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
+        d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
+    }
+
+    return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8);
+}
+
+// ---------------------------------------------------------------------------
+// Reference: mul_mat_vec_q_moe<GGML_TYPE_Q2_K, 2> (mmvq.cu:608), verbatim up
+// to Q2_K constant folding (qk 256, qi 16, vdr 1) and fastmodulo -> % (index
+// math only).  GB10 (sm_121 >= Ada): launch bounds MMVQ_MAX_BATCH_SIZE(8)*32.
+// ---------------------------------------------------------------------------
+template <int c_rows_per_block>
+__launch_bounds__(8*32, 1)
+static __global__ void ref_mul_mat_vec_q_moe(
+        const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ ids,
+        float * __restrict__ dst,
+        const uint32_t ncols_x, const uint32_t nchannels_y, const uint32_t nrows_x,
+        const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
+        const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
+        const uint32_t ncols_dst, const uint32_t ids_stride) {
+
+    constexpr int qk  = QK_K;
+    constexpr int qi  = QI2_K;
+    constexpr int vdr = 1;                 // VDR_Q2_K_Q8_1_MMVQ
+    constexpr int warp_size = 32;
+
+    const uint32_t token_idx   = threadIdx.y;
+    const int      row0        = c_rows_per_block*blockIdx.x;
+    const int      blocks_per_row_x = ncols_x / qk;
+    constexpr int  blocks_per_iter  = vdr * warp_size / qi;   // 2
+
+    const uint32_t channel_dst = blockIdx.y;
+
+    if (token_idx >= ncols_dst) {
+        return;
+    }
+
+    const int32_t  id_raw     = ids[channel_dst + token_idx * ids_stride];
+    const bool     invalid_id = id_raw < 0;
+    const uint32_t channel_x  = invalid_id ? 0u : (uint32_t)id_raw;
+    const uint32_t channel_y  = channel_dst % nchannels_y;
+
+    const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y;
+    const int kbx_offset  = channel_x*stride_channel_x + row0*stride_row_x;
+
+    float tmp[c_rows_per_block] = {0.0f};
+
+    for (int kbx = threadIdx.x / (qi/vdr); !invalid_id && kbx < blocks_per_row_x; kbx += blocks_per_iter) {
+        const int kby = kbx * (qk/QK8_1);
+        const int kqs = vdr * (threadIdx.x % (qi/vdr));
+
+#pragma unroll
+        for (int i = 0; i < c_rows_per_block; ++i) {
+            tmp[i] += vec_dot_q2_K_q8_1(vx, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs);
+        }
+    }
+
+#pragma unroll
+    for (int i = 0; i < c_rows_per_block; ++i) {
+        tmp[i] = warp_reduce_sum<warp_size>(tmp[i]);
+    }
+
+    if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) {
+        dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = tmp[threadIdx.x];
+    }
+}
+
+// ---------------------------------------------------------------------------
+// V1: plain SoA twin.  Same skeleton, same lane mapping, same accumulation
+// order; only the three weight streams are re-sourced from aligned sections.
+// ---------------------------------------------------------------------------
+template <int c_rows_per_block>
+__launch_bounds__(8*32, 1)
+static __global__ void v1_soa_moe_kernel(
+        const __half2 * __restrict__ dm_soa,       // [nblk]
+        const uint8_t * __restrict__ scales_soa,   // [nblk*16]
+        const int     * __restrict__ qs_soa,       // [nblk*16]
+        const void * __restrict__ vy, const int32_t * __restrict__ ids,
+        float * __restrict__ dst,
+        const uint32_t ncols_x, const uint32_t nchannels_y, const uint32_t nrows_x,
+        const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
+        const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
+        const uint32_t ncols_dst, const uint32_t ids_stride) {
+
+    constexpr int qk  = QK_K;
+    constexpr int qi  = QI2_K;
+    constexpr int vdr = 1;
+    constexpr int warp_size = 32;
+
+    const uint32_t token_idx   = threadIdx.y;
+    const int      row0        = c_rows_per_block*blockIdx.x;
+    const int      blocks_per_row_x = ncols_x / qk;
+    constexpr int  blocks_per_iter  = vdr * warp_size / qi;
+
+    const uint32_t channel_dst = blockIdx.y;
+
+    if (token_idx >= ncols_dst) {
+        return;
+    }
+
+    const int32_t  id_raw     = ids[channel_dst + token_idx * ids_stride];
+    const bool     invalid_id = id_raw < 0;
+    const uint32_t channel_x  = invalid_id ? 0u : (uint32_t)id_raw;
+    const uint32_t channel_y  = channel_dst % nchannels_y;
+
+    const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y;
+    const int kbx_offset  = channel_x*stride_channel_x + row0*stride_row_x;
+
+    float tmp[c_rows_per_block] = {0.0f};
+
+    for (int kbx = threadIdx.x / (qi/vdr); !invalid_id && kbx < blocks_per_row_x; kbx += blocks_per_iter) {
+        const int kby = kbx * (qk/QK8_1);
+        const int kqs = vdr * (threadIdx.x % (qi/vdr));
+
+        const int iqs = kqs;
+        const int bq8_offset = QR2_K * (iqs / QI8_1);
+        const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2);
+        const block_q8_1 * bq8_1 = &y[kby];
+
+        int    u[QR2_K];
+        float d8[QR2_K];
+#pragma unroll
+        for (int i = 0; i < QR2_K; ++i) {
+            u[i]  = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
+            d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
+        }
+
+#pragma unroll
+        for (int i = 0; i < c_rows_per_block; ++i) {
+            const size_t blk = (size_t)(kbx_offset + i*stride_row_x + kbx);
+            const int v = qs_soa[blk*16u + (unsigned)iqs];
+            tmp[i] += vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales_soa + blk*16u + (unsigned)scale_offset,
+                                                  dm_soa[blk], d8);
+        }
+    }
+
+#pragma unroll
+    for (int i = 0; i < c_rows_per_block; ++i) {
+        tmp[i] = warp_reduce_sum<warp_size>(tmp[i]);
+    }
+
+    if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) {
+        dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = tmp[threadIdx.x];
+    }
+}
+
+// ---------------------------------------------------------------------------
+// V2: row-pair SoA twin (rows_per_block == 2 baked into the layout).
+// Per lane-iter: one uint2 qs load (rows 0/1), one int4 scales-window load
+// (both rows' 8B window for this lane's iqs half), one uint2 dm load.  Scale
+// BYTES, qs ints, dm halves and the float chain are value-identical to the
+// reference; only load shapes change.
+// ---------------------------------------------------------------------------
+__launch_bounds__(8*32, 1)
+static __global__ void v2_pair_soa_moe_kernel(
+        const uint2 * __restrict__ dm2_soa,   // [npair*nb]      {row0 half2, row1 half2}
+        const int4  * __restrict__ sc4_soa,   // [npair*nb*2]    {row0 w0,w1, row1 w0,w1} per half
+        const uint2 * __restrict__ qs2_soa,   // [npair*nb*16]   {row0 int, row1 int}
+        const void * __restrict__ vy, const int32_t * __restrict__ ids,
+        float * __restrict__ dst,
+        const uint32_t ncols_x, const uint32_t nchannels_y, const uint32_t nrows_x,
+        const uint32_t stride_col_y, const uint32_t stride_col_dst,
+        const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
+        const uint32_t ncols_dst, const uint32_t ids_stride) {
+
+    constexpr int qk  = QK_K;
+    constexpr int qi  = QI2_K;
+    constexpr int vdr = 1;
+    constexpr int warp_size = 32;
+
+    const uint32_t token_idx   = threadIdx.y;
+    const int      row0        = 2*blockIdx.x;
+    const int      blocks_per_row_x = ncols_x / qk;
+    constexpr int  blocks_per_iter  = vdr * warp_size / qi;
+
+    const uint32_t channel_dst = blockIdx.y;
+
+    if (token_idx >= ncols_dst) {
+        return;
+    }
+
+    const int32_t  id_raw     = ids[channel_dst + token_idx * ids_stride];
+    const bool     invalid_id = id_raw < 0;
+    const uint32_t channel_x  = invalid_id ? 0u : (uint32_t)id_raw;
+    const uint32_t channel_y  = channel_dst % nchannels_y;
+
+    const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y;
+    // pair-block base for (expert, row pair): pairs are laid out expert-major,
+    // pair-major, block-minor -- same linear order as raw blocks per pair.
+    const size_t pair_base = ((size_t)channel_x * (nrows_x/2u) + (size_t)blockIdx.x) * (size_t)blocks_per_row_x;
+
+    float tmp[2] = {0.0f, 0.0f};
+
+    for (int kbx = threadIdx.x / (qi/vdr); !invalid_id && kbx < blocks_per_row_x; kbx += blocks_per_iter) {
+        const int kby = kbx * (qk/QK8_1);
+        const int kqs = vdr * (threadIdx.x % (qi/vdr));
+
+        const int iqs = kqs;
+        const int bq8_offset = QR2_K * (iqs / QI8_1);
+        const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2);
+        const int whalf = iqs / QI8_1;            // 0: window at scales[0..7], 1: at [8..15]
+        const int lo    = scale_offset - 8*whalf; // 0 or 1 within the window
+        const block_q8_1 * bq8_1 = &y[kby];
+
+        int    u[QR2_K];
+        float d8[QR2_K];
+#pragma unroll
+        for (int i = 0; i < QR2_K; ++i) {
+            u[i]  = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
+            d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
+        }
+
+        const size_t pblk = pair_base + (size_t)kbx;
+        const uint2 v2  = qs2_soa[pblk*16u + (unsigned)iqs];
+        const uint2 dmw = dm2_soa[pblk];
+        const int4  scw = sc4_soa[pblk*2u + (unsigned)whalf];
+        const __half2 dm0 = *(const __half2 *)&dmw.x;
+        const __half2 dm1 = *(const __half2 *)&dmw.y;
+
+        tmp[0] += vec_dot_q2_K_q8_1_impl_mmvq_w((int)v2.x, u, (uint32_t)scw.x, (uint32_t)scw.y, lo, dm0, d8);
+        tmp[1] += vec_dot_q2_K_q8_1_impl_mmvq_w((int)v2.y, u, (uint32_t)scw.z, (uint32_t)scw.w, lo, dm1, d8);
+    }
+
+#pragma unroll
+    for (int i = 0; i < 2; ++i) {
+        tmp[i] = warp_reduce_sum<warp_size>(tmp[i]);
+    }
+
+    if (threadIdx.x < 2 && uint32_t(row0 + threadIdx.x) < nrows_x) {
+        dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = tmp[threadIdx.x];
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Host helpers.
+// ---------------------------------------------------------------------------
+static const int M   = 4096;   // out rows (down out_dim)
+static const int K   = 2048;   // in cols (expert_mid_dim)
+static const int NB  = K / QK_K;    // 8 superblocks per row
+static const int NYB = K / QK8_1;   // 64 q8_1 blocks per activation row
+static const int NEXP = 256;
+static const int NUSED = 6;
+
+static size_t nblk_layer() { return (size_t)NEXP * M * NB; }
+
+static void fill_raw_layer(std::vector<block_q2_K> &w, std::mt19937 &rng) {
+    std::uniform_int_distribution<int> byte(0, 255);
+    std::uniform_real_distribution<float> mag(0.001f, 0.05f);
+    for (block_q2_K &b : w) {
+        for (int i = 0; i < 16; i++) b.scales[i] = (uint8_t)byte(rng);
+        for (int i = 0; i < 64; i++) b.qs[i] = (uint8_t)byte(rng);
+        b.dm = __floats2half2_rn(mag(rng), mag(rng));
+    }
+}
+
+static void quantize_q8_1(const float *x, block_q8_1 *y, int n) {
+    for (int ib = 0; ib < n / QK8_1; ib++) {
+        const float *xb = x + ib*QK8_1;
+        float amax = 0.0f;
+        for (int j = 0; j < QK8_1; j++) { float a = fabsf(xb[j]); if (a > amax) amax = a; }
+        const float d = amax / 127.0f;
+        const float idf = d != 0.0f ? 1.0f/d : 0.0f;
+        int sum = 0;
+        for (int j = 0; j < QK8_1; j++) {
+            int q = (int)roundf(xb[j] * idf);
+            if (q > 127) q = 127;
+            if (q < -127) q = -127;
+            y[ib].qs[j] = (int8_t)q;
+            sum += q;
+        }
+        y[ib].ds = __floats2half2_rn(d, d * (float)sum);
+    }
+}
+
+// V1 repack: [dm nblk*4][pad64][scales nblk*16][pad64][qs nblk*64]
+static void repack_v1(const block_q2_K *raw, size_t nblk,
+                      __half2 *dm, uint8_t *scales, int *qs) {
+    for (size_t b = 0; b < nblk; b++) {
+        dm[b] = raw[b].dm;
+        memcpy(scales + b*16u, raw[b].scales, 16);
+        memcpy(qs + b*16u, raw[b].qs, 64);
+    }
+}
+
+// V2 repack: pair p covers raw blocks of rows (2p, 2p+1) within an expert.
+// Raw linear block index of (expert e, row r, block b) = (e*M + r)*NB + b.
+static void repack_v2(const block_q2_K *raw,
+                      uint2 *dm2, int4 *sc4, uint2 *qs2) {
+    const size_t npair_blk = (size_t)NEXP * (M/2) * NB;
+    for (size_t pb = 0; pb < npair_blk; pb++) {
+        const size_t b   = pb % NB;
+        const size_t pr  = (pb / NB) % (M/2);
+        const size_t e   = pb / ((size_t)NB * (M/2));
+        const block_q2_K *r0 = raw + ((e*M + 2*pr + 0)*NB + b);
+        const block_q2_K *r1 = raw + ((e*M + 2*pr + 1)*NB + b);
+        uint2 dmw;
+        memcpy(&dmw.x, &r0->dm, 4);
+        memcpy(&dmw.y, &r1->dm, 4);
+        dm2[pb] = dmw;
+        // sc4[pb*2 + h] = {row0 w(h)[0:4], row0 w(h)[4:8], row1 w(h)[0:4], row1 w(h)[4:8]}
+        for (int h = 0; h < 2; h++) {
+            int4 s;
+            memcpy(&s.x, r0->scales + 8*h + 0, 4);
+            memcpy(&s.y, r0->scales + 8*h + 4, 4);
+            memcpy(&s.z, r1->scales + 8*h + 0, 4);
+            memcpy(&s.w, r1->scales + 8*h + 4, 4);
+            sc4[pb*2 + h] = s;
+        }
+        for (int i = 0; i < 16; i++) {
+            uint2 q;
+            memcpy(&q.x, r0->qs + 4*i, 4);
+            memcpy(&q.y, r1->qs + 4*i, 4);
+            qs2[pb*16 + i] = q;
+        }
+    }
+}
+
+struct dev_layer {
+    block_q2_K *raw = nullptr;
+    __half2 *v1_dm = nullptr; uint8_t *v1_sc = nullptr; int *v1_qs = nullptr;
+    uint2 *v2_dm = nullptr; int4 *v2_sc = nullptr; uint2 *v2_qs = nullptr;
+};
+
+// Production launch config for the down leg (routed_moe_launch vec tier):
+// ncols_dst = n_assignments, nchannels_dst = 1, ids_stride = 1,
+// stride_col_dst = M, stride_channel_dst = M, stride_col_y = NYB.
+static void launch_ref(const dev_layer &L, const block_q8_1 *y, const int32_t *ids,
+                       float *dst, int ncols, cudaStream_t s) {
+    dim3 grid(M/2, 1);
+    dim3 block(32, ncols);
+    ref_mul_mat_vec_q_moe<2><<<grid, block, 0, s>>>(
+        L.raw, y, ids, dst,
+        (uint32_t)K, 1u, (uint32_t)M,
+        (uint32_t)NB, (uint32_t)NYB, (uint32_t)M,
+        (uint32_t)(M*NB), (uint32_t)NYB, (uint32_t)M,
+        (uint32_t)ncols, 1u);
+}
+
+static void launch_v1(const dev_layer &L, const block_q8_1 *y, const int32_t *ids,
+                      float *dst, int ncols, cudaStream_t s) {
+    dim3 grid(M/2, 1);
+    dim3 block(32, ncols);
+    v1_soa_moe_kernel<2><<<grid, block, 0, s>>>(
+        L.v1_dm, L.v1_sc, L.v1_qs, y, ids, dst,
+        (uint32_t)K, 1u, (uint32_t)M,
+        (uint32_t)NB, (uint32_t)NYB, (uint32_t)M,
+        (uint32_t)(M*NB), (uint32_t)NYB, (uint32_t)M,
+        (uint32_t)ncols, 1u);
+}
+
+static void launch_v2(const dev_layer &L, const block_q8_1 *y, const int32_t *ids,
+                      float *dst, int ncols, cudaStream_t s) {
+    dim3 grid(M/2, 1);
+    dim3 block(32, ncols);
+    v2_pair_soa_moe_kernel<<<grid, block, 0, s>>>(
+        L.v2_dm, L.v2_sc, L.v2_qs, y, ids, dst,
+        (uint32_t)K, 1u, (uint32_t)M,
+        (uint32_t)NYB, (uint32_t)M,
+        (uint32_t)NYB, (uint32_t)M,
+        (uint32_t)ncols, 1u);
+}
+
+int main(int argc, char **argv) {
+    int n_layers = 6;
+    int timing_iters = 300;
+    if (argc > 1) n_layers = atoi(argv[1]);
+    if (argc > 2) timing_iters = atoi(argv[2]);
+
+    const size_t nblk = nblk_layer();
+    const size_t raw_bytes = nblk * sizeof(block_q2_K);
+    printf("proto_m2_q2k: M=%d K=%d E=%d nb=%d layers=%d raw=%.1f MiB/layer\n",
+           M, K, NEXP, NB, n_layers, raw_bytes / 1048576.0);
+
+    std::mt19937 rng(0xC0FFEE);
+
+    // Host-side raw layer template + repacks (one host copy reused per layer
+    // with a per-layer byte perturbation so layers differ without n_layers
+    // full random generations).
+    std::vector<block_q2_K> hraw(nblk);
+    fill_raw_layer(hraw, rng);
+
+    std::vector<__half2> h_v1dm(nblk);
+    std::vector<uint8_t> h_v1sc(nblk * 16u);
+    std::vector<int>     h_v1qs(nblk * 16u);
+    std::vector<uint2>   h_v2dm(nblk / 2u);
+    std::vector<int4>    h_v2sc(nblk);            // nblk/2 pairs * 2
+    std::vector<uint2>   h_v2qs((nblk / 2u) * 16u);
+
+    std::vector<dev_layer> layers(n_layers);
+    for (int l = 0; l < n_layers; l++) {
+        if (l > 0) {
+            // cheap per-layer variation: rotate scales/qs bytes by l
+            for (size_t b = 0; b < nblk; b += 977u) {
+                hraw[b].qs[(l*7) & 63] ^= (uint8_t)(0x5A + l);
+                hraw[b].scales[(l*3) & 15] ^= (uint8_t)(0xA5 - l);
+            }
+        }
+        repack_v1(hraw.data(), nblk, h_v1dm.data(), h_v1sc.data(), h_v1qs.data());
+        repack_v2(hraw.data(), h_v2dm.data(), h_v2sc.data(), h_v2qs.data());
+
+        dev_layer &L = layers[l];
+        CHECK(cudaMalloc(&L.raw, raw_bytes));
+        CHECK(cudaMalloc(&L.v1_dm, nblk * 4u));
+        CHECK(cudaMalloc(&L.v1_sc, nblk * 16u));
+        CHECK(cudaMalloc(&L.v1_qs, nblk * 64u));
+        CHECK(cudaMalloc(&L.v2_dm, (nblk/2u) * 8u));
+        CHECK(cudaMalloc(&L.v2_sc, nblk * 16u));
+        CHECK(cudaMalloc(&L.v2_qs, (nblk/2u) * 16u * 8u));
+        CHECK(cudaMemcpy(L.raw, hraw.data(), raw_bytes, cudaMemcpyHostToDevice));
+        CHECK(cudaMemcpy(L.v1_dm, h_v1dm.data(), nblk * 4u, cudaMemcpyHostToDevice));
+        CHECK(cudaMemcpy(L.v1_sc, h_v1sc.data(), nblk * 16u, cudaMemcpyHostToDevice));
+        CHECK(cudaMemcpy(L.v1_qs, h_v1qs.data(), nblk * 64u, cudaMemcpyHostToDevice));
+        CHECK(cudaMemcpy(L.v2_dm, h_v2dm.data(), (nblk/2u) * 8u, cudaMemcpyHostToDevice));
+        CHECK(cudaMemcpy(L.v2_sc, h_v2sc.data(), nblk * 16u, cudaMemcpyHostToDevice));
+        CHECK(cudaMemcpy(L.v2_qs, h_v2qs.data(), (nblk/2u) * 16u * 8u, cudaMemcpyHostToDevice));
+    }
+
+    // Activations: MAX_COLS q8_1 rows.
+    const int MAX_COLS = 8;
+    std::vector<float> hx(MAX_COLS * K);
+    std::vector<block_q8_1> hy(MAX_COLS * NYB);
+    std::uniform_real_distribution<float> xf(-1.0f, 1.0f);
+    for (float &v : hx) v = xf(rng);
+    for (int c = 0; c < MAX_COLS; c++) quantize_q8_1(hx.data() + c*K, hy.data() + c*NYB, K);
+    block_q8_1 *dy = nullptr;
+    CHECK(cudaMalloc(&dy, hy.size() * sizeof(block_q8_1)));
+    CHECK(cudaMemcpy(dy, hy.data(), hy.size() * sizeof(block_q8_1), cudaMemcpyHostToDevice));
+
+    int32_t *dids = nullptr;
+    CHECK(cudaMalloc(&dids, MAX_COLS * sizeof(int32_t)));
+    float *dref = nullptr, *dv1 = nullptr, *dv2 = nullptr;
+    const size_t out_bytes = (size_t)MAX_COLS * M * sizeof(float);
+    CHECK(cudaMalloc(&dref, out_bytes));
+    CHECK(cudaMalloc(&dv1, out_bytes));
+    CHECK(cudaMalloc(&dv2, out_bytes));
+
+    cudaStream_t stream;
+    CHECK(cudaStreamCreate(&stream));
+
+    // ------------------------------------------------------------------
+    // Parity: seeds x id patterns x ncols on rotating layers.
+    // ------------------------------------------------------------------
+    std::vector<float> href(MAX_COLS * M), hv1(MAX_COLS * M), hv2(MAX_COLS * M);
+    const int ncols_cases[3] = {1, 6, 8};
+    int checks = 0, fails = 0;
+    std::uniform_int_distribution<int> expd(0, NEXP - 1);
+    for (int seed = 0; seed < 10; seed++) {
+        for (int pat = 0; pat < 4; pat++) {
+            for (int nc = 0; nc < 3; nc++) {
+                const int ncols = ncols_cases[nc];
+                const dev_layer &L = layers[(seed + pat + nc) % n_layers];
+                int32_t ids[MAX_COLS];
+                for (int i = 0; i < ncols; i++) {
+                    switch (pat) {
+                        case 0: ids[i] = expd(rng); break;                    // random
+                        case 1: ids[i] = (i < 2) ? expd(rng) : ids[i % 2];    // duplicates
+                                break;
+                        case 2: ids[i] = (i % 3 == 1) ? -1 : expd(rng); break; // sparse -1
+                        case 3: ids[i] = -1; break;                            // all -1
+                    }
+                }
+                CHECK(cudaMemcpy(dids, ids, ncols * sizeof(int32_t), cudaMemcpyHostToDevice));
+                CHECK(cudaMemset(dref, 0xCB, out_bytes));
+                CHECK(cudaMemset(dv1, 0xCB, out_bytes));
+                CHECK(cudaMemset(dv2, 0xCB, out_bytes));
+                launch_ref(L, dy, dids, dref, ncols, stream);
+                launch_v1(L, dy, dids, dv1, ncols, stream);
+                launch_v2(L, dy, dids, dv2, ncols, stream);
+                CHECK(cudaStreamSynchronize(stream));
+                CHECK(cudaGetLastError());
+                const size_t cmp = (size_t)ncols * M * sizeof(float);
+                CHECK(cudaMemcpy(href.data(), dref, cmp, cudaMemcpyDeviceToHost));
+                CHECK(cudaMemcpy(hv1.data(), dv1, cmp, cudaMemcpyDeviceToHost));
+                CHECK(cudaMemcpy(hv2.data(), dv2, cmp, cudaMemcpyDeviceToHost));
+                checks += 2;
+                if (memcmp(href.data(), hv1.data(), cmp) != 0) {
+                    fails++;
+                    printf("PARITY FAIL V1 seed=%d pat=%d ncols=%d\n", seed, pat, ncols);
+                }
+                if (memcmp(href.data(), hv2.data(), cmp) != 0) {
+                    fails++;
+                    printf("PARITY FAIL V2 seed=%d pat=%d ncols=%d\n", seed, pat, ncols);
+                }
+            }
+        }
+    }
+    printf("parity: %d/%d bit-identical%s\n", checks - fails, checks, fails ? "  <-- FAIL" : "");
+
+    // ------------------------------------------------------------------
+    // Capture/replay parity: capture each variant once, replay, compare to
+    // the eager output just validated.
+    // ------------------------------------------------------------------
+    {
+        int32_t ids[MAX_COLS];
+        for (int i = 0; i < NUSED; i++) ids[i] = expd(rng);
+        ids[3] = -1;
+        CHECK(cudaMemcpy(dids, ids, NUSED * sizeof(int32_t), cudaMemcpyHostToDevice));
+        const dev_layer &L = layers[0];
+        const size_t cmp = (size_t)NUSED * M * sizeof(float);
+
+        launch_ref(L, dy, dids, dref, NUSED, stream);
+        CHECK(cudaStreamSynchronize(stream));
+        CHECK(cudaMemcpy(href.data(), dref, cmp, cudaMemcpyDeviceToHost));
+
+        int replay_fails = 0;
+        for (int variant = 0; variant < 2; variant++) {
+            float *dout = variant == 0 ? dv1 : dv2;
+            CHECK(cudaMemset(dout, 0xCB, out_bytes));
+            cudaGraph_t graph;
+            cudaGraphExec_t exec;
+            CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));
+            if (variant == 0) launch_v1(L, dy, dids, dout, NUSED, stream);
+            else              launch_v2(L, dy, dids, dout, NUSED, stream);
+            CHECK(cudaStreamEndCapture(stream, &graph));
+            CHECK(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0));
+            CHECK(cudaGraphLaunch(exec, stream));
+            CHECK(cudaStreamSynchronize(stream));
+            CHECK(cudaMemcpy(hv1.data(), dout, cmp, cudaMemcpyDeviceToHost));
+            if (memcmp(href.data(), hv1.data(), cmp) != 0) {
+                replay_fails++;
+                printf("REPLAY FAIL V%d\n", variant + 1);
+            }
+            CHECK(cudaGraphExecDestroy(exec));
+            CHECK(cudaGraphDestroy(graph));
+        }
+        printf("capture/replay: %d/2 bit-identical%s\n", 2 - replay_fails, replay_fails ? "  <-- FAIL" : "");
+        fails += replay_fails;
+    }
+
+    // ------------------------------------------------------------------
+    // Timing: rotating layers, 6 distinct experts per launch (production
+    // decode shape), events around the whole loop.
+    // ------------------------------------------------------------------
+    {
+        // Pre-generate rotating id sets on device (one buffer per iteration
+        // slot would perturb the graph; production re-writes ids in place, so
+        // do the same: ids uploaded once per timing pass, kernels rotate
+        // layers only -- weight reads dominate).
+        int32_t ids[MAX_COLS];
+        bool used[NEXP] = {false};
+        int n = 0;
+        while (n < NUSED) {
+            int e = expd(rng);
+            if (used[e]) continue;
+            used[e] = true;
+            ids[n++] = e;
+        }
+        CHECK(cudaMemcpy(dids, ids, NUSED * sizeof(int32_t), cudaMemcpyHostToDevice));
+
+        const double gib = (double)NUSED * M * NB * 84.0;  // distinct-expert weight bytes/launch
+        cudaEvent_t e0, e1;
+        CHECK(cudaEventCreate(&e0));
+        CHECK(cudaEventCreate(&e1));
+
+        for (int variant = 0; variant < 3; variant++) {
+            const char *name = variant == 0 ? "ref-raw" : (variant == 1 ? "v1-soa " : "v2-pair");
+            float *dout = variant == 0 ? dref : (variant == 1 ? dv1 : dv2);
+            // warm
+            for (int it = 0; it < 30; it++) {
+                const dev_layer &L = layers[it % n_layers];
+                if (variant == 0) launch_ref(L, dy, dids, dout, NUSED, stream);
+                else if (variant == 1) launch_v1(L, dy, dids, dout, NUSED, stream);
+                else launch_v2(L, dy, dids, dout, NUSED, stream);
+            }
+            CHECK(cudaStreamSynchronize(stream));
+            CHECK(cudaEventRecord(e0, stream));
+            for (int it = 0; it < timing_iters; it++) {
+                const dev_layer &L = layers[it % n_layers];
+                if (variant == 0) launch_ref(L, dy, dids, dout, NUSED, stream);
+                else if (variant == 1) launch_v1(L, dy, dids, dout, NUSED, stream);
+                else launch_v2(L, dy, dids, dout, NUSED, stream);
+            }
+            CHECK(cudaEventRecord(e1, stream));
+            CHECK(cudaStreamSynchronize(stream));
+            float ms = 0.0f;
+            CHECK(cudaEventElapsedTime(&ms, e0, e1));
+            const double us = ms * 1000.0 / timing_iters;
+            printf("%s: %8.2f us/launch  %7.1f GB/s  (x43 = %.3f ms/step)\n",
+                   name, us, gib / (us * 1e-6) / 1e9, us * 43.0 / 1000.0);
+        }
+        CHECK(cudaEventDestroy(e0));
+        CHECK(cudaEventDestroy(e1));
+    }
+
+    printf(fails ? "RESULT: FAIL\n" : "RESULT: PASS\n");
+    return fails ? 1 : 0;
+}

cuda/mmq/test/proto_m2_qkv.cu

@@ -0,0 +1,604 @@
+/* proto_m2_qkv.cu — M2-Inc2 prototypes: QKV-post chain fusions.
+ *
+ * Three kill-switch-ready fusions over the decode QKV-post pool, each proven
+ * BIT-EXACT here against verbatim copies of the production kernels:
+ *
+ *   2a  dsv4_qkv_rms_norm_rows + q8_1 emission of the q row (kills the q_b
+ *       mmvq quantize prelude, quantize_q8_1 grid=4).
+ *   2b  head_rms_norm + rope_tail(q) fused, device-scalars pos source (the
+ *       existing fused kernel is dead code because it bakes host pos0 into
+ *       captures; this is its capture-safe twin).
+ *   2c  rope_tail(kv) + fp8_kv_quantize + store_raw_kv fused into ONE
+ *       8-block kernel: rope touches only the 64-float rotary tail, fp8
+ *       covers only the 448 nope elems (disjoint), and the production fp8
+ *       kernel's 5.9 us is 7 SEQUENTIAL 64-elem groups in a single 64-thread
+ *       block — the groups are independent, so they become 7 parallel blocks
+ *       (+1 rope/tail block).  Raw-store f16 round-trip preserved.
+ *
+ * Build: nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 proto_m2_qkv.cu -o proto_qkv
+ */
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <cstdint>
+#include <cmath>
+
+#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
+    printf("CUDA error %s at %s:%d\n", cudaGetErrorString(e_), __FILE__, __LINE__); exit(1); } } while (0)
+
+#define N_HEAD   64u
+#define HEAD_DIM 512u
+#define N_ROT    64u
+#define Q_RANK   1024u
+#define KV_DIM   512u
+#define RAW_CAP  1024u
+#define RMS_EPS  1e-6f
+
+/* mirror of the production decode-scalars substrate (only pos0/raw_row read) */
+struct ds4_decode_scalars {
+    uint32_t pos0, raw_row, raw_start, n_raw, n_comp, emit_phase,
+             comp_row, index_row, flags, token;
+};
+
+typedef struct { __half2 ds; int8_t qs[32]; } block_q8_1_t;
+
+/* ============ verbatim production device helpers ============ */
+__device__ static float rope_yarn_ramp_dev(float low, float high, int i0) {
+    float y = ((float)(i0 / 2) - low) / fmaxf(0.001f, high - low);
+    return 1.0f - fminf(1.0f, fmaxf(0.0f, y));
+}
+__device__ static float dsv4_e4m3fn_value_dev(int i) {
+    int exp = (i >> 3) & 15;
+    int mant = i & 7;
+    if (exp == 0) return (float)mant * 0.001953125f;
+    return (1.0f + (float)mant * 0.125f) * exp2f((float)exp - 7.0f);
+}
+__device__ static float dsv4_e4m3fn_dequant_dev(float x) {
+    float sign = x < 0.0f ? -1.0f : 1.0f;
+    float ax = fminf(fabsf(x), 448.0f);
+    int lo = 0, hi = 126;
+    while (lo < hi) {
+        int mid = (lo + hi + 1) >> 1;
+        if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
+        else hi = mid - 1;
+    }
+    int best = lo;
+    if (best < 126) {
+        float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
+        float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
+        if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
+    }
+    return sign * dsv4_e4m3fn_value_dev(best);
+}
+
+/* ============ verbatim production kernels (references) ============ */
+__global__ static void dsv4_qkv_rms_norm_rows_kernel(
+        float *q_out, const float *q, const float *q_w, uint32_t q_n,
+        float *kv_out, const float *kv, const float *kv_w, uint32_t kv_n,
+        uint32_t rows, float eps) {
+    const uint32_t row = blockIdx.x;
+    const uint32_t which = blockIdx.y;
+    if (row >= rows || which > 1u) return;
+    const uint32_t n = which == 0u ? q_n : kv_n;
+    const float *xr = (which == 0u ? q : kv) + (uint64_t)row * n;
+    float *orow = (which == 0u ? q_out : kv_out) + (uint64_t)row * n;
+    const float *w = which == 0u ? q_w : kv_w;
+    float sum = 0.0f;
+    for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
+        const float v = xr[i];
+        sum += v * v;
+    }
+    __shared__ float partial[256];
+    partial[threadIdx.x] = sum;
+    __syncthreads();
+    for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
+        if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
+        __syncthreads();
+    }
+    const float scale = rsqrtf(partial[0] / (float)n + eps);
+    for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
+        orow[i] = xr[i] * scale * w[i];
+    }
+}
+
+__global__ static void head_rms_norm_kernel(float *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, float eps) {
+    uint32_t row = blockIdx.x;
+    if (row >= n_tok * n_head) return;
+    float *xr = x + (uint64_t)row * head_dim;
+    float sum = 0.0f;
+    for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) {
+        float v = xr[i];
+        sum += v * v;
+    }
+    __shared__ float partial[256];
+    partial[threadIdx.x] = sum;
+    __syncthreads();
+    for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
+        if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
+        __syncthreads();
+    }
+    float scale = rsqrtf(partial[0] / (float)head_dim + eps);
+    for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) xr[i] *= scale;
+}
+
+__global__ static void rope_tail_scalars_kernel(
+        float *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, uint32_t n_rot,
+        const struct ds4_decode_scalars * __restrict__ scalars,
+        int32_t pos_offset, uint32_t pos_stride, uint32_t n_ctx_orig, int inverse,
+        float freq_base, float freq_scale, float ext_factor, float attn_factor,
+        float beta_fast, float beta_slow) {
+    uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
+    uint32_t pairs = n_tok * n_head * (n_rot / 2);
+    if (gid >= pairs) return;
+    uint32_t pair = gid % (n_rot / 2);
+    uint32_t tmp = gid / (n_rot / 2);
+    uint32_t h = tmp % n_head;
+    uint32_t t = tmp / n_head;
+    uint32_t n_nope = head_dim - n_rot;
+    uint32_t i = pair * 2;
+    const uint32_t pos0 = (uint32_t)((int32_t)scalars->pos0 + pos_offset);
+    float corr0 = 0.0f, corr1 = 0.0f;
+    if (ext_factor != 0.0f) {
+        float denom = 2.0f * logf(freq_base);
+        corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
+        corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
+        corr0 = fmaxf(0.0f, corr0);
+        corr1 = fminf((float)(n_rot - 1), corr1);
+    }
+    float theta_extrap = (float)(pos0 + t * pos_stride) * powf(freq_base, -((float)i) / (float)n_rot);
+    float theta_interp = freq_scale * theta_extrap;
+    float theta = theta_interp;
+    float mscale = attn_factor;
+    if (ext_factor != 0.0f) {
+        float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
+        theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
+        mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
+    }
+    float c = cosf(theta) * mscale;
+    float s = sinf(theta) * mscale;
+    if (inverse) s = -s;
+    float *tail = x + ((uint64_t)t * n_head + h) * head_dim + n_nope;
+    float x0 = tail[i];
+    float x1 = tail[i + 1];
+    tail[i] = x0 * c - x1 * s;
+    tail[i + 1] = x0 * s + x1 * c;
+}
+
+__global__ static void fp8_kv_quantize_kernel(
+        float *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot,
+        unsigned char * __restrict__ codes_base, float * __restrict__ scale_base) {
+    uint32_t row = blockIdx.x;
+    uint32_t tid = threadIdx.x;
+    uint32_t n_nope = head_dim - n_rot;
+    const uint64_t row_stride = (uint64_t)n_nope + (uint64_t)n_rot * sizeof(float);
+    const uint64_t scale_stride = (uint64_t)(n_nope / 64u);
+    float *xr = x + (uint64_t)row * head_dim;
+    unsigned char *codes_row = codes_base
+        ? codes_base + (uint64_t)row * row_stride : (unsigned char *)0;
+    float *scale_row = scale_base
+        ? scale_base + (uint64_t)row * scale_stride : (float *)0;
+    __shared__ float scratch[64];
+    for (uint32_t off = 0; off < n_nope; off += 64) {
+        float v = 0.0f;
+        if (off + tid < n_nope) v = xr[off + tid];
+        scratch[tid] = off + tid < n_nope ? fabsf(v) : 0.0f;
+        __syncthreads();
+        for (uint32_t stride = 32; stride > 0; stride >>= 1) {
+            if (tid < stride) scratch[tid] = fmaxf(scratch[tid], scratch[tid + stride]);
+            __syncthreads();
+        }
+        float scale = exp2f(ceilf(log2f(fmaxf(scratch[0], 1.0e-4f) / 448.0f)));
+        if (off + tid < n_nope) {
+            float clamp = fminf(448.0f, fmaxf(-448.0f, v / scale));
+            xr[off + tid] = dsv4_e4m3fn_dequant_dev(clamp) * scale;
+            if (codes_row) codes_row[off + tid] = 0; /* unused in decode path */
+        }
+        if (scale_row && tid == 0) scale_row[off / 64u] = scale;
+        __syncthreads();
+    }
+    (void)row_stride;
+}
+
+__global__ static void store_raw_kv_batch_kernel(
+        float *raw, const float *kv,
+        uint32_t raw_cap, uint32_t pos0, uint32_t n_tokens, uint32_t head_dim,
+        const int32_t * __restrict__ positions, const int32_t * __restrict__ seq_id,
+        const struct ds4_decode_scalars * __restrict__ s_override) {
+    uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    uint64_t n = (uint64_t)n_tokens * head_dim;
+    if (gid >= n) return;
+    uint32_t d = gid % head_dim;
+    uint32_t t = gid / head_dim;
+    uint32_t row;
+    if (s_override != NULL && n_tokens == 1u) {
+        row = s_override->raw_row;
+    } else {
+        uint32_t pos = positions ? (uint32_t)positions[t] : pos0 + t;
+        uint32_t seq_base = seq_id ? (uint32_t)seq_id[t] * raw_cap : 0u;
+        row = seq_base + (pos % raw_cap);
+    }
+    raw[(uint64_t)row * head_dim + d] = __half2float(__float2half(kv[(uint64_t)t * head_dim + d]));
+}
+
+/* q8_1 reference (vendored quantize_q8_1 semantics, contiguous case) */
+__global__ static void ref_quantize_q8_1_kernel(block_q8_1_t *y, const float *x, uint32_t n) {
+    const uint32_t i = blockIdx.x * 32u + threadIdx.x;
+    if (i >= n) return;
+    const float xi = x[i];
+    float amax = fabsf(xi);
+    float sum = xi;
+    for (uint32_t off = 16u; off > 0u; off >>= 1u) {
+        amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off, 32));
+        sum += __shfl_xor_sync(0xffffffffu, sum, off, 32);
+    }
+    const float d = amax / 127.0f;
+    const int8_t q = amax == 0.0f ? (int8_t)0 : (int8_t)roundf(xi / d);
+    y[blockIdx.x].qs[threadIdx.x] = q;
+    if (threadIdx.x == 0) y[blockIdx.x].ds = __floats2half2_rn(d, sum);
+}
+
+/* ============ 2a: qkv_rms_norm_rows + q8_1(q row) ============ */
+__global__ static void dsv4_qkv_rms_norm_rows_q81_kernel(
+        float *q_out, const float *q, const float *q_w, uint32_t q_n,
+        float *kv_out, const float *kv, const float *kv_w, uint32_t kv_n,
+        uint32_t rows, float eps, block_q8_1_t *q81) {
+    const uint32_t row = blockIdx.x;
+    const uint32_t which = blockIdx.y;
+    if (row >= rows || which > 1u) return;
+    const uint32_t n = which == 0u ? q_n : kv_n;
+    const float *xr = (which == 0u ? q : kv) + (uint64_t)row * n;
+    float *orow = (which == 0u ? q_out : kv_out) + (uint64_t)row * n;
+    const float *w = which == 0u ? q_w : kv_w;
+    float sum = 0.0f;
+    for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
+        const float v = xr[i];
+        sum += v * v;
+    }
+    __shared__ float partial[256];
+    partial[threadIdx.x] = sum;
+    __syncthreads();
+    for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
+        if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
+        __syncthreads();
+    }
+    const float scale = rsqrtf(partial[0] / (float)n + eps);
+    for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
+        orow[i] = xr[i] * scale * w[i];
+    }
+    /* q8_1 emission of the q row (row 0 decode): re-read the just-written
+     * orow values (identical bits) in warp-contiguous granules. */
+    if (which == 0u && q81 != NULL && row == 0u && (n & 31u) == 0u) {
+        __syncthreads();
+        const uint32_t lane = threadIdx.x & 31u;
+        const uint32_t warp = threadIdx.x >> 5u;
+        const uint32_t nwarp = blockDim.x >> 5u;
+        for (uint32_t qb = warp; qb < n / 32u; qb += nwarp) {
+            const float v = orow[qb * 32u + lane];
+            float amax = fabsf(v);
+            float s = v;
+            for (uint32_t off = 16u; off > 0u; off >>= 1u) {
+                amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off, 32));
+                s += __shfl_xor_sync(0xffffffffu, s, off, 32);
+            }
+            const float d1 = amax / 127.0f;
+            const int8_t qv = amax == 0.0f ? (int8_t)0 : (int8_t)roundf(v / d1);
+            q81[qb].qs[lane] = qv;
+            if (lane == 0u) q81[qb].ds = __floats2half2_rn(d1, s);
+        }
+    }
+}
+
+/* ============ 2b: head_rms + rope_tail(q), scalars pos ============ */
+__global__ static void head_rms_norm_rope_tail_scalars_kernel(
+        float *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, uint32_t n_rot,
+        const struct ds4_decode_scalars * __restrict__ scalars,
+        int32_t pos_offset, uint32_t pos_stride, uint32_t n_ctx_orig, int inverse,
+        float freq_base, float freq_scale, float ext_factor, float attn_factor,
+        float beta_fast, float beta_slow, float eps) {
+    uint32_t row = blockIdx.x;
+    if (row >= n_tok * n_head) return;
+    uint32_t t = row / n_head;
+    float *xr = x + (uint64_t)row * head_dim;
+    float sum = 0.0f;
+    for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) {
+        float v = xr[i];
+        sum += v * v;
+    }
+    __shared__ float partial[256];
+    partial[threadIdx.x] = sum;
+    __syncthreads();
+    for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
+        if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
+        __syncthreads();
+    }
+    const float scale = rsqrtf(partial[0] / (float)head_dim + eps);
+    const uint32_t n_nope = head_dim - n_rot;
+    for (uint32_t i = threadIdx.x; i < n_nope; i += blockDim.x) {
+        xr[i] *= scale;
+    }
+    const uint32_t pos0 = (uint32_t)((int32_t)scalars->pos0 + pos_offset);
+    float corr0 = 0.0f, corr1 = 0.0f;
+    if (ext_factor != 0.0f) {
+        float denom = 2.0f * logf(freq_base);
+        corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
+        corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
+        corr0 = fmaxf(0.0f, corr0);
+        corr1 = fminf((float)(n_rot - 1), corr1);
+    }
+    for (uint32_t pair = threadIdx.x; pair < n_rot / 2; pair += blockDim.x) {
+        uint32_t i = pair * 2u;
+        float theta_extrap = (float)(pos0 + t * pos_stride) * powf(freq_base, -((float)i) / (float)n_rot);
+        float theta_interp = freq_scale * theta_extrap;
+        float theta = theta_interp;
+        float mscale = attn_factor;
+        if (ext_factor != 0.0f) {
+            float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
+            theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
+            mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
+        }
+        float c = cosf(theta) * mscale;
+        float s = sinf(theta) * mscale;
+        if (inverse) s = -s;
+        float *tail = xr + n_nope;
+        float x0 = tail[i] * scale;
+        float x1 = tail[i + 1] * scale;
+        tail[i] = x0 * c - x1 * s;
+        tail[i + 1] = x0 * s + x1 * c;
+    }
+}
+
+/* ============ 2c: rope_tail(kv) + fp8 quantize + raw store ============ */
+/* Grid: (n_nope/64 + 1) blocks x 64 threads.  Blocks [0, n_nope/64) each own
+ * one independent fp8 group (identical 64-wide shared-tree reduce as the
+ * production kernel, so scales and rounded values are bit-exact) and store
+ * their post-quant region to the raw row.  The last block ropes the rotary
+ * tail (32 pairs, verbatim rope_tail_scalars math at n_head=1/t=0) and
+ * stores the rotated tail.  Disjoint regions -> no cross-block ordering. */
+__global__ static void kv_rope_fp8_store_scalars_kernel(
+        float *kv, float *raw, uint32_t raw_cap, uint32_t head_dim, uint32_t n_rot,
+        const struct ds4_decode_scalars * __restrict__ scalars,
+        int32_t pos_offset, uint32_t n_ctx_orig, int inverse,
+        float freq_base, float freq_scale, float ext_factor, float attn_factor,
+        float beta_fast, float beta_slow) {
+    const uint32_t n_nope = head_dim - n_rot;
+    const uint32_t b = blockIdx.x;
+    const uint32_t tid = threadIdx.x;
+    const uint32_t row = scalars->raw_row;
+    float *raw_row = raw + (uint64_t)(row % raw_cap) * head_dim;
+    if (b < n_nope / 64u) {
+        /* fp8 group b: verbatim group math from fp8_kv_quantize_kernel */
+        const uint32_t off = b * 64u;
+        float v = 0.0f;
+        if (off + tid < n_nope) v = kv[off + tid];
+        __shared__ float scratch[64];
+        scratch[tid] = off + tid < n_nope ? fabsf(v) : 0.0f;
+        __syncthreads();
+        for (uint32_t stride = 32; stride > 0; stride >>= 1) {
+            if (tid < stride) scratch[tid] = fmaxf(scratch[tid], scratch[tid + stride]);
+            __syncthreads();
+        }
+        float scale = exp2f(ceilf(log2f(fmaxf(scratch[0], 1.0e-4f) / 448.0f)));
+        if (off + tid < n_nope) {
+            float clamp = fminf(448.0f, fmaxf(-448.0f, v / scale));
+            const float qv = dsv4_e4m3fn_dequant_dev(clamp) * scale;
+            kv[off + tid] = qv;
+            raw_row[off + tid] = __half2float(__float2half(qv));
+        }
+    } else {
+        /* rope the rotary tail then store it */
+        const uint32_t pos0 = (uint32_t)((int32_t)scalars->pos0 + pos_offset);
+        float corr0 = 0.0f, corr1 = 0.0f;
+        if (ext_factor != 0.0f) {
+            float denom = 2.0f * logf(freq_base);
+            corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
+            corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
+            corr0 = fmaxf(0.0f, corr0);
+            corr1 = fminf((float)(n_rot - 1), corr1);
+        }
+        float *tail = kv + n_nope;
+        if (tid < n_rot / 2u) {
+            const uint32_t i = tid * 2u;
+            float theta_extrap = (float)pos0 * powf(freq_base, -((float)i) / (float)n_rot);
+            float theta_interp = freq_scale * theta_extrap;
+            float theta = theta_interp;
+            float mscale = attn_factor;
+            if (ext_factor != 0.0f) {
+                float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
+                theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
+                mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
+            }
+            float c = cosf(theta) * mscale;
+            float s = sinf(theta) * mscale;
+            if (inverse) s = -s;
+            float x0 = tail[i];
+            float x1 = tail[i + 1];
+            const float r0 = x0 * c - x1 * s;
+            const float r1 = x0 * s + x1 * c;
+            tail[i] = r0;
+            tail[i + 1] = r1;
+            raw_row[n_nope + i] = __half2float(__float2half(r0));
+            raw_row[n_nope + i + 1] = __half2float(__float2half(r1));
+        }
+    }
+}
+
+/* ============ harness ============ */
+static float *dmalloc_f(uint64_t n) { float *p; CK(cudaMalloc(&p, n * sizeof(float))); return p; }
+static void fill_rand(float *d, uint64_t n, unsigned seed) {
+    float *h = (float *)malloc(n * sizeof(float));
+    srand(seed);
+    for (uint64_t i = 0; i < n; i++) h[i] = ((float)rand() / RAND_MAX - 0.5f) * 4.0f;
+    CK(cudaMemcpy(d, h, n * sizeof(float), cudaMemcpyHostToDevice));
+    free(h);
+}
+static int cmp_dev(const char *what, const float *a, const float *b, uint64_t n) {
+    float *ha = (float *)malloc(n * sizeof(float)), *hb = (float *)malloc(n * sizeof(float));
+    CK(cudaMemcpy(ha, a, n * sizeof(float), cudaMemcpyDeviceToHost));
+    CK(cudaMemcpy(hb, b, n * sizeof(float), cudaMemcpyDeviceToHost));
+    int bad = memcmp(ha, hb, n * sizeof(float)) != 0;
+    if (bad) {
+        uint64_t k = 0;
+        for (; k < n; k++) if (ha[k] != hb[k]) break;
+        printf("  %s: MISMATCH at %llu (%a vs %a)\n", what, (unsigned long long)k, ha[k], hb[k]);
+    } else printf("  %s: bit-identical (%llu floats)\n", what, (unsigned long long)n);
+    free(ha); free(hb);
+    return bad;
+}
+
+int main() {
+    cudaStream_t st;
+    CK(cudaStreamCreate(&st));
+    int total_bad = 0;
+
+    /* rope parameter sets: plain and yarn-extended */
+    struct RP { float fb, fs, ef, af, bf, bs; uint32_t ctxo; const char *name; };
+    RP rps[2] = {
+        { 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 0u, "plain" },
+        { 10000.0f, 0.25f, 1.0f, 1.0f, 32.0f, 1.0f, 4096u, "yarn" },
+    };
+    uint32_t poss[3] = { 0u, 999u, 32767u };
+
+    ds4_decode_scalars hs;
+    memset(&hs, 0, sizeof(hs));
+    ds4_decode_scalars *ds;
+    CK(cudaMalloc(&ds, sizeof(hs)));
+
+    /* ---- 2a ---- */
+    {
+        float *q_in = dmalloc_f(Q_RANK), *kv_in = dmalloc_f(KV_DIM);
+        float *qw = dmalloc_f(Q_RANK), *kvw = dmalloc_f(KV_DIM);
+        float *qo_r = dmalloc_f(Q_RANK), *kvo_r = dmalloc_f(KV_DIM);
+        float *qo_f = dmalloc_f(Q_RANK), *kvo_f = dmalloc_f(KV_DIM);
+        block_q8_1_t *q81_r, *q81_f;
+        CK(cudaMalloc(&q81_r, (Q_RANK / 32) * sizeof(block_q8_1_t)));
+        CK(cudaMalloc(&q81_f, (Q_RANK / 32) * sizeof(block_q8_1_t)));
+        fill_rand(q_in, Q_RANK, 11); fill_rand(kv_in, KV_DIM, 12);
+        fill_rand(qw, Q_RANK, 13); fill_rand(kvw, KV_DIM, 14);
+        dim3 g(1, 2, 1);
+        dsv4_qkv_rms_norm_rows_kernel<<<g, 256, 0, st>>>(qo_r, q_in, qw, Q_RANK, kvo_r, kv_in, kvw, KV_DIM, 1, RMS_EPS);
+        ref_quantize_q8_1_kernel<<<Q_RANK / 32, 32, 0, st>>>(q81_r, qo_r, Q_RANK);
+        dsv4_qkv_rms_norm_rows_q81_kernel<<<g, 256, 0, st>>>(qo_f, q_in, qw, Q_RANK, kvo_f, kv_in, kvw, KV_DIM, 1, RMS_EPS, q81_f);
+        CK(cudaStreamSynchronize(st));
+        printf("2a qkv_rms + q8_1 emission:\n");
+        total_bad += cmp_dev("q_out", qo_r, qo_f, Q_RANK);
+        total_bad += cmp_dev("kv_out", kvo_r, kvo_f, KV_DIM);
+        char *h1 = (char *)malloc((Q_RANK / 32) * sizeof(block_q8_1_t));
+        char *h2 = (char *)malloc((Q_RANK / 32) * sizeof(block_q8_1_t));
+        CK(cudaMemcpy(h1, q81_r, (Q_RANK / 32) * sizeof(block_q8_1_t), cudaMemcpyDeviceToHost));
+        CK(cudaMemcpy(h2, q81_f, (Q_RANK / 32) * sizeof(block_q8_1_t), cudaMemcpyDeviceToHost));
+        int bad = memcmp(h1, h2, (Q_RANK / 32) * sizeof(block_q8_1_t)) != 0;
+        printf("  q8_1 codes: %s\n", bad ? "MISMATCH" : "bit-identical");
+        total_bad += bad;
+        free(h1); free(h2);
+    }
+
+    /* ---- 2b ---- */
+    for (int r = 0; r < 2; r++) for (int p = 0; p < 3; p++) {
+        RP rp = rps[r];
+        hs.pos0 = poss[p];
+        CK(cudaMemcpy(ds, &hs, sizeof(hs), cudaMemcpyHostToDevice));
+        const uint64_t qn = (uint64_t)N_HEAD * HEAD_DIM;
+        float *qa = dmalloc_f(qn), *qb = dmalloc_f(qn);
+        fill_rand(qa, qn, 100 + r * 10 + p);
+        CK(cudaMemcpy(qb, qa, qn * sizeof(float), cudaMemcpyDeviceToDevice));
+        head_rms_norm_kernel<<<N_HEAD, 256, 0, st>>>(qa, 1, N_HEAD, HEAD_DIM, RMS_EPS);
+        uint32_t pairs = N_HEAD * (N_ROT / 2);
+        rope_tail_scalars_kernel<<<(pairs + 255) / 256, 256, 0, st>>>(qa, 1, N_HEAD, HEAD_DIM, N_ROT,
+                ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
+        head_rms_norm_rope_tail_scalars_kernel<<<N_HEAD, 256, 0, st>>>(qb, 1, N_HEAD, HEAD_DIM, N_ROT,
+                ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs, RMS_EPS);
+        CK(cudaStreamSynchronize(st));
+        printf("2b head_rms+rope (%s, pos=%u):\n", rp.name, poss[p]);
+        total_bad += cmp_dev("q", qa, qb, qn);
+        CK(cudaFree(qa)); CK(cudaFree(qb));
+    }
+
+    /* ---- 2c ---- */
+    for (int r = 0; r < 2; r++) for (int p = 0; p < 3; p++) {
+        RP rp = rps[r];
+        hs.pos0 = poss[p];
+        hs.raw_row = poss[p] % RAW_CAP;
+        CK(cudaMemcpy(ds, &hs, sizeof(hs), cudaMemcpyHostToDevice));
+        float *kva = dmalloc_f(KV_DIM), *kvb = dmalloc_f(KV_DIM);
+        float *rawa = dmalloc_f((uint64_t)RAW_CAP * HEAD_DIM), *rawb = dmalloc_f((uint64_t)RAW_CAP * HEAD_DIM);
+        CK(cudaMemset(rawa, 0, (uint64_t)RAW_CAP * HEAD_DIM * sizeof(float)));
+        CK(cudaMemset(rawb, 0, (uint64_t)RAW_CAP * HEAD_DIM * sizeof(float)));
+        fill_rand(kva, KV_DIM, 200 + r * 10 + p);
+        CK(cudaMemcpy(kvb, kva, KV_DIM * sizeof(float), cudaMemcpyDeviceToDevice));
+        /* reference chain: rope(kv) -> fp8 -> store */
+        uint32_t pairs = 1 * (N_ROT / 2);
+        rope_tail_scalars_kernel<<<(pairs + 255) / 256, 256, 0, st>>>(kva, 1, 1, HEAD_DIM, N_ROT,
+                ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
+        fp8_kv_quantize_kernel<<<1, 64, 0, st>>>(kva, 1, HEAD_DIM, N_ROT, NULL, NULL);
+        store_raw_kv_batch_kernel<<<(HEAD_DIM + 255) / 256, 256, 0, st>>>(rawa, kva, RAW_CAP, 0, 1, HEAD_DIM, NULL, NULL, ds);
+        /* fused */
+        kv_rope_fp8_store_scalars_kernel<<<(HEAD_DIM - N_ROT) / 64 + 1, 64, 0, st>>>(kvb, rawb, RAW_CAP, HEAD_DIM, N_ROT,
+                ds, 0, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
+        CK(cudaStreamSynchronize(st));
+        printf("2c kv rope+fp8+store (%s, pos=%u, row=%u):\n", rp.name, poss[p], hs.raw_row);
+        total_bad += cmp_dev("kv", kva, kvb, KV_DIM);
+        total_bad += cmp_dev("raw row", rawa + (uint64_t)hs.raw_row * HEAD_DIM,
+                             rawb + (uint64_t)hs.raw_row * HEAD_DIM, HEAD_DIM);
+        CK(cudaFree(kva)); CK(cudaFree(kvb)); CK(cudaFree(rawa)); CK(cudaFree(rawb));
+    }
+
+    /* ---- timing: reference chains vs fused, 43-layer sweeps ---- */
+    {
+        const int REPS = 500, LAYERS = 43;
+        hs.pos0 = 1000u; hs.raw_row = 1000u % RAW_CAP;
+        CK(cudaMemcpy(ds, &hs, sizeof(hs), cudaMemcpyHostToDevice));
+        RP rp = rps[1];
+        const uint64_t qn = (uint64_t)N_HEAD * HEAD_DIM;
+        float *q = dmalloc_f(qn), *kv = dmalloc_f(KV_DIM);
+        float *raw = dmalloc_f((uint64_t)RAW_CAP * HEAD_DIM);
+        float *qi = dmalloc_f(Q_RANK), *kvi = dmalloc_f(KV_DIM);
+        float *qw = dmalloc_f(Q_RANK), *kvw = dmalloc_f(KV_DIM);
+        float *qo = dmalloc_f(Q_RANK), *kvo = dmalloc_f(KV_DIM);
+        block_q8_1_t *q81; CK(cudaMalloc(&q81, (Q_RANK / 32) * sizeof(block_q8_1_t)));
+        fill_rand(q, qn, 31); fill_rand(kv, KV_DIM, 32);
+        fill_rand(qi, Q_RANK, 33); fill_rand(kvi, KV_DIM, 34);
+        fill_rand(qw, Q_RANK, 35); fill_rand(kvw, KV_DIM, 36);
+        cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
+        uint32_t pairs_q = N_HEAD * (N_ROT / 2), pairs_kv = N_ROT / 2;
+        dim3 g2(1, 2, 1);
+        float ms;
+        /* reference chain */
+        for (int w = 0; w < 2; w++) {
+            if (w) CK(cudaEventRecord(e0, st));
+            for (int rr = 0; rr < (w ? REPS : 10); rr++) for (int l = 0; l < LAYERS; l++) {
+                dsv4_qkv_rms_norm_rows_kernel<<<g2, 256, 0, st>>>(qo, qi, qw, Q_RANK, kvo, kvi, kvw, KV_DIM, 1, RMS_EPS);
+                ref_quantize_q8_1_kernel<<<Q_RANK / 32, 32, 0, st>>>(q81, qo, Q_RANK);
+                head_rms_norm_kernel<<<N_HEAD, 256, 0, st>>>(q, 1, N_HEAD, HEAD_DIM, RMS_EPS);
+                rope_tail_scalars_kernel<<<(pairs_q + 255) / 256, 256, 0, st>>>(q, 1, N_HEAD, HEAD_DIM, N_ROT, ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
+                rope_tail_scalars_kernel<<<1, 256, 0, st>>>(kv, 1, 1, HEAD_DIM, N_ROT, ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
+                fp8_kv_quantize_kernel<<<1, 64, 0, st>>>(kv, 1, HEAD_DIM, N_ROT, NULL, NULL);
+                store_raw_kv_batch_kernel<<<(HEAD_DIM + 255) / 256, 256, 0, st>>>(raw, kv, RAW_CAP, 0, 1, HEAD_DIM, NULL, NULL, ds);
+                (void)pairs_kv;
+            }
+            if (w) { CK(cudaEventRecord(e1, st)); CK(cudaStreamSynchronize(st)); }
+            else CK(cudaStreamSynchronize(st));
+        }
+        CK(cudaEventElapsedTime(&ms, e0, e1));
+        printf("timing reference chain (7 launches): %.2f us/layer (x43 = %.3f ms/step)\n",
+               1000.0f * ms / (REPS * LAYERS), ms / REPS);
+        /* fused chain */
+        for (int w = 0; w < 2; w++) {
+            if (w) CK(cudaEventRecord(e0, st));
+            for (int rr = 0; rr < (w ? REPS : 10); rr++) for (int l = 0; l < LAYERS; l++) {
+                dsv4_qkv_rms_norm_rows_q81_kernel<<<g2, 256, 0, st>>>(qo, qi, qw, Q_RANK, kvo, kvi, kvw, KV_DIM, 1, RMS_EPS, q81);
+                head_rms_norm_rope_tail_scalars_kernel<<<N_HEAD, 256, 0, st>>>(q, 1, N_HEAD, HEAD_DIM, N_ROT, ds, 0, 1u, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs, RMS_EPS);
+                kv_rope_fp8_store_scalars_kernel<<<(HEAD_DIM - N_ROT) / 64 + 1, 64, 0, st>>>(kv, raw, RAW_CAP, HEAD_DIM, N_ROT, ds, 0, rp.ctxo, 0, rp.fb, rp.fs, rp.ef, rp.af, rp.bf, rp.bs);
+            }
+            if (w) { CK(cudaEventRecord(e1, st)); CK(cudaStreamSynchronize(st)); }
+            else CK(cudaStreamSynchronize(st));
+        }
+        CK(cudaEventElapsedTime(&ms, e0, e1));
+        printf("timing fused chain (3 launches):   %.2f us/layer (x43 = %.3f ms/step)\n",
+               1000.0f * ms / (REPS * LAYERS), ms / REPS);
+    }
+
+    printf(total_bad ? "PROTO_M2_QKV FAILED (%d)\n" : "PROTO_M2_QKV_DONE all bit-exact\n", total_bad);
+    return total_bad != 0;
+}

cuda/mmq/test/proto_m2_router.cu

@@ -0,0 +1,939 @@
+/* proto_m2_router.cu — M2-Inc3 prototype: fused router
+ * (f16 split-K logits matmul + combine + top-6 select) as ONE cooperative
+ * kernel, replacing the production 3-kernel chain
+ *   matmul_f16_splitk_kernel (grid 256x8, 256 thr)
+ * + matmul_f16_splitk_combine_kernel (1 blk, 256 thr)
+ * + router_select_warp_topk_kernel (1 blk, 32x4)
+ * at the router shape in_dim=4096, out_dim=256 (=n_expert), n_tok=1, ksplit=8.
+ *
+ * Gate: all four outputs (logits, probs, selected, weights) BIT-IDENTICAL to
+ * the reference chain, in topk-bias / no-bias / hash modes, incl. exact-tie
+ * logits; capture==eager bit-identity; then per-layer timing vs the chain.
+ *
+ * Bit-exactness argument for the matmul phase: the ref (row,kseg) block gives
+ * each thread tid<64 exactly one 8-half chunk (seg 512 = 64 chunks); its
+ * block_reduce_sum_256 reduces warp 0 (chunks 0..31) and warp 1 (chunks
+ * 32..63) with warp_sum_f32, then the cross-warp tree only adds those two
+ * warp sums and exact zeros.  The fused warp-per-tile form computes the SAME
+ * chunk values at the SAME lanes, applies the SAME warp_sum_f32 trees, and
+ * adds s0+s1 -- identical up to x+0.0f==x (breaks only for a partial that is
+ * exactly -0.0f, i.e. an all-zero activation segment; unreachable for
+ * rms-normed activations, and select-invariant even then).
+ *
+ * Build (GB10):
+ *   nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 proto_m2_router.cu -o proto_m2_router
+ */
+#include <cooperative_groups.h>
+#include <cuda_fp16.h>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <cstdint>
+#include <cmath>
+
+#define CHECK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
+    fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(e_), __FILE__, __LINE__); exit(1); } } while (0)
+
+#define N_IN     4096u
+#define N_EXPERT 256u
+#define KSPLIT   8u
+#define SEG      512u
+#define N_TOP    6u
+
+/* ------------------------------------------------------------------ */
+/* Vendored production device functions + kernels (ds4_cuda.cu, verbatim) */
+
+struct ds4_decode_scalars {
+    uint32_t pos0;
+    uint32_t raw_row;
+    uint32_t raw_start;
+    uint32_t n_raw;
+    uint32_t n_comp;
+    uint32_t emit_phase;
+    uint32_t comp_row;
+    uint32_t index_row;
+    uint32_t flags;
+    uint32_t token;
+};
+
+__device__ static float warp_sum_f32(float v) {
+    for (int offset = 16; offset > 0; offset >>= 1) {
+        v += __shfl_down_sync(0xffffffffu, v, offset);
+    }
+    return v;
+}
+
+__device__ static float block_reduce_sum_256(float v) {
+    v = warp_sum_f32(v);
+    __shared__ float s[8];
+    const uint32_t lane = threadIdx.x & 31u;
+    const uint32_t warp = threadIdx.x >> 5u;
+    if (lane == 0u) s[warp] = v;
+    __syncthreads();
+    float r = 0.0f;
+    if (warp == 0u) {
+        r = (lane < 8u) ? s[lane] : 0.0f;
+        r += __shfl_down_sync(0xffffffffu, r, 4);
+        r += __shfl_down_sync(0xffffffffu, r, 2);
+        r += __shfl_down_sync(0xffffffffu, r, 1);
+    }
+    return r;
+}
+
+__global__ static void matmul_f16_splitk_kernel(
+        float *partial,
+        const __half *w,
+        const float *x,
+        uint64_t in_dim,
+        uint64_t out_dim,
+        uint32_t ksplit,
+        int vec_ok) {
+    const uint64_t row = (uint64_t)blockIdx.x;
+    const uint32_t kseg = blockIdx.y;
+    if (row >= out_dim) return;
+    uint64_t seg = (in_dim + (uint64_t)ksplit - 1u) / (uint64_t)ksplit;
+    seg = (seg + 7u) & ~(uint64_t)7u;
+    const uint64_t k0 = (uint64_t)kseg * seg;
+    if (k0 >= in_dim) {
+        if (threadIdx.x == 0u) partial[row * (uint64_t)ksplit + (uint64_t)kseg] = 0.0f;
+        return;
+    }
+    uint64_t k1 = k0 + seg;
+    if (k1 > in_dim) k1 = in_dim;
+    const __half *wr = w + row * in_dim;
+    const uint32_t tid = threadIdx.x;
+    const uint32_t nt = blockDim.x;
+    float sum = 0.0f;
+    if (vec_ok) {
+        const uint64_t vend = k0 + ((k1 - k0) & ~(uint64_t)7u);
+        for (uint64_t i = k0 + (uint64_t)tid * 8u; i < vend; i += (uint64_t)nt * 8u) {
+            const uint4 wv = *reinterpret_cast<const uint4 *>(wr + i);
+            const float4 xa = *reinterpret_cast<const float4 *>(x + i);
+            const float4 xb = *reinterpret_cast<const float4 *>(x + i + 4);
+            const __half2 *h = reinterpret_cast<const __half2 *>(&wv);
+            const float2 w0 = __half22float2(h[0]);
+            const float2 w1 = __half22float2(h[1]);
+            const float2 w2 = __half22float2(h[2]);
+            const float2 w3 = __half22float2(h[3]);
+            sum += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+                 + w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
+        }
+        for (uint64_t i = vend + tid; i < k1; i += nt) sum += __half2float(wr[i]) * x[i];
+    } else {
+        for (uint64_t i = k0 + tid; i < k1; i += nt) sum += __half2float(wr[i]) * x[i];
+    }
+    sum = block_reduce_sum_256(sum);
+    if (threadIdx.x == 0u) partial[row * (uint64_t)ksplit + (uint64_t)kseg] = sum;
+}
+
+__global__ static void matmul_f16_splitk_combine_kernel(
+        float *out,
+        const float *partial,
+        uint64_t out_dim,
+        uint32_t ksplit) {
+    const uint64_t row = (uint64_t)blockIdx.x * (uint64_t)blockDim.x + (uint64_t)threadIdx.x;
+    if (row >= out_dim) return;
+    const float *p = partial + row * (uint64_t)ksplit;
+    float s = 0.0f;
+    for (uint32_t k = 0u; k < ksplit; k++) s += p[k];
+    out[row] = s;
+}
+
+__device__ static float softplus_dev(float x) {
+    if (x > 20.0f) return x;
+    if (x < -20.0f) return expf(x);
+    return log1pf(expf(x));
+}
+
+__device__ __forceinline__ static bool router_score_better(float av, uint32_t ai, float bv, uint32_t bi) {
+    return av > bv || (av == bv && ai < bi);
+}
+
+__global__ static void router_select_warp_topk_kernel(
+        int32_t *selected,
+        float *weights,
+        float *probs,
+        const float *bias,
+        const int32_t *hash,
+        const float *logits,
+        const int32_t *tokens,
+        int32_t token_scalar,
+        uint32_t hash_rows,
+        uint32_t n_tokens,
+        int has_bias,
+        int hash_mode,
+        const struct ds4_decode_scalars *scalars) {
+    const uint32_t lane = threadIdx.x;
+    const uint32_t row_in_block = threadIdx.y;
+    const uint32_t t = blockIdx.x * blockDim.y + row_in_block;
+    if (t >= n_tokens || lane >= 32u) return;
+
+    const float *log = logits + (uint64_t)t * 256u;
+    float *prob = probs + (uint64_t)t * 256u;
+    int32_t *sel = selected + (uint64_t)t * 6u;
+    float *w = weights + (uint64_t)t * 6u;
+    __shared__ float sprob[4][256];
+    float local_prob[8];
+    float local_score[8];
+
+    #pragma unroll
+    for (uint32_t j = 0; j < 8u; j++) {
+        const uint32_t e = lane + j * 32u;
+        const float p = sqrtf(softplus_dev(log[e]));
+        local_prob[j] = p;
+        local_score[j] = p + (has_bias ? bias[e] : 0.0f);
+        sprob[row_in_block][e] = p;
+        prob[e] = p;
+    }
+    __syncwarp();
+
+    if (hash_mode) {
+        if (lane == 0) {
+            int32_t tok = tokens ? tokens[t]
+                                 : (scalars ? (int32_t)scalars->token : token_scalar);
+            if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0;
+            const int32_t *row = hash + (uint64_t)tok * 6u;
+            float sum = 0.0f;
+            #pragma unroll
+            for (uint32_t j = 0; j < 6u; j++) {
+                const int32_t e = row[j];
+                sel[j] = e;
+                const float v = (e >= 0 && e < 256) ? sprob[row_in_block][(uint32_t)e] : 0.0f;
+                w[j] = v;
+                sum += v;
+            }
+            sum = fmaxf(sum, 6.103515625e-5f);
+            #pragma unroll
+            for (uint32_t j = 0; j < 6u; j++) w[j] = w[j] / sum * 1.5f;
+        }
+        return;
+    }
+
+    float out_prob[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+    uint32_t out_idx[6] = {0, 0, 0, 0, 0, 0};
+    #pragma unroll
+    for (uint32_t k = 0; k < 6u; k++) {
+        float best_score = -INFINITY;
+        float best_prob = 0.0f;
+        uint32_t best_idx = UINT32_MAX;
+        #pragma unroll
+        for (uint32_t j = 0; j < 8u; j++) {
+            const uint32_t e = lane + j * 32u;
+            const float s = local_score[j];
+            if (router_score_better(s, e, best_score, best_idx)) {
+                best_score = s;
+                best_prob = local_prob[j];
+                best_idx = e;
+            }
+        }
+        #pragma unroll
+        for (uint32_t mask = 16u; mask > 0u; mask >>= 1u) {
+            const float other_score = __shfl_xor_sync(0xffffffffu, best_score, mask);
+            const float other_prob = __shfl_xor_sync(0xffffffffu, best_prob, mask);
+            const uint32_t other_idx = __shfl_xor_sync(0xffffffffu, best_idx, mask);
+            if (router_score_better(other_score, other_idx, best_score, best_idx)) {
+                best_score = other_score;
+                best_prob = other_prob;
+                best_idx = other_idx;
+            }
+        }
+        #pragma unroll
+        for (uint32_t j = 0; j < 8u; j++) {
+            const uint32_t e = lane + j * 32u;
+            if (e == best_idx) local_score[j] = -INFINITY;
+        }
+        if (lane == 0) {
+            out_idx[k] = best_idx;
+            out_prob[k] = best_prob;
+        }
+    }
+
+    if (lane == 0) {
+        float sum = 0.0f;
+        #pragma unroll
+        for (uint32_t j = 0; j < 6u; j++) {
+            sel[j] = (int32_t)out_idx[j];
+            w[j] = out_prob[j];
+            sum += out_prob[j];
+        }
+        sum = fmaxf(sum, 6.103515625e-5f);
+        #pragma unroll
+        for (uint32_t j = 0; j < 6u; j++) w[j] = w[j] / sum * 1.5f;
+    }
+}
+
+/* ------------------------------------------------------------------ */
+/* FUSED: one cooperative kernel for the whole router stage. */
+
+__global__ static void router_fused_coop_kernel(
+        float *logits,            /* [256] */
+        float *partials,          /* [256*8] dedicated scratch */
+        int32_t *selected,        /* [6] */
+        float *weights,           /* [6] */
+        float *probs,             /* [256] */
+        const __half *w,          /* [256][4096] f16 */
+        const float *x,           /* [4096] */
+        const float *bias,        /* [256] or NULL */
+        const int32_t *hash,      /* [hash_rows][6] or NULL */
+        const int32_t *tokens,    /* decode: NULL */
+        int32_t token_scalar,
+        uint32_t hash_rows,
+        int has_bias,
+        int hash_mode,
+        const struct ds4_decode_scalars *scalars) {
+    __shared__ float sbias[256];
+    __shared__ float sprob[256];
+    __shared__ int32_t shash[6];
+    const uint32_t lane = threadIdx.x & 31u;
+    const uint32_t warp = threadIdx.x >> 5u;
+
+    /* Phase 0 (block 0): prefetch the tail's small model-map reads so their
+     * latency hides behind the matmul phase.  Verbatim copies -- the select
+     * below consumes identical values. */
+    if (blockIdx.x == 0u) {
+        if (has_bias) sbias[threadIdx.x] = bias[threadIdx.x];
+        if (hash_mode && threadIdx.x == 0u) {
+            int32_t tok = tokens ? tokens[0]
+                                 : (scalars ? (int32_t)scalars->token : token_scalar);
+            if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0;
+            const int32_t *row = hash + (uint64_t)tok * 6u;
+            #pragma unroll
+            for (uint32_t j = 0; j < 6u; j++) shash[j] = row[j];
+        }
+    }
+
+    /* Phase 1: split-K logits matmul, one warp per (row, kseg) tile.
+     * Chunk values, lane placement, and both warp_sum_f32 trees replicate the
+     * ref block's warps 0/1 exactly; s0+s1 matches its cross-warp tree, which
+     * otherwise only adds exact zeros. */
+    const uint32_t nwarp = gridDim.x * (blockDim.x >> 5u);
+    for (uint32_t t = blockIdx.x * (blockDim.x >> 5u) + warp;
+         t < N_EXPERT * KSPLIT; t += nwarp) {
+        const uint32_t row = t >> 3u;
+        const uint32_t kseg = t & 7u;
+        const uint64_t k0 = (uint64_t)kseg * SEG;
+        const __half *wr = w + (uint64_t)row * N_IN;
+        float sum0 = 0.0f;
+        float sum1 = 0.0f;
+        {
+            const uint64_t i = k0 + (uint64_t)lane * 8u;
+            const uint4 wv = *reinterpret_cast<const uint4 *>(wr + i);
+            const float4 xa = *reinterpret_cast<const float4 *>(x + i);
+            const float4 xb = *reinterpret_cast<const float4 *>(x + i + 4);
+            const __half2 *h = reinterpret_cast<const __half2 *>(&wv);
+            const float2 w0 = __half22float2(h[0]);
+            const float2 w1 = __half22float2(h[1]);
+            const float2 w2 = __half22float2(h[2]);
+            const float2 w3 = __half22float2(h[3]);
+            sum0 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+                  + w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
+        }
+        {
+            const uint64_t i = k0 + ((uint64_t)lane + 32u) * 8u;
+            const uint4 wv = *reinterpret_cast<const uint4 *>(wr + i);
+            const float4 xa = *reinterpret_cast<const float4 *>(x + i);
+            const float4 xb = *reinterpret_cast<const float4 *>(x + i + 4);
+            const __half2 *h = reinterpret_cast<const __half2 *>(&wv);
+            const float2 w0 = __half22float2(h[0]);
+            const float2 w1 = __half22float2(h[1]);
+            const float2 w2 = __half22float2(h[2]);
+            const float2 w3 = __half22float2(h[3]);
+            sum1 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+                  + w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
+        }
+        const float s0 = warp_sum_f32(sum0);
+        const float s1 = warp_sum_f32(sum1);
+        if (lane == 0u) partials[(uint64_t)row * KSPLIT + kseg] = s0 + s1;
+    }
+
+    cooperative_groups::this_grid().sync();
+    if (blockIdx.x != 0u) return;
+
+    /* Phase 2: combine partials in fixed kseg order (verbatim combine kernel;
+     * out_dim 256 == one 256-thread block). */
+    {
+        const uint64_t row = (uint64_t)threadIdx.x;
+        const float *p = partials + row * KSPLIT;
+        float s = 0.0f;
+        for (uint32_t k = 0u; k < KSPLIT; k++) s += p[k];
+        logits[row] = s;
+    }
+    __syncthreads();
+
+    /* Phase 3: verbatim router_select_warp_topk_kernel warp body (t = 0,
+     * row_in_block = 0); bias/hash reads swapped for the prefetched copies. */
+    if (warp != 0u) return;
+    const float *log = logits;
+    float local_prob[8];
+    float local_score[8];
+    #pragma unroll
+    for (uint32_t j = 0; j < 8u; j++) {
+        const uint32_t e = lane + j * 32u;
+        const float p = sqrtf(softplus_dev(log[e]));
+        local_prob[j] = p;
+        local_score[j] = p + (has_bias ? sbias[e] : 0.0f);
+        sprob[e] = p;
+        probs[e] = p;
+    }
+    __syncwarp();
+
+    if (hash_mode) {
+        if (lane == 0u) {
+            float sum = 0.0f;
+            #pragma unroll
+            for (uint32_t j = 0; j < 6u; j++) {
+                const int32_t e = shash[j];
+                selected[j] = e;
+                const float v = (e >= 0 && e < 256) ? sprob[(uint32_t)e] : 0.0f;
+                weights[j] = v;
+                sum += v;
+            }
+            sum = fmaxf(sum, 6.103515625e-5f);
+            #pragma unroll
+            for (uint32_t j = 0; j < 6u; j++) weights[j] = weights[j] / sum * 1.5f;
+        }
+        return;
+    }
+
+    float out_prob[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+    uint32_t out_idx[6] = {0, 0, 0, 0, 0, 0};
+    #pragma unroll
+    for (uint32_t k = 0; k < 6u; k++) {
+        float best_score = -INFINITY;
+        float best_prob = 0.0f;
+        uint32_t best_idx = UINT32_MAX;
+        #pragma unroll
+        for (uint32_t j = 0; j < 8u; j++) {
+            const uint32_t e = lane + j * 32u;
+            const float s = local_score[j];
+            if (router_score_better(s, e, best_score, best_idx)) {
+                best_score = s;
+                best_prob = local_prob[j];
+                best_idx = e;
+            }
+        }
+        #pragma unroll
+        for (uint32_t mask = 16u; mask > 0u; mask >>= 1u) {
+            const float other_score = __shfl_xor_sync(0xffffffffu, best_score, mask);
+            const float other_prob = __shfl_xor_sync(0xffffffffu, best_prob, mask);
+            const uint32_t other_idx = __shfl_xor_sync(0xffffffffu, best_idx, mask);
+            if (router_score_better(other_score, other_idx, best_score, best_idx)) {
+                best_score = other_score;
+                best_prob = other_prob;
+                best_idx = other_idx;
+            }
+        }
+        #pragma unroll
+        for (uint32_t j = 0; j < 8u; j++) {
+            const uint32_t e = lane + j * 32u;
+            if (e == best_idx) local_score[j] = -INFINITY;
+        }
+        if (lane == 0) {
+            out_idx[k] = best_idx;
+            out_prob[k] = best_prob;
+        }
+    }
+
+    if (lane == 0) {
+        float sum = 0.0f;
+        #pragma unroll
+        for (uint32_t j = 0; j < 6u; j++) {
+            selected[j] = (int32_t)out_idx[j];
+            weights[j] = out_prob[j];
+            sum += out_prob[j];
+        }
+        sum = fmaxf(sum, 6.103515625e-5f);
+        #pragma unroll
+        for (uint32_t j = 0; j < 6u; j++) weights[j] = weights[j] / sum * 1.5f;
+    }
+}
+
+/* V2: identical numerics, but each warp owns TWO tiles per pass with all four
+ * DRAM-latency w-vector loads hoisted before any math, doubling per-warp
+ * memory-level parallelism (x reads stay L1-resident after warmup). */
+__global__ static void router_fused_coop_kernel_v2(
+        float *logits,
+        float *partials,
+        int32_t *selected,
+        float *weights,
+        float *probs,
+        const __half *w,
+        const float *x,
+        const float *bias,
+        const int32_t *hash,
+        const int32_t *tokens,
+        int32_t token_scalar,
+        uint32_t hash_rows,
+        int has_bias,
+        int hash_mode,
+        const struct ds4_decode_scalars *scalars) {
+    __shared__ float sbias[256];
+    __shared__ float sprob[256];
+    __shared__ int32_t shash[6];
+    const uint32_t lane = threadIdx.x & 31u;
+    const uint32_t warp = threadIdx.x >> 5u;
+
+    if (blockIdx.x == 0u) {
+        if (has_bias) sbias[threadIdx.x] = bias[threadIdx.x];
+        if (hash_mode && threadIdx.x == 0u) {
+            int32_t tok = tokens ? tokens[0]
+                                 : (scalars ? (int32_t)scalars->token : token_scalar);
+            if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0;
+            const int32_t *row = hash + (uint64_t)tok * 6u;
+            #pragma unroll
+            for (uint32_t j = 0; j < 6u; j++) shash[j] = row[j];
+        }
+    }
+
+    const uint32_t nwarp = gridDim.x * (blockDim.x >> 5u);
+    const uint32_t gwarp = blockIdx.x * (blockDim.x >> 5u) + warp;
+    for (uint32_t t0 = gwarp; t0 < N_EXPERT * KSPLIT; t0 += 2u * nwarp) {
+        const uint32_t t1 = t0 + nwarp;
+        const int have1 = t1 < N_EXPERT * KSPLIT;
+        const uint64_t k0a = (uint64_t)(t0 & 7u) * SEG;
+        const __half *wra = w + (uint64_t)(t0 >> 3u) * N_IN;
+        const uint64_t ia0 = k0a + (uint64_t)lane * 8u;
+        const uint64_t ia1 = k0a + ((uint64_t)lane + 32u) * 8u;
+        const uint64_t k0b = have1 ? (uint64_t)(t1 & 7u) * SEG : k0a;
+        const __half *wrb = have1 ? w + (uint64_t)(t1 >> 3u) * N_IN : wra;
+        const uint64_t ib0 = k0b + (uint64_t)lane * 8u;
+        const uint64_t ib1 = k0b + ((uint64_t)lane + 32u) * 8u;
+        /* all four w vectors in flight before any dependent math */
+        const uint4 wva0 = *reinterpret_cast<const uint4 *>(wra + ia0);
+        const uint4 wva1 = *reinterpret_cast<const uint4 *>(wra + ia1);
+        const uint4 wvb0 = *reinterpret_cast<const uint4 *>(wrb + ib0);
+        const uint4 wvb1 = *reinterpret_cast<const uint4 *>(wrb + ib1);
+        float s0, s1;
+        {
+            float sum0 = 0.0f;
+            float sum1 = 0.0f;
+            {
+                const float4 xa = *reinterpret_cast<const float4 *>(x + ia0);
+                const float4 xb = *reinterpret_cast<const float4 *>(x + ia0 + 4);
+                const __half2 *h = reinterpret_cast<const __half2 *>(&wva0);
+                const float2 w0 = __half22float2(h[0]);
+                const float2 w1 = __half22float2(h[1]);
+                const float2 w2 = __half22float2(h[2]);
+                const float2 w3 = __half22float2(h[3]);
+                sum0 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+                      + w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
+            }
+            {
+                const float4 xa = *reinterpret_cast<const float4 *>(x + ia1);
+                const float4 xb = *reinterpret_cast<const float4 *>(x + ia1 + 4);
+                const __half2 *h = reinterpret_cast<const __half2 *>(&wva1);
+                const float2 w0 = __half22float2(h[0]);
+                const float2 w1 = __half22float2(h[1]);
+                const float2 w2 = __half22float2(h[2]);
+                const float2 w3 = __half22float2(h[3]);
+                sum1 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+                      + w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
+            }
+            s0 = warp_sum_f32(sum0);
+            s1 = warp_sum_f32(sum1);
+        }
+        float s2, s3;
+        if (have1) {
+            float sum0 = 0.0f;
+            float sum1 = 0.0f;
+            {
+                const float4 xa = *reinterpret_cast<const float4 *>(x + ib0);
+                const float4 xb = *reinterpret_cast<const float4 *>(x + ib0 + 4);
+                const __half2 *h = reinterpret_cast<const __half2 *>(&wvb0);
+                const float2 w0 = __half22float2(h[0]);
+                const float2 w1 = __half22float2(h[1]);
+                const float2 w2 = __half22float2(h[2]);
+                const float2 w3 = __half22float2(h[3]);
+                sum0 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+                      + w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
+            }
+            {
+                const float4 xa = *reinterpret_cast<const float4 *>(x + ib1);
+                const float4 xb = *reinterpret_cast<const float4 *>(x + ib1 + 4);
+                const __half2 *h = reinterpret_cast<const __half2 *>(&wvb1);
+                const float2 w0 = __half22float2(h[0]);
+                const float2 w1 = __half22float2(h[1]);
+                const float2 w2 = __half22float2(h[2]);
+                const float2 w3 = __half22float2(h[3]);
+                sum1 += w0.x * xa.x + w0.y * xa.y + w1.x * xa.z + w1.y * xa.w
+                      + w2.x * xb.x + w2.y * xb.y + w3.x * xb.z + w3.y * xb.w;
+            }
+            s2 = warp_sum_f32(sum0);
+            s3 = warp_sum_f32(sum1);
+        }
+        if (lane == 0u) {
+            partials[(uint64_t)(t0 >> 3u) * KSPLIT + (t0 & 7u)] = s0 + s1;
+            if (have1) partials[(uint64_t)(t1 >> 3u) * KSPLIT + (t1 & 7u)] = s2 + s3;
+        }
+    }
+
+    cooperative_groups::this_grid().sync();
+    if (blockIdx.x != 0u) return;
+
+    {
+        const uint64_t row = (uint64_t)threadIdx.x;
+        const float *p = partials + row * KSPLIT;
+        float s = 0.0f;
+        for (uint32_t k = 0u; k < KSPLIT; k++) s += p[k];
+        logits[row] = s;
+    }
+    __syncthreads();
+
+    if (warp != 0u) return;
+    const float *log = logits;
+    float local_prob[8];
+    float local_score[8];
+    #pragma unroll
+    for (uint32_t j = 0; j < 8u; j++) {
+        const uint32_t e = lane + j * 32u;
+        const float p = sqrtf(softplus_dev(log[e]));
+        local_prob[j] = p;
+        local_score[j] = p + (has_bias ? sbias[e] : 0.0f);
+        sprob[e] = p;
+        probs[e] = p;
+    }
+    __syncwarp();
+
+    if (hash_mode) {
+        if (lane == 0u) {
+            float sum = 0.0f;
+            #pragma unroll
+            for (uint32_t j = 0; j < 6u; j++) {
+                const int32_t e = shash[j];
+                selected[j] = e;
+                const float v = (e >= 0 && e < 256) ? sprob[(uint32_t)e] : 0.0f;
+                weights[j] = v;
+                sum += v;
+            }
+            sum = fmaxf(sum, 6.103515625e-5f);
+            #pragma unroll
+            for (uint32_t j = 0; j < 6u; j++) weights[j] = weights[j] / sum * 1.5f;
+        }
+        return;
+    }
+
+    float out_prob[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+    uint32_t out_idx[6] = {0, 0, 0, 0, 0, 0};
+    #pragma unroll
+    for (uint32_t k = 0; k < 6u; k++) {
+        float best_score = -INFINITY;
+        float best_prob = 0.0f;
+        uint32_t best_idx = UINT32_MAX;
+        #pragma unroll
+        for (uint32_t j = 0; j < 8u; j++) {
+            const uint32_t e = lane + j * 32u;
+            const float s = local_score[j];
+            if (router_score_better(s, e, best_score, best_idx)) {
+                best_score = s;
+                best_prob = local_prob[j];
+                best_idx = e;
+            }
+        }
+        #pragma unroll
+        for (uint32_t mask = 16u; mask > 0u; mask >>= 1u) {
+            const float other_score = __shfl_xor_sync(0xffffffffu, best_score, mask);
+            const float other_prob = __shfl_xor_sync(0xffffffffu, best_prob, mask);
+            const uint32_t other_idx = __shfl_xor_sync(0xffffffffu, best_idx, mask);
+            if (router_score_better(other_score, other_idx, best_score, best_idx)) {
+                best_score = other_score;
+                best_prob = other_prob;
+                best_idx = other_idx;
+            }
+        }
+        #pragma unroll
+        for (uint32_t j = 0; j < 8u; j++) {
+            const uint32_t e = lane + j * 32u;
+            if (e == best_idx) local_score[j] = -INFINITY;
+        }
+        if (lane == 0) {
+            out_idx[k] = best_idx;
+            out_prob[k] = best_prob;
+        }
+    }
+
+    if (lane == 0) {
+        float sum = 0.0f;
+        #pragma unroll
+        for (uint32_t j = 0; j < 6u; j++) {
+            selected[j] = (int32_t)out_idx[j];
+            weights[j] = out_prob[j];
+            sum += out_prob[j];
+        }
+        sum = fmaxf(sum, 6.103515625e-5f);
+        #pragma unroll
+        for (uint32_t j = 0; j < 6u; j++) weights[j] = weights[j] / sum * 1.5f;
+    }
+}
+
+/* ------------------------------------------------------------------ */
+/* Host harness */
+
+struct outs {
+    float logits[N_EXPERT];
+    float probs[N_EXPERT];
+    int32_t sel[N_TOP];
+    float wgt[N_TOP];
+};
+
+static void fill_rand(float *p, size_t n, unsigned seed, float scale) {
+    unsigned s = seed;
+    for (size_t i = 0; i < n; i++) {
+        s = s * 1664525u + 1013904223u;
+        p[i] = ((float)(s >> 8) / (float)(1u << 24) - 0.5f) * 2.0f * scale;
+    }
+}
+
+struct dev_bufs {
+    __half *w;
+    float *x, *bias, *logits, *probs, *partials_ref, *partials_fu, *wgt;
+    int32_t *hash, *sel;
+    ds4_decode_scalars *scalars;
+    uint32_t hash_rows;
+};
+
+static void run_ref(const dev_bufs &b, int has_bias, int hash_mode, int32_t token_via_scalars) {
+    dim3 sgrid(N_EXPERT, KSPLIT, 1);
+    matmul_f16_splitk_kernel<<<sgrid, 256>>>(b.partials_ref, b.w, b.x, N_IN, N_EXPERT, KSPLIT, 1);
+    matmul_f16_splitk_combine_kernel<<<1, 256>>>(b.logits, b.partials_ref, N_EXPERT, KSPLIT);
+    dim3 block(32, 4, 1);
+    router_select_warp_topk_kernel<<<1, block>>>(b.sel, b.wgt, b.probs,
+            has_bias && !hash_mode ? b.bias : NULL,
+            hash_mode ? b.hash : NULL,
+            b.logits, NULL, /*token_scalar=*/-1, b.hash_rows, 1,
+            has_bias && !hash_mode, hash_mode,
+            token_via_scalars ? b.scalars : NULL);
+    CHECK(cudaGetLastError());
+}
+
+static int launch_fused_k(void *kernel, const dev_bufs &b, int has_bias, int hash_mode,
+                          int32_t token_via_scalars, unsigned grid_blocks, cudaStream_t stream) {
+    float *logits = b.logits, *partials = b.partials_fu, *probs = b.probs, *wgt = b.wgt, *x = b.x;
+    const __half *w = b.w;
+    const float *bias = has_bias && !hash_mode ? b.bias : NULL;
+    const int32_t *hash = hash_mode ? b.hash : NULL;
+    int32_t *sel = b.sel;
+    const int32_t *tokens = NULL;
+    int32_t token_scalar = -1;
+    uint32_t hash_rows = b.hash_rows;
+    int hb = has_bias && !hash_mode;
+    int hm = hash_mode;
+    const ds4_decode_scalars *scalars = token_via_scalars ? b.scalars : NULL;
+    void *args[] = {
+        (void *)&logits, (void *)&partials, (void *)&sel, (void *)&wgt, (void *)&probs,
+        (void *)&w, (void *)&x, (void *)&bias, (void *)&hash, (void *)&tokens,
+        (void *)&token_scalar, (void *)&hash_rows, (void *)&hb, (void *)&hm, (void *)&scalars };
+    cudaError_t err = cudaLaunchCooperativeKernel(
+            kernel,
+            dim3(grid_blocks, 1, 1), dim3(256, 1, 1), args, 0, stream);
+    if (err != cudaSuccess) {
+        fprintf(stderr, "coop launch failed (%u blk): %s\n", grid_blocks, cudaGetErrorString(err));
+        return 0;
+    }
+    return 1;
+}
+
+static int launch_fused(const dev_bufs &b, int has_bias, int hash_mode, int32_t token_via_scalars,
+                        unsigned grid_blocks, cudaStream_t stream) {
+    return launch_fused_k((void *)router_fused_coop_kernel, b, has_bias, hash_mode,
+                          token_via_scalars, grid_blocks, stream);
+}
+
+static void read_outs(const dev_bufs &b, outs *o) {
+    CHECK(cudaMemcpy(o->logits, b.logits, sizeof(o->logits), cudaMemcpyDeviceToHost));
+    CHECK(cudaMemcpy(o->probs, b.probs, sizeof(o->probs), cudaMemcpyDeviceToHost));
+    CHECK(cudaMemcpy(o->sel, b.sel, sizeof(o->sel), cudaMemcpyDeviceToHost));
+    CHECK(cudaMemcpy(o->wgt, b.wgt, sizeof(o->wgt), cudaMemcpyDeviceToHost));
+}
+
+static int cmp_outs(const outs &a, const outs &bb, const char *tag) {
+    int bad = 0;
+    if (memcmp(a.logits, bb.logits, sizeof(a.logits))) { printf("  [%s] logits MISMATCH\n", tag); bad = 1; }
+    if (memcmp(a.probs, bb.probs, sizeof(a.probs)))    { printf("  [%s] probs MISMATCH\n", tag); bad = 1; }
+    if (memcmp(a.sel, bb.sel, sizeof(a.sel)))          { printf("  [%s] selected MISMATCH\n", tag); bad = 1; }
+    if (memcmp(a.wgt, bb.wgt, sizeof(a.wgt)))          { printf("  [%s] weights MISMATCH\n", tag); bad = 1; }
+    if (bad) {
+        for (int i = 0; i < 8; i++) printf("    logit[%d] ref=%.9g fused=%.9g\n", i, a.logits[i], bb.logits[i]);
+        for (int i = 0; i < 6; i++) printf("    sel[%d] ref=%d fused=%d  w ref=%.9g fused=%.9g\n",
+                                           i, a.sel[i], bb.sel[i], a.wgt[i], bb.wgt[i]);
+    }
+    return bad;
+}
+
+int main() {
+    int dev = 0, coop = 0, n_sm = 0, bps = 0;
+    CHECK(cudaGetDevice(&dev));
+    CHECK(cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, dev));
+    CHECK(cudaDeviceGetAttribute(&n_sm, cudaDevAttrMultiProcessorCount, dev));
+    CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&bps, router_fused_coop_kernel, 256, 0));
+    const unsigned max_blocks = (unsigned)(bps * n_sm);
+    printf("coop=%d n_sm=%d blocks/sm=%d max co-resident=%u\n", coop, n_sm, bps, max_blocks);
+    if (!coop) { printf("NO cooperative launch support\n"); return 1; }
+
+    dev_bufs b = {};
+    b.hash_rows = 1000;
+    CHECK(cudaMalloc(&b.w, (size_t)N_EXPERT * N_IN * sizeof(__half)));
+    CHECK(cudaMalloc(&b.x, N_IN * sizeof(float)));
+    CHECK(cudaMalloc(&b.bias, N_EXPERT * sizeof(float)));
+    CHECK(cudaMalloc(&b.logits, N_EXPERT * sizeof(float)));
+    CHECK(cudaMalloc(&b.probs, N_EXPERT * sizeof(float)));
+    CHECK(cudaMalloc(&b.partials_ref, N_EXPERT * KSPLIT * sizeof(float)));
+    CHECK(cudaMalloc(&b.partials_fu, N_EXPERT * KSPLIT * sizeof(float)));
+    CHECK(cudaMalloc(&b.sel, N_TOP * sizeof(int32_t)));
+    CHECK(cudaMalloc(&b.wgt, N_TOP * sizeof(float)));
+    CHECK(cudaMalloc(&b.hash, (size_t)b.hash_rows * N_TOP * sizeof(int32_t)));
+    CHECK(cudaMalloc(&b.scalars, sizeof(ds4_decode_scalars)));
+
+    float *hw = (float *)malloc((size_t)N_EXPERT * N_IN * sizeof(float));
+    __half *hwh = (__half *)malloc((size_t)N_EXPERT * N_IN * sizeof(__half));
+    float *hx = (float *)malloc(N_IN * sizeof(float));
+    float *hb = (float *)malloc(N_EXPERT * sizeof(float));
+    int32_t *hh = (int32_t *)malloc((size_t)b.hash_rows * N_TOP * sizeof(int32_t));
+
+    /* -------- parity sweep -------- */
+    int total_bad = 0, n_cases = 0;
+    const unsigned grids[] = { 48u, 96u, 128u, 192u, 256u };
+    for (unsigned seed = 1; seed <= 10; seed++) {
+        fill_rand(hw, (size_t)N_EXPERT * N_IN, seed * 7919u, 0.05f);
+        if (seed == 3) {
+            /* exact-tie case: duplicate rows -> equal logits (and equal bias) ->
+             * tie-break by lower index must survive fusion */
+            for (int r = 0; r < 6; r++)
+                memcpy(hw + (size_t)(100 + r) * N_IN, hw + (size_t)(10 + r) * N_IN, N_IN * sizeof(float));
+        }
+        for (size_t i = 0; i < (size_t)N_EXPERT * N_IN; i++) hwh[i] = __float2half(hw[i]);
+        fill_rand(hx, N_IN, seed * 104729u, 1.0f);
+        if (seed == 5) for (size_t i = 0; i < N_IN; i++) hx[i] = 0.0f;  /* all-equal logits */
+        fill_rand(hb, N_EXPERT, seed * 31u, 0.3f);
+        if (seed == 3) for (int r = 0; r < 6; r++) hb[100 + r] = hb[10 + r];
+        unsigned s = seed * 2654435761u;
+        for (size_t i = 0; i < (size_t)b.hash_rows * N_TOP; i++) {
+            s = s * 1664525u + 1013904223u;
+            hh[i] = (int32_t)(s % 300u) - 20;  /* includes out-of-range ids */
+        }
+        CHECK(cudaMemcpy(b.w, hwh, (size_t)N_EXPERT * N_IN * sizeof(__half), cudaMemcpyHostToDevice));
+        CHECK(cudaMemcpy(b.x, hx, N_IN * sizeof(float), cudaMemcpyHostToDevice));
+        CHECK(cudaMemcpy(b.bias, hb, N_EXPERT * sizeof(float), cudaMemcpyHostToDevice));
+        CHECK(cudaMemcpy(b.hash, hh, (size_t)b.hash_rows * N_TOP * sizeof(int32_t), cudaMemcpyHostToDevice));
+        ds4_decode_scalars sc = {};
+        sc.token = (seed * 37u) % 1200u;   /* sometimes >= hash_rows -> clamp path */
+        CHECK(cudaMemcpy(b.scalars, &sc, sizeof(sc), cudaMemcpyHostToDevice));
+
+        struct { int has_bias, hash_mode, via_scalars; const char *name; } cases[] = {
+            { 1, 0, 0, "bias" }, { 0, 0, 0, "nobias" }, { 0, 1, 1, "hash" },
+        };
+        for (auto &c : cases) {
+            outs oref, ofu;
+            run_ref(b, c.has_bias, c.hash_mode, c.via_scalars);
+            CHECK(cudaDeviceSynchronize());
+            read_outs(b, &oref);
+            for (unsigned g : grids) {
+                if (g > max_blocks) continue;
+                /* scrub outputs so a silent non-write is caught */
+                CHECK(cudaMemset(b.logits, 0xCB, N_EXPERT * sizeof(float)));
+                CHECK(cudaMemset(b.probs, 0xCB, N_EXPERT * sizeof(float)));
+                CHECK(cudaMemset(b.sel, 0xCB, N_TOP * sizeof(int32_t)));
+                CHECK(cudaMemset(b.wgt, 0xCB, N_TOP * sizeof(float)));
+                if (!launch_fused(b, c.has_bias, c.hash_mode, c.via_scalars, g, 0)) return 1;
+                CHECK(cudaDeviceSynchronize());
+                read_outs(b, &ofu);
+                char tag[64];
+                snprintf(tag, sizeof(tag), "seed%u %s g%u", seed, c.name, g);
+                total_bad += cmp_outs(oref, ofu, tag);
+                n_cases++;
+                CHECK(cudaMemset(b.logits, 0xCB, N_EXPERT * sizeof(float)));
+                CHECK(cudaMemset(b.probs, 0xCB, N_EXPERT * sizeof(float)));
+                CHECK(cudaMemset(b.sel, 0xCB, N_TOP * sizeof(int32_t)));
+                CHECK(cudaMemset(b.wgt, 0xCB, N_TOP * sizeof(float)));
+                if (!launch_fused_k((void *)router_fused_coop_kernel_v2,
+                                    b, c.has_bias, c.hash_mode, c.via_scalars, g, 0)) return 1;
+                CHECK(cudaDeviceSynchronize());
+                read_outs(b, &ofu);
+                snprintf(tag, sizeof(tag), "V2 seed%u %s g%u", seed, c.name, g);
+                total_bad += cmp_outs(oref, ofu, tag);
+                n_cases++;
+            }
+        }
+    }
+    printf("parity: %d/%d cases bit-identical\n", n_cases - total_bad, n_cases);
+    if (total_bad) return 1;
+
+    /* -------- capture/replay bit-identity -------- */
+    {
+        const unsigned g = max_blocks < 256u ? max_blocks : 256u;
+        outs oeager, oreplay;
+        if (!launch_fused(b, 1, 0, 0, g, 0)) return 1;
+        CHECK(cudaDeviceSynchronize());
+        read_outs(b, &oeager);
+        cudaStream_t cs;
+        CHECK(cudaStreamCreate(&cs));
+        cudaGraph_t graph;
+        cudaGraphExec_t gexec;
+        CHECK(cudaStreamBeginCapture(cs, cudaStreamCaptureModeGlobal));
+        if (!launch_fused(b, 1, 0, 0, g, cs)) return 1;
+        CHECK(cudaStreamEndCapture(cs, &graph));
+        CHECK(cudaGraphInstantiate(&gexec, graph, NULL, NULL, 0));
+        for (int r = 0; r < 3; r++) {
+            CHECK(cudaMemset(b.logits, 0xCB, N_EXPERT * sizeof(float)));
+            CHECK(cudaGraphLaunch(gexec, cs));
+            CHECK(cudaStreamSynchronize(cs));
+            read_outs(b, &oreplay);
+            if (cmp_outs(oeager, oreplay, "capture-replay")) return 1;
+        }
+        printf("capture/replay: bit-identical (3 replays, %u blk)\n", g);
+        CHECK(cudaGraphExecDestroy(gexec));
+        CHECK(cudaGraphDestroy(graph));
+        CHECK(cudaStreamDestroy(cs));
+    }
+
+    /* -------- timing: rotating 68-layer weight set (L2-defeated) -------- */
+    {
+        const int NL = 68;
+        __half *wl[NL];
+        for (int l = 0; l < NL; l++) {
+            CHECK(cudaMalloc(&wl[l], (size_t)N_EXPERT * N_IN * sizeof(__half)));
+            CHECK(cudaMemcpy(wl[l], b.w, (size_t)N_EXPERT * N_IN * sizeof(__half), cudaMemcpyDeviceToDevice));
+        }
+        const int ITERS = 3000, WARM = 300;
+        cudaEvent_t e0, e1;
+        CHECK(cudaEventCreate(&e0));
+        CHECK(cudaEventCreate(&e1));
+        float ms;
+
+        /* ref chain */
+        dev_bufs bl = b;
+        for (int it = 0; it < WARM; it++) { bl.w = wl[it % NL]; run_ref(bl, 1, 0, 0); }
+        CHECK(cudaDeviceSynchronize());
+        CHECK(cudaEventRecord(e0));
+        for (int it = 0; it < ITERS; it++) { bl.w = wl[it % NL]; run_ref(bl, 1, 0, 0); }
+        CHECK(cudaEventRecord(e1));
+        CHECK(cudaDeviceSynchronize());
+        CHECK(cudaEventElapsedTime(&ms, e0, e1));
+        printf("ref chain   : %8.3f us/layer\n", ms * 1000.0f / ITERS);
+
+        int bps2 = 0;
+        CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&bps2, router_fused_coop_kernel_v2, 256, 0));
+        const unsigned max_blocks2 = (unsigned)(bps2 * n_sm);
+        printf("V2 blocks/sm=%d max co-resident=%u\n", bps2, max_blocks2);
+        struct { void *k; const char *name; unsigned cap; } variants[] = {
+            { (void *)router_fused_coop_kernel, "V1", max_blocks },
+            { (void *)router_fused_coop_kernel_v2, "V2", max_blocks2 },
+        };
+        for (auto &v : variants) {
+            for (unsigned g : grids) {
+                if (g > v.cap) continue;
+                for (int it = 0; it < WARM; it++) { bl.w = wl[it % NL]; launch_fused_k(v.k, bl, 1, 0, 0, g, 0); }
+                CHECK(cudaDeviceSynchronize());
+                CHECK(cudaEventRecord(e0));
+                for (int it = 0; it < ITERS; it++) { bl.w = wl[it % NL]; launch_fused_k(v.k, bl, 1, 0, 0, g, 0); }
+                CHECK(cudaEventRecord(e1));
+                CHECK(cudaDeviceSynchronize());
+                CHECK(cudaEventElapsedTime(&ms, e0, e1));
+                printf("fused %s g=%3u : %8.3f us/layer\n", v.name, g, ms * 1000.0f / ITERS);
+            }
+        }
+        for (int l = 0; l < NL; l++) CHECK(cudaFree(wl[l]));
+    }
+
+    printf("PROTO_M2_ROUTER OK\n");
+    return 0;
+}

cuda/mmq/test/proto_mm_ids.cu

@@ -0,0 +1,475 @@
+/* proto_mm_ids.cu — mm_ids_helper parity harness.
+ *
+ * P4 Inc1 (landed 86c5d4d): mm_ids_helper<1> fast path for the routed-MoE
+ * down matmul (ds4_mmq.cu ds4_mmq_moe_impl: n_expert_used=1, si1=1, sis1=1,
+ * nchannels_y=1, n_experts=256).  Without a case-1 dispatch the down matmul
+ * fell to the generic <0> template: one warp per expert scanning all
+ * assignment rows with a single active lane — 22.5 ms/launch at W4096.
+ *
+ * P5 (this extension): mm_ids_helper_global — the large-n variant with NO
+ * shared-memory staging.  The smem kernel stages the compact (it, iex_used)
+ * list in n_tokens*4 B of dynamic shared memory, capping n_tokens at smpbo/4
+ * (~25k on GB10); the down matmul passes n_tokens = assignment rows = 6x the
+ * forward width, so 8192-row prefill chunks hit 48384 "tokens", the mmq
+ * entries refused, and the WHOLE MoE block fell back to the pre-mmq
+ * expert-tile kernels (the W8192 prefill cliff, trace cmtp5w8.sqlite).  The
+ * global variant runs the same per-expert scan twice: pass 0 counts
+ * (nex_prev, bucket size), pass 1 re-scans and writes ids_src1/ids_dst
+ * directly at final offsets.
+ *
+ * Gate: ids_src1 / ids_dst / expert_bounds BIT-IDENTICAL across
+ *   (a) smem <0> vs smem <1>/<6>          (the landed Inc1 claim),
+ *   (b) smem vs global at every sub-cap shape,
+ *   (c) global vs a host reference at EVERY shape, including past-cap legs
+ *       (48384/49152 assignments = the W8192 down shape; a >cap
+ *       n_expert_used=6 leg; the cap boundary +/-1)
+ * with the production pre-zero of both id maps (dropped -1 router ids leave
+ * slots unwritten; expert_bounds[0] counts them, so the zeros sit below
+ * expert 0's bucket and are never consumed).
+ *
+ * Bit-exactness argument for the global variant: identical (it, iex_used)
+ * tuples in identical token order and identical output expressions; the
+ * (lossless) 22/10-bit store round-trip is simply removed.  Both variants
+ * are undefined if one token lists the same expert twice (router top-k is
+ * without replacement); the generators here sample distinct experts.
+ *
+ * Build (GB10):
+ *   nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 proto_mm_ids.cu -o proto_mm_ids
+ */
+#include <cuda_runtime.h>
+#include <climits>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <cstdint>
+#include <cassert>
+#include <vector>
+
+#define CHECK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
+    fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(e_), __FILE__, __LINE__); exit(1); } } while (0)
+
+static constexpr int WARP_SIZE = 32;
+
+/* ------------------------------------------------------------------ */
+/* Vendored from cuda/mmq/mmid.cu + common.cuh (verbatim semantics)     */
+
+template <int width = WARP_SIZE>
+static __device__ __forceinline__ int warp_reduce_any(int x) {
+    if (width == WARP_SIZE) {
+        return __any_sync(0xffffffff, x);
+    } else {
+#pragma unroll
+        for (int offset = width/2; offset > 0; offset >>= 1) {
+            x = __shfl_xor_sync(0xffffffff, x, offset, width) || x;
+        }
+        return x;
+    }
+}
+
+template <int width = WARP_SIZE>
+static __device__ __forceinline__ int warp_reduce_sum(int x) {
+#pragma unroll
+    for (int offset = width/2; offset > 0; offset >>= 1) {
+        x += __shfl_xor_sync(0xffffffff, x, offset, width);
+    }
+    return x;
+}
+
+struct mm_ids_helper_store {
+    uint32_t data;
+    __device__ mm_ids_helper_store(const uint32_t it, const uint32_t iex_used) {
+        data = (it & 0x003FFFFF) | (iex_used << 22);
+    }
+    __device__ uint32_t it() const { return data & 0x003FFFFF; }
+    __device__ uint32_t iex_used() const { return data >> 22; }
+};
+static_assert(sizeof(mm_ids_helper_store) == 4, "unexpected size for mm_ids_helper_store");
+
+template <int n_expert_used_template>
+__launch_bounds__(WARP_SIZE, 1)
+static __global__ void mm_ids_helper(
+        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
+        const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1) {
+    constexpr int warp_size = WARP_SIZE;
+    const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
+    const int expert = blockIdx.x;
+
+    extern __shared__ char data_mm_ids_helper[];
+    mm_ids_helper_store * store = (mm_ids_helper_store *) data_mm_ids_helper;
+
+    int nex_prev   = 0;
+    int it_compact = 0;
+
+    if constexpr (n_expert_used_template == 0) {
+        for (int it = 0; it < n_tokens; ++it) {
+            int iex_used = -1;
+            for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
+                const int expert_used = ids[it*si1 + iex];
+                nex_prev += expert_used < expert;
+                if (expert_used == expert) {
+                    iex_used = iex;
+                }
+            }
+            if (iex_used != -1) {
+                store[it_compact] = mm_ids_helper_store(it, iex_used);
+            }
+            if (warp_reduce_any<warp_size>(iex_used != -1)) {
+                it_compact++;
+            }
+        }
+    } else {
+        static_assert(n_expert_used_template == 6 || warp_size % n_expert_used_template == 0, "bad n_expert_used");
+        constexpr int neu_padded = n_expert_used_template == 6 ? 8 : n_expert_used_template;
+        for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) {
+            const int it = it0 + threadIdx.x / neu_padded;
+
+            const int iex = threadIdx.x % neu_padded;
+            const int expert_used = (neu_padded == n_expert_used_template || iex < n_expert_used_template) && it < n_tokens ?
+                ids[it*si1 + iex] : INT_MAX;
+            const int iex_used = expert_used == expert ? iex : -1;
+            nex_prev += expert_used < expert;
+
+            const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);
+
+            int it_compact_add_lower = 0;
+#pragma unroll
+            for (int offset = neu_padded; offset < warp_size; offset += neu_padded) {
+                const int tmp = __shfl_up_sync(0xFFFFFFFF, it_compact_add_self, offset, warp_size);
+                if (threadIdx.x >= static_cast<unsigned int>(offset)) {
+                    it_compact_add_lower += tmp;
+                }
+            }
+
+            if (iex_used != -1) {
+                store[it_compact + it_compact_add_lower] = mm_ids_helper_store(it, iex_used);
+            }
+
+            it_compact += __shfl_sync(0xFFFFFFFF, it_compact_add_lower + it_compact_add_self, warp_size - 1, warp_size);
+        }
+    }
+    nex_prev = warp_reduce_sum<warp_size>(nex_prev);
+
+    __syncwarp();
+
+    for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) {
+        const mm_ids_helper_store store_it = store[itc];
+        const int it       = store_it.it();
+        const int iex_used = store_it.iex_used();
+        ids_src1[nex_prev + itc] = it*sis1          + iex_used % nchannels_y;
+        ids_dst [nex_prev + itc] = it*n_expert_used + iex_used;
+    }
+
+    if (threadIdx.x != 0) {
+        return;
+    }
+    expert_bounds[expert] = nex_prev;
+    if (expert < static_cast<int>(gridDim.x) - 1) {
+        return;
+    }
+    expert_bounds[gridDim.x] = nex_prev + it_compact;
+}
+
+/* Vendored mm_ids_helper_global (P5 large-n variant, mmid.cu) */
+template <int n_expert_used_template>
+__launch_bounds__(WARP_SIZE, 1)
+static __global__ void mm_ids_helper_global(
+        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
+        const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1) {
+    constexpr int warp_size = WARP_SIZE;
+    const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
+    const int expert = blockIdx.x;
+
+    int nex_prev         = 0;
+    int it_compact_count = 0;
+
+#pragma unroll 1
+    for (int pass = 0; pass < 2; ++pass) {
+        int it_compact = 0;
+
+        if constexpr (n_expert_used_template == 0) {
+            for (int it = 0; it < n_tokens; ++it) {
+                int iex_used = -1;
+                for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
+                    const int expert_used = ids[it*si1 + iex];
+                    if (pass == 0) {
+                        nex_prev += expert_used < expert;
+                    }
+                    if (expert_used == expert) {
+                        iex_used = iex;
+                    }
+                }
+
+                if (pass == 1 && iex_used != -1) {
+                    ids_src1[nex_prev + it_compact] = it*sis1          + iex_used % nchannels_y;
+                    ids_dst [nex_prev + it_compact] = it*n_expert_used + iex_used;
+                }
+
+                if (warp_reduce_any<warp_size>(iex_used != -1)) {
+                    it_compact++;
+                }
+            }
+        } else {
+            static_assert(n_expert_used_template == 6 || warp_size % n_expert_used_template == 0, "bad n_expert_used");
+            constexpr int neu_padded = n_expert_used_template == 6 ? 8 : n_expert_used_template;
+            for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) {
+                const int it = it0 + threadIdx.x / neu_padded;
+
+                const int iex = threadIdx.x % neu_padded;
+                const int expert_used = (neu_padded == n_expert_used_template || iex < n_expert_used_template) && it < n_tokens ?
+                    ids[it*si1 + iex] : INT_MAX;
+                const int iex_used = expert_used == expert ? iex : -1;
+                if (pass == 0) {
+                    nex_prev += expert_used < expert;
+                }
+
+                const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);
+
+                int it_compact_add_lower = 0;
+#pragma unroll
+                for (int offset = neu_padded; offset < warp_size; offset += neu_padded) {
+                    const int tmp = __shfl_up_sync(0xFFFFFFFF, it_compact_add_self, offset, warp_size);
+                    if (threadIdx.x >= static_cast<unsigned int>(offset)) {
+                        it_compact_add_lower += tmp;
+                    }
+                }
+
+                if (pass == 1 && iex_used != -1) {
+                    const int itc = it_compact + it_compact_add_lower;
+                    ids_src1[nex_prev + itc] = it*sis1          + iex_used % nchannels_y;
+                    ids_dst [nex_prev + itc] = it*n_expert_used + iex_used;
+                }
+
+                it_compact += __shfl_sync(0xFFFFFFFF, it_compact_add_lower + it_compact_add_self, warp_size - 1, warp_size);
+            }
+        }
+
+        if (pass == 0) {
+            it_compact_count = it_compact;
+            nex_prev = warp_reduce_sum<warp_size>(nex_prev);
+        }
+    }
+
+    if (threadIdx.x != 0) {
+        return;
+    }
+
+    expert_bounds[expert] = nex_prev;
+
+    if (expert < static_cast<int>(gridDim.x) - 1) {
+        return;
+    }
+
+    expert_bounds[gridDim.x] = nex_prev + it_compact_count;
+}
+
+/* NOTE: the proto instantiations replace the vendored kernels' runtime-const
+ * neu_padded with a constexpr derived from the template arg (identical value
+ * per instantiation; the vendored code relies on the same constant folding). */
+
+/* ------------------------------------------------------------------ */
+
+static constexpr int N_EXPERTS = 256;
+
+/* Host reference implementing the mm_ids_helper output contract:
+ *   expert_bounds[e] = #{assignment values < e}  (signed: -1 counts for all
+ *   experts, so invalid rows occupy never-consumed slots below expert 0);
+ *   bucket e = token-ascending (it, iex) pairs with ids[it*si1+iex] == e;
+ *   ids_src1[pos] = it*sis1 + iex % nchannels_y; ids_dst[pos] = it*neu + iex.
+ * src1/dst must be pre-zeroed by the caller (production memset semantics). */
+static void cpu_ref(const int32_t *ids, int nt, int neu, int nchannels_y, int si1, int sis1,
+                    int32_t *src1, int32_t *dst, int32_t *eb) {
+    std::vector<int32_t> h(N_EXPERTS + 2, 0); // h[b+1] counts value b, b in [-1, N_EXPERTS]
+    for (int it = 0; it < nt; ++it) {
+        for (int iex = 0; iex < neu; ++iex) {
+            int v = ids[it*si1 + iex];
+            if (v < -1) v = -1;
+            if (v > N_EXPERTS) v = N_EXPERTS;
+            h[v + 1]++;
+        }
+    }
+    std::vector<int32_t> below(N_EXPERTS + 1, 0); // below[e] = #{v < e}
+    int32_t acc = h[0];
+    for (int e = 0; e <= N_EXPERTS; ++e) {
+        below[e] = acc;
+        if (e < N_EXPERTS) acc += h[e + 1];
+    }
+    for (int e = 0; e <= N_EXPERTS; ++e) {
+        eb[e] = below[e];
+    }
+    std::vector<int32_t> fill(N_EXPERTS, 0);
+    for (int it = 0; it < nt; ++it) {
+        for (int iex = 0; iex < neu; ++iex) {
+            const int e = ids[it*si1 + iex];
+            if (e < 0 || e >= N_EXPERTS) continue;
+            const int64_t pos = (int64_t)below[e] + fill[e]++;
+            src1[pos] = it*sis1 + iex % nchannels_y;
+            dst [pos] = it*neu  + iex;
+        }
+    }
+}
+
+struct leg {
+    const char *name;
+    int n_tokens;      // assignment rows (neu=1) or tokens (neu=6)
+    int neu;           // n_expert_used: 1 = down shape, 6 = gate/up shape
+    int invalid_pct;   // % of slots with -1 (router NaN path)
+    bool skew;         // hot-expert distribution
+};
+
+int main() {
+    int smpbo = 0;
+    CHECK(cudaDeviceGetAttribute(&smpbo, cudaDevAttrMaxSharedMemoryPerBlockOptin, 0));
+    CHECK(cudaFuncSetAttribute((const void*)mm_ids_helper<0>, cudaFuncAttributeMaxDynamicSharedMemorySize, smpbo));
+    CHECK(cudaFuncSetAttribute((const void*)mm_ids_helper<1>, cudaFuncAttributeMaxDynamicSharedMemorySize, smpbo));
+    CHECK(cudaFuncSetAttribute((const void*)mm_ids_helper<6>, cudaFuncAttributeMaxDynamicSharedMemorySize, smpbo));
+    const int cap_tokens = smpbo / 4;
+    printf("smpbo=%d (cap %d tokens)\n", smpbo, cap_tokens);
+
+    std::vector<leg> legs = {
+        // down shape (neu=1): n_tokens = rows*6 assignments
+        {"W4096x6dn uniform",   4096*6, 1, 0, false},
+        {"W4096x6dn 2%invalid", 4096*6, 1, 2, false},
+        {"W4096x6dn skewed",    4096*6, 1, 1, true},
+        {"W2048x6dn uniform",   2048*6, 1, 2, false},
+        {"decode w1x6",         6,      1, 0, false},
+        {"decode w8x6",         48,     1, 0, false},
+        // past-cap down shapes (the P5 cliff: W8192 -> 8064*6 = 48384 traced, 8192*6 = 49152 max)
+        {"W8192x6dn uniform",   48384,  1, 0, false},
+        {"W8192x6dn 2%invalid", 48384,  1, 2, false},
+        {"W8192x6dn skewed",    48384,  1, 1, true},
+        {"W8192x6dn max",       49152,  1, 2, false},
+        // gate/up shape (neu=6): n_tokens = rows, si1 = 6, distinct experts/token
+        {"W8192gu uniform",     8064,   6, 0, false},
+        {"W8192gu 2%invalid",   8064,   6, 2, false},
+        {"gu past-cap",         30000,  6, 2, false},
+    };
+    // cap-boundary seam legs (runtime-sized)
+    leg seam_at   = {"cap-boundary",  cap_tokens,     1, 2, false};
+    leg seam_over = {"cap-boundary+1", cap_tokens + 1, 1, 2, false};
+    legs.push_back(seam_at);
+    legs.push_back(seam_over);
+
+    int failures = 0;
+
+    for (const leg &L : legs) {
+        const int nt  = L.n_tokens;
+        const int neu = L.neu;
+        const int si1 = neu, sis1 = 1, nchannels_y = 1;
+        const int64_t n_ids  = (int64_t)nt * neu;
+        const size_t  smem   = (size_t)nt * 4;
+        const bool smem_fits = smem <= (size_t)smpbo;
+
+        int32_t *h_ids = (int32_t*)malloc(n_ids * sizeof(int32_t));
+        srand(20260708);
+        if (neu == 1) {
+            for (int i = 0; i < nt; i++) {
+                if (L.invalid_pct && rand() % 100 < L.invalid_pct) { h_ids[i] = -1; continue; }
+                h_ids[i] = L.skew ? (rand() % 100 < 60 ? rand() % 8 : rand() % N_EXPERTS)
+                                  : rand() % N_EXPERTS;
+            }
+        } else {
+            // router top-k semantics: distinct experts per token
+            for (int it = 0; it < nt; it++) {
+                int chosen[8]; int c = 0;
+                while (c < neu) {
+                    const int e = L.skew ? (rand() % 100 < 60 ? rand() % 8 : rand() % N_EXPERTS)
+                                         : rand() % N_EXPERTS;
+                    bool dup = false;
+                    for (int j = 0; j < c; j++) dup = dup || chosen[j] == e;
+                    if (!dup) chosen[c++] = e;
+                }
+                for (int iex = 0; iex < neu; iex++) {
+                    h_ids[(int64_t)it*neu + iex] =
+                        (L.invalid_pct && rand() % 100 < L.invalid_pct) ? -1 : chosen[iex];
+                }
+            }
+        }
+
+        // host reference
+        int32_t *r_src1 = (int32_t*)calloc(n_ids, sizeof(int32_t));
+        int32_t *r_dst  = (int32_t*)calloc(n_ids, sizeof(int32_t));
+        int32_t r_eb[N_EXPERTS+1];
+        cpu_ref(h_ids, nt, neu, nchannels_y, si1, sis1, r_src1, r_dst, r_eb);
+
+        int32_t *d_ids, *d_src1, *d_dst, *d_eb;
+        CHECK(cudaMalloc(&d_ids,  n_ids * sizeof(int32_t)));
+        CHECK(cudaMalloc(&d_src1, n_ids * sizeof(int32_t)));
+        CHECK(cudaMalloc(&d_dst,  n_ids * sizeof(int32_t)));
+        CHECK(cudaMalloc(&d_eb,   (N_EXPERTS+1) * sizeof(int32_t)));
+        CHECK(cudaMemcpy(d_ids, h_ids, n_ids * sizeof(int32_t), cudaMemcpyHostToDevice));
+
+        const dim3 grid(N_EXPERTS, 1, 1), block(WARP_SIZE, 1, 1);
+        int32_t *h_src1 = (int32_t*)malloc(n_ids * sizeof(int32_t));
+        int32_t *h_dst  = (int32_t*)malloc(n_ids * sizeof(int32_t));
+        int32_t h_eb[N_EXPERTS+1];
+
+        // candidate kernels for this leg: smem generic, smem fast, global fast
+        struct cand { const char *name; int variant; }; // 0=<0> smem, 1=<neu> smem, 2=<neu> global
+        std::vector<cand> cands;
+        if (smem_fits) { cands.push_back({"smem<0>", 0}); cands.push_back({"smem<t>", 1}); }
+        cands.push_back({"global<t>", 2});
+
+        bool ok = true;
+        float ms[3] = {0, 0, 0};
+        for (const cand &C : cands) {
+            // production pre-zero of both id maps
+            CHECK(cudaMemset(d_src1, 0, n_ids * sizeof(int32_t)));
+            CHECK(cudaMemset(d_dst,  0, n_ids * sizeof(int32_t)));
+            CHECK(cudaMemset(d_eb, 0xFF, (N_EXPERTS+1) * sizeof(int32_t)));
+
+            auto launch = [&](void) {
+                switch (C.variant) {
+                    case 0: mm_ids_helper<0><<<grid, block, smem>>>(d_ids, d_src1, d_dst, d_eb, nt, neu, nchannels_y, si1, sis1); break;
+                    case 1:
+                        if (neu == 1) mm_ids_helper<1><<<grid, block, smem>>>(d_ids, d_src1, d_dst, d_eb, nt, neu, nchannels_y, si1, sis1);
+                        else          mm_ids_helper<6><<<grid, block, smem>>>(d_ids, d_src1, d_dst, d_eb, nt, neu, nchannels_y, si1, sis1);
+                        break;
+                    case 2:
+                        if (neu == 1) mm_ids_helper_global<1><<<grid, block, 0>>>(d_ids, d_src1, d_dst, d_eb, nt, neu, nchannels_y, si1, sis1);
+                        else          mm_ids_helper_global<6><<<grid, block, 0>>>(d_ids, d_src1, d_dst, d_eb, nt, neu, nchannels_y, si1, sis1);
+                        break;
+                }
+            };
+            launch();
+            CHECK(cudaGetLastError());
+            CHECK(cudaDeviceSynchronize());
+
+            CHECK(cudaMemcpy(h_src1, d_src1, n_ids*4, cudaMemcpyDeviceToHost));
+            CHECK(cudaMemcpy(h_dst,  d_dst,  n_ids*4, cudaMemcpyDeviceToHost));
+            CHECK(cudaMemcpy(h_eb,   d_eb,   (N_EXPERTS+1)*4, cudaMemcpyDeviceToHost));
+            const bool m1 = memcmp(h_src1, r_src1, n_ids*4) == 0;
+            const bool m2 = memcmp(h_dst,  r_dst,  n_ids*4) == 0;
+            const bool m3 = memcmp(h_eb,   r_eb,   (N_EXPERTS+1)*4) == 0;
+            if (!(m1 && m2 && m3)) {
+                ok = false;
+                printf("  %-18s %-9s MISMATCH src1=%d dst=%d eb=%d\n", L.name, C.name, m1, m2, m3);
+            }
+
+            // timing (50 reps after 5 warmups)
+            cudaEvent_t t0, t1;
+            CHECK(cudaEventCreate(&t0)); CHECK(cudaEventCreate(&t1));
+            for (int w = 0; w < 5; w++) launch();
+            CHECK(cudaEventRecord(t0));
+            for (int r = 0; r < 50; r++) launch();
+            CHECK(cudaEventRecord(t1)); CHECK(cudaEventSynchronize(t1));
+            float m = 0; CHECK(cudaEventElapsedTime(&m, t0, t1));
+            ms[C.variant] = m / 50;
+            CHECK(cudaEventDestroy(t0)); CHECK(cudaEventDestroy(t1));
+        }
+
+        if (smem_fits) {
+            printf("%-20s nt=%-6d neu=%d %s  smem<0> %8.3f ms  smem<t> %8.3f ms  global<t> %8.3f ms\n",
+                   L.name, nt, neu, ok ? "PARITY-OK " : "MISMATCH!!", ms[0], ms[1], ms[2]);
+        } else {
+            printf("%-20s nt=%-6d neu=%d %s  (past cap)                              global<t> %8.3f ms\n",
+                   L.name, nt, neu, ok ? "PARITY-OK " : "MISMATCH!!", ms[2]);
+        }
+        if (!ok) failures++;
+
+        free(h_ids); free(h_src1); free(h_dst); free(r_src1); free(r_dst);
+        cudaFree(d_ids); cudaFree(d_src1); cudaFree(d_dst); cudaFree(d_eb);
+    }
+
+    printf(failures ? "PROTO FAIL (%d legs)\n" : "PROTO PASS\n", failures);
+    return failures ? 1 : 0;
+}

cuda/mmq/test/proto_q8_aligned.cu

@@ -0,0 +1,296 @@
+// proto_q8_aligned.cu — M1-Inc3 prototype (megakernel program).
+//
+// HYPOTHESIS: the dense Q8_0 decode GEMVs run below the ~200-202 GB/s kernel
+// ceiling (attrib15b: attn_q_b 1024x32768 @164, 2048x4096 @157, head
+// 4096x129280 @146) because block_q8_0 is 34 bytes, so the int8 code stream
+// is only 2-byte aligned (same misalignment class as IQ2_XXS, proven +12%
+// by proto_iq2_aligned).  Repacking into an aligned SoA layout (d[] halves
+// separate, qs[] 64B-aligned 32B per block) allows int4 (16B) aligned loads
+// at identical byte count.
+//
+// A/B at the three production decode shapes (K = in_dim, M = out rows):
+//   attn_q_b : K=1024,  M=32768   (34.0 MiB/call, 164 GB/s baseline)
+//   mid      : K=2048,  M=4096    ( 8.5 MiB/call, 157 GB/s baseline)
+//   head     : K=4096,  M=129280  (536.6 MiB/call, 146 GB/s baseline)
+//   baseline = ds4_mmq_q8_0_dense_vec (production decode entry, quantizes X
+//              internally; bare-kernel aligned numbers exclude that ~small
+//              cost — Inc1 lesson: judge the KERNEL here, entry design later)
+//
+// Build (box):
+//   nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 -I.. \
+//        proto_q8_aligned.cu ../*.o -lcudart -lcuda -o proto_q8_aligned
+
+#include "ds4_mmq.h"
+
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <random>
+#include <vector>
+
+#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
+    printf("CUDA ERR %s @%d: %s\n", #x, __LINE__, cudaGetErrorString(e_)); exit(1); } } while (0)
+
+namespace {
+
+constexpr int QK = 32;               // values per Q8_0 block
+constexpr int BLOCK_BYTES = 34;      // sizeof(block_q8_0) = 2 + 32
+
+// ---------------------------------------------------------------- host q8_1
+struct q8_1_block { __half2 ds; int8_t qs[32]; };   // mirrors block_q8_1
+
+static void quantize_q8_1_host(const float *x, q8_1_block *out, int k) {
+    for (int b = 0; b < k / 32; b++) {
+        float amax = 0.0f;
+        for (int i = 0; i < 32; i++) amax = fmaxf(amax, fabsf(x[b * 32 + i]));
+        const float d = amax / 127.0f;
+        const float id = d != 0.0f ? 1.0f / d : 0.0f;
+        int sum = 0;
+        for (int i = 0; i < 32; i++) {
+            int q = (int)roundf(x[b * 32 + i] * id);
+            if (q > 127) q = 127;
+            if (q < -127) q = -127;
+            out[b].qs[i] = (int8_t)q;
+            sum += q;
+        }
+        out[b].ds = __floats2half2_rn(d, d * (float)sum);
+    }
+}
+
+// ------------------------------------------------------------ device kernel
+// Aligned SoA weights: dq[blk] __half scale; qs stream 64B-aligned, 32B per
+// block (int4 x2).  Warp-per-row GEMV: lane covers one block per pass.
+__global__ void q8_aligned_dense_vec_kernel(
+        float             *out,        // [M]
+        const int4        *qs,         // aligned codes, 2 int4 per block
+        const __half      *dq,         // block scales
+        const q8_1_block  *x8,         // [K/32] quantized activation
+        int                M,
+        int                nb)         // blocks per row = K/32
+{
+    const int row  = blockIdx.x;
+    const int lane = threadIdx.x;
+    const long long rbase = (long long)row * nb;
+
+    float acc = 0.0f;
+    for (int b0 = 0; b0 < nb; b0 += 32) {
+        const int b = b0 + lane;
+        const int4 w0 = qs[(rbase + b) * 2 + 0];   // aligned 16B
+        const int4 w1 = qs[(rbase + b) * 2 + 1];
+        const int *u = (const int *)x8[b].qs;
+        int sumi = 0;
+        sumi = __dp4a(w0.x, u[0], sumi);
+        sumi = __dp4a(w0.y, u[1], sumi);
+        sumi = __dp4a(w0.z, u[2], sumi);
+        sumi = __dp4a(w0.w, u[3], sumi);
+        sumi = __dp4a(w1.x, u[4], sumi);
+        sumi = __dp4a(w1.y, u[5], sumi);
+        sumi = __dp4a(w1.z, u[6], sumi);
+        sumi = __dp4a(w1.w, u[7], sumi);
+        acc += __half2float(dq[rbase + b]) * __low2float(x8[b].ds) * (float)sumi;
+    }
+#pragma unroll
+    for (int off = 16; off > 0; off >>= 1)
+        acc += __shfl_down_sync(0xffffffffu, acc, off);
+    if (lane == 0) out[row] = acc;
+}
+
+// Variant R2: two rows per 64-thread block (one warp each) — probes whether
+// higher occupancy per SM changes the DRAM rate at skinny K.
+__global__ void q8_aligned_dense_vec_r2_kernel(
+        float             *out,
+        const int4        *qs,
+        const __half      *dq,
+        const q8_1_block  *x8,
+        int                M,
+        int                nb)
+{
+    const int row  = blockIdx.x * 2 + (threadIdx.x >> 5);
+    const int lane = threadIdx.x & 31;
+    if (row >= M) return;
+    const long long rbase = (long long)row * nb;
+    float acc = 0.0f;
+    for (int b0 = 0; b0 < nb; b0 += 32) {
+        const int b = b0 + lane;
+        const int4 w0 = qs[(rbase + b) * 2 + 0];
+        const int4 w1 = qs[(rbase + b) * 2 + 1];
+        const int *u = (const int *)x8[b].qs;
+        int sumi = 0;
+        sumi = __dp4a(w0.x, u[0], sumi);
+        sumi = __dp4a(w0.y, u[1], sumi);
+        sumi = __dp4a(w0.z, u[2], sumi);
+        sumi = __dp4a(w0.w, u[3], sumi);
+        sumi = __dp4a(w1.x, u[4], sumi);
+        sumi = __dp4a(w1.y, u[5], sumi);
+        sumi = __dp4a(w1.z, u[6], sumi);
+        sumi = __dp4a(w1.w, u[7], sumi);
+        acc += __half2float(dq[rbase + b]) * __low2float(x8[b].ds) * (float)sumi;
+    }
+#pragma unroll
+    for (int off = 16; off > 0; off >>= 1)
+        acc += __shfl_down_sync(0xffffffffu, acc, off);
+    if (lane == 0) out[row] = acc;
+}
+
+static void run_shape(const char *label, int M, int K, int iters, uint32_t seed) {
+    const int nb = K / QK;
+    const long long row_bytes = (long long)nb * BLOCK_BYTES;
+    const double call_gb = (double)M * row_bytes / 1e9;
+
+    std::mt19937 rng(seed);
+
+    // Random bytes are valid Q8_0 content; keep block scales sane.
+    std::vector<uint8_t> W((size_t)M * row_bytes);
+    for (auto &b : W) b = (uint8_t)(rng() & 0xff);
+    const long long nblk = (long long)M * nb;
+    for (long long blk = 0; blk < nblk; blk++) {
+        uint16_t h = (uint16_t)(0x2c00 | (rng() & 0xff));
+        memcpy(&W[blk * BLOCK_BYTES], &h, 2);
+    }
+
+    std::vector<float> X(K);
+    for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
+    std::vector<q8_1_block> X8(K / 32);
+    quantize_q8_1_host(X.data(), X8.data(), K);
+
+    // Aligned SoA repack: [half dq[nblk]][pad to 64B][int8 qs[nblk*32]]
+    const uint64_t dq_bytes = ((uint64_t)nblk * 2u + 63u) & ~63ull;
+    std::vector<uint8_t> ART(dq_bytes + (uint64_t)nblk * 32u);
+    for (long long blk = 0; blk < nblk; blk++) {
+        memcpy(&ART[blk * 2], &W[blk * BLOCK_BYTES], 2);
+        memcpy(&ART[dq_bytes + (uint64_t)blk * 32u], &W[blk * BLOCK_BYTES + 2], 32);
+    }
+
+    // L2 trap (proven on the iq2 pair test AND the first run of this proto:
+    // "mid" printed 931 GB/s, 4x the DRAM ceiling): a weight buffer smaller
+    // than ~L2 stays resident across timing iters.  Rotate through enough
+    // copies that consecutive iters never re-hit cache (~256 MB footprint).
+    const int n_copies = (int)((256ull * 1024 * 1024) / (W.size() + 1)) + 1;
+    uint8_t *dW, *dArt; float *dX, *dOutBase, *dOutAl; q8_1_block *dX8;
+    CK(cudaMalloc(&dW, W.size() * n_copies));
+    CK(cudaMalloc(&dArt, ART.size() * n_copies));
+    CK(cudaMalloc(&dX, sizeof(float) * K));
+    CK(cudaMalloc(&dOutBase, sizeof(float) * M));
+    CK(cudaMalloc(&dOutAl, sizeof(float) * M));
+    CK(cudaMalloc(&dX8, X8.size() * sizeof(q8_1_block)));
+    for (int c = 0; c < n_copies; c++) {
+        CK(cudaMemcpy(dW + (size_t)c * W.size(), W.data(), W.size(), cudaMemcpyHostToDevice));
+        CK(cudaMemcpy(dArt + (size_t)c * ART.size(), ART.data(), ART.size(), cudaMemcpyHostToDevice));
+    }
+    CK(cudaMemcpy(dX, X.data(), sizeof(float) * K, cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dX8, X8.data(), X8.size() * sizeof(q8_1_block), cudaMemcpyHostToDevice));
+
+    cudaStream_t stream; CK(cudaStreamCreate(&stream));
+    cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
+
+    // ---- baseline: production dense vec entry ------------------------------
+    int rc = ds4_mmq_q8_0_dense_vec(dW, dX, dOutBase, M, 1, K, stream);
+    if (rc != 0) { printf("%s: baseline dense_vec rc=%d\n", label, rc); exit(1); }
+    CK(cudaStreamSynchronize(stream));
+    CK(cudaEventRecord(e0, stream));
+    for (int i = 0; i < iters; i++)
+        (void)ds4_mmq_q8_0_dense_vec(dW + (size_t)(i % n_copies) * W.size(),
+                                     dX, dOutBase, M, 1, K, stream);
+    CK(cudaEventRecord(e1, stream));
+    CK(cudaStreamSynchronize(stream));
+    float ms_base = 0.0f; CK(cudaEventElapsedTime(&ms_base, e0, e1));
+    ms_base /= iters;
+
+    // ---- aligned prototype (warp per row) ----------------------------------
+    const int4 *qs = (const int4 *)(dArt + dq_bytes);
+    const __half *dq = (const __half *)dArt;
+    q8_aligned_dense_vec_kernel<<<M, 32, 0, stream>>>(dOutAl, qs, dq, dX8, M, nb);
+    CK(cudaGetLastError());
+    CK(cudaStreamSynchronize(stream));
+    CK(cudaEventRecord(e0, stream));
+    for (int i = 0; i < iters; i++) {
+        const uint8_t *art_i = dArt + (size_t)(i % n_copies) * ART.size();
+        q8_aligned_dense_vec_kernel<<<M, 32, 0, stream>>>(
+            dOutAl, (const int4 *)(art_i + dq_bytes), (const __half *)art_i, dX8, M, nb);
+    }
+    CK(cudaEventRecord(e1, stream));
+    CK(cudaStreamSynchronize(stream));
+    float ms_al = 0.0f; CK(cudaEventElapsedTime(&ms_al, e0, e1));
+    ms_al /= iters;
+
+    // ---- variant R2 ---------------------------------------------------------
+    q8_aligned_dense_vec_r2_kernel<<<(M + 1) / 2, 64, 0, stream>>>(dOutAl, qs, dq, dX8, M, nb);
+    CK(cudaGetLastError());
+    CK(cudaStreamSynchronize(stream));
+    CK(cudaEventRecord(e0, stream));
+    for (int i = 0; i < iters; i++) {
+        const uint8_t *art_i = dArt + (size_t)(i % n_copies) * ART.size();
+        q8_aligned_dense_vec_r2_kernel<<<(M + 1) / 2, 64, 0, stream>>>(
+            dOutAl, (const int4 *)(art_i + dq_bytes), (const __half *)art_i, dX8, M, nb);
+    }
+    CK(cudaEventRecord(e1, stream));
+    CK(cudaStreamSynchronize(stream));
+    float ms_r2 = 0.0f; CK(cudaEventElapsedTime(&ms_r2, e0, e1));
+    ms_r2 /= iters;
+
+    // ---- correctness: both kernels vs a DOUBLE host reference --------------
+    // The per-block integer dot is exact on both paths; only the float
+    // accumulation order differs (256 block terms at K=8192), so
+    // kernel-vs-kernel tolerance false-fails on cancellation-heavy rows.
+    // Judge each kernel by its distance from the double reference instead.
+    (void)ds4_mmq_q8_0_dense_vec(dW, dX, dOutBase, M, 1, K, stream);
+    q8_aligned_dense_vec_kernel<<<M, 32, 0, stream>>>(dOutAl, qs, dq, dX8, M, nb);
+    CK(cudaStreamSynchronize(stream));
+    std::vector<float> ob(M), oa(M);
+    CK(cudaMemcpy(ob.data(), dOutBase, sizeof(float) * M, cudaMemcpyDeviceToHost));
+    CK(cudaMemcpy(oa.data(), dOutAl, sizeof(float) * M, cudaMemcpyDeviceToHost));
+    double err_b = 0.0, err_a = 0.0; int bad = 0;
+    for (int i = 0; i < M; i++) {
+        double ref = 0.0;
+        const uint8_t *rw = &W[(size_t)i * row_bytes];
+        for (int b = 0; b < nb; b++) {
+            const uint8_t *blk = rw + (size_t)b * BLOCK_BYTES;
+            uint16_t h; memcpy(&h, blk, 2);
+            const float dw = __half2float(__ushort_as_half(h));
+            const int8_t *q = (const int8_t *)(blk + 2);
+            const int8_t *u = X8[b].qs;
+            int s = 0;
+            for (int j = 0; j < 32; j++) s += (int)q[j] * (int)u[j];
+            ref += (double)dw * (double)__half2float(__low2half(X8[b].ds)) * (double)s;
+        }
+        const double eb = fabs((double)ob[i] - ref);
+        const double ea = fabs((double)oa[i] - ref);
+        if (eb > err_b) err_b = eb;
+        if (ea > err_a) err_a = ea;
+        // aligned must not be meaningfully further from ref than baseline
+        if (ea > eb * 4.0 + 1e-2) bad++;
+    }
+
+    printf("%-9s K=%-5d M=%-7d (%.1f MB/call, iters=%d)\n", label, K, M, call_gb * 1000.0, iters);
+    printf("  baseline dense_vec : %.4f ms -> %6.1f GB/s\n", ms_base, call_gb / (ms_base / 1e3));
+    printf("  aligned  warp/row  : %.4f ms -> %6.1f GB/s  (%+.1f%%)\n",
+           ms_al, call_gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0));
+    printf("  aligned  2row/blk  : %.4f ms -> %6.1f GB/s  (%+.1f%%)\n",
+           ms_r2, call_gb / (ms_r2 / 1e3), 100.0 * (ms_base / ms_r2 - 1.0));
+    printf("  parity vs double ref: max_abs_err base=%.3e aligned=%.3e bad=%d -> %s\n",
+           err_b, err_a, bad, bad == 0 ? "PASS" : "FAIL");
+    if (bad != 0) exit(2);
+
+    CK(cudaFree(dW)); CK(cudaFree(dArt)); CK(cudaFree(dX));
+    CK(cudaFree(dOutBase)); CK(cudaFree(dOutAl)); CK(cudaFree(dX8));
+    CK(cudaEventDestroy(e0)); CK(cudaEventDestroy(e1));
+    CK(cudaStreamDestroy(stream));
+}
+
+} // namespace
+
+int main(int argc, char **argv) {
+    const uint32_t seed = argc > 1 ? (uint32_t)atoi(argv[1]) : 1234u;
+    if (ds4_mmq_init(0) != 0) { printf("ds4_mmq_init failed\n"); return 1; }
+    printf("PROTO_Q8_ALIGNED (block_q8_0 34B -> aligned SoA [dq][pad64][qs 32B])\n");
+    run_shape("attn_q_b", 32768, 1024, 2000, seed);
+    run_shape("mid",      4096,  2048, 2000, seed + 1);
+    run_shape("out_a",    4096,  8192, 2000, seed + 3);
+    run_shape("head",     129280, 4096, 200, seed + 2);
+    return 0;
+}

cuda/mmq/test/proto_q8_warp8.cu

@@ -0,0 +1,665 @@
+// proto_q8_warp8.cu — Inc4 prototype (megakernel program).
+//
+// The nsys truth table (nsys33) put the three CUSTOM warp8 Q8_0 decode
+// kernels at 21.2 ms/step combined (hc_expand 9.09 / grouped_a 6.98 /
+// pair 5.13) running at 150-195 GB/s on the raw 34-byte block layout.
+// proto_q8_aligned proved the aligned SoA layout ([half dq[nblk]][pad64]
+// [int8 qs[nblk*32]], derived kind 5) reaches 235-245 GB/s at these K.
+//
+// This proto A/Bs each warp8 kernel VERBATIM against an aligned variant
+// with the SAME warp8 structure, prequant inputs (xq int8 + xscale float,
+// quantize_q8_0_f32 semantics), accumulation order, and epilogue — only
+// the weight load path changes (unaligned 34B-stride assembles -> two
+// aligned int4 per block).  Production Flash shapes:
+//   hc_expand : out=4096 in=8192 (attn_output_b)  + out=4096 in=2048 (shared_down)
+//   pair      : in=4096 out0=1024 out1=512 (q_a+kv) + out0=out1=2048 (shared gate+up)
+//   grouped   : group_dim=4096 rank=1024 n_groups=8 (attn_output_a)
+//
+// L2-defeating rotation (>=256 MB of weight copies) + double host reference
+// parity, per the proto_q8_aligned method rules.
+//
+// Build (box):
+//   nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 -I.. \
+//        proto_q8_warp8.cu -lcudart -lcuda -o proto_q8_warp8
+
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <random>
+#include <vector>
+
+#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
+    printf("CUDA ERR %s @%d: %s\n", #x, __LINE__, cudaGetErrorString(e_)); exit(1); } } while (0)
+
+namespace {
+
+// ---- device helpers copied verbatim from ds4_cuda.cu ----------------------
+__device__ __forceinline__ int32_t load_i8x4_i32_unaligned(const int8_t *p) {
+    const uint8_t *u = (const uint8_t *)p;
+    return (int32_t)((uint32_t)u[0] |
+                     ((uint32_t)u[1] << 8) |
+                     ((uint32_t)u[2] << 16) |
+                     ((uint32_t)u[3] << 24));
+}
+__device__ __forceinline__ int32_t load_i8x4_i32_aligned(const int8_t *p) {
+    return *(const int32_t *)p;
+}
+__device__ __forceinline__ int32_t dot_i8x32_dp4a(const int8_t *a, const int8_t *b) {
+    int32_t dot = 0;
+#pragma unroll
+    for (uint32_t i = 0; i < 32u; i += 4u) {
+        dot = __dp4a(load_i8x4_i32_unaligned(a + i), load_i8x4_i32_aligned(b + i), dot);
+    }
+    return dot;
+}
+__device__ __forceinline__ int32_t dot_i8_block(const int8_t *a, const int8_t *b, uint64_t n, int use_dp4a) {
+    if (use_dp4a && n == 32u) return dot_i8x32_dp4a(a, b);
+    int32_t dot = 0;
+    for (uint64_t i = 0; i < n; i++) dot += (int32_t)a[i] * (int32_t)b[i];
+    return dot;
+}
+__device__ float warp_sum_f32(float v) {
+    for (int offset = 16; offset > 0; offset >>= 1) {
+        v += __shfl_down_sync(0xffffffffu, v, offset);
+    }
+    return v;
+}
+
+// Aligned per-block dot: two 16B loads + 8 dp4a, same integer result.
+__device__ __forceinline__ int32_t dot_i8x32_aligned(const int8_t *qs32, const int8_t *xqb) {
+    const int4 w0 = *(const int4 *)qs32;
+    const int4 w1 = *(const int4 *)(qs32 + 16);
+    const int *u = (const int *)xqb;
+    int sumi = 0;
+    sumi = __dp4a(w0.x, u[0], sumi);
+    sumi = __dp4a(w0.y, u[1], sumi);
+    sumi = __dp4a(w0.z, u[2], sumi);
+    sumi = __dp4a(w0.w, u[3], sumi);
+    sumi = __dp4a(w1.x, u[4], sumi);
+    sumi = __dp4a(w1.y, u[5], sumi);
+    sumi = __dp4a(w1.z, u[6], sumi);
+    sumi = __dp4a(w1.w, u[7], sumi);
+    return sumi;
+}
+
+// ==== 1. hc_expand ==========================================================
+// Baseline: verbatim matmul_q8_0_hc_expand_preq_warp8_kernel (ds4_cuda.cu).
+__global__ void hc_expand_base_kernel(
+        float *out_hc, float *block_out, const float *block_add,
+        const float *residual_hc, const float *split, const unsigned char *w,
+        const int8_t *xq, const float *xscale,
+        uint64_t in_dim, uint64_t out_dim, uint32_t n_embd, uint32_t n_hc,
+        uint64_t blocks, int has_add, int use_dp4a) {
+    const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
+    const uint32_t lane = threadIdx.x & 31u;
+    if (row >= out_dim) return;
+    const unsigned char *wr = w + row * blocks * 34;
+    float acc = 0.0f;
+    for (uint64_t b = lane; b < blocks; b += 32u) {
+        const uint64_t i0 = b * 32;
+        const uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
+        const __half *scale_h = (const __half *)(wr + b * 34);
+        const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
+        const int8_t *xqb = xq + b * 32;
+        int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
+        acc += __half2float(*scale_h) * xscale[b] * (float)dot;
+    }
+    acc = warp_sum_f32(acc);
+    if (lane == 0) {
+        const uint32_t d = (uint32_t)row;
+        block_out[d] = acc;
+        float block_v = acc;
+        if (has_add) block_v += block_add[d];
+        const float *post = split + n_hc;
+        const float *comb = split + 2u * n_hc;
+        for (uint32_t dst_hc = 0; dst_hc < n_hc; dst_hc++) {
+            float hc_acc = block_v * post[dst_hc];
+            for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
+                const float comb_v = comb[dst_hc + (uint64_t)src_hc * n_hc];
+                const float res_v = residual_hc[(uint64_t)src_hc * n_embd + d];
+                hc_acc += comb_v * res_v;
+            }
+            out_hc[(uint64_t)dst_hc * n_embd + d] = hc_acc;
+        }
+    }
+}
+
+// Aligned variant: identical loop/epilogue; weight loads via [dq|qs] artifact.
+__global__ void hc_expand_aligned_kernel(
+        float *out_hc, float *block_out, const float *block_add,
+        const float *residual_hc, const float *split,
+        const int8_t *qs_art, const __half *dq_art,
+        const int8_t *xq, const float *xscale,
+        uint64_t in_dim, uint64_t out_dim, uint32_t n_embd, uint32_t n_hc,
+        uint64_t blocks, int has_add) {
+    const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
+    const uint32_t lane = threadIdx.x & 31u;
+    if (row >= out_dim) return;
+    const uint64_t rbase = row * blocks;
+    float acc = 0.0f;
+    for (uint64_t b = lane; b < blocks; b += 32u) {
+        const int sumi = dot_i8x32_aligned(qs_art + (rbase + b) * 32u, xq + b * 32u);
+        acc += __half2float(dq_art[rbase + b]) * xscale[b] * (float)sumi;
+    }
+    acc = warp_sum_f32(acc);
+    if (lane == 0) {
+        const uint32_t d = (uint32_t)row;
+        block_out[d] = acc;
+        float block_v = acc;
+        if (has_add) block_v += block_add[d];
+        const float *post = split + n_hc;
+        const float *comb = split + 2u * n_hc;
+        for (uint32_t dst_hc = 0; dst_hc < n_hc; dst_hc++) {
+            float hc_acc = block_v * post[dst_hc];
+            for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
+                const float comb_v = comb[dst_hc + (uint64_t)src_hc * n_hc];
+                const float res_v = residual_hc[(uint64_t)src_hc * n_embd + d];
+                hc_acc += comb_v * res_v;
+            }
+            out_hc[(uint64_t)dst_hc * n_embd + d] = hc_acc;
+        }
+    }
+}
+
+// ==== 2. pair ===============================================================
+__global__ void pair_base_kernel(
+        float *out0, float *out1, const unsigned char *w0, const unsigned char *w1,
+        const int8_t *xq, const float *xscale,
+        uint64_t in_dim, uint64_t out0_dim, uint64_t out1_dim, uint64_t blocks,
+        int use_dp4a) {
+    uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
+    uint32_t lane = threadIdx.x & 31u;
+    if (row >= out0_dim && row >= out1_dim) return;
+    float acc0 = 0.0f;
+    float acc1 = 0.0f;
+    const unsigned char *wr0 = row < out0_dim ? w0 + row * blocks * 34 : NULL;
+    const unsigned char *wr1 = row < out1_dim ? w1 + row * blocks * 34 : NULL;
+    for (uint64_t b = lane; b < blocks; b += 32u) {
+        uint64_t i0 = b * 32;
+        uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
+        const int8_t *xqb = xq + b * 32;
+        const float xs = xscale[b];
+        if (wr0) {
+            const __half *scale_h = (const __half *)(wr0 + b * 34);
+            const int8_t *qs = (const int8_t *)(wr0 + b * 34 + 2);
+            int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
+            acc0 += __half2float(*scale_h) * xs * (float)dot;
+        }
+        if (wr1) {
+            const __half *scale_h = (const __half *)(wr1 + b * 34);
+            const int8_t *qs = (const int8_t *)(wr1 + b * 34 + 2);
+            int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
+            acc1 += __half2float(*scale_h) * xs * (float)dot;
+        }
+    }
+    acc0 = warp_sum_f32(acc0);
+    acc1 = warp_sum_f32(acc1);
+    if (lane == 0) {
+        if (row < out0_dim) out0[row] = acc0;
+        if (row < out1_dim) out1[row] = acc1;
+    }
+}
+
+__global__ void pair_aligned_kernel(
+        float *out0, float *out1,
+        const int8_t *qs0, const __half *dq0,
+        const int8_t *qs1, const __half *dq1,
+        const int8_t *xq, const float *xscale,
+        uint64_t in_dim, uint64_t out0_dim, uint64_t out1_dim, uint64_t blocks) {
+    uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
+    uint32_t lane = threadIdx.x & 31u;
+    if (row >= out0_dim && row >= out1_dim) return;
+    const uint64_t rbase = row * blocks;
+    const int in0 = row < out0_dim;
+    const int in1 = row < out1_dim;
+    float acc0 = 0.0f;
+    float acc1 = 0.0f;
+    for (uint64_t b = lane; b < blocks; b += 32u) {
+        const int8_t *xqb = xq + b * 32;
+        const float xs = xscale[b];
+        if (in0) {
+            const int sumi = dot_i8x32_aligned(qs0 + (rbase + b) * 32u, xqb);
+            acc0 += __half2float(dq0[rbase + b]) * xs * (float)sumi;
+        }
+        if (in1) {
+            const int sumi = dot_i8x32_aligned(qs1 + (rbase + b) * 32u, xqb);
+            acc1 += __half2float(dq1[rbase + b]) * xs * (float)sumi;
+        }
+    }
+    acc0 = warp_sum_f32(acc0);
+    acc1 = warp_sum_f32(acc1);
+    if (lane == 0) {
+        if (row < out0_dim) out0[row] = acc0;
+        if (row < out1_dim) out1[row] = acc1;
+    }
+}
+
+// ==== 3. grouped ============================================================
+__global__ void grouped_base_kernel(
+        float *low, const unsigned char *w, const int8_t *xq, const float *xscale,
+        uint64_t group_dim, uint64_t rank, uint32_t n_groups, uint32_t n_tokens,
+        uint64_t blocks, int use_dp4a) {
+    const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
+    const uint64_t tok = (uint64_t)blockIdx.y;
+    const uint32_t lane = threadIdx.x & 31u;
+    const uint64_t low_dim = (uint64_t)n_groups * rank;
+    if (row >= low_dim || tok >= n_tokens) return;
+
+    const uint64_t group = row / rank;
+    const uint64_t row_in_group = row - group * rank;
+    const unsigned char *wr = w + (group * rank + row_in_group) * blocks * 34;
+    const uint64_t xrow = tok * (uint64_t)n_groups + group;
+    const int8_t *xqr = xq + xrow * blocks * 32;
+    const float *xsr = xscale + xrow * blocks;
+    float acc = 0.0f;
+
+    for (uint64_t b = lane; b < blocks; b += 32u) {
+        const uint64_t i0 = b * 32;
+        const uint64_t bn = group_dim - i0 < 32 ? group_dim - i0 : 32;
+        const __half *scale_h = (const __half *)(wr + b * 34);
+        const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
+        const int8_t *xqb = xqr + b * 32;
+        int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
+        acc += __half2float(*scale_h) * xsr[b] * (float)dot;
+    }
+    acc = warp_sum_f32(acc);
+    if (lane == 0) low[tok * low_dim + row] = acc;
+}
+
+__global__ void grouped_aligned_kernel(
+        float *low, const int8_t *qs_art, const __half *dq_art,
+        const int8_t *xq, const float *xscale,
+        uint64_t group_dim, uint64_t rank, uint32_t n_groups, uint32_t n_tokens,
+        uint64_t blocks) {
+    const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
+    const uint64_t tok = (uint64_t)blockIdx.y;
+    const uint32_t lane = threadIdx.x & 31u;
+    const uint64_t low_dim = (uint64_t)n_groups * rank;
+    if (row >= low_dim || tok >= n_tokens) return;
+
+    const uint64_t group = row / rank;
+    const uint64_t rbase = row * blocks;
+    const uint64_t xrow = tok * (uint64_t)n_groups + group;
+    const int8_t *xqr = xq + xrow * blocks * 32;
+    const float *xsr = xscale + xrow * blocks;
+    float acc = 0.0f;
+
+    for (uint64_t b = lane; b < blocks; b += 32u) {
+        const int sumi = dot_i8x32_aligned(qs_art + (rbase + b) * 32u, xqr + b * 32u);
+        acc += __half2float(dq_art[rbase + b]) * xsr[b] * (float)sumi;
+    }
+    acc = warp_sum_f32(acc);
+    if (lane == 0) low[tok * low_dim + row] = acc;
+}
+
+// ==== host helpers ==========================================================
+// Mirror quantize_q8_0_f32_kernel semantics (lrintf, clamp 127/-128).
+static void quantize_q8_0_host(const float *x, int8_t *xq, float *xscale, int K) {
+    for (int b = 0; b < K / 32; b++) {
+        float amax = 0.0f;
+        for (int i = 0; i < 32; i++) amax = fmaxf(amax, fabsf(x[b * 32 + i]));
+        const float d = amax / 127.0f;
+        const float id = d != 0.0f ? 1.0f / d : 0.0f;
+        xscale[b] = d;
+        for (int i = 0; i < 32; i++) {
+            int v = (int)lrintf(x[b * 32 + i] * id);
+            v = v > 127 ? 127 : (v < -128 ? -128 : v);
+            xq[b * 32 + i] = (int8_t)v;
+        }
+    }
+}
+
+struct RawArt {
+    std::vector<uint8_t> W;    // raw 34B-block rows
+    std::vector<uint8_t> ART;  // [dq halves][pad64][qs 32B]
+    uint64_t dq_bytes;
+};
+
+static RawArt make_weights(std::mt19937 &rng, long long rows, int blocks) {
+    RawArt r;
+    const long long nblk = rows * blocks;
+    r.W.resize((size_t)nblk * 34);
+    for (auto &b : r.W) b = (uint8_t)(rng() & 0xff);
+    for (long long blk = 0; blk < nblk; blk++) {
+        uint16_t h = (uint16_t)(0x2c00 | (rng() & 0xff));
+        memcpy(&r.W[blk * 34], &h, 2);
+    }
+    r.dq_bytes = ((uint64_t)nblk * 2u + 63u) & ~63ull;
+    r.ART.resize(r.dq_bytes + (uint64_t)nblk * 32u);
+    for (long long blk = 0; blk < nblk; blk++) {
+        memcpy(&r.ART[blk * 2], &r.W[blk * 34], 2);
+        memcpy(&r.ART[r.dq_bytes + (uint64_t)blk * 32u], &r.W[blk * 34 + 2], 32);
+    }
+    return r;
+}
+
+// Double reference GEMV row: sum_b dq[b] * xscale[b] * (int dot)
+static double ref_row(const uint8_t *wrow, const int8_t *xq, const float *xs, int blocks) {
+    double ref = 0.0;
+    for (int b = 0; b < blocks; b++) {
+        const uint8_t *blk = wrow + (size_t)b * 34;
+        uint16_t h; memcpy(&h, blk, 2);
+        const float dw = __half2float(__ushort_as_half(h));
+        const int8_t *q = (const int8_t *)(blk + 2);
+        const int8_t *u = xq + (size_t)b * 32;
+        int s = 0;
+        for (int j = 0; j < 32; j++) s += (int)q[j] * (int)u[j];
+        ref += (double)dw * (double)xs[b] * (double)s;
+    }
+    return ref;
+}
+
+struct DevCopies {
+    uint8_t *dW, *dArt;
+    int n_copies;
+    size_t w_sz, art_sz;
+    const uint8_t *w(int i) const { return dW + (size_t)(i % n_copies) * w_sz; }
+    const uint8_t *art(int i) const { return dArt + (size_t)(i % n_copies) * art_sz; }
+    const int8_t *art_qs(int i, uint64_t dqb) const { return (const int8_t *)(art(i) + dqb); }
+    const __half *art_dq(int i) const { return (const __half *)art(i); }
+};
+
+static DevCopies upload_rotating(const RawArt &r, size_t min_footprint = 256ull << 20) {
+    DevCopies d;
+    d.w_sz = r.W.size(); d.art_sz = r.ART.size();
+    d.n_copies = (int)(min_footprint / (d.w_sz + 1)) + 1;
+    CK(cudaMalloc(&d.dW, d.w_sz * d.n_copies));
+    CK(cudaMalloc(&d.dArt, d.art_sz * d.n_copies));
+    for (int c = 0; c < d.n_copies; c++) {
+        CK(cudaMemcpy(d.dW + (size_t)c * d.w_sz, r.W.data(), d.w_sz, cudaMemcpyHostToDevice));
+        CK(cudaMemcpy(d.dArt + (size_t)c * d.art_sz, r.ART.data(), d.art_sz, cudaMemcpyHostToDevice));
+    }
+    return d;
+}
+
+static float time_loop(cudaStream_t stream, cudaEvent_t e0, cudaEvent_t e1, int iters,
+                       void (*launch)(int, void *), void *ud) {
+    launch(0, ud);
+    CK(cudaGetLastError());
+    CK(cudaStreamSynchronize(stream));
+    CK(cudaEventRecord(e0, stream));
+    for (int i = 0; i < iters; i++) launch(i, ud);
+    CK(cudaEventRecord(e1, stream));
+    CK(cudaStreamSynchronize(stream));
+    float ms = 0.0f; CK(cudaEventElapsedTime(&ms, e0, e1));
+    return ms / iters;
+}
+
+static int check_pair_tol(const std::vector<float> &base, const std::vector<float> &al,
+                          const std::vector<double> &ref, const char *tag) {
+    double eb_max = 0.0, ea_max = 0.0; int bad = 0;
+    for (size_t i = 0; i < ref.size(); i++) {
+        const double eb = fabs((double)base[i] - ref[i]);
+        const double ea = fabs((double)al[i] - ref[i]);
+        if (eb > eb_max) eb_max = eb;
+        if (ea > ea_max) ea_max = ea;
+        if (ea > eb * 4.0 + 1e-2) bad++;
+    }
+    printf("  parity(%s): max_abs_err base=%.3e aligned=%.3e bad=%d -> %s\n",
+           tag, eb_max, ea_max, bad, bad == 0 ? "PASS" : "FAIL");
+    return bad;
+}
+
+// ==== section runners =======================================================
+struct HcCtx {
+    cudaStream_t stream; DevCopies dev; uint64_t dq_bytes;
+    float *out_hc, *block_out; const float *block_add, *residual, *split;
+    const int8_t *xq; const float *xs;
+    uint64_t in_dim, out_dim, blocks; uint32_t n_embd, n_hc; int base;
+};
+static void hc_launch(int i, void *p) {
+    HcCtx *c = (HcCtx *)p;
+    const unsigned grid = ((unsigned)c->out_dim + 7u) / 8u;
+    if (c->base) {
+        hc_expand_base_kernel<<<grid, 256, 0, c->stream>>>(
+            c->out_hc, c->block_out, c->block_add, c->residual, c->split,
+            c->dev.w(i), c->xq, c->xs, c->in_dim, c->out_dim, c->n_embd, c->n_hc,
+            c->blocks, 1, 1);
+    } else {
+        hc_expand_aligned_kernel<<<grid, 256, 0, c->stream>>>(
+            c->out_hc, c->block_out, c->block_add, c->residual, c->split,
+            c->dev.art_qs(i, c->dq_bytes), c->dev.art_dq(i), c->xq, c->xs,
+            c->in_dim, c->out_dim, c->n_embd, c->n_hc, c->blocks, 1);
+    }
+}
+
+static int run_hc_expand(const char *label, int out_dim, int in_dim, int iters, uint32_t seed) {
+    const int blocks = in_dim / 32, n_hc = 4;
+    std::mt19937 rng(seed);
+    RawArt r = make_weights(rng, out_dim, blocks);
+    std::vector<float> X(in_dim);
+    for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
+    std::vector<int8_t> XQ(in_dim); std::vector<float> XS(blocks);
+    quantize_q8_0_host(X.data(), XQ.data(), XS.data(), in_dim);
+    std::vector<float> ADD(out_dim), RES((size_t)n_hc * out_dim), SPL(2 * n_hc + n_hc * n_hc);
+    for (auto &v : ADD) v = ((float)(rng() % 2000) - 1000.0f) / 1000.0f;
+    for (auto &v : RES) v = ((float)(rng() % 2000) - 1000.0f) / 1000.0f;
+    for (auto &v : SPL) v = ((float)(rng() % 2000) - 1000.0f) / 2000.0f;
+
+    DevCopies dev = upload_rotating(r);
+    float *dOutHc, *dBlk, *dAdd, *dRes, *dSpl, *dXs; int8_t *dXq;
+    CK(cudaMalloc(&dOutHc, sizeof(float) * n_hc * out_dim));
+    CK(cudaMalloc(&dBlk, sizeof(float) * out_dim));
+    CK(cudaMalloc(&dAdd, sizeof(float) * out_dim));
+    CK(cudaMalloc(&dRes, sizeof(float) * RES.size()));
+    CK(cudaMalloc(&dSpl, sizeof(float) * SPL.size()));
+    CK(cudaMalloc(&dXq, XQ.size()));
+    CK(cudaMalloc(&dXs, sizeof(float) * XS.size()));
+    CK(cudaMemcpy(dAdd, ADD.data(), sizeof(float) * ADD.size(), cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dRes, RES.data(), sizeof(float) * RES.size(), cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dSpl, SPL.data(), sizeof(float) * SPL.size(), cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dXq, XQ.data(), XQ.size(), cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dXs, XS.data(), sizeof(float) * XS.size(), cudaMemcpyHostToDevice));
+
+    cudaStream_t stream; CK(cudaStreamCreate(&stream));
+    cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
+
+    HcCtx c = {stream, dev, r.dq_bytes, dOutHc, dBlk, dAdd, dRes, dSpl,
+               dXq, dXs, (uint64_t)in_dim, (uint64_t)out_dim, (uint64_t)blocks,
+               (uint32_t)out_dim, (uint32_t)n_hc, 1};
+    const float ms_base = time_loop(stream, e0, e1, iters, hc_launch, &c);
+    std::vector<float> bBlk(out_dim), bHc((size_t)n_hc * out_dim);
+    CK(cudaMemcpy(bBlk.data(), dBlk, sizeof(float) * out_dim, cudaMemcpyDeviceToHost));
+    CK(cudaMemcpy(bHc.data(), dOutHc, sizeof(float) * bHc.size(), cudaMemcpyDeviceToHost));
+    c.base = 0;
+    const float ms_al = time_loop(stream, e0, e1, iters, hc_launch, &c);
+    std::vector<float> aBlk(out_dim), aHc((size_t)n_hc * out_dim);
+    CK(cudaMemcpy(aBlk.data(), dBlk, sizeof(float) * out_dim, cudaMemcpyDeviceToHost));
+    CK(cudaMemcpy(aHc.data(), dOutHc, sizeof(float) * aHc.size(), cudaMemcpyDeviceToHost));
+
+    // double refs
+    std::vector<double> rBlk(out_dim), rHc((size_t)n_hc * out_dim);
+    for (int d = 0; d < out_dim; d++) {
+        rBlk[d] = ref_row(&r.W[(size_t)d * blocks * 34], XQ.data(), XS.data(), blocks);
+        const double bv = rBlk[d] + (double)ADD[d];
+        for (int dst = 0; dst < n_hc; dst++) {
+            double hc = bv * (double)SPL[n_hc + dst];
+            for (int src = 0; src < n_hc; src++)
+                hc += (double)SPL[2 * n_hc + dst + src * n_hc] * (double)RES[(size_t)src * out_dim + d];
+            rHc[(size_t)dst * out_dim + d] = hc;
+        }
+    }
+    const double gb = (double)out_dim * blocks * 34 / 1e9;
+    printf("%-22s out=%d in=%d (%.1f MB/call, iters=%d)\n", label, out_dim, in_dim, gb * 1000, iters);
+    printf("  baseline warp8     : %.4f ms -> %6.1f GB/s\n", ms_base, gb / (ms_base / 1e3));
+    printf("  aligned  warp8     : %.4f ms -> %6.1f GB/s  (%+.1f%%)\n",
+           ms_al, gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0));
+    int bad = check_pair_tol(bBlk, aBlk, rBlk, "block_out");
+    bad += check_pair_tol(bHc, aHc, rHc, "out_hc");
+
+    CK(cudaFree(dev.dW)); CK(cudaFree(dev.dArt));
+    CK(cudaFree(dOutHc)); CK(cudaFree(dBlk)); CK(cudaFree(dAdd)); CK(cudaFree(dRes));
+    CK(cudaFree(dSpl)); CK(cudaFree(dXq)); CK(cudaFree(dXs));
+    CK(cudaEventDestroy(e0)); CK(cudaEventDestroy(e1)); CK(cudaStreamDestroy(stream));
+    return bad;
+}
+
+struct PairCtx {
+    cudaStream_t stream; DevCopies d0, d1; uint64_t dqb0, dqb1;
+    float *out0, *out1; const int8_t *xq; const float *xs;
+    uint64_t in_dim, out0_dim, out1_dim, blocks; int base;
+};
+static void pair_launch(int i, void *p) {
+    PairCtx *c = (PairCtx *)p;
+    const uint64_t max_out = c->out0_dim > c->out1_dim ? c->out0_dim : c->out1_dim;
+    const unsigned grid = ((unsigned)max_out + 7u) / 8u;
+    if (c->base) {
+        pair_base_kernel<<<grid, 256, 0, c->stream>>>(
+            c->out0, c->out1, c->d0.w(i), c->d1.w(i), c->xq, c->xs,
+            c->in_dim, c->out0_dim, c->out1_dim, c->blocks, 1);
+    } else {
+        pair_aligned_kernel<<<grid, 256, 0, c->stream>>>(
+            c->out0, c->out1,
+            c->d0.art_qs(i, c->dqb0), c->d0.art_dq(i),
+            c->d1.art_qs(i, c->dqb1), c->d1.art_dq(i),
+            c->xq, c->xs, c->in_dim, c->out0_dim, c->out1_dim, c->blocks);
+    }
+}
+
+static int run_pair(const char *label, int out0, int out1, int in_dim, int iters, uint32_t seed) {
+    const int blocks = in_dim / 32;
+    std::mt19937 rng(seed);
+    RawArt r0 = make_weights(rng, out0, blocks);
+    RawArt r1 = make_weights(rng, out1, blocks);
+    std::vector<float> X(in_dim);
+    for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
+    std::vector<int8_t> XQ(in_dim); std::vector<float> XS(blocks);
+    quantize_q8_0_host(X.data(), XQ.data(), XS.data(), in_dim);
+
+    DevCopies d0 = upload_rotating(r0, 128ull << 20);
+    DevCopies d1 = upload_rotating(r1, 128ull << 20);
+    float *dOut0, *dOut1, *dXs; int8_t *dXq;
+    CK(cudaMalloc(&dOut0, sizeof(float) * out0));
+    CK(cudaMalloc(&dOut1, sizeof(float) * out1));
+    CK(cudaMalloc(&dXq, XQ.size()));
+    CK(cudaMalloc(&dXs, sizeof(float) * XS.size()));
+    CK(cudaMemcpy(dXq, XQ.data(), XQ.size(), cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dXs, XS.data(), sizeof(float) * XS.size(), cudaMemcpyHostToDevice));
+
+    cudaStream_t stream; CK(cudaStreamCreate(&stream));
+    cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
+
+    PairCtx c = {stream, d0, d1, r0.dq_bytes, r1.dq_bytes, dOut0, dOut1,
+                 dXq, dXs, (uint64_t)in_dim, (uint64_t)out0, (uint64_t)out1, (uint64_t)blocks, 1};
+    const float ms_base = time_loop(stream, e0, e1, iters, pair_launch, &c);
+    std::vector<float> b0(out0), b1(out1);
+    CK(cudaMemcpy(b0.data(), dOut0, sizeof(float) * out0, cudaMemcpyDeviceToHost));
+    CK(cudaMemcpy(b1.data(), dOut1, sizeof(float) * out1, cudaMemcpyDeviceToHost));
+    c.base = 0;
+    const float ms_al = time_loop(stream, e0, e1, iters, pair_launch, &c);
+    std::vector<float> a0(out0), a1(out1);
+    CK(cudaMemcpy(a0.data(), dOut0, sizeof(float) * out0, cudaMemcpyDeviceToHost));
+    CK(cudaMemcpy(a1.data(), dOut1, sizeof(float) * out1, cudaMemcpyDeviceToHost));
+
+    std::vector<double> ref0(out0), ref1(out1);
+    for (int d = 0; d < out0; d++) ref0[d] = ref_row(&r0.W[(size_t)d * blocks * 34], XQ.data(), XS.data(), blocks);
+    for (int d = 0; d < out1; d++) ref1[d] = ref_row(&r1.W[(size_t)d * blocks * 34], XQ.data(), XS.data(), blocks);
+
+    const double gb = ((double)out0 + out1) * blocks * 34 / 1e9;
+    printf("%-22s out0=%d out1=%d in=%d (%.1f MB/call, iters=%d)\n", label, out0, out1, in_dim, gb * 1000, iters);
+    printf("  baseline warp8     : %.4f ms -> %6.1f GB/s\n", ms_base, gb / (ms_base / 1e3));
+    printf("  aligned  warp8     : %.4f ms -> %6.1f GB/s  (%+.1f%%)\n",
+           ms_al, gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0));
+    int bad = check_pair_tol(b0, a0, ref0, "out0");
+    bad += check_pair_tol(b1, a1, ref1, "out1");
+
+    CK(cudaFree(d0.dW)); CK(cudaFree(d0.dArt)); CK(cudaFree(d1.dW)); CK(cudaFree(d1.dArt));
+    CK(cudaFree(dOut0)); CK(cudaFree(dOut1)); CK(cudaFree(dXq)); CK(cudaFree(dXs));
+    CK(cudaEventDestroy(e0)); CK(cudaEventDestroy(e1)); CK(cudaStreamDestroy(stream));
+    return bad;
+}
+
+struct GroupCtx {
+    cudaStream_t stream; DevCopies dev; uint64_t dqb;
+    float *low; const int8_t *xq; const float *xs;
+    uint64_t group_dim, rank; uint32_t n_groups; uint64_t blocks; int base;
+};
+static void grouped_launch(int i, void *p) {
+    GroupCtx *c = (GroupCtx *)p;
+    const uint64_t low_dim = (uint64_t)c->n_groups * c->rank;
+    const unsigned grid = ((unsigned)low_dim + 7u) / 8u;
+    if (c->base) {
+        grouped_base_kernel<<<grid, 256, 0, c->stream>>>(
+            c->low, c->dev.w(i), c->xq, c->xs, c->group_dim, c->rank,
+            c->n_groups, 1u, c->blocks, 1);
+    } else {
+        grouped_aligned_kernel<<<grid, 256, 0, c->stream>>>(
+            c->low, c->dev.art_qs(i, c->dqb), c->dev.art_dq(i), c->xq, c->xs,
+            c->group_dim, c->rank, c->n_groups, 1u, c->blocks);
+    }
+}
+
+static int run_grouped(const char *label, int group_dim, int rank, int n_groups, int iters, uint32_t seed) {
+    const int blocks = group_dim / 32;
+    const long long low_dim = (long long)n_groups * rank;
+    std::mt19937 rng(seed);
+    RawArt r = make_weights(rng, low_dim, blocks);
+    // per-group activation rows
+    std::vector<float> X((size_t)n_groups * group_dim);
+    for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
+    std::vector<int8_t> XQ(X.size()); std::vector<float> XS((size_t)n_groups * blocks);
+    for (int g = 0; g < n_groups; g++)
+        quantize_q8_0_host(&X[(size_t)g * group_dim], &XQ[(size_t)g * group_dim],
+                           &XS[(size_t)g * blocks], group_dim);
+
+    DevCopies dev = upload_rotating(r);
+    float *dLow, *dXs; int8_t *dXq;
+    CK(cudaMalloc(&dLow, sizeof(float) * low_dim));
+    CK(cudaMalloc(&dXq, XQ.size()));
+    CK(cudaMalloc(&dXs, sizeof(float) * XS.size()));
+    CK(cudaMemcpy(dXq, XQ.data(), XQ.size(), cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dXs, XS.data(), sizeof(float) * XS.size(), cudaMemcpyHostToDevice));
+
+    cudaStream_t stream; CK(cudaStreamCreate(&stream));
+    cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
+
+    GroupCtx c = {stream, dev, r.dq_bytes, dLow, dXq, dXs,
+                  (uint64_t)group_dim, (uint64_t)rank, (uint32_t)n_groups, (uint64_t)blocks, 1};
+    const float ms_base = time_loop(stream, e0, e1, iters, grouped_launch, &c);
+    std::vector<float> bl(low_dim);
+    CK(cudaMemcpy(bl.data(), dLow, sizeof(float) * low_dim, cudaMemcpyDeviceToHost));
+    c.base = 0;
+    const float ms_al = time_loop(stream, e0, e1, iters, grouped_launch, &c);
+    std::vector<float> al(low_dim);
+    CK(cudaMemcpy(al.data(), dLow, sizeof(float) * low_dim, cudaMemcpyDeviceToHost));
+
+    std::vector<double> ref(low_dim);
+    for (long long d = 0; d < low_dim; d++) {
+        const int g = (int)(d / rank);
+        ref[d] = ref_row(&r.W[(size_t)d * blocks * 34], &XQ[(size_t)g * group_dim],
+                         &XS[(size_t)g * blocks], blocks);
+    }
+
+    const double gb = (double)low_dim * blocks * 34 / 1e9;
+    printf("%-22s gdim=%d rank=%d groups=%d (%.1f MB/call, iters=%d)\n",
+           label, group_dim, rank, n_groups, gb * 1000, iters);
+    printf("  baseline warp8     : %.4f ms -> %6.1f GB/s\n", ms_base, gb / (ms_base / 1e3));
+    printf("  aligned  warp8     : %.4f ms -> %6.1f GB/s  (%+.1f%%)\n",
+           ms_al, gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0));
+    int bad = check_pair_tol(bl, al, ref, "low");
+
+    CK(cudaFree(dev.dW)); CK(cudaFree(dev.dArt));
+    CK(cudaFree(dLow)); CK(cudaFree(dXq)); CK(cudaFree(dXs));
+    CK(cudaEventDestroy(e0)); CK(cudaEventDestroy(e1)); CK(cudaStreamDestroy(stream));
+    return bad;
+}
+
+} // namespace
+
+int main(int argc, char **argv) {
+    const uint32_t seed = argc > 1 ? (uint32_t)atoi(argv[1]) : 1234u;
+    printf("PROTO_Q8_WARP8 (Inc4: aligned-SoA variants of the custom warp8 q8_0 kernels)\n");
+    int bad = 0;
+    bad += run_hc_expand("hc_expand(output_b)", 4096, 8192, 1500, seed);
+    bad += run_hc_expand("hc_expand(shr_down)", 4096, 2048, 2000, seed + 1);
+    bad += run_pair("pair(q_a+kv)",       1024, 512, 4096, 2000, seed + 2);
+    bad += run_pair("pair(shr_gate+up)",  2048, 2048, 4096, 2000, seed + 3);
+    bad += run_grouped("grouped(output_a)", 4096, 1024, 8, 1500, seed + 4);
+    printf(bad == 0 ? "ALL PASS\n" : "FAILURES: %d\n", bad);
+    return bad == 0 ? 0 : 2;
+}

cuda/mmq/test/test_f16_gemm_lt.cu

@@ -0,0 +1,295 @@
+/* P3.3 proto (2026-07-07, CLOSED "cuBLAS knows best"): are the cutlass_80_*
+ * (Ampere-tile) f16 GEMMs on sm_121 leaving throughput on the table?  NO:
+ * - The 6.25 ms cutlass_80_64x256 launch (129x/12k admission) is the
+ *   ATTENTION OUTPUT projection (M=8192=8 groups x 1024 rank, K=4096,
+ *   N=4096 -> 47 TFLOPS, near roof).  The "11 TFLOPS" that chartered this
+ *   recon assumed it was the indexer q_b (K=1024) -- wrong dims.  The real
+ *   indexer q_b GEMM is producer-gated (n_comp > top-k; last chunk only at
+ *   12k) and runs cutlass_80_128x256 at 1.3 ms / 53 TFLOPS.
+ * - cublasLt heuristics beat GemmEx by only ~6-10% on these shapes (not
+ *   worth the plumbing); TF32 math mode, operand alignment (A/B/C +16/+32B),
+ *   L2 warmth (8-set rotation) and host-vs-device A were all bracketed:
+ *   only host-mapped A costs real time (2x), and production weights already
+ *   resolve to device-resident copies (HBM cache / WS import) -- verified
+ *   in-server via DS4_F16_PTR_PROBE (recon-only, not in tree).
+ * - hc mix (m=24) is latency-floor at any tile; comp_proj m=256 is fine.
+ *
+ * Shapes measured (production call form: A=weights f16 OP_T lda=k, B=acts
+ * f16 OP_N ldb=k, C f32 ldc=m, computeType F32, 32 MiB handle workspace):
+ *
+ * Build (on a GB10 box):
+ *   /usr/local/cuda/bin/nvcc -O3 -arch=sm_121 -std=c++17 \
+ *     -o /home/ent/dspark_work/test_f16_gemm_lt cuda/mmq/test/test_f16_gemm_lt.cu \
+ *     -lcublasLt -lcublas
+ * Run:  ./test_f16_gemm_lt            (all shapes, 100 iters each)
+ * Kernel names: nsys profile -o /tmp/lt ./test_f16_gemm_lt && inspect sqlite.
+ */
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <cuda_fp16.h>
+#include <cublas_v2.h>
+#include <cublasLt.h>
+
+#define CK(x) do { auto _e = (x); if (_e != cudaSuccess) { \
+    fprintf(stderr, "CUDA ERR %s @%d: %s\n", #x, __LINE__, cudaGetErrorString(_e)); exit(1); } } while (0)
+#define CB(x) do { auto _e = (x); if (_e != CUBLAS_STATUS_SUCCESS) { \
+    fprintf(stderr, "CUBLAS ERR %s @%d: %d\n", #x, __LINE__, (int)_e); exit(1); } } while (0)
+
+static void fill_half(__half *p, size_t nelem, unsigned seed) {
+    __half *h = (__half *)malloc(nelem * sizeof(__half));
+    unsigned s = seed;
+    for (size_t i = 0; i < nelem; i++) {
+        s = s * 1664525u + 1013904223u;
+        h[i] = __float2half(((float)(s >> 8) / (float)(1u << 24) - 0.5f) * 0.25f);
+    }
+    CK(cudaMemcpy(p, h, nelem * sizeof(__half), cudaMemcpyHostToDevice));
+    free(h);
+}
+
+static float time_iters(cudaStream_t st, int iters, void (*fn)(void *), void *ud) {
+    cudaEvent_t e0, e1;
+    CK(cudaEventCreate(&e0));
+    CK(cudaEventCreate(&e1));
+    for (int i = 0; i < 10; i++) fn(ud);       /* warmup */
+    CK(cudaStreamSynchronize(st));
+    CK(cudaEventRecord(e0, st));
+    for (int i = 0; i < iters; i++) fn(ud);
+    CK(cudaEventRecord(e1, st));
+    CK(cudaEventSynchronize(e1));
+    float ms = 0.f;
+    CK(cudaEventElapsedTime(&ms, e0, e1));
+    cudaEventDestroy(e0);
+    cudaEventDestroy(e1);
+    return ms / iters;
+}
+
+struct GemmCase {
+    int m, n, k;
+    const __half *A, *B;
+    float *C;
+    cublasHandle_t bl;
+    cublasLtHandle_t lt;
+    void *ws;
+    size_t ws_bytes;
+    cublasLtMatmulDesc_t op;
+    cublasLtMatrixLayout_t la, lb, lc;
+    cublasLtMatmulAlgo_t algo;
+    int has_algo;
+};
+
+static void run_gemmex(void *ud) {
+    GemmCase *g = (GemmCase *)ud;
+    const float alpha = 1.f, beta = 0.f;
+    CB(cublasGemmEx(g->bl, CUBLAS_OP_T, CUBLAS_OP_N, g->m, g->n, g->k,
+                    &alpha, g->A, CUDA_R_16F, g->k, g->B, CUDA_R_16F, g->k,
+                    &beta, g->C, CUDA_R_32F, g->m, CUDA_R_32F, CUBLAS_GEMM_DEFAULT));
+}
+
+static void run_lt(void *ud) {
+    GemmCase *g = (GemmCase *)ud;
+    const float alpha = 1.f, beta = 0.f;
+    CB(cublasLtMatmul(g->lt, g->op, &alpha, g->A, g->la, g->B, g->lb,
+                      &beta, g->C, g->lc, g->C, g->lc,
+                      g->has_algo ? &g->algo : NULL, g->ws, g->ws_bytes, 0));
+}
+
+int main(void) {
+    struct { const char *name; int m, n, k; } shapes[] = {
+        { "indexer_q_b m=8192 n=4096 k=1024", 8192, 4096, 1024 },
+        { "hc_mix      m=24   n=4096 k=16384", 24, 4096, 16384 },
+        { "comp_proj   m=256  n=4096 k=4096", 256, 4096, 4096 },
+    };
+    cublasHandle_t bl;
+    CB(cublasCreate(&bl));
+    void *hws = NULL;
+    const size_t hws_bytes = 32u * 1024u * 1024u;
+    CK(cudaMalloc(&hws, hws_bytes));
+    CB(cublasSetWorkspace(bl, hws, hws_bytes));
+    cublasLtHandle_t lt;
+    CB(cublasLtCreate(&lt));
+    void *ltws = NULL;
+    const size_t ltws_bytes = 32u * 1024u * 1024u;
+    CK(cudaMalloc(&ltws, ltws_bytes));
+
+    for (auto &s : shapes) {
+        __half *A, *B;
+        float *C;
+        CK(cudaMalloc(&A, (size_t)s.m * s.k * sizeof(__half)));
+        CK(cudaMalloc(&B, (size_t)s.n * s.k * sizeof(__half)));
+        CK(cudaMalloc(&C, (size_t)s.m * s.n * sizeof(float)));
+        fill_half(A, (size_t)s.m * s.k, 1234);
+        fill_half(B, (size_t)s.n * s.k, 5678);
+
+        GemmCase g = {};
+        g.m = s.m; g.n = s.n; g.k = s.k;
+        g.A = A; g.B = B; g.C = C;
+        g.bl = bl; g.lt = lt; g.ws = ltws; g.ws_bytes = ltws_bytes;
+
+        const double flops = 2.0 * s.m * s.n * s.k;
+        const float ms_ex = time_iters(0, 100, run_gemmex, &g);
+        printf("%-36s GemmEx DEFAULT      : %8.3f ms  %6.1f TFLOPS\n",
+               s.name, ms_ex, flops / (ms_ex * 1e9));
+        /* Production handle state (ds4_gpu_init): deprecated TF32 math mode.
+         * On CUDA 13 this may shunt the heuristic onto a legacy kernel table
+         * -- the suspected source of the 64x256 cutlass_80 pick. */
+        CB(cublasSetMathMode(bl, CUBLAS_TF32_TENSOR_OP_MATH));
+        const float ms_tf = time_iters(0, 100, run_gemmex, &g);
+        printf("%-36s GemmEx TF32-mode    : %8.3f ms  %6.1f TFLOPS\n",
+               s.name, ms_tf, flops / (ms_tf * 1e9));
+        CB(cublasSetMathMode(bl, CUBLAS_DEFAULT_MATH));
+        /* Production pointer classes: A = gguf tensor offset inside a mapped
+         * model (32B alignment class, not cudaMalloc's 256B); B = cuda_tmp
+         * suballocation.  Heuristics gate kernels on operand alignment. */
+        {
+            __half *A32, *B32;
+            CK(cudaMalloc(&A32, (size_t)s.m * s.k * sizeof(__half) + 256));
+            CK(cudaMalloc(&B32, (size_t)s.n * s.k * sizeof(__half) + 256));
+            CK(cudaMemcpy((char *)A32 + 32, A, (size_t)s.m * s.k * sizeof(__half), cudaMemcpyDeviceToDevice));
+            CK(cudaMemcpy((char *)B32 + 32, B, (size_t)s.n * s.k * sizeof(__half), cudaMemcpyDeviceToDevice));
+            GemmCase g2 = g;
+            g2.A = (const __half *)((const char *)A32 + 32);
+            const float ms_a = time_iters(0, 100, run_gemmex, &g2);
+            printf("%-36s GemmEx A+32B        : %8.3f ms  %6.1f TFLOPS\n",
+                   s.name, ms_a, flops / (ms_a * 1e9));
+            GemmCase g3 = g;
+            g3.B = (const __half *)((const char *)B32 + 32);
+            const float ms_b = time_iters(0, 100, run_gemmex, &g3);
+            printf("%-36s GemmEx B+32B        : %8.3f ms  %6.1f TFLOPS\n",
+                   s.name, ms_b, flops / (ms_b * 1e9));
+            GemmCase g4 = g2;
+            g4.B = g3.B;
+            const float ms_ab = time_iters(0, 100, run_gemmex, &g4);
+            printf("%-36s GemmEx A+32B B+32B  : %8.3f ms  %6.1f TFLOPS\n",
+                   s.name, ms_ab, flops / (ms_ab * 1e9));
+            float *C32;
+            CK(cudaMalloc(&C32, (size_t)s.m * s.n * sizeof(float) + 256));
+            GemmCase g5 = g;
+            g5.C = (float *)((char *)C32 + 32);
+            const float ms_c = time_iters(0, 100, run_gemmex, &g5);
+            printf("%-36s GemmEx C+32B        : %8.3f ms  %6.1f TFLOPS\n",
+                   s.name, ms_c, flops / (ms_c * 1e9));
+            GemmCase g6 = g;
+            g6.C = (float *)((char *)C32 + 16);
+            const float ms_c16 = time_iters(0, 100, run_gemmex, &g6);
+            printf("%-36s GemmEx C+16B        : %8.3f ms  %6.1f TFLOPS\n",
+                   s.name, ms_c16, flops / (ms_c16 * 1e9));
+            CK(cudaFree(C32));
+            CK(cudaFree(A32));
+            CK(cudaFree(B32));
+        }
+
+        /* Lt setup: same TN f16/f16/f32-compute layout */
+        CB(cublasLtMatmulDescCreate(&g.op, CUBLAS_COMPUTE_32F, CUDA_R_32F));
+        cublasOperation_t ta = CUBLAS_OP_T, tb = CUBLAS_OP_N;
+        CB(cublasLtMatmulDescSetAttribute(g.op, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)));
+        CB(cublasLtMatmulDescSetAttribute(g.op, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)));
+        CB(cublasLtMatrixLayoutCreate(&g.la, CUDA_R_16F, s.k, s.m, s.k)); /* A: k x m, lda=k (OP_T) */
+        CB(cublasLtMatrixLayoutCreate(&g.lb, CUDA_R_16F, s.k, s.n, s.k));
+        CB(cublasLtMatrixLayoutCreate(&g.lc, CUDA_R_32F, s.m, s.n, s.m));
+
+        /* Substrate x L2-warmth matrix.  Production weight pointers are
+         * host-mapped (client mmap-registered / WS IPC import -- cmtp3d
+         * measured both at the same speed); the proto default is cudaMalloc
+         * A reused every iteration (L2-warm).  8-set rotation defeats L2
+         * reuse; host-alloc A reproduces the mapped-weight class. */
+        {
+            enum { NSETS = 8 };
+            static __half *As[NSETS], *Bs[NSETS];
+            static float *Cs[NSETS];
+            static GemmCase gsets;
+            static int set_idx;
+            struct Rot {
+                static void run(void *ud) {
+                    GemmCase *g = (GemmCase *)ud;
+                    (void)g;
+                    GemmCase tmp = gsets;
+                    tmp.A = As[set_idx % NSETS];
+                    tmp.B = Bs[set_idx % NSETS];
+                    tmp.C = Cs[set_idx % NSETS];
+                    set_idx++;
+                    run_gemmex(&tmp);
+                }
+            };
+            gsets = g;
+            /* dev-cold */
+            for (int i = 0; i < NSETS; i++) {
+                CK(cudaMalloc(&As[i], (size_t)s.m * s.k * sizeof(__half)));
+                CK(cudaMalloc(&Bs[i], (size_t)s.n * s.k * sizeof(__half)));
+                CK(cudaMalloc(&Cs[i], (size_t)s.m * s.n * sizeof(float)));
+                fill_half(As[i], (size_t)s.m * s.k, 100 + i);
+                fill_half(Bs[i], (size_t)s.n * s.k, 200 + i);
+            }
+            set_idx = 0;
+            const float ms_dc = time_iters(0, 100, Rot::run, &g);
+            printf("%-36s GemmEx devA cold    : %8.3f ms  %6.1f TFLOPS\n",
+                   s.name, ms_dc, flops / (ms_dc * 1e9));
+            for (int i = 0; i < NSETS; i++) { CK(cudaFree(As[i])); CK(cudaFree(Bs[i])); CK(cudaFree(Cs[i])); }
+            /* hostA warm */
+            __half *Ah;
+            CK(cudaHostAlloc(&Ah, (size_t)s.m * s.k * sizeof(__half), cudaHostAllocMapped));
+            {
+                unsigned sd = 42;
+                for (size_t i = 0; i < (size_t)s.m * s.k; i++) {
+                    sd = sd * 1664525u + 1013904223u;
+                    Ah[i] = __float2half(((float)(sd >> 8) / (float)(1u << 24) - 0.5f) * 0.25f);
+                }
+            }
+            GemmCase gh = g;
+            gh.A = Ah;
+            const float ms_hw = time_iters(0, 100, run_gemmex, &gh);
+            printf("%-36s GemmEx hostA warm   : %8.3f ms  %6.1f TFLOPS\n",
+                   s.name, ms_hw, flops / (ms_hw * 1e9));
+            /* hostA cold: rotate device B/C, host A stays (weights are the
+             * suspect; hostA likely bypasses L2 anyway) */
+            for (int i = 0; i < NSETS; i++) {
+                CK(cudaMalloc(&Bs[i], (size_t)s.n * s.k * sizeof(__half)));
+                CK(cudaMalloc(&Cs[i], (size_t)s.m * s.n * sizeof(float)));
+                fill_half(Bs[i], (size_t)s.n * s.k, 300 + i);
+                As[i] = Ah;
+            }
+            gsets = gh;
+            set_idx = 0;
+            const float ms_hc = time_iters(0, 100, Rot::run, &g);
+            printf("%-36s GemmEx hostA cold   : %8.3f ms  %6.1f TFLOPS\n",
+                   s.name, ms_hc, flops / (ms_hc * 1e9));
+            for (int i = 0; i < NSETS; i++) { CK(cudaFree(Bs[i])); CK(cudaFree(Cs[i])); }
+            CK(cudaFreeHost(Ah));
+        }
+        cublasLtMatmulPreference_t pref;
+        CB(cublasLtMatmulPreferenceCreate(&pref));
+        CB(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+                                                &ltws_bytes, sizeof(ltws_bytes)));
+        cublasLtMatmulHeuristicResult_t res[8];
+        int nres = 0;
+        CB(cublasLtMatmulAlgoGetHeuristic(lt, g.op, g.la, g.lb, g.lc, g.lc, pref, 8, res, &nres));
+        printf("%-36s Lt heuristics: %d candidates\n", s.name, nres);
+        float best_ms = 1e30f;
+        int best_i = -1;
+        for (int i = 0; i < nres; i++) {
+            g.algo = res[i].algo;
+            g.has_algo = 1;
+            const float ms = time_iters(0, 100, run_lt, &g);
+            int algo_id = -1, tile = -1, splitk = -1;
+            size_t sz;
+            cublasLtMatmulAlgoConfigGetAttribute(&g.algo, CUBLASLT_ALGO_CONFIG_ID, &algo_id, sizeof(algo_id), &sz);
+            cublasLtMatmulAlgoConfigGetAttribute(&g.algo, CUBLASLT_ALGO_CONFIG_TILE_ID, &tile, sizeof(tile), &sz);
+            cublasLtMatmulAlgoConfigGetAttribute(&g.algo, CUBLASLT_ALGO_CONFIG_SPLITK_NUM, &splitk, sizeof(splitk), &sz);
+            printf("%-36s   Lt[%d] algo=%d tile=%d splitk=%d ws=%zu: %8.3f ms  %6.1f TFLOPS\n",
+                   s.name, i, algo_id, tile, splitk, res[i].workspaceSize,
+                   ms, flops / (ms * 1e9));
+            if (ms < best_ms) { best_ms = ms; best_i = i; }
+        }
+        if (best_i >= 0)
+            printf("%-36s Lt BEST[%d]          : %8.3f ms  %6.1f TFLOPS  (GemmEx/Lt = %.2fx)\n",
+                   s.name, best_i, best_ms, flops / (best_ms * 1e9), ms_ex / best_ms);
+        cublasLtMatmulPreferenceDestroy(pref);
+        cublasLtMatrixLayoutDestroy(g.la);
+        cublasLtMatrixLayoutDestroy(g.lb);
+        cublasLtMatrixLayoutDestroy(g.lc);
+        cublasLtMatmulDescDestroy(g.op);
+        CK(cudaFree(A)); CK(cudaFree(B)); CK(cudaFree(C));
+    }
+    printf("LT-GEMM DONE\n");
+    return 0;
+}

cuda/mmq/test/test_iq2_aligned_entry.cu

@@ -0,0 +1,329 @@
+// test_iq2_aligned_entry.cu — M1-Inc1 integration parity gate.
+//
+// A/B of the PRODUCTION aligned entry (ds4_mmq_iq2_xxs_aligned_moe_vec, fed
+// the artifact layout the weight server builds under --repack-iq2-aligned)
+// against the production raw-layout vec entry (ds4_mmq_iq2_xxs_moe_vec) at
+// the decode shape.  Both entries quantize the activation through
+// quantize_row_q8_1_cuda, so only the float accumulation order differs.
+//
+// The kernel-level A/B with hand-rolled quantize lives in
+// proto_iq2_aligned.cu (the original +12% proof); this test locks the
+// integrated entry + artifact layout contract.
+//
+// Build (box):
+//   nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 -I.. \
+//        test_iq2_aligned_entry.cu ../*.o -lcudart -lcuda -o test_iq2_aligned_entry
+
+#include "ds4_mmq.h"
+
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <random>
+#include <vector>
+
+#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \
+    printf("CUDA ERR %s @%d: %s\n", #x, __LINE__, cudaGetErrorString(e_)); exit(1); } } while (0)
+
+int main(int argc, char **argv) {
+    const int M = 2048, K = 4096;
+    const int n_experts = 256;
+    const int n_slots = 6;
+    const int n_idsets = 32;
+    const int iters = 2000;
+    const int nb = K / 256;
+    const long long row_bytes = (long long)nb * 66;
+    const double slot_gb = (double)n_slots * M * row_bytes / 1e9;
+
+    if (ds4_mmq_init(0) != 0) { printf("ds4_mmq_init failed\n"); return 1; }
+
+    std::mt19937 rng(argc > 1 ? (uint32_t)atoi(argv[1]) : 1234u);
+
+    // Random bytes are valid IQ2_XXS content; keep block scales sane.
+    std::vector<uint8_t> W((size_t)n_experts * M * row_bytes);
+    for (auto &b : W) b = (uint8_t)(rng() & 0xff);
+    const long long nblk = (long long)n_experts * M * nb;
+    for (long long blk = 0; blk < nblk; blk++) {
+        uint16_t h = (uint16_t)(0x2c00 | (rng() & 0xff));
+        memcpy(&W[blk * 66], &h, 2);
+    }
+
+    std::vector<float> X(K);
+    for (auto &v : X) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
+
+    std::vector<int32_t> ids(n_idsets * n_slots);
+    for (int s = 0; s < n_idsets; s++)
+        for (int j = 0; j < n_slots; j++)
+            ids[s * n_slots + j] = (s * n_slots + j) % n_experts;
+
+    // Artifact layout, exactly as the weight server repack kernel builds it:
+    //   [__half dq[nblk]][pad to 64B][uint2 qs[nblk*8]]
+    const uint64_t expect_bytes = ds4_mmq_iq2_xxs_aligned_bytes(M, K, n_experts);
+    const uint64_t dq_bytes = ((uint64_t)nblk * 2u + 63u) & ~63ull;
+    const uint64_t art_bytes = dq_bytes + (uint64_t)nblk * 64u;
+    if (expect_bytes != art_bytes) {
+        printf("aligned_bytes mismatch: helper=%llu local=%llu\n",
+               (unsigned long long)expect_bytes, (unsigned long long)art_bytes);
+        return 1;
+    }
+    std::vector<uint8_t> ART(art_bytes);
+    for (long long blk = 0; blk < nblk; blk++) {
+        memcpy(&ART[blk * 2], &W[blk * 66], 2);
+        memcpy(&ART[dq_bytes + (uint64_t)blk * 64u], &W[blk * 66 + 2], 64);
+    }
+
+    // dArt2: identical bytes at a DIFFERENT address, used as the "up" stream
+    // in the pair/fused entries.  Passing dArt twice lets the second stream's
+    // reads ride the L2 lines the first stream just fetched, inflating the
+    // effective rate past the DRAM ceiling; production gate/up are distinct
+    // 553 MB artifacts.
+    uint8_t *dW, *dArt, *dArt2; float *dX, *dOutBase, *dOutAl; int32_t *dIds;
+    CK(cudaMalloc(&dW, W.size()));
+    CK(cudaMalloc(&dArt, ART.size()));
+    CK(cudaMalloc(&dArt2, ART.size()));
+    CK(cudaMalloc(&dX, sizeof(float) * K));
+    CK(cudaMalloc(&dOutBase, sizeof(float) * n_slots * M));
+    CK(cudaMalloc(&dOutAl, sizeof(float) * n_slots * M));
+    CK(cudaMalloc(&dIds, sizeof(int32_t) * ids.size()));
+    CK(cudaMemcpy(dW, W.data(), W.size(), cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dArt, ART.data(), ART.size(), cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dArt2, dArt, ART.size(), cudaMemcpyDeviceToDevice));
+    CK(cudaMemcpy(dX, X.data(), sizeof(float) * K, cudaMemcpyHostToDevice));
+    CK(cudaMemcpy(dIds, ids.data(), sizeof(int32_t) * ids.size(), cudaMemcpyHostToDevice));
+
+    cudaStream_t stream; CK(cudaStreamCreate(&stream));
+    cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1));
+
+    // ---- de-repack round-trip (M1-Inc2b): artifact -> raw must be byte-exact
+    {
+        uint8_t *dRaw2; CK(cudaMalloc(&dRaw2, W.size()));
+        CK(cudaMemset(dRaw2, 0xAB, W.size()));
+        const int drc = ds4_mmq_iq2_xxs_aligned_derepack(dArt, dRaw2, M, K, n_experts, stream);
+        if (drc != 0) { printf("derepack rc=%d\n", drc); return 1; }
+        CK(cudaStreamSynchronize(stream));
+        std::vector<uint8_t> RT(W.size());
+        CK(cudaMemcpy(RT.data(), dRaw2, W.size(), cudaMemcpyDeviceToHost));
+        if (memcmp(RT.data(), W.data(), W.size()) != 0) {
+            size_t off = 0; while (off < W.size() && RT[off] == W[off]) off++;
+            printf("derepack round-trip MISMATCH at byte %zu\n", off);
+            return 1;
+        }
+        CK(cudaEventRecord(e0, stream));
+        const int dr_iters = 50;
+        for (int i = 0; i < dr_iters; i++)
+            (void)ds4_mmq_iq2_xxs_aligned_derepack(dArt, dRaw2, M, K, n_experts, stream);
+        CK(cudaEventRecord(e1, stream));
+        CK(cudaStreamSynchronize(stream));
+        float ms_dr = 0.0f; CK(cudaEventElapsedTime(&ms_dr, e0, e1));
+        ms_dr /= dr_iters;
+        printf("  derepack round-trip: byte-exact PASS, %.3f ms/fill -> %.1f GB/s (%.1f MB tensor)\n",
+               ms_dr, (double)(ART.size() + W.size()) / ms_dr / 1e6,
+               (double)W.size() / 1e6);
+        CK(cudaFree(dRaw2));
+    }
+
+    // ---- baseline: raw-layout production vec entry -------------------------
+    int rc = ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds, dOutBase, M, K, 1, n_experts, n_slots, stream);
+    if (rc != 0) { printf("baseline moe_vec rc=%d\n", rc); return 1; }
+    CK(cudaStreamSynchronize(stream));
+    CK(cudaEventRecord(e0, stream));
+    for (int i = 0; i < iters; i++)
+        (void)ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds + (i % n_idsets) * n_slots, dOutBase,
+                                      M, K, 1, n_experts, n_slots, stream);
+    CK(cudaEventRecord(e1, stream));
+    CK(cudaStreamSynchronize(stream));
+    float ms_base = 0.0f; CK(cudaEventElapsedTime(&ms_base, e0, e1));
+    ms_base /= iters;
+
+    // ---- aligned production entry ------------------------------------------
+    rc = ds4_mmq_iq2_xxs_aligned_moe_vec(dArt, dX, dIds, dOutAl, M, K, 1, n_experts, n_slots, stream);
+    if (rc != 0) { printf("aligned entry rc=%d\n", rc); return 1; }
+    // width > 16 (beyond the vec-tier envelope) must be rejected so the
+    // caller can fall back; 1..16 are accepted since the multi-token fix.
+    if (ds4_mmq_iq2_xxs_aligned_moe_vec(dArt, dX, dIds, dOutAl, M, K, 17, n_experts, n_slots, stream) == 0) {
+        printf("aligned entry accepted n_tokens=17 (must reject)\n");
+        return 1;
+    }
+    CK(cudaStreamSynchronize(stream));
+    CK(cudaEventRecord(e0, stream));
+    for (int i = 0; i < iters; i++)
+        (void)ds4_mmq_iq2_xxs_aligned_moe_vec(dArt, dX, dIds + (i % n_idsets) * n_slots, dOutAl,
+                                              M, K, 1, n_experts, n_slots, stream);
+    CK(cudaEventRecord(e1, stream));
+    CK(cudaStreamSynchronize(stream));
+    float ms_al = 0.0f; CK(cudaEventElapsedTime(&ms_al, e0, e1));
+    ms_al /= iters;
+
+    // ---- Inc2 variant P: pair (one quantize, one z=2 launch) ----------------
+    float *dOutUpP; CK(cudaMalloc(&dOutUpP, sizeof(float) * n_slots * M));
+    rc = ds4_mmq_iq2_xxs_aligned_moe_pair_vec(dArt, dArt2, dX, dIds, dOutAl, dOutUpP,
+                                              M, K, 1, n_experts, n_slots, stream);
+    if (rc != 0) { printf("pair entry rc=%d\n", rc); return 1; }
+    CK(cudaStreamSynchronize(stream));
+    CK(cudaEventRecord(e0, stream));
+    for (int i = 0; i < iters; i++)
+        (void)ds4_mmq_iq2_xxs_aligned_moe_pair_vec(dArt, dArt2, dX,
+                                                   dIds + (i % n_idsets) * n_slots,
+                                                   dOutAl, dOutUpP, M, K, 1, n_experts, n_slots, stream);
+    CK(cudaEventRecord(e1, stream));
+    CK(cudaStreamSynchronize(stream));
+    float ms_pair = 0.0f; CK(cudaEventElapsedTime(&ms_pair, e0, e1));
+    ms_pair /= iters;
+
+    // ---- Inc2 variant F: fused gate+up+swiglu -> mid ------------------------
+    float *dMid, *dSlotW; CK(cudaMalloc(&dMid, sizeof(float) * n_slots * M));
+    CK(cudaMalloc(&dSlotW, sizeof(float) * n_slots));
+    std::vector<float> slotw(n_slots);
+    for (int j = 0; j < n_slots; j++) slotw[j] = 0.5f + 0.1f * j;
+    CK(cudaMemcpy(dSlotW, slotw.data(), sizeof(float) * n_slots, cudaMemcpyHostToDevice));
+    const float clampv = 10.0f;
+    rc = ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(dArt, dArt2, dX, dIds, dSlotW, dMid,
+                                                     M, K, 1, n_experts, n_slots, clampv, stream);
+    if (rc != 0) { printf("fused entry rc=%d\n", rc); return 1; }
+    CK(cudaStreamSynchronize(stream));
+    CK(cudaEventRecord(e0, stream));
+    for (int i = 0; i < iters; i++)
+        (void)ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(dArt, dArt2, dX,
+                                                          dIds + (i % n_idsets) * n_slots,
+                                                          dSlotW, dMid, M, K, 1, n_experts, n_slots,
+                                                          clampv, stream);
+    CK(cudaEventRecord(e1, stream));
+    CK(cudaStreamSynchronize(stream));
+    float ms_fused = 0.0f; CK(cudaEventElapsedTime(&ms_fused, e0, e1));
+    ms_fused /= iters;
+
+    // NOTE: pair/fused per-iter weight bytes are 2x slot_gb (gate + up); the
+    // baseline-equivalent cost of one iter is 2x the single-entry time.
+
+    // ---- parity on idset 0 --------------------------------------------------
+    (void)ds4_mmq_iq2_xxs_moe_vec(dW, dX, dIds, dOutBase, M, K, 1, n_experts, n_slots, stream);
+    (void)ds4_mmq_iq2_xxs_aligned_moe_vec(dArt, dX, dIds, dOutAl, M, K, 1, n_experts, n_slots, stream);
+    float *dPairG; CK(cudaMalloc(&dPairG, sizeof(float) * n_slots * M));
+    (void)ds4_mmq_iq2_xxs_aligned_moe_pair_vec(dArt, dArt2, dX, dIds, dPairG, dOutUpP,
+                                               M, K, 1, n_experts, n_slots, stream);
+    (void)ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(dArt, dArt2, dX, dIds, dSlotW, dMid,
+                                                      M, K, 1, n_experts, n_slots, clampv, stream);
+    CK(cudaStreamSynchronize(stream));
+    std::vector<float> ob(n_slots * M), oa(n_slots * M), opg(n_slots * M), opu(n_slots * M), om(n_slots * M);
+    CK(cudaMemcpy(ob.data(), dOutBase, sizeof(float) * ob.size(), cudaMemcpyDeviceToHost));
+    CK(cudaMemcpy(oa.data(), dOutAl, sizeof(float) * oa.size(), cudaMemcpyDeviceToHost));
+    CK(cudaMemcpy(opg.data(), dPairG, sizeof(float) * opg.size(), cudaMemcpyDeviceToHost));
+    CK(cudaMemcpy(opu.data(), dOutUpP, sizeof(float) * opu.size(), cudaMemcpyDeviceToHost));
+    CK(cudaMemcpy(om.data(), dMid, sizeof(float) * om.size(), cudaMemcpyDeviceToHost));
+
+    auto count_bad = [](const std::vector<float> &ref, const std::vector<float> &got,
+                        double *mr, double *ma) {
+        double max_rel = 0.0, max_abs = 0.0; int bad = 0;
+        for (size_t i = 0; i < ref.size(); i++) {
+            const double a = ref[i], b = got[i];
+            const double ad = fabs(a - b);
+            const double rd = ad / (fabs(a) > 1e-3 ? fabs(a) : 1e-3);
+            if (ad > max_abs) max_abs = ad;
+            if (rd > max_rel) max_rel = rd;
+            if (rd > 2e-2 && ad > 1e-2) bad++;
+        }
+        if (mr) *mr = max_rel;
+        if (ma) *ma = max_abs;
+        return bad;
+    };
+    double max_rel = 0.0, max_abs = 0.0;
+    int bad = count_bad(ob, oa, &max_rel, &max_abs);
+    // pair: gate and up both computed from dArt, so both must match the
+    // aligned single-entry output exactly (same kernel math, same inputs).
+    int bad_pair = count_bad(ob, opg, nullptr, nullptr) + count_bad(ob, opu, nullptr, nullptr);
+    // fused: expected mid from the baseline outputs + the exact epilogue.
+    std::vector<float> em(n_slots * M);
+    for (int s = 0; s < n_slots; s++) {
+        for (int r = 0; r < M; r++) {
+            float g = ob[(size_t)s * M + r], u = g;
+            if (!std::isfinite(g)) g = 0.0f;
+            if (!std::isfinite(u)) u = 0.0f;
+            if (clampv > 1.0e-6f) {
+                if (g > clampv) g = clampv;
+                if (u > clampv) u = clampv;
+                if (u < -clampv) u = -clampv;
+            }
+            const float silu = g / (1.0f + expf(-g));
+            em[(size_t)s * M + r] = silu * u * slotw[s];
+        }
+    }
+    int bad_mid = count_bad(em, om, nullptr, nullptr);
+
+    // ---- multi-token parity (batch_eval-hang fix, 2026-07-03): aligned vs
+    // raw production entries at n_tokens=4, flat assignment layout
+    // [tok*n_slots+slot], including one -1 router id (NaN-path guard).
+    int bad4 = 0, bad4_mid = 0;
+    {
+        const int n_tok = 4, n_asg = n_tok * n_slots;
+        std::vector<float> X4((size_t)n_tok * K);
+        for (auto &v : X4) v = ((float)(rng() % 2000) - 1000.0f) / 500.0f;
+        std::vector<int32_t> ids4(n_asg);
+        for (int a = 0; a < n_asg; a++) ids4[a] = (a * 7 + 3) % n_experts;
+        ids4[13] = -1;  // router NaN path: both entries must write a clean 0 row
+        std::vector<float> w4(n_asg);
+        for (int a = 0; a < n_asg; a++) w4[a] = 0.25f + 0.05f * a;
+        float *dX4, *dOut4B, *dOut4A, *dMid4, *dW4; int32_t *dIds4;
+        CK(cudaMalloc(&dX4, sizeof(float) * X4.size()));
+        CK(cudaMalloc(&dOut4B, sizeof(float) * n_asg * M));
+        CK(cudaMalloc(&dOut4A, sizeof(float) * n_asg * M));
+        CK(cudaMalloc(&dMid4, sizeof(float) * n_asg * M));
+        CK(cudaMalloc(&dW4, sizeof(float) * n_asg));
+        CK(cudaMalloc(&dIds4, sizeof(int32_t) * n_asg));
+        CK(cudaMemcpy(dX4, X4.data(), sizeof(float) * X4.size(), cudaMemcpyHostToDevice));
+        CK(cudaMemcpy(dW4, w4.data(), sizeof(float) * n_asg, cudaMemcpyHostToDevice));
+        CK(cudaMemcpy(dIds4, ids4.data(), sizeof(int32_t) * n_asg, cudaMemcpyHostToDevice));
+        int r1 = ds4_mmq_iq2_xxs_moe_vec(dW, dX4, dIds4, dOut4B, M, K, n_tok, n_experts, n_slots, stream);
+        int r2 = ds4_mmq_iq2_xxs_aligned_moe_vec(dArt, dX4, dIds4, dOut4A, M, K, n_tok, n_experts, n_slots, stream);
+        int r3 = ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(dArt, dArt2, dX4, dIds4, dW4, dMid4,
+                                                             M, K, n_tok, n_experts, n_slots, clampv, stream);
+        if (r1 != 0 || r2 != 0 || r3 != 0) {
+            printf("multi-token entries rc=%d/%d/%d\n", r1, r2, r3);
+            return 1;
+        }
+        CK(cudaStreamSynchronize(stream));
+        std::vector<float> ob4((size_t)n_asg * M), oa4((size_t)n_asg * M), om4((size_t)n_asg * M);
+        CK(cudaMemcpy(ob4.data(), dOut4B, sizeof(float) * ob4.size(), cudaMemcpyDeviceToHost));
+        CK(cudaMemcpy(oa4.data(), dOut4A, sizeof(float) * oa4.size(), cudaMemcpyDeviceToHost));
+        CK(cudaMemcpy(om4.data(), dMid4, sizeof(float) * om4.size(), cudaMemcpyDeviceToHost));
+        bad4 = count_bad(ob4, oa4, nullptr, nullptr);
+        std::vector<float> em4((size_t)n_asg * M);
+        for (int a = 0; a < n_asg; a++) {
+            for (int r = 0; r < M; r++) {
+                float g = ob4[(size_t)a * M + r], u = g;
+                if (!std::isfinite(g)) g = 0.0f;
+                if (!std::isfinite(u)) u = 0.0f;
+                if (clampv > 1.0e-6f) {
+                    if (g > clampv) g = clampv;
+                    if (u > clampv) u = clampv;
+                    if (u < -clampv) u = -clampv;
+                }
+                const float silu = g / (1.0f + expf(-g));
+                em4[(size_t)a * M + r] = silu * u * w4[a];
+            }
+        }
+        bad4_mid = count_bad(em4, om4, nullptr, nullptr);
+        CK(cudaFree(dX4)); CK(cudaFree(dOut4B)); CK(cudaFree(dOut4A));
+        CK(cudaFree(dMid4)); CK(cudaFree(dW4)); CK(cudaFree(dIds4));
+    }
+
+    printf("TEST_IQ2_ALIGNED_ENTRY  M=%d K=%d slots=%d experts=%d iters=%d (weights/iter %.1f MB)\n",
+           M, K, n_slots, n_experts, iters, slot_gb * 1000.0);
+    printf("  baseline moe_vec : %.4f ms  -> %6.1f GB/s   (x2 for gate+up: %.4f ms)\n",
+           ms_base, slot_gb / (ms_base / 1e3), 2.0f * ms_base);
+    printf("  aligned  entry   : %.4f ms  -> %6.1f GB/s   (%+.1f%%; x2: %.4f ms)\n",
+           ms_al, slot_gb / (ms_al / 1e3), 100.0 * (ms_base / ms_al - 1.0), 2.0f * ms_al);
+    printf("  pair     entry   : %.4f ms  -> %6.1f GB/s   (vs 2x aligned %+.1f%%)\n",
+           ms_pair, 2.0 * slot_gb / (ms_pair / 1e3), 100.0 * (2.0f * ms_al / ms_pair - 1.0));
+    printf("  fused    entry   : %.4f ms  -> %6.1f GB/s   (vs 2x aligned %+.1f%%)\n",
+           ms_fused, 2.0 * slot_gb / (ms_fused / 1e3), 100.0 * (2.0f * ms_al / ms_fused - 1.0));
+    printf("  parity: aligned max_rel=%.3e max_abs=%.3e bad=%d | pair bad=%d | fused-mid bad=%d"
+           " | ntok4 bad=%d ntok4-mid bad=%d -> %s\n",
+           max_rel, max_abs, bad, bad_pair, bad_mid, bad4, bad4_mid,
+           (bad == 0 && bad_pair == 0 && bad_mid == 0 && bad4 == 0 && bad4_mid == 0) ? "PASS" : "FAIL");
+    return (bad == 0 && bad_pair == 0 && bad_mid == 0 && bad4 == 0 && bad4_mid == 0) ? 0 : 2;
+}

cuda/mmq/test/test_mmq_soa_tiles.cu

@@ -0,0 +1,285 @@
+// SPDX-License-Identifier: MIT
+// test_mmq_soa_tiles.cu - P4 Inc3 decision gate: SoA-direct mmq tile loaders
+// vs the raw-block path, at the production V4-Flash MoE shapes.
+//
+// Method: generate a RANDOM aligned-SoA artifact (the weight-server format),
+// derive the raw block stream from it with the production derepack kernel
+// (ds4_mmq_*_aligned_derepack -- byte-exactly what the scratch path serves),
+// then run the same MoE matmul through both entries and require bit-identical
+// outputs:
+//
+//   Q2_K down leg:      ds4_mmq_q2_K_moe(raw)          vs ds4_mmq_q2_K_moe_soa(soa)
+//   IQ2_XXS gate/up leg: ds4_mmq_iq2_xxs_moe_pair(raw)  vs ds4_mmq_iq2_xxs_moe_pair_soa(soa)
+//
+// Shapes are the V4-Flash production calls at a W4096 prefill chunk:
+//   down:    M=4096 (n_embd), K=2048 (n_ff_exp), 24576 assignments x 1 expert
+//   gate/up: M=2048,          K=4096,            4096 tokens x 6 experts
+// with 256 experts and uniform routing (~96 columns/expert, the trace shape).
+//
+// Timing: cudaEvent over REPS calls of each entry (identical quantize +
+// mm_ids overhead on both sides, so the delta is the tile-load effect), plus
+// the standalone derepack kernel cost (what the SoA path deletes per layer
+// switch).
+//
+// Build (on the box, after `make -j6 cuda-spark` built the mmq objects):
+//   ar rcs libds4mmq.a cuda/mmq/ds4_ggml_stubs.o cuda/mmq/ds4_mmq.o \
+//          cuda/mmq/quantize.o cuda/mmq/mmid.o cuda/mmq/mmvq.o
+//   /usr/local/cuda/bin/nvcc -O3 --use_fast_math -std=c++17 -arch=sm_121 \
+//        -Icuda/mmq cuda/mmq/test/test_mmq_soa_tiles.cu libds4mmq.a \
+//        -lcudart -lcublas -lcuda -o test_mmq_soa_tiles
+
+#include "ds4_mmq.h"
+
+#include <cuda_runtime.h>
+
+// libds4mmq.a references this ds4_cuda.cu symbol from the q8-fold vec paths
+// (C3 Inc4); the entries under test never reach it, so a "no fold available"
+// stub satisfies the link.
+extern "C" int ds4_cuda_q8_fold_take_q81(const void *src, uint64_t in_dim, void *out) {
+    (void)src; (void)in_dim; (void)out;
+    return 0;
+}
+
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <random>
+#include <vector>
+
+#define CHECK(call) do { \
+    cudaError_t err_ = (call); \
+    if (err_ != cudaSuccess) { \
+        fprintf(stderr, "CUDA error %s at %s:%d: %s\n", #call, __FILE__, __LINE__, \
+                cudaGetErrorString(err_)); \
+        exit(1); \
+    } \
+} while (0)
+
+namespace {
+
+int g_failures = 0;
+
+// Random bytes with half-precision words nudged into a sane exponent range
+// (bit-parity would hold for any bits -- both paths read the same words --
+// but normal-range scales keep the timing legs representative).
+uint16_t sane_half(std::mt19937 &rng) {
+    const uint16_t mant = (uint16_t)(rng() & 0x03FFu);
+    const uint16_t sign = (uint16_t)((rng() & 1u) << 15);
+    return (uint16_t)(sign | (14u << 10) | mant);   // +/- [0.25, 0.5)
+}
+
+void fill_random(std::mt19937 &rng, void *dst, size_t bytes) {
+    uint32_t *p = (uint32_t *)dst;
+    for (size_t i = 0; i < bytes / 4; ++i) p[i] = rng();
+}
+
+float *device_random_floats(std::mt19937 &rng, size_t n) {
+    std::vector<float> h(n);
+    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
+    for (size_t i = 0; i < n; ++i) h[i] = dist(rng);
+    float *d = nullptr;
+    CHECK(cudaMalloc(&d, n * sizeof(float)));
+    CHECK(cudaMemcpy(d, h.data(), n * sizeof(float), cudaMemcpyHostToDevice));
+    return d;
+}
+
+bool compare_outputs(const char *what, const float *d_a, const float *d_b, size_t n) {
+    std::vector<float> a(n), b(n);
+    CHECK(cudaMemcpy(a.data(), d_a, n * sizeof(float), cudaMemcpyDeviceToHost));
+    CHECK(cudaMemcpy(b.data(), d_b, n * sizeof(float), cudaMemcpyDeviceToHost));
+    if (memcmp(a.data(), b.data(), n * sizeof(float)) == 0) {
+        printf("  %-28s PARITY-OK (bit-identical, %zu floats)\n", what, n);
+        return true;
+    }
+    size_t first = 0, count = 0;
+    for (size_t i = 0; i < n; ++i) {
+        if (memcmp(&a[i], &b[i], 4) != 0) { if (!count) first = i; ++count; }
+    }
+    printf("  %-28s PARITY-FAIL: %zu/%zu differ, first at %zu (raw=%.9g soa=%.9g)\n",
+           what, count, n, first, a[first], b[first]);
+    ++g_failures;
+    return false;
+}
+
+template <typename F>
+float time_ms(F &&fn, int reps) {
+    cudaEvent_t e0, e1;
+    CHECK(cudaEventCreate(&e0));
+    CHECK(cudaEventCreate(&e1));
+    fn();                                   // warmup
+    CHECK(cudaDeviceSynchronize());
+    CHECK(cudaEventRecord(e0));
+    for (int r = 0; r < reps; ++r) fn();
+    CHECK(cudaEventRecord(e1));
+    CHECK(cudaEventSynchronize(e1));
+    float ms = 0.0f;
+    CHECK(cudaEventElapsedTime(&ms, e0, e1));
+    CHECK(cudaEventDestroy(e0));
+    CHECK(cudaEventDestroy(e1));
+    return ms / reps;
+}
+
+// ---------------------------------------------------------------------------
+// Q2_K down leg.
+// ---------------------------------------------------------------------------
+void run_q2k_leg(std::mt19937 &rng, int reps) {
+    const int E = 256, M = 4096, K = 2048;
+    const int n_tokens = 4096 * 6;          // assignments, n_expert_used == 1
+    const uint64_t nb_row = (uint64_t)K / 256;
+    const uint64_t nblk   = (uint64_t)E * M * nb_row;
+    const uint64_t npair  = nblk / 2;
+    const uint64_t soa_bytes = ds4_mmq_q2_k_aligned_bytes(M, K, E);
+    const uint64_t raw_bytes = nblk * 84ull;
+    printf("Q2_K down leg: E=%d M=%d K=%d cols=%d (soa %.1f MiB, raw %.1f MiB)\n",
+           E, M, K, n_tokens, soa_bytes / 1048576.0, raw_bytes / 1048576.0);
+
+    // Random SoA artifact; dm halves get sane exponents.
+    std::vector<uint8_t> h_soa(soa_bytes);
+    fill_random(rng, h_soa.data(), soa_bytes);
+    uint16_t *dm = (uint16_t *)h_soa.data();
+    for (uint64_t i = 0; i < npair * 4; ++i) dm[i] = sane_half(rng);
+
+    void *d_soa = nullptr, *d_raw = nullptr;
+    CHECK(cudaMalloc(&d_soa, soa_bytes));
+    CHECK(cudaMalloc(&d_raw, raw_bytes));
+    CHECK(cudaMemcpy(d_soa, h_soa.data(), soa_bytes, cudaMemcpyHostToDevice));
+
+    // Raw twin = production derepack of the artifact.
+    if (ds4_mmq_q2_K_aligned_derepack(d_soa, d_raw, M, K, E, 0) != 0) {
+        fprintf(stderr, "q2k derepack failed\n"); exit(1);
+    }
+    CHECK(cudaDeviceSynchronize());
+
+    float *d_x = device_random_floats(rng, (size_t)n_tokens * K);
+    std::vector<int32_t> h_ids(n_tokens);
+    for (int i = 0; i < n_tokens; ++i) h_ids[i] = (int32_t)(rng() % E);
+    int32_t *d_ids = nullptr;
+    CHECK(cudaMalloc(&d_ids, n_tokens * sizeof(int32_t)));
+    CHECK(cudaMemcpy(d_ids, h_ids.data(), n_tokens * sizeof(int32_t), cudaMemcpyHostToDevice));
+
+    const size_t out_n = (size_t)M * n_tokens;
+    float *d_out_raw = nullptr, *d_out_soa = nullptr;
+    CHECK(cudaMalloc(&d_out_raw, out_n * sizeof(float)));
+    CHECK(cudaMalloc(&d_out_soa, out_n * sizeof(float)));
+
+    if (ds4_mmq_q2_K_moe(d_raw, d_x, d_ids, d_out_raw, M, K, n_tokens, E, 1, 0) != 0) {
+        fprintf(stderr, "q2k raw moe failed\n"); exit(1);
+    }
+    if (ds4_mmq_q2_K_moe_soa(d_soa, d_x, d_ids, d_out_soa, M, K, n_tokens, E, 1, 0) != 0) {
+        fprintf(stderr, "q2k soa moe failed\n"); exit(1);
+    }
+    CHECK(cudaDeviceSynchronize());
+    compare_outputs("q2k down W4096", d_out_raw, d_out_soa, out_n);
+
+    const float ms_raw = time_ms([&] {
+        ds4_mmq_q2_K_moe(d_raw, d_x, d_ids, d_out_raw, M, K, n_tokens, E, 1, 0);
+    }, reps);
+    const float ms_soa = time_ms([&] {
+        ds4_mmq_q2_K_moe_soa(d_soa, d_x, d_ids, d_out_soa, M, K, n_tokens, E, 1, 0);
+    }, reps);
+    const float ms_dr = time_ms([&] {
+        ds4_mmq_q2_K_aligned_derepack(d_soa, d_raw, M, K, E, 0);
+    }, reps);
+    printf("  q2k down: raw %.3f ms  soa %.3f ms (%+.1f%%)  [derepack alone %.3f ms/launch]\n",
+           ms_raw, ms_soa, 100.0 * (ms_soa - ms_raw) / ms_raw, ms_dr);
+
+    CHECK(cudaFree(d_soa));  CHECK(cudaFree(d_raw));
+    CHECK(cudaFree(d_x));    CHECK(cudaFree(d_ids));
+    CHECK(cudaFree(d_out_raw)); CHECK(cudaFree(d_out_soa));
+}
+
+// ---------------------------------------------------------------------------
+// IQ2_XXS gate/up pair leg.
+// ---------------------------------------------------------------------------
+void run_iq2_leg(std::mt19937 &rng, int reps) {
+    const int E = 256, M = 2048, K = 4096;
+    const int n_tokens = 4096, n_used = 6;
+    const uint64_t nb_row = (uint64_t)K / 256;
+    const uint64_t nblk   = (uint64_t)E * M * nb_row;
+    const uint64_t soa_bytes = ds4_mmq_iq2_xxs_aligned_bytes(M, K, E);
+    const uint64_t raw_bytes = nblk * 66ull;
+    printf("IQ2_XXS gate/up leg: E=%d M=%d K=%d tokens=%dx%d (soa %.1f MiB x2, raw %.1f MiB x2)\n",
+           E, M, K, n_tokens, n_used, soa_bytes / 1048576.0, raw_bytes / 1048576.0);
+
+    void *d_soa[2] = {nullptr, nullptr}, *d_raw[2] = {nullptr, nullptr};
+    for (int t = 0; t < 2; ++t) {
+        std::vector<uint8_t> h_soa(soa_bytes);
+        fill_random(rng, h_soa.data(), soa_bytes);
+        uint16_t *dq = (uint16_t *)h_soa.data();
+        for (uint64_t i = 0; i < nblk; ++i) dq[i] = sane_half(rng);
+        CHECK(cudaMalloc(&d_soa[t], soa_bytes));
+        CHECK(cudaMalloc(&d_raw[t], raw_bytes));
+        CHECK(cudaMemcpy(d_soa[t], h_soa.data(), soa_bytes, cudaMemcpyHostToDevice));
+        if (ds4_mmq_iq2_xxs_aligned_derepack(d_soa[t], d_raw[t], M, K, E, 0) != 0) {
+            fprintf(stderr, "iq2 derepack failed\n"); exit(1);
+        }
+    }
+    CHECK(cudaDeviceSynchronize());
+
+    float *d_x = device_random_floats(rng, (size_t)n_tokens * K);
+    // Production routing: n_used DISTINCT experts per token.
+    std::vector<int32_t> h_ids((size_t)n_tokens * n_used);
+    for (int i = 0; i < n_tokens; ++i) {
+        int picked[6];
+        for (int j = 0; j < n_used; ++j) {
+            int e;
+            bool dup;
+            do {
+                e = (int)(rng() % E);
+                dup = false;
+                for (int q = 0; q < j; ++q) dup |= (picked[q] == e);
+            } while (dup);
+            picked[j] = e;
+            h_ids[(size_t)i * n_used + j] = e;
+        }
+    }
+    int32_t *d_ids = nullptr;
+    CHECK(cudaMalloc(&d_ids, h_ids.size() * sizeof(int32_t)));
+    CHECK(cudaMemcpy(d_ids, h_ids.data(), h_ids.size() * sizeof(int32_t), cudaMemcpyHostToDevice));
+
+    const size_t out_n = (size_t)M * n_tokens * n_used;
+    float *d_out[4] = {nullptr, nullptr, nullptr, nullptr};   // raw a/b, soa a/b
+    for (int i = 0; i < 4; ++i) CHECK(cudaMalloc(&d_out[i], out_n * sizeof(float)));
+
+    if (ds4_mmq_iq2_xxs_moe_pair(d_raw[0], d_raw[1], d_x, d_ids, d_out[0], d_out[1],
+                                 M, K, n_tokens, E, n_used, 0) != 0) {
+        fprintf(stderr, "iq2 raw pair failed\n"); exit(1);
+    }
+    if (ds4_mmq_iq2_xxs_moe_pair_soa(d_soa[0], d_soa[1], d_x, d_ids, d_out[2], d_out[3],
+                                     M, K, n_tokens, E, n_used, 0) != 0) {
+        fprintf(stderr, "iq2 soa pair failed\n"); exit(1);
+    }
+    CHECK(cudaDeviceSynchronize());
+    compare_outputs("iq2 gate W4096", d_out[0], d_out[2], out_n);
+    compare_outputs("iq2 up   W4096", d_out[1], d_out[3], out_n);
+
+    const float ms_raw = time_ms([&] {
+        ds4_mmq_iq2_xxs_moe_pair(d_raw[0], d_raw[1], d_x, d_ids, d_out[0], d_out[1],
+                                 M, K, n_tokens, E, n_used, 0);
+    }, reps);
+    const float ms_soa = time_ms([&] {
+        ds4_mmq_iq2_xxs_moe_pair_soa(d_soa[0], d_soa[1], d_x, d_ids, d_out[2], d_out[3],
+                                     M, K, n_tokens, E, n_used, 0);
+    }, reps);
+    const float ms_dr = time_ms([&] {
+        ds4_mmq_iq2_xxs_aligned_derepack(d_soa[0], d_raw[0], M, K, E, 0);
+    }, reps);
+    printf("  iq2 pair: raw %.3f ms  soa %.3f ms (%+.1f%%)  [derepack alone %.3f ms/launch, x2 tensors]\n",
+           ms_raw, ms_soa, 100.0 * (ms_soa - ms_raw) / ms_raw, ms_dr);
+
+    for (int t = 0; t < 2; ++t) { CHECK(cudaFree(d_soa[t])); CHECK(cudaFree(d_raw[t])); }
+    CHECK(cudaFree(d_x)); CHECK(cudaFree(d_ids));
+    for (int i = 0; i < 4; ++i) CHECK(cudaFree(d_out[i]));
+}
+
+} // namespace
+
+int main(int argc, char **argv) {
+    const int reps = (argc > 1) ? atoi(argv[1]) : 20;
+    std::mt19937 rng(0x5A0Au);
+    run_q2k_leg(rng, reps);
+    run_iq2_leg(rng, reps);
+    printf(g_failures ? "SOA-TILES FAIL (%d)\n" : "SOA-TILES PASS\n", g_failures);
+    return g_failures ? 1 : 0;
+}

cuda/mmq/unary.cuh

@@ -0,0 +1,114 @@
+#pragma once
+#include "common.cuh"
+
+#define CUDA_NEG_BLOCK_SIZE 256
+#define CUDA_STEP_BLOCK_SIZE 256
+#define CUDA_GELU_BLOCK_SIZE 256
+#define CUDA_SILU_BLOCK_SIZE 256
+#define CUDA_SILU_BACK_BLOCK_SIZE 256
+#define CUDA_TANH_BLOCK_SIZE 256
+#define CUDA_RELU_BLOCK_SIZE 256
+#define CUDA_SIGMOID_BLOCK_SIZE 256
+#define CUDA_HARDSIGMOID_BLOCK_SIZE 256
+#define CUDA_EXP_BLOCK_SIZE 256
+#define CUDA_HARDSWISH_BLOCK_SIZE 256
+#define CUDA_SQR_BLOCK_SIZE 256
+#define CUDA_SQRT_BLOCK_SIZE 256
+#define CUDA_SIN_BLOCK_SIZE 256
+#define CUDA_COS_BLOCK_SIZE 256
+#define CUDA_GLU_BLOCK_SIZE 256
+#define CUDA_XIELU_BLOCK_SIZE 256
+
+void ggml_cuda_op_abs(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_sgn(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_neg(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_step(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_gelu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_silu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_silu_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_gelu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_gelu_quick(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_tanh(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_relu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_sigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_hardsigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_exp(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_hardswish(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_leaky_relu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_sqrt(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_sin(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_cos(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_log(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_expm1(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_softplus(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_elu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_floor(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_ceil(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_round(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_trunc(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_reglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_swiglu_oai(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_geglu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_geglu_quick(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_xielu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node);
+
+void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_node, ggml_tensor * sqr_node);
+
+__device__ __forceinline__ float ggml_cuda_op_silu_single(float x) {
+    return x / (1.0f + expf(-x));
+}
+
+__device__ __forceinline__ float ggml_cuda_op_gelu_single(float x) {
+    const float GELU_COEF_A    = 0.044715f;
+    const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
+
+    return 0.5f * x * (1.0f + tanhf(SQRT_2_OVER_PI * x * (1.0f + GELU_COEF_A * x * x)));
+}
+
+__device__ __forceinline__ float ggml_cuda_op_swiglu_oai_single(float x, float g, float alpha = 1.702f, float limit = 7.0f) {
+    x = fminf(x, limit);
+    g = fmaxf(fminf(g, limit), -limit);
+
+    float out_glu = x / (1.0f + expf(-x * alpha));
+    out_glu = out_glu * (1.0f + g);
+    return out_glu;
+}

cuda/mmq/vendors/cuda.h

@@ -0,0 +1,28 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cuda.h>
+#include <cublas_v2.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#ifdef GGML_USE_NCCL
+#include <nccl.h>
+#endif // GGML_USE_NCCL
+
+#if CUDART_VERSION >= 11080
+#include <cuda_fp8.h>
+#define FP8_AVAILABLE
+#endif // CUDART_VERSION >= 11080
+
+#if CUDART_VERSION >= 12080
+#include <cuda_fp4.h>
+#endif // CUDART_VERSION >= 12080
+
+#if CUDART_VERSION < 11020
+#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED
+#define CUBLAS_TF32_TENSOR_OP_MATH CUBLAS_TENSOR_OP_MATH
+#define CUBLAS_COMPUTE_16F CUDA_R_16F
+#define CUBLAS_COMPUTE_32F CUDA_R_32F
+#define cublasComputeType_t cudaDataType_t
+#endif // CUDART_VERSION < 11020

dir-steering/README.md

@@ -1,8 +1,9 @@
 # Directional Steering
 
 Directional steering is a runtime activation edit for DS4. A steering file is a
-flat `f32` matrix with one normalized 4096-wide direction per layer. During
-inference, ds4 can apply the edit after attention outputs, FFN outputs, or both:
+flat `f32` matrix with one normalized hidden-width direction per normal
+transformer layer. During inference, ds4 can apply the edit after attention
+outputs, FFN outputs, or both:
 
 ```text
 y = y - scale * direction[layer] * dot(direction[layer], y)
@@ -11,10 +12,17 @@ y = y - scale * direction[layer] * dot(direction[layer], y)
 Positive scale removes the represented direction. Negative scale amplifies it.
 With no steering file or zero scales, ds4 follows the normal inference path.
 
+The file shape depends on the model:
+
+- DeepSeek V4 Flash: `43 x 4096`.
+- GLM 5.3 Flash: `45 x 4096`. The separate MTP predictor layer is omitted.
+
+GLM 5.2 steering is not implemented.
+
 ## Runtime Options
 
 ```text
---dir-steering-file FILE   load a 43 x 4096 f32 direction file
+--dir-steering-file FILE   load one f32 direction per normal model layer
 --dir-steering-ffn F       apply steering after FFN outputs; default is 1 when a file is provided
 --dir-steering-attn F      apply steering after attention outputs; default is 0
 ```
@@ -23,6 +31,27 @@ The FFN output is usually the best first target because it is late enough in
 each layer to represent behavior, style, and topic signals. Attention steering
 is available for experiments, but it can be more fragile.
 
+## GLM 5.3 Example
+
+Build a GLM 5.3 direction from paired target and control prompt lists:
+
+```sh
+python3 dir-steering/tools/build_direction.py \
+  --profile glm-5.3-flash \
+  --ds4 ./ds4 \
+  --model gguf/GLM-5.3-Flash-Q2.gguf \
+  --good-file /path/to/target-prompts.txt \
+  --bad-file /path/to/control-prompts.txt \
+  --out dir-steering/out/glm53-direction.json \
+  --component ffn_out \
+  --ctx 512
+```
+
+Generated `.f32` vectors are local artifacts and are not stored in the
+repository. GLM 5.3 steering works with `--mtp`, `ds4-server`, native session
+batching, and two-Mac tensor parallelism. For tensor parallelism, pass the same
+steering file and scales to both the worker and coordinator.
+
 ## Verbosity Example
 
 The bundled example builds a style direction from 100 paired prompts. Each pair
@@ -39,6 +68,7 @@ Build the vector:
 
 ```sh
 python3 dir-steering/tools/build_direction.py \
+  --profile deepseek-v4-flash \
   --ds4 ./ds4 \
   --model ds4flash.gguf \
   --good-file dir-steering/examples/succinct.txt \

dir-steering/tools/build_direction.py

@@ -1,8 +1,8 @@
 #!/usr/bin/env python3
 """Build a DS4 directional-steering vector from paired prompt sets.
 
-The extractor asks ds4 to dump one 4096-wide activation row per layer, averages
-the target and control rows, and writes a flat f32 file with 43 layer vectors.
+The extractor asks ds4 to dump one hidden activation row per normal transformer
+layer, averages the target and control rows, and writes a flat f32 file.
 At runtime ds4 applies:
 
     y = y - scale * direction[layer] * dot(direction[layer], y)
@@ -20,15 +20,9 @@ import tempfile
 from pathlib import Path
 
 
-N_LAYER = 43
-N_EMBD = 4096
-
-SPECIALS = {
-    "bos": "<｜begin▁of▁sentence｜>",
-    "user": "<｜User｜>",
-    "assistant": "<｜Assistant｜>",
-    "think": "<think>",
-    "nothink": "</think>",
+MODEL_PROFILES = {
+    "deepseek-v4-flash": (43, 4096),
+    "glm-5.3-flash": (45, 4096),
 }
 
 
@@ -45,20 +39,6 @@ def read_prompt_file(path: Path) -> list[str]:
     return prompts
 
 
-def render_ds4_prompt(system: str, user: str, think: bool) -> str:
-    """Render the minimal DS4 chat prefix used for activation capture."""
-    pieces = [SPECIALS["bos"]]
-    if system:
-        pieces.append(system)
-    pieces += [
-        SPECIALS["user"],
-        user,
-        SPECIALS["assistant"],
-        SPECIALS["think"] if think else SPECIALS["nothink"],
-    ]
-    return "".join(pieces)
-
-
 def normalize(v: list[float]) -> list[float]:
     n2 = sum(x * x for x in v)
     if n2 <= 0.0:
@@ -79,11 +59,13 @@ def run_capture(
     think: bool,
     ctx: int,
     component: str,
+    n_layer: int,
+    n_embd: int,
     work: Path,
 ) -> list[list[float]]:
     """Run ds4 once and return the last prompt-row dump for every layer."""
     prompt_path = work / "prompt.txt"
-    prompt_path.write_text(render_ds4_prompt(system, prompt, think), encoding="utf-8")
+    prompt_path.write_text(prompt, encoding="utf-8")
     dump_prefix = work / "dump"
 
     env = os.environ.copy()
@@ -98,23 +80,29 @@ def run_capture(
         "--prompt-file", str(prompt_path),
         "-n", "1",
     ]
-    subprocess.run(cmd, cwd=ds4.parent, env=env, check=True,
-                   stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
+    if system:
+        cmd += ["--system", system]
+    cmd.append("--think-high" if think else "--nothink")
+    result = subprocess.run(cmd, cwd=ds4.parent, env=env, check=False,
+                            stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
+    if result.returncode != 0:
+        error = result.stderr.decode("utf-8", errors="replace")
+        raise RuntimeError(f"ds4 activation capture failed:\n{error}")
 
     rows: list[list[float]] = []
-    for layer in range(N_LAYER):
+    for layer in range(n_layer):
         path = work / f"dump_{component}-{layer}_pos0.bin"
         data = array.array("f")
         with path.open("rb") as f:
             data.fromfile(f, path.stat().st_size // 4)
-        if len(data) < N_EMBD or len(data) % N_EMBD != 0:
+        if len(data) < n_embd or len(data) % n_embd != 0:
             raise RuntimeError(f"bad dump shape for {path}: {len(data)} floats")
-        rows.append(list(data[-N_EMBD:]))
+        rows.append(list(data[-n_embd:]))
     return rows
 
 
-def add_rows(total: list[list[float]], rows: list[list[float]]) -> None:
-    for layer in range(N_LAYER):
+def add_rows(total: list[list[float]], rows: list[list[float]], n_layer: int) -> None:
+    for layer in range(n_layer):
         dst = total[layer]
         src = rows[layer]
         for i, value in enumerate(src):
@@ -125,12 +113,17 @@ def main() -> None:
     ap = argparse.ArgumentParser()
     ap.add_argument("--ds4", default="./ds4", help="path to the ds4 CLI")
     ap.add_argument("--model", default="ds4flash.gguf", help="GGUF model path")
+    ap.add_argument("--profile", choices=MODEL_PROFILES,
+                    default="deepseek-v4-flash",
+                    help="model shape and chat-template profile")
     ap.add_argument("--good-file", required=True,
                     help="desired/target prompts, one per line")
     ap.add_argument("--bad-file", required=True,
                     help="contrast/control prompts, one per line")
     ap.add_argument("--out", default="dir-steering/out/direction.json",
                     help="metadata JSON path; .f32 is written next to it")
+    ap.add_argument("--source", default="",
+                    help="optional corpus source or provenance note")
     ap.add_argument("--ctx", type=int, default=512)
     ap.add_argument("--system", default="You are a helpful assistant.")
     ap.add_argument("--component", default="ffn_out",
@@ -145,16 +138,18 @@ def main() -> None:
     args = ap.parse_args()
 
     ds4 = Path(args.ds4).resolve()
-    model = Path(args.model).resolve()
+    model_arg = Path(args.model)
+    model = model_arg.resolve()
+    n_layer, n_embd = MODEL_PROFILES[args.profile]
     good_prompts = read_prompt_file(Path(args.good_file))
     bad_prompts = read_prompt_file(Path(args.bad_file))
     n = min(len(good_prompts), len(bad_prompts))
     good_prompts = good_prompts[:n]
     bad_prompts = bad_prompts[:n]
 
-    good_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
-    bad_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
-    pair_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
+    good_sum = [[0.0] * n_embd for _ in range(n_layer)]
+    bad_sum = [[0.0] * n_embd for _ in range(n_layer)]
+    pair_sum = [[0.0] * n_embd for _ in range(n_layer)]
 
     with tempfile.TemporaryDirectory(prefix="ds4-dir-steer-") as td:
         root = Path(td)
@@ -165,22 +160,22 @@ def main() -> None:
             gw.mkdir()
             bw.mkdir()
             good_rows = run_capture(ds4, model, good, args.system, args.think,
-                                    args.ctx, args.component, gw)
+                                    args.ctx, args.component, n_layer, n_embd, gw)
             bad_rows = run_capture(ds4, model, bad, args.system, args.think,
-                                   args.ctx, args.component, bw)
-            add_rows(good_sum, good_rows)
-            add_rows(bad_sum, bad_rows)
+                                   args.ctx, args.component, n_layer, n_embd, bw)
+            add_rows(good_sum, good_rows, n_layer)
+            add_rows(bad_sum, bad_rows, n_layer)
             if args.pair_normalize:
-                for layer in range(N_LAYER):
+                for layer in range(n_layer):
                     diff = normalize([
                         good_rows[layer][j] - bad_rows[layer][j]
-                        for j in range(N_EMBD)
+                        for j in range(n_embd)
                     ])
                     for j, value in enumerate(diff):
                         pair_sum[layer][j] += value
 
     layers = []
-    for layer in range(N_LAYER):
+    for layer in range(n_layer):
         good_mean = [x / n for x in good_sum[layer]]
         bad_mean = [x / n for x in bad_sum[layer]]
         if args.pair_normalize:
@@ -188,14 +183,14 @@ def main() -> None:
         else:
             direction = normalize([
                 good_mean[i] - bad_mean[i]
-                for i in range(N_EMBD)
+                for i in range(n_embd)
             ])
         if not args.no_orthogonalize:
             base = normalize(bad_mean)
             projection = dot(direction, base)
             direction = normalize([
                 direction[i] - projection * base[i]
-                for i in range(N_EMBD)
+                for i in range(n_embd)
             ])
         layers.append(direction)
 
@@ -203,14 +198,16 @@ def main() -> None:
     out.parent.mkdir(parents=True, exist_ok=True)
     payload = {
         "format": "ds4-directional-steering-v1",
-        "shape": [N_LAYER, N_EMBD],
+        "shape": [n_layer, n_embd],
+        "profile": args.profile,
         "component": args.component,
         "thinking": bool(args.think),
         "pair_normalize": bool(args.pair_normalize),
         "orthogonalize_control_mean": not args.no_orthogonalize,
         "good_file": str(Path(args.good_file)),
         "bad_file": str(Path(args.bad_file)),
-        "model": str(model),
+        "model": str(model_arg),
+        "source": args.source,
         "note": "runtime positive scale suppresses this direction; negative scale amplifies it",
     }
     out.write_text(json.dumps(payload, indent=2), encoding="utf-8")

download_model.sh

@@ -1,12 +1,28 @@
 #!/bin/sh
 set -e
 
+GLM_UNSLOTH_REPO="unsloth/GLM-5.2-GGUF"
+GLM_ANTIREZ_REPO="antirez/GLM-5.2-GGUF"
+GLM53_REPO="antirez/glm-5.3-flash-gguf"
 REPO="antirez/deepseek-v4-gguf"
-Q2_FILE="DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2.gguf"
-Q2_IMATRIX_FILE="DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf"
-Q4_FILE="DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2.gguf"
-Q4_IMATRIX_FILE="DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-imatrix.gguf"
-MTP_FILE="DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.gguf"
+DS4F_Q2_FILE="DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf"
+DS4F_Q4_FILE="DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-imatrix-0731.gguf"
+DS4F_MXFP4_FILE="DeepSeek-V4-Flash-MXFP4Experts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-mxfp4-0731.gguf"
+DS4F_Q2_Q4_FILE="DeepSeek-V4-Flash-Layers37-42Q4KExperts-OtherExpertLayersIQ2XXSGateUp-Q2KDown-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-fixed-0731.gguf"
+PRO_Q2_IMATRIX_FILE="DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix-0813.gguf"
+PRO_Q4_LAYERS00_30_FILE="DeepSeek-V4-Pro-Q4K-Layers00-30.gguf"
+PRO_Q4_LAYERS31_OUTPUT_FILE="DeepSeek-V4-Pro-Q4K-Layers-31-output.gguf"
+DS4F_DSPARK_FILE="DeepSeek-V4-Flash-DSpark-support-0731.gguf"
+GLM_UNSLOTH_Q4_REMOTE_BASE="UD-Q4_K_XL/GLM-5.2-UD-Q4_K_XL"
+GLM_UNSLOTH_Q4_LOCAL_BASE="GLM-5.2-UD-Q4_K_XL"
+GLM_UNSLOTH_Q4_FIRST_FILE="$GLM_UNSLOTH_Q4_LOCAL_BASE-00001-of-00011.gguf"
+GLM_ANTIREZ_IQ2XXS_FILE="GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf"
+GLM_ANTIREZ_Q2_FILE="GLM-5.2-UD-Q2_K_RoutedQ2K.gguf"
+GLM_ANTIREZ_Q4_FILE="GLM-5.2-UD-Q4_K_RoutedQ4K.gguf"
+GLM53_Q2_FILE="GLM-5.3-Flash-Q2.gguf"
+GLM53_Q4_FILE="GLM-5.3-Flash-Q4_K.gguf"
+GLM53_FP8_FILE="GLM-5.3-Flash-FP8.gguf"
+GLM53_VISION_FILE="GLM-5.3-Flash-Vision-Encoder.gguf"
 
 ROOT=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
 OUT_DIR=${DS4_GGUF_DIR:-"$ROOT/gguf"}
@@ -18,39 +34,101 @@ TOKEN=${HF_TOKEN:-}
 
 usage() {
     cat <<EOF
-DeepSeek V4 Flash GGUF downloader
+DwarfStar GGUF downloader
 
 Usage:
-  ./download_model.sh q2-imatrix [--token TOKEN]
-  ./download_model.sh q4-imatrix [--token TOKEN]
-  ./download_model.sh q2 [--token TOKEN]
-  ./download_model.sh q4 [--token TOKEN]
-  ./download_model.sh mtp [--token TOKEN]
+  ./download_model.sh ds4f-q2 [--token TOKEN]
+  ./download_model.sh ds4f-q2-q4 [--token TOKEN]
+  ./download_model.sh ds4f-q4 [--token TOKEN]
+  ./download_model.sh ds4f-mxfp4 [--token TOKEN]
+  ./download_model.sh ds4f-dspark [--token TOKEN]
+  ./download_model.sh pro-q2-imatrix [--token TOKEN]
+  ./download_model.sh pro-q4-layers00-30 [--token TOKEN]
+  ./download_model.sh pro-q4-layers31-output [--token TOKEN]
+  ./download_model.sh pro-q4-split [--token TOKEN]
+  ./download_model.sh glm-unsloth-q4 [--token TOKEN]
+  ./download_model.sh glm-antirez-iq2xxs [--token TOKEN]
+  ./download_model.sh glm-antirez-q2 [--token TOKEN]
+  ./download_model.sh glm-antirez-q4 [--token TOKEN]
+  ./download_model.sh glm53-q2 [--token TOKEN]
+  ./download_model.sh glm53-q4 [--token TOKEN]
+  ./download_model.sh glm53-fp8 [--token TOKEN]
+  ./download_model.sh glm53-vision [--token TOKEN]
 
 Targets:
-  *** PREFERRED GGUF FILES: USE THE IMATRIX VERSIONS BELOW ***
 
-  q2-imatrix
+  ds4f-q2
        2-bit routed experts, about 81 GB on disk.
        Recommended model for 96 and 128 GB RAM machines.
 
-  q4-imatrix
+  ds4f-q2-q4
+       Mixed Flash quant: mostly q2 routed experts, with the last 6 layers
+       using q4 routed experts. About 98 GB on disk. Good for higher
+       quality inference for 128 GB MacBooks. Works on DGX Spark but loading
+       may struggle compared to ds4f-q2.
+
+  ds4f-q4
        4-bit routed experts, about 153 GB on disk.
        Recommended model for machines with 256 GB RAM or more.
 
-  Legacy GGUF files:
+  ds4f-mxfp4
+       Native DeepSeek V4 Flash MXFP4 routed experts, about 156 GB on disk.
+       Supported by Metal and CUDA; Blackwell uses FP4 tensor cores for batched
+       expert work, while CUDA decode keeps Q8 activations.
 
-  q2   2-bit routed experts, about 81 GB on disk.
-       Older non-imatrix model for 96 and 128 GB RAM machines. Prefer
-       q2-imatrix unless you specifically need the legacy quant.
+  ds4f-dspark
+       Optional DSpark speculative decoding support GGUF for Flash 0731, about
+       6 GB. Enable it with --dspark and --mtp-model when running ds4 or ds4-server.
 
-  q4   4-bit routed experts, about 153 GB on disk.
-       Older non-imatrix model for machines with 256 GB RAM or more. Prefer
-       q4-imatrix unless you specifically need the legacy quant.
+  pro-q2-imatrix
+       DeepSeek V4 PRO 0813 q2 imatrix quant, as a single GGUF file. About
+       430 GB on disk; intended for 512 GB RAM machines.
 
-  mtp  Optional speculative decoding component, about 3.5 GB on disk.
-       It is useful with q2-imatrix, q4-imatrix, q2, and q4, but must be
-       enabled explicitly with --mtp when running ds4 or ds4-server.
+  pro-q4-layers00-30
+       First half of the DeepSeek V4 PRO Q4 routed-expert quant, layers 0..30.
+       Use on the coordinator in a two-Mac-Studio distributed run. About 426 GB.
+
+  pro-q4-layers31-output
+       Second half of the DeepSeek V4 PRO Q4 routed-expert quant, layers
+       31..output. Use on the worker in a two-Mac-Studio distributed run.
+       About 412 GB.
+
+  pro-q4-split
+       Downloads both PRO Q4 split files into the download directory. About
+       838 GB total. This target does not update ./ds4flash.gguf.
+
+  glm-unsloth-q4
+       GLM 5.2 Unsloth UD-Q4_K_XL quant from unsloth/GLM-5.2-GGUF.
+       Downloads all 11 shards and links ./ds4flash.gguf to the first shard.
+
+  glm-antirez-iq2xxs
+       GLM 5.2 antirez routed IQ2_XXS GGUF from antirez/GLM-5.2-GGUF.
+       Includes Q2_K block 78 and is intended for reduced-memory testing.
+
+  glm-antirez-q2
+       GLM 5.2 antirez routed Q2_K GGUF from antirez/GLM-5.2-GGUF.
+       About 262 GB on disk.
+
+  glm-antirez-q4
+       GLM 5.2 antirez routed Q4_K GGUF from antirez/GLM-5.2-GGUF.
+       About 434 GB on disk.
+
+  glm53-q2
+       GLM 5.3 Flash imatrix Q2 GGUF, about 90 GiB on disk. Intended for
+       resident inference on 128 GB Macs.
+
+  glm53-q4
+       GLM 5.3 Flash Q4_K GGUF, about 178 GiB on disk. Intended for two-Mac
+       tensor parallelism or single-Mac SSD streaming.
+
+  glm53-fp8
+       Text-only GLM 5.3 Flash native FP8 GGUF, about 305 GiB on disk. It
+       preserves the released weights without requantization. DwarfStar
+       inference support for this paired FP8-code/scale format is pending.
+
+  glm53-vision
+       GLM 5.3 Flash vision encoder, about 1.1 GB on disk. Load it separately
+       with --vision; this target does not update ./ds4flash.gguf.
 
 Options:
   --token TOKEN  Hugging Face token. Otherwise HF_TOKEN or the local HF token
@@ -60,15 +138,19 @@ Environment:
   DS4_GGUF_DIR   Directory used for downloaded GGUF files.
                  Default: ./gguf
 
-After q2-imatrix/q4-imatrix/q2/q4 downloads the script updates:
+After main-model downloads the script updates:
   ./ds4flash.gguf -> <download directory>/<selected model>
 
 Then the default commands work:
   ./ds4 -p "Hello"
   ./ds4-server --ctx 100000
 
-After downloading mtp, enable it explicitly, for example:
-  ./ds4 --mtp <download directory>/$MTP_FILE --mtp-draft 2
+After downloading DSpark support, enable it explicitly in greedy mode:
+  ./ds4 --dspark --mtp-model <download directory>/$DS4F_DSPARK_FILE --temp 0
+
+PRO and GLM files are downloaded with the official Hugging Face downloader
+because they are too large, sharded, or nested for the curl path used by the
+smaller DeepSeek Flash GGUF files.
 EOF
 }
 
@@ -79,13 +161,71 @@ fi
 
 MODEL=$1
 shift
+MODEL_FILES=
+LINK_MODEL=1
+FORCE_HF_DOWNLOAD=0
+FLATTEN_DOWNLOADS=0
 
 case "$MODEL" in
-    q2-imatrix) MODEL_FILE=$Q2_IMATRIX_FILE ;;
-    q4-imatrix) MODEL_FILE=$Q4_IMATRIX_FILE ;;
-    q2) MODEL_FILE=$Q2_FILE ;;
-    q4) MODEL_FILE=$Q4_FILE ;;
-    mtp) MODEL_FILE=$MTP_FILE ;;
+    ds4f-q2) MODEL_FILE=$DS4F_Q2_FILE ;;
+    ds4f-q2-q4) MODEL_FILE=$DS4F_Q2_Q4_FILE ;;
+    ds4f-q4) MODEL_FILE=$DS4F_Q4_FILE ;;
+    ds4f-mxfp4) MODEL_FILE=$DS4F_MXFP4_FILE; FORCE_HF_DOWNLOAD=1 ;;
+    ds4f-dspark) MODEL_FILE=$DS4F_DSPARK_FILE; LINK_MODEL=0 ;;
+    pro-q2-imatrix) MODEL_FILE=$PRO_Q2_IMATRIX_FILE ;;
+    pro-q4-layers00-30) MODEL_FILE=$PRO_Q4_LAYERS00_30_FILE; LINK_MODEL=0 ;;
+    pro-q4-layers31-output) MODEL_FILE=$PRO_Q4_LAYERS31_OUTPUT_FILE; LINK_MODEL=0 ;;
+    pro-q4-split)
+        MODEL_FILES="$PRO_Q4_LAYERS00_30_FILE $PRO_Q4_LAYERS31_OUTPUT_FILE"
+        LINK_MODEL=0
+        ;;
+    glm-unsloth-q4)
+        REPO=$GLM_UNSLOTH_REPO
+        MODEL_FILE=$GLM_UNSLOTH_Q4_FIRST_FILE
+        MODEL_FILES=
+        for part in 00001 00002 00003 00004 00005 00006 00007 00008 00009 00010 00011; do
+            MODEL_FILES="$MODEL_FILES $GLM_UNSLOTH_Q4_REMOTE_BASE-${part}-of-00011.gguf"
+        done
+        FORCE_HF_DOWNLOAD=1
+        FLATTEN_DOWNLOADS=1
+        ;;
+    glm-antirez-q2)
+        REPO=$GLM_ANTIREZ_REPO
+        MODEL_FILE=$GLM_ANTIREZ_Q2_FILE
+        FORCE_HF_DOWNLOAD=1
+        ;;
+    glm-antirez-iq2xxs)
+        REPO=$GLM_ANTIREZ_REPO
+        MODEL_FILE=$GLM_ANTIREZ_IQ2XXS_FILE
+        FORCE_HF_DOWNLOAD=1
+        ;;
+    glm-antirez-q4)
+        REPO=$GLM_ANTIREZ_REPO
+        MODEL_FILE=$GLM_ANTIREZ_Q4_FILE
+        FORCE_HF_DOWNLOAD=1
+        ;;
+    glm53-q2)
+        REPO=$GLM53_REPO
+        MODEL_FILE=$GLM53_Q2_FILE
+        FORCE_HF_DOWNLOAD=1
+        ;;
+    glm53-q4)
+        REPO=$GLM53_REPO
+        MODEL_FILE=$GLM53_Q4_FILE
+        FORCE_HF_DOWNLOAD=1
+        ;;
+    glm53-fp8)
+        REPO=$GLM53_REPO
+        MODEL_FILE=$GLM53_FP8_FILE
+        FORCE_HF_DOWNLOAD=1
+        LINK_MODEL=0
+        ;;
+    glm53-vision)
+        REPO=$GLM53_REPO
+        MODEL_FILE=$GLM53_VISION_FILE
+        FORCE_HF_DOWNLOAD=1
+        LINK_MODEL=0
+        ;;
     -h|--help|help)
         usage
         exit 0
@@ -120,14 +260,107 @@ if [ -z "$TOKEN" ] && [ -s "$HOME/.cache/huggingface/token" ]; then
     TOKEN=$(cat "$HOME/.cache/huggingface/token")
 fi
 
+needs_hf_download() {
+    if [ "${FORCE_HF_DOWNLOAD:-0}" -eq 1 ]; then
+        return 0
+    fi
+    case "$1" in
+        "$PRO_Q2_IMATRIX_FILE"|"$PRO_Q4_LAYERS00_30_FILE"|"$PRO_Q4_LAYERS31_OUTPUT_FILE")
+            return 0
+            ;;
+        *)
+            return 1
+            ;;
+    esac
+}
+
+find_hf_command() {
+    if command -v hf >/dev/null 2>&1; then
+        printf '%s\n' hf
+        return 0
+    fi
+    for dir in "$HOME"/Library/Python/*/bin "$HOME"/.local/bin; do
+        if [ -x "$dir/hf" ]; then
+            printf '%s\n' "$dir/hf"
+            return 0
+        fi
+    done
+    return 1
+}
+
+local_download_name() {
+    if [ "${FLATTEN_DOWNLOADS:-0}" -eq 1 ]; then
+        basename "$1"
+    else
+        printf '%s\n' "$1"
+    fi
+}
+
+download_one_hf() {
+    file=$1
+    local_file=$(local_download_name "$file")
+    out="$OUT_DIR/$local_file"
+    hf_out="$OUT_DIR/$file"
+    part="$out.part"
+
+    mkdir -p "$(dirname "$out")"
+
+    if [ -s "$out" ]; then
+        echo "Already downloaded: $out"
+        return
+    fi
+
+    if [ -e "$part" ]; then
+        echo "Found curl partial download: $part" >&2
+        echo "The Hugging Face downloader cannot resume curl .part files." >&2
+        echo "Move or remove that partial download before retrying this target." >&2
+        exit 1
+    fi
+
+    HF_CMD=$(find_hf_command || true)
+    if [ -z "$HF_CMD" ]; then
+        echo "Large GGUF downloads require the official Hugging Face CLI." >&2
+        echo "Install it with:" >&2
+        echo "  python3 -m pip install -U huggingface_hub hf_xet" >&2
+        exit 1
+    fi
+
+    echo "Downloading $file"
+    echo "from https://huggingface.co/$REPO"
+    echo "using $HF_CMD download"
+    echo "If the download stops, run the same command again to resume it."
+
+    if [ -n "$TOKEN" ]; then
+        "$HF_CMD" download "$REPO" "$file" --repo-type model --local-dir "$OUT_DIR" --token "$TOKEN"
+    else
+        "$HF_CMD" download "$REPO" "$file" --repo-type model --local-dir "$OUT_DIR"
+    fi
+
+    if [ "$hf_out" != "$out" ] && [ -s "$hf_out" ]; then
+        mv "$hf_out" "$out"
+        rmdir "$(dirname "$hf_out")" 2>/dev/null || true
+    fi
+
+    if [ ! -s "$out" ]; then
+        echo "Hugging Face download finished but expected file is missing: $out" >&2
+        exit 1
+    fi
+}
+
 download_one() {
     file=$1
-    out="$OUT_DIR/$file"
+    local_file=$(local_download_name "$file")
+    out="$OUT_DIR/$local_file"
     part="$out.part"
     aria2_part="$out.aria2"
     url="https://huggingface.co/$REPO/resolve/main/$file"
 
-    mkdir -p "$OUT_DIR"
+    if needs_hf_download "$file"; then
+        download_one_hf "$file"
+        return
+    fi
+
+    mkdir -p "$(dirname "$out")"
 
     if [ -e "$aria2_part" ]; then
         echo "Found incomplete aria2 download sidecar: $aria2_part" >&2
@@ -153,14 +386,23 @@ download_one() {
     mv "$part" "$out"
 }
 
-download_one "$MODEL_FILE"
-
-if [ "$MODEL" = "mtp" ]; then
-    echo
-    echo "MTP is an optional component for q2-imatrix, q4-imatrix, q2, and q4."
-    echo "Enable it explicitly, for example:"
-    echo "  ./ds4 --mtp $OUT_DIR/$MTP_FILE --mtp-draft 2"
+if [ -n "$MODEL_FILES" ]; then
+    for file in $MODEL_FILES; do
+        download_one "$file"
+    done
 else
+    download_one "$MODEL_FILE"
+fi
+
+if [ "$MODEL" = "ds4f-dspark" ]; then
+    echo
+    echo "DSpark support downloaded. Enable it explicitly in greedy mode:"
+    echo "  ./ds4 --dspark -m ./ds4flash.gguf --mtp-model $OUT_DIR/$DS4F_DSPARK_FILE --temp 0"
+elif [ "$MODEL" = "pro-q4-layers00-30" ] || [ "$MODEL" = "pro-q4-layers31-output" ] || [ "$MODEL" = "pro-q4-split" ]; then
+    echo
+    echo "Downloaded PRO Q4 distributed split file(s). Use them with --layers,"
+    echo "for example coordinator layers 0:30 and worker layers 31:output."
+elif [ "$LINK_MODEL" -eq 1 ]; then
     cd "$ROOT"
     ln -sfn "$OUT_DIR/$MODEL_FILE" ds4flash.gguf
     echo "Linked ./ds4flash.gguf -> $OUT_DIR/$MODEL_FILE"

ds4.h

@@ -6,6 +6,8 @@
 #include <stdint.h>
 #include <stdio.h>
 
+#include "ds4_ssd.h"
+
 /* Public engine boundary.
  *
  * The CLI and server should treat ds4_engine as the loaded model and
@@ -50,25 +52,143 @@ typedef struct {
     float logprob;
 } ds4_token_score;
 
+#define DS4_DEFAULT_TEMPERATURE 1.0f
+#define DS4_DEFAULT_TOP_P 1.0f
+#define DS4_DEFAULT_MIN_P 0.05f
+
 typedef struct ds4_engine ds4_engine;
 typedef struct ds4_session ds4_session;
 
 typedef void (*ds4_session_progress_fn)(void *ud, const char *event, int current, int total);
+typedef bool (*ds4_session_cancel_fn)(void *ud);
+
+#define DS4_SESSION_SYNC_INTERRUPTED 2
+
+typedef enum {
+    DS4_DISTRIBUTED_NONE = 0,
+    DS4_DISTRIBUTED_COORDINATOR,
+    DS4_DISTRIBUTED_WORKER,
+} ds4_distributed_role;
+
+typedef struct {
+    uint32_t start;
+    uint32_t end;
+    bool has_output;
+    bool set;
+} ds4_distributed_layers;
+
+typedef struct {
+    ds4_distributed_role role;
+    ds4_distributed_layers layers;
+    const char *listen_host;
+    int listen_port;
+    const char *coordinator_host;
+    int coordinator_port;
+    uint32_t prefill_chunk;
+    uint32_t prefill_window;
+    uint32_t activation_bits;
+    bool replay_check;
+    bool debug;
+} ds4_distributed_options;
+
+/* Tensor parallelism: two identical machines run the model in lockstep and
+ * split the heavy per-layer matvecs, exchanging partial sums at gates inside
+ * the graph (see misc/METAL_TENSOR_PARALLELISM.md).  Each rank keeps one
+ * contiguous half of the routed experts resident; dense and shared weights
+ * remain replicated.  The leader owns prompt/sampling and listens; the worker
+ * dials in and mirrors every session sync/eval. */
+typedef enum {
+    DS4_TP_NONE = 0,
+    DS4_TP_LEADER,
+    DS4_TP_WORKER,
+} ds4_tp_role;
+
+typedef enum {
+    DS4_TP_TRANSPORT_AUTO = 0,
+    DS4_TP_TRANSPORT_RDMA,
+    DS4_TP_TRANSPORT_TCP,
+} ds4_tp_transport;
+
+typedef struct {
+    ds4_tp_role role;
+    bool requested;             /* --tensor-parallel with shared role options */
+    const char *listen_host;    /* leader listens here for the worker */
+    int listen_port;
+    const char *leader_host;    /* worker dials the leader */
+    int leader_port;
+    ds4_tp_transport transport;
+    const char *rdma_device;
+    int rdma_gid_index;
+    bool rdma_gid_index_set;
+    bool glm_token_prefill;
+    int debug_hash;             /* cross-check hidden state every N tokens */
+} ds4_tp_options;
 
 typedef struct {
     const char *model_path;
     const char *mtp_path;
+    const char *vision_path;
     ds4_backend backend;
     int n_threads;
+    int context_size;
+    uint32_t prefill_chunk;
     int mtp_draft_tokens;
     float mtp_margin;
+    float dspark_confidence_threshold;
     const char *directional_steering_file;
+    const char *expert_profile_path;
     float directional_steering_attn;
     float directional_steering_ffn;
+    int power_percent;
+    uint32_t ssd_streaming_cache_experts;
+    uint64_t ssd_streaming_cache_bytes;
+    uint32_t ssd_streaming_full_layers;
+    uint32_t ssd_streaming_preload_experts;
+    uint64_t simulate_used_memory_bytes;
     bool warm_weights;
     bool quality;
+    bool glm_mtp;
+    bool glm_mtp_timing;
+    bool dspark;
+    bool dspark_strict;
+    bool dspark_exact_sampling;
+    bool dspark_confidence_threshold_set;
+    bool cuda_tensor_parallel;
+    bool ssd_streaming;
+    bool ssd_streaming_cold;
+    bool ssd_streaming_full_layers_set;
+    bool inspect_only;
+    /* Multi-GPU placement uses this to price per-layer KV storage. */
+    int placement_ctx_hint;
+    /* Number of independently allocated session graphs/caches to reserve. */
+    int placement_session_count_hint;
+    /* Server batch mode serializes execution and can share prefill scratch. */
+    bool share_session_prefill_workspace;
+    bool first_token_test;
+    bool metal_graph_test;
+    bool load_slice;
+    uint32_t load_layer_start;
+    uint32_t load_layer_end;
+    bool load_output;
+    ds4_distributed_options distributed;
+    ds4_tp_options tp;
 } ds4_engine_options;
 
+typedef struct {
+    float *data;
+    uint32_t token_count;
+    uint32_t width;
+    uint32_t height;
+    uint32_t content_width;
+    uint32_t content_height;
+    uint8_t fingerprint[32];
+} ds4_vision_embedding;
+
+typedef struct {
+    uint32_t token_start;
+    ds4_vision_embedding embedding;
+} ds4_vision_span;
+
 typedef void (*ds4_token_emit_fn)(void *ud, int token);
 typedef void (*ds4_generation_done_fn)(void *ud);
 
@@ -88,16 +208,117 @@ typedef struct {
     uint64_t cap;
 } ds4_session_snapshot;
 
+typedef struct {
+    char *path;
+    uint64_t bytes;
+} ds4_session_payload_file;
+
 int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt);
+
+/* Multi-GPU pipeline-parallel entry point (wave 2).
+ *
+ * Accepts an optional ds4_gpu_config (defined in ds4_gpu_mgpu.h) that
+ * lets callers describe a multi-GPU placement target. Passing NULL is
+ * back-compatible with ds4_engine_open and produces identical engine
+ * state — bit-equivalent execution at runtime.
+ *
+ * When a non-NULL config is supplied AND the computed placement spans
+ * more than one tier (either multiple GPUs or any CPU-spill), this
+ * wave-2 implementation prints the layout and refuses to open: full
+ * multi-tier execution wiring lands in a follow-up task
+ * (mgpu-graph-session-execution). Callers receive a non-zero return
+ * and a documented stderr notice. */
+/* ds4_gpu_config is declared in ds4_gpu_mgpu.h, which callers should
+ * include separately. We forward-declare it here so this header can be
+ * used as-is (callers passing NULL don't need the struct definition). */
+struct ds4_gpu_config;
+int ds4_engine_create_with_gpu_config(ds4_engine **out,
+                                       const ds4_engine_options *opt,
+                                       const struct ds4_gpu_config *gpu_cfg);
 void ds4_engine_close(ds4_engine *e);
 void ds4_engine_summary(ds4_engine *e);
+int ds4_engine_vocab_size(ds4_engine *e);
+uint32_t ds4_engine_prefill_chunk(ds4_engine *e);
+int ds4_engine_power(ds4_engine *e);
+int ds4_engine_set_power(ds4_engine *e, int power_percent);
+const char *ds4_engine_model_name(ds4_engine *e);
+int ds4_engine_layer_count(ds4_engine *e);
+/* Decode gate schedule for the TP transport; see ds4_tp_identity. */
+enum { DS4_TP_GATE_MASK_WORDS = 3 };
+void ds4_engine_tp_gate_schedule(ds4_engine *e,
+                                 uint32_t *start,
+                                 uint32_t *step,
+                                 uint32_t *per_token,
+                                 uint64_t mask[DS4_TP_GATE_MASK_WORDS]);
+uint32_t ds4_engine_layer_compress_ratio(ds4_engine *e, uint32_t layer);
+uint64_t ds4_engine_hidden_f32_values(ds4_engine *e);
+int ds4_engine_embd_dim(ds4_engine *e);
+uint64_t ds4_engine_model_bytes(ds4_engine *e);
+bool ds4_engine_has_vision(ds4_engine *e);
+int ds4_engine_vision_encode_file(ds4_engine *e,
+                                  const char *path,
+                                  ds4_vision_embedding *out,
+                                  char *error,
+                                  size_t error_cap);
+int ds4_engine_vision_encode_memory(ds4_engine *e,
+                                    const uint8_t *encoded,
+                                    size_t encoded_len,
+                                    ds4_vision_embedding *out,
+                                    char *error,
+                                    size_t error_cap);
+void ds4_vision_embedding_free(ds4_vision_embedding *embedding);
+int ds4_prompt_append_vision(ds4_engine *e,
+                             ds4_tokens *tokens,
+                             ds4_vision_span *span,
+                             ds4_vision_embedding *embedding,
+                             char *error,
+                             size_t error_cap);
+/* Append one user or tool message whose text parts alternate with images.
+ * text_parts must contain image_count + 1 entries. On success ownership of
+ * each embedding is transferred to the corresponding output span. */
+int ds4_chat_append_multimodal_message(ds4_engine *e,
+                                       ds4_tokens *tokens,
+                                       const char *role,
+                                       const char *const *text_parts,
+                                       ds4_vision_embedding *embeddings,
+                                       size_t image_count,
+                                       ds4_vision_span *spans,
+                                       char *error,
+                                       size_t error_cap);
+int ds4_engine_tp_vocab_split(ds4_engine *e);
+bool ds4_engine_glm_layer_payload_bytes(ds4_engine *e,
+                                        uint32_t layer,
+                                        uint32_t full_live,
+                                        uint32_t key_dim,
+                                        uint32_t value_dim,
+                                        uint32_t compact_live,
+                                        uint32_t index_live,
+                                        uint64_t *out);
+/* Stable id for cache compatibility.  0 is the original Flash shape, so old
+ * KV files with the previously-zero reserved byte remain Flash-compatible;
+ * Pro and later shapes must use nonzero ids. */
+int ds4_engine_model_id(ds4_engine *e);
+bool ds4_engine_is_glm_dsa(ds4_engine *e);
+bool ds4_engine_is_glm53(ds4_engine *e);
 const char *ds4_backend_name(ds4_backend backend);
 bool ds4_think_mode_enabled(ds4_think_mode mode);
 const char *ds4_think_mode_name(ds4_think_mode mode);
 const char *ds4_think_max_prefix(void);
+const char *ds4_glm_reasoning_effort_text(ds4_think_mode mode);
 uint32_t ds4_think_max_min_context(void);
 ds4_think_mode ds4_think_mode_for_context(ds4_think_mode mode, int ctx_size);
+/* Uses the active model shape selected by ds4_engine_open(); call after opening
+ * the GGUF so Flash/Pro dimensions are known. */
 ds4_context_memory ds4_context_memory_estimate(ds4_backend backend, int ctx_size);
+ds4_context_memory ds4_context_memory_estimate_with_prefill(
+        ds4_backend backend,
+        int ctx_size,
+        uint32_t prefill_chunk);
+ds4_context_memory ds4_context_memory_estimate_with_prefill_mode(
+        ds4_backend backend,
+        int ctx_size,
+        uint32_t prefill_chunk,
+        bool ssd_streaming);
 bool ds4_log_is_tty(FILE *fp);
 void ds4_log(FILE *fp, ds4_log_type type, const char *fmt, ...);
 int ds4_engine_generate_argmax(ds4_engine *e, const ds4_tokens *prompt,
@@ -115,7 +336,13 @@ int ds4_engine_collect_imatrix(ds4_engine *e,
                                int max_tokens);
 void ds4_engine_dump_tokens(ds4_engine *e, const ds4_tokens *tokens);
 int ds4_dump_text_tokenization(const char *model_path, const char *text, FILE *fp);
+int ds4_dump_chat_tokenization(const char *model_path,
+                               const char *system,
+                               const char *prompt,
+                               ds4_think_mode think_mode,
+                               FILE *fp);
 int ds4_engine_head_test(ds4_engine *e, const ds4_tokens *prompt);
+bool ds4_engine_is_glm_dsa(ds4_engine *e);
 int ds4_engine_first_token_test(ds4_engine *e, const ds4_tokens *prompt);
 int ds4_engine_metal_graph_test(ds4_engine *e, const ds4_tokens *prompt);
 int ds4_engine_metal_graph_full_test(ds4_engine *e, const ds4_tokens *prompt);
@@ -141,10 +368,42 @@ void ds4_chat_append_assistant_prefix(ds4_engine *e, ds4_tokens *tokens, ds4_thi
 
 char *ds4_token_text(ds4_engine *e, int token, size_t *len);
 int ds4_token_eos(ds4_engine *e);
+bool ds4_token_is_stop(ds4_engine *e, int token);
+bool ds4_token_is_thinking_control(ds4_engine *e, int token);
+bool ds4_token_is_stop_for_think_mode(ds4_engine *e,
+                                      int token,
+                                      ds4_think_mode mode);
+int ds4_token_user(ds4_engine *e);
+int ds4_token_assistant(ds4_engine *e);
+
+/* Tensor-parallel binding: allocates the GPU gate slab, registers it with
+ * the transport and arms the per-layer gate machinery.  Call once, after
+ * ds4_tp_create() and before any session work.  Transport lifecycle stays
+ * with the caller. */
+struct ds4_tp;
+int ds4_engine_tp_bind(ds4_engine *e, struct ds4_tp *tp, char *err, size_t errlen);
 
 int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size);
 void ds4_session_free(ds4_session *s);
+int ds4_session_power(ds4_session *s);
+int ds4_session_set_power(ds4_session *s, int power_percent);
+float ds4_session_directional_steering_ffn(ds4_session *s);
+/* Change steering for future evaluation without rebuilding the existing KV
+ * state. Live changes are currently limited to non-distributed sessions. */
+int ds4_session_set_directional_steering_ffn(ds4_session *s, float scale);
+bool ds4_session_is_distributed(ds4_session *s);
 void ds4_session_set_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud);
+/* UI-only progress. It may report fine-grained progress inside a prefill chunk;
+ * callers must not treat it as a durable KV checkpoint boundary. */
+void ds4_session_set_display_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud);
+/* Optional cooperative cancellation.  ds4_session_sync() checks it only at
+ * safe boundaries where the live checkpoint is either unchanged or represents a
+ * valid token prefix, and returns DS4_SESSION_SYNC_INTERRUPTED when it stops. */
+void ds4_session_set_cancel(ds4_session *s, ds4_session_cancel_fn fn, void *ud);
+void ds4_session_report_progress(ds4_session *s, const char *event, int current, int total);
+/* Distributed coordinator sessions return 1 when the full layer route is
+ * available, 0 when it is still incomplete, and -1 for a local API error. */
+int ds4_session_distributed_route_ready(ds4_session *s, char *err, size_t errlen);
 
 typedef enum {
     DS4_SESSION_REWRITE_ERROR = -1,
@@ -157,7 +416,14 @@ typedef enum {
 /* Synchronize the live session to a full prompt token prefix.  If the current
  * checkpoint is a prefix, only the suffix is evaluated; otherwise the backend
  * state is refilled from scratch. */
+#define DS4_SESSION_SYNC_INTERRUPTED 2
 int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t errlen);
+int ds4_session_sync_multimodal(ds4_session *s,
+                                const ds4_tokens *prompt,
+                                const ds4_vision_span *images,
+                                size_t image_count,
+                                char *err,
+                                size_t errlen);
 bool ds4_session_rewrite_requires_rebuild(int live_len, int canonical_len, int common);
 ds4_session_rewrite_result ds4_session_rewrite_from_common(
         ds4_session *s, const ds4_tokens *prompt, int common,
@@ -165,30 +431,149 @@ ds4_session_rewrite_result ds4_session_rewrite_from_common(
 int ds4_session_common_prefix(ds4_session *s, const ds4_tokens *prompt);
 int ds4_session_argmax(ds4_session *s);
 int ds4_session_argmax_excluding(ds4_session *s, int excluded_id);
+int ds4_sample_logits(const float *logits, int n_vocab, float temperature,
+                      int top_k, float top_p, float min_p, uint64_t *rng);
 int ds4_session_sample(ds4_session *s, float temperature, int top_k, float top_p, float min_p, uint64_t *rng);
+#ifdef DS4_TEST_HOOKS
+int ds4_test_sample_logits(const float *logits, uint32_t n_vocab,
+                           float temperature, int top_k,
+                           float top_p, float min_p, uint64_t *rng,
+                           float *prob_scratch);
+int ds4_test_sampling_probabilities(const float *logits, uint32_t n_vocab,
+                                    float temperature, int top_k,
+                                    float top_p, float min_p, float *probs);
+int ds4_test_speculative_sample(const float *target_logits,
+                                const float *draft_logits,
+                                uint32_t n_vocab,
+                                float temperature,
+                                int top_k,
+                                float top_p,
+                                float min_p,
+                                uint64_t *rng,
+                                float *target_probs,
+                                float *draft_probs);
+int ds4_test_speculative_delta_sample(const float *target_logits,
+                                      uint32_t n_vocab,
+                                      int draft_token,
+                                      float temperature,
+                                      int top_k,
+                                      float top_p,
+                                      float min_p,
+                                      uint64_t *rng,
+                                      float *target_probs);
+int ds4_test_argmax_excluding_logits(const float *logits, uint32_t n_vocab,
+                                     int excluded_id);
+uint64_t ds4_test_mixed_native_count(void);
+#endif
 int ds4_session_top_logprobs(ds4_session *s, ds4_token_score *out, int k);
 int ds4_session_token_logprob(ds4_session *s, int token, ds4_token_score *out);
+int ds4_session_copy_logits(ds4_session *s, float *out, int cap);
+int ds4_session_set_logits(ds4_session *s, const float *logits, int n);
+/* Pay the one-time first-submission GPU cost outside any measured window;
+ * used by the TP worker right after session create (no-op on CPU/GLM). */
+void ds4_session_gpu_warmup(ds4_session *s);
 int ds4_session_eval(ds4_session *s, int token, char *err, size_t errlen);
+
+typedef struct {
+    ds4_session *session;
+    int token;
+} ds4_decode_item;
+
+/* Advance independent sessions by one token each. Batch size one is exactly
+ * ds4_session_eval(). Backends without native batching use a correctness-first
+ * sequential fallback. */
+int ds4_sessions_eval_batch(ds4_decode_item *items, int count,
+                            char *err, size_t errlen);
+/* Advance one resumed prefill suffix and an independent decode batch as one
+ * scheduling step. Unsupported combinations use the ordinary serialized
+ * session operations. */
+int ds4_sessions_eval_batch_with_prefill(
+        ds4_decode_item *items, int count,
+        ds4_session *prefill_session, const ds4_tokens *prefill_prompt,
+        char *err, size_t errlen);
 int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token,
                                         int max_tokens, int eos_token,
                                         int *accepted, int accepted_cap,
                                         char *err, size_t errlen);
+/* Evaluate one already-sampled target token and speculatively extend it.
+ * Positive-temperature DSpark normally commits greedily verified draft
+ * tokens; dspark_exact_sampling selects exact stochastic p/q acceptance for
+ * DSpark or an internal GLM MTP block. */
+int ds4_session_eval_speculative(ds4_session *s, int first_token,
+                                 int max_tokens, int eos_token,
+                                 float temperature, int top_k,
+                                 float top_p, float min_p, uint64_t *rng,
+                                 int *accepted, int accepted_cap,
+                                 char *err, size_t errlen);
+/* TP worker side of a mirrored speculative-verify block: run its half of the
+ * batch verify for KV side effects, then obey the leader's commit frame
+ * (keep, or roll back and replay). Only called from ds4_tp_worker_run. */
+int ds4_session_tp_spec_cycle(ds4_session *s, const int *drafts, int draft_n,
+                              char *err, size_t errlen);
 void ds4_session_invalidate(ds4_session *s);
 void ds4_session_rewind(ds4_session *s, int pos);
 int ds4_session_pos(ds4_session *s);
 int ds4_session_ctx(ds4_session *s);
+int ds4_session_prefill_cap(ds4_session *s);
 int ds4_engine_routed_quant_bits(ds4_engine *e);
+bool ds4_engine_has_output_head(ds4_engine *e);
 bool ds4_engine_has_mtp(ds4_engine *e);
 int ds4_engine_mtp_draft_tokens(ds4_engine *e);
 const ds4_tokens *ds4_session_tokens(ds4_session *s);
 
-/* Disk KV cache payload helpers.  The server owns the outer file header and
- * policy; the engine owns the DS4-specific serialized graph state. */
+/* Low-level graph slice entry points used by distributed inference.  The
+ * transport/session routing logic lives in ds4_distributed.c. */
+int ds4_session_layer_slice_reset(ds4_session *s, char *err, size_t errlen);
+int ds4_session_eval_layer_slice(ds4_session *s,
+                                 const int *tokens,
+                                 uint32_t n_tokens,
+                                 uint32_t pos0,
+                                 uint32_t layer_start,
+                                 uint32_t layer_end,
+                                 const float *input_hc,
+                                 float *output_hc,
+                                 bool output_logits,
+                                 float *logits,
+                                 char *err,
+                                 size_t errlen);
+int ds4_session_eval_output_head_from_hc(ds4_session *s,
+                                         const float *hidden_hc,
+                                         uint32_t n_tokens,
+                                         float *logits,
+                                         char *err,
+                                         size_t errlen);
+
+/* Disk KV payload helpers.  HTTP/agent code owns the outer file header and
+ * persistence policy; the engine owns the DS4-specific serialized graph state. */
+#define DS4_SESSION_PAYLOAD_MAGIC UINT32_C(0x34565344) /* "DSV4" */
+#define DS4_SESSION_PAYLOAD_VERSION UINT32_C(2)
+#define DS4_SESSION_PAYLOAD_U32_FIELDS 13u
+#define DS4_SESSION_LAYER_PAYLOAD_MAGIC UINT32_C(0x4c565344) /* "DSVL" */
+#define DS4_SESSION_LAYER_PAYLOAD_VERSION UINT32_C(1)
+#define DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS 14u
+
 uint64_t ds4_session_payload_bytes(ds4_session *s);
+int ds4_session_stage_payload(ds4_session *s, ds4_session_payload_file *out,
+                              char *err, size_t errlen);
+int ds4_session_write_staged_payload(const ds4_session_payload_file *payload,
+                                     FILE *fp, char *err, size_t errlen);
+void ds4_session_payload_file_free(ds4_session_payload_file *payload);
 int ds4_session_save_payload(ds4_session *s, FILE *fp, char *err, size_t errlen);
 int ds4_session_load_payload(ds4_session *s, FILE *fp, uint64_t payload_bytes, char *err, size_t errlen);
 int ds4_session_save_snapshot(ds4_session *s, ds4_session_snapshot *snap, char *err, size_t errlen);
 int ds4_session_load_snapshot(ds4_session *s, const ds4_session_snapshot *snap, char *err, size_t errlen);
 void ds4_session_snapshot_free(ds4_session_snapshot *snap);
 
+uint64_t ds4_session_layer_payload_bytes(ds4_session *s,
+                                         uint32_t layer_start,
+                                         uint32_t layer_end);
+int ds4_session_save_layer_payload(ds4_session *s, FILE *fp,
+                                   uint32_t layer_start, uint32_t layer_end,
+                                   char *err, size_t errlen);
+int ds4_session_load_layer_payload(ds4_session *s, FILE *fp,
+                                   uint64_t payload_bytes,
+                                   const int *tokens, uint32_t n_tokens,
+                                   uint32_t layer_start, uint32_t layer_end,
+                                   char *err, size_t errlen);
+
 #endif

ds4_bench.c

@@ -1,13 +1,17 @@
 #include "ds4.h"
+#include "ds4_distributed.h"
+#include "ds4_gpu_args.h"
+#include "ds4_help.h"
+#include "ds4_tp.h"
 
 /* Purpose-built throughput benchmark.
  *
  * The benchmark walks one fixed token sequence to configurable context
  * frontiers, measuring only the newest prefill interval at each frontier.  It
- * then snapshots the live session in memory, performs a fixed greedy decode
- * run without allowing EOS, restores the snapshot, and continues to the next
- * frontier.  Snapshot save/restore time is intentionally outside both timing
- * windows.
+ * then snapshots the live session in memory when the payload is small enough,
+ * performs a fixed greedy decode run without allowing EOS, restores the
+ * snapshot or replays the prefix, and continues to the next frontier.  Snapshot
+ * save/restore time is intentionally outside both timing windows.
  */
 
 #include <errno.h>
@@ -20,12 +24,25 @@
 #include <string.h>
 #include <time.h>
 
+/* CUDA builds resolve these weak symbols through libcudart; other builds
+ * remain independent of CUDA headers and libraries. */
+#if defined(__GNUC__) && !defined(__APPLE__)
+#define DS4_BENCH_HAVE_CUDA_PROFILER 1
+extern int cudaProfilerStart(void) __attribute__((weak));
+extern int cudaProfilerStop(void) __attribute__((weak));
+#endif
+
+#define DS4_BENCH_DEFAULT_SNAPSHOT_MAX_BYTES (UINT64_C(1) << 30)
+
 typedef struct {
     const char *model_path;
     const char *prompt_path;
     const char *chat_prompt_path;
     const char *system;
     const char *csv_path;
+    const char *expert_profile_path;
+    const char *gpu_vram_arg;
+    const char *gpu_devices_arg;
     ds4_backend backend;
     int threads;
     int ctx_start;
@@ -33,9 +50,25 @@ typedef struct {
     int ctx_alloc;
     int step_incr;
     int gen_tokens;
+    int power_percent;
+    uint32_t prefill_chunk;
+    uint32_t ssd_streaming_cache_experts;
+    uint64_t ssd_streaming_cache_bytes;
+    uint32_t ssd_streaming_full_layers;
+    uint32_t ssd_streaming_preload_experts;
+    uint64_t simulate_used_memory_bytes;
     double step_mul;
+    const char *dump_frontier_logits_dir;
+    ds4_dist_options dist;
+    ds4_tp_options tp;
     bool warm_weights;
     bool quality;
+    bool ssd_streaming;
+    bool ssd_streaming_cold;
+    bool ssd_streaming_full_layers_set;
+    bool cuda_tensor_parallel;
+    bool show_output;
+    bool teacher_forced_decode;
 } bench_config;
 
 static double bench_now_sec(void) {
@@ -44,42 +77,31 @@ static double bench_now_sec(void) {
     return (double)ts.tv_sec + (double)ts.tv_nsec / 1000000000.0;
 }
 
-static void usage(FILE *fp) {
-    fprintf(fp,
-        "Usage: ds4-bench --prompt-file FILE [options]\n"
-        "\n"
-        "Benchmarks instantaneous prefill and generation throughput at context\n"
-        "frontiers such as 2048, 4096, 6144, ... . Generation is always greedy,\n"
-        "runs for exactly --gen-tokens tokens, and skips EOS so every row is\n"
-        "comparable.\n"
-        "\n"
-        "Input:\n"
-        "  --prompt-file FILE\n"
-        "      Raw benchmark text. The fixed token sequence is sliced at each frontier.\n"
-        "  --chat-prompt-file FILE\n"
-        "      Render FILE as one no-thinking chat user message, then slice that sequence.\n"
-        "  -sys, --system TEXT\n"
-        "      System prompt used only with --chat-prompt-file.\n"
-        "\n"
-        "Model and backend:\n"
-        "  -m, --model FILE       GGUF model path. Default: ds4flash.gguf\n"
-        "  --metal | --cuda | --cpu | --backend NAME\n"
-        "      Select backend explicitly. Defaults to Metal on macOS, CUDA elsewhere.\n"
-        "  -t, --threads N        CPU helper threads.\n"
-        "  --quality              Prefer exact kernels where applicable.\n"
-        "  --warm-weights         Touch mapped tensor pages before benchmarking.\n"
-        "\n"
-        "Sweep:\n"
-        "  --ctx-start N          First measured frontier. Default: 2048\n"
-        "  --ctx-max N            Last measured frontier. Default: 32768\n"
-        "  --ctx-alloc N          Allocated context. Default: ctx-max + gen-tokens + 1\n"
-        "  --step-mul F           Multiplicative step. Default: 1\n"
-        "  --step-incr N          Linear step when --step-mul is 1. Default: 2048\n"
-        "  --gen-tokens N         Greedy decode tokens per frontier. Default: 128\n"
-        "\n"
-        "Output:\n"
-        "  --csv FILE             Write CSV there instead of stdout.\n"
-        "  -h, --help             Show this help.\n");
+static uint64_t bench_snapshot_max_bytes(void) {
+    const char *env = getenv("DS4_BENCH_SNAPSHOT_MAX_BYTES");
+    if (!env || env[0] == '\0') return DS4_BENCH_DEFAULT_SNAPSHOT_MAX_BYTES;
+    if (!strcmp(env, "unlimited") || !strcmp(env, "UNLIMITED") ||
+        !strcmp(env, "inf") || !strcmp(env, "INF")) {
+        return UINT64_MAX;
+    }
+    char *end = NULL;
+    unsigned long long v = strtoull(env, &end, 10);
+    if (env[0] == '\0' || !end || *end != '\0') {
+        fprintf(stderr,
+                "ds4-bench: invalid DS4_BENCH_SNAPSHOT_MAX_BYTES=%s; using default %llu\n",
+                env,
+                (unsigned long long)DS4_BENCH_DEFAULT_SNAPSHOT_MAX_BYTES);
+        return DS4_BENCH_DEFAULT_SNAPSHOT_MAX_BYTES;
+    }
+    return (uint64_t)v;
+}
+
+static double bytes_to_gib(uint64_t bytes) {
+    return (double)bytes / (1024.0 * 1024.0 * 1024.0);
+}
+
+static void usage(FILE *fp, const char *topic) {
+    ds4_help_print(fp, DS4_HELP_BENCH, topic);
 }
 
 static int parse_int(const char *s, const char *opt) {
@@ -92,6 +114,16 @@ static int parse_int(const char *s, const char *opt) {
     return (int)v;
 }
 
+static int parse_nonnegative_int(const char *s, const char *opt) {
+    char *end = NULL;
+    long v = strtol(s, &end, 10);
+    if (s[0] == '\0' || *end != '\0' || v < 0 || v > INT_MAX) {
+        fprintf(stderr, "ds4-bench: invalid value for %s: %s\n", opt, s);
+        exit(2);
+    }
+    return (int)v;
+}
+
 static double parse_double_arg(const char *s, const char *opt) {
     char *end = NULL;
     double v = strtod(s, &end);
@@ -112,10 +144,18 @@ static const char *need_arg(int *i, int argc, char **argv, const char *opt) {
 
 static ds4_backend parse_backend(const char *s, const char *opt) {
     if (!strcmp(s, "metal")) return DS4_BACKEND_METAL;
+#ifdef DS4_ROCM_BUILD
+    if (!strcmp(s, "rocm")) return DS4_BACKEND_CUDA;
+#else
     if (!strcmp(s, "cuda")) return DS4_BACKEND_CUDA;
+#endif
     if (!strcmp(s, "cpu")) return DS4_BACKEND_CPU;
     fprintf(stderr, "ds4-bench: invalid value for %s: %s\n", opt, s);
+#ifdef DS4_ROCM_BUILD
+    fprintf(stderr, "ds4-bench: valid backends are: metal, rocm, cpu\n");
+#else
     fprintf(stderr, "ds4-bench: valid backends are: metal, cuda, cpu\n");
+#endif
     exit(2);
 }
 
@@ -183,9 +223,46 @@ static bench_config parse_options(int argc, char **argv) {
     for (int i = 1; i < argc; i++) {
         const char *arg = argv[i];
         if (!strcmp(arg, "-h") || !strcmp(arg, "--help")) {
-            usage(stdout);
+            const char *topic = (i + 1 < argc && argv[i + 1][0] != '-') ?
+                argv[i + 1] : NULL;
+            usage(stdout, topic);
             exit(0);
-        } else if (!strcmp(arg, "-m") || !strcmp(arg, "--model")) {
+        }
+        char dist_parse_err[256] = {0};
+        ds4_dist_cli_parse_result dist_parse =
+            ds4_dist_parse_cli_arg(arg,
+                                   &i,
+                                   argc,
+                                   argv,
+                                   &c.dist,
+                                   dist_parse_err,
+                                   sizeof(dist_parse_err));
+        if (dist_parse == DS4_DIST_CLI_ERROR) {
+            fprintf(stderr,
+                    "ds4-bench: %s\n",
+                    dist_parse_err[0] ? dist_parse_err : "invalid distributed option");
+            exit(2);
+        }
+        if (dist_parse == DS4_DIST_CLI_MATCHED) continue;
+
+        char tp_parse_err[256] = {0};
+        ds4_tp_cli_parse_result tp_parse =
+            ds4_tp_parse_cli_arg(arg,
+                                 &i,
+                                 argc,
+                                 argv,
+                                 &c.tp,
+                                 tp_parse_err,
+                                 sizeof(tp_parse_err));
+        if (tp_parse == DS4_TP_CLI_ERROR) {
+            fprintf(stderr,
+                    "ds4-bench: %s\n",
+                    tp_parse_err[0] ? tp_parse_err : "invalid tensor-parallel option");
+            exit(2);
+        }
+        if (tp_parse == DS4_TP_CLI_MATCHED) continue;
+
+        if (!strcmp(arg, "-m") || !strcmp(arg, "--model")) {
             c.model_path = need_arg(&i, argc, argv, arg);
         } else if (!strcmp(arg, "--prompt-file")) {
             c.prompt_path = need_arg(&i, argc, argv, arg);
@@ -204,26 +281,86 @@ static bench_config parse_options(int argc, char **argv) {
         } else if (!strcmp(arg, "--step-mul")) {
             c.step_mul = parse_double_arg(need_arg(&i, argc, argv, arg), arg);
         } else if (!strcmp(arg, "--gen-tokens") || !strcmp(arg, "--tokens") || !strcmp(arg, "-n")) {
-            c.gen_tokens = parse_int(need_arg(&i, argc, argv, arg), arg);
+            c.gen_tokens = parse_nonnegative_int(need_arg(&i, argc, argv, arg), arg);
         } else if (!strcmp(arg, "--csv")) {
             c.csv_path = need_arg(&i, argc, argv, arg);
+        } else if (!strcmp(arg, "--dump-frontier-logits-dir")) {
+            c.dump_frontier_logits_dir = need_arg(&i, argc, argv, arg);
+        } else if (!strcmp(arg, "--expert-profile")) {
+            c.expert_profile_path = need_arg(&i, argc, argv, arg);
         } else if (!strcmp(arg, "-t") || !strcmp(arg, "--threads")) {
             c.threads = parse_int(need_arg(&i, argc, argv, arg), arg);
         } else if (!strcmp(arg, "--backend")) {
             c.backend = parse_backend(need_arg(&i, argc, argv, arg), arg);
         } else if (!strcmp(arg, "--metal")) {
             c.backend = DS4_BACKEND_METAL;
+#ifdef DS4_ROCM_BUILD
+        } else if (!strcmp(arg, "--rocm")) {
+            c.backend = DS4_BACKEND_CUDA;
+#else
         } else if (!strcmp(arg, "--cuda")) {
             c.backend = DS4_BACKEND_CUDA;
+#endif
+        } else if (!strcmp(arg, "--gpu-vram")) {
+            c.gpu_vram_arg = need_arg(&i, argc, argv, arg);
+        } else if (!strcmp(arg, "--gpu-devices")) {
+            c.gpu_devices_arg = need_arg(&i, argc, argv, arg);
+        } else if (!strcmp(arg, "--cuda-tensor-parallel")) {
+            c.cuda_tensor_parallel = true;
         } else if (!strcmp(arg, "--cpu")) {
             c.backend = DS4_BACKEND_CPU;
         } else if (!strcmp(arg, "--quality")) {
             c.quality = true;
+        } else if (!strcmp(arg, "--ssd-streaming")) {
+            c.ssd_streaming = true;
+        } else if (!strcmp(arg, "--ssd-streaming-cold")) {
+            c.ssd_streaming_cold = true;
+        } else if (!strcmp(arg, "--ssd-streaming-cache-experts")) {
+            uint32_t experts = 0;
+            uint64_t bytes = 0;
+            if (!ds4_parse_streaming_cache_experts_arg(
+                    need_arg(&i, argc, argv, arg), &experts, &bytes)) {
+                fprintf(stderr,
+                        "ds4-bench: --ssd-streaming-cache-experts must be a positive count or <number>GB\n");
+                exit(2);
+            }
+            c.ssd_streaming_cache_experts = experts;
+            c.ssd_streaming_cache_bytes = bytes;
+        } else if (!strcmp(arg, "--ssd-streaming-full-layers")) {
+            int v = parse_nonnegative_int(need_arg(&i, argc, argv, arg), arg);
+            c.ssd_streaming_full_layers = (uint32_t)v;
+            c.ssd_streaming_full_layers_set = true;
+        } else if (!strcmp(arg, "--ssd-streaming-preload-experts")) {
+            int v = parse_int(need_arg(&i, argc, argv, arg), arg);
+            if (v <= 0) {
+                fprintf(stderr, "ds4-bench: --ssd-streaming-preload-experts must be positive\n");
+                exit(2);
+            }
+            c.ssd_streaming_preload_experts = (uint32_t)v;
+        } else if (!strcmp(arg, "--simulate-used-memory")) {
+            if (!ds4_parse_gib_arg(need_arg(&i, argc, argv, arg),
+                                   &c.simulate_used_memory_bytes)) {
+                fprintf(stderr,
+                        "ds4-bench: --simulate-used-memory must be a positive GiB value, e.g. 64GB\n");
+                exit(2);
+            }
+        } else if (!strcmp(arg, "--prefill-chunk")) {
+            c.prefill_chunk = (uint32_t)parse_int(need_arg(&i, argc, argv, arg), arg);
+        } else if (!strcmp(arg, "--power")) {
+            c.power_percent = parse_int(need_arg(&i, argc, argv, arg), arg);
+            if (c.power_percent < 1 || c.power_percent > 100) {
+                fprintf(stderr, "ds4-bench: --power must be between 1 and 100\n");
+                exit(2);
+            }
         } else if (!strcmp(arg, "--warm-weights")) {
             c.warm_weights = true;
+        } else if (!strcmp(arg, "--show-output")) {
+            c.show_output = true;
+        } else if (!strcmp(arg, "--teacher-forced-decode")) {
+            c.teacher_forced_decode = true;
         } else {
             fprintf(stderr, "ds4-bench: unknown option: %s\n", arg);
-            usage(stderr);
+            usage(stderr, NULL);
             exit(2);
         }
     }
@@ -253,9 +390,122 @@ static bench_config parse_options(int argc, char **argv) {
         fprintf(stderr, "ds4-bench: --ctx-alloc must be greater than ctx-max + gen-tokens\n");
         exit(2);
     }
+    char tp_err[256];
+    if (!ds4_tp_adopt_distributed_options(&c.tp,
+                                          &c.dist,
+                                          tp_err,
+                                          sizeof(tp_err))) {
+        fprintf(stderr, "ds4-bench: %s\n", tp_err);
+        exit(2);
+    }
+    char dist_err[256];
+    if (ds4_dist_prepare_engine_options(&c.dist, NULL, dist_err, sizeof(dist_err)) != 0) {
+        fprintf(stderr, "ds4-bench: %s\n", dist_err);
+        exit(2);
+    }
+    if (c.dist.role == DS4_DISTRIBUTED_WORKER || c.tp.role == DS4_TP_WORKER) {
+        fprintf(stderr, "ds4-bench: --role worker is a serving mode; start workers with ./ds4\n");
+        exit(2);
+    }
     return c;
 }
 
+static void json_write_string(FILE *fp, const char *s) {
+    fputc('"', fp);
+    if (s) {
+        for (const unsigned char *p = (const unsigned char *)s; *p; p++) {
+            switch (*p) {
+            case '"':  fputs("\\\"", fp); break;
+            case '\\': fputs("\\\\", fp); break;
+            case '\b': fputs("\\b", fp); break;
+            case '\f': fputs("\\f", fp); break;
+            case '\n': fputs("\\n", fp); break;
+            case '\r': fputs("\\r", fp); break;
+            case '\t': fputs("\\t", fp); break;
+            default:
+                if (*p < 0x20) fprintf(fp, "\\u%04x", (unsigned)*p);
+                else fputc((char)*p, fp);
+                break;
+            }
+        }
+    }
+    fputc('"', fp);
+}
+
+static int write_frontier_logits_json(
+        const bench_config *cfg,
+        ds4_engine         *engine,
+        ds4_session        *session,
+        int                 frontier,
+        int                 previous) {
+    if (!cfg->dump_frontier_logits_dir) return 0;
+
+    const int vocab = ds4_engine_vocab_size(engine);
+    float *logits = malloc((size_t)vocab * sizeof(logits[0]));
+    if (!logits) {
+        fprintf(stderr, "ds4-bench: out of memory copying frontier logits\n");
+        return 1;
+    }
+    if (ds4_session_copy_logits(session, logits, vocab) != vocab) {
+        fprintf(stderr, "ds4-bench: failed to copy frontier logits at %d\n", frontier);
+        free(logits);
+        return 1;
+    }
+
+    char path[PATH_MAX];
+    const int n = snprintf(path,
+                           sizeof(path),
+                           "%s/frontier_%06d.logits.json",
+                           cfg->dump_frontier_logits_dir,
+                           frontier);
+    if (n <= 0 || (size_t)n >= sizeof(path)) {
+        fprintf(stderr, "ds4-bench: frontier logits path is too long\n");
+        free(logits);
+        return 1;
+    }
+
+    FILE *fp = fopen(path, "wb");
+    if (!fp) {
+        fprintf(stderr, "ds4-bench: failed to open %s: %s\n", path, strerror(errno));
+        free(logits);
+        return 1;
+    }
+
+    const int argmax = ds4_session_argmax(session);
+    fprintf(fp, "{\n  \"source\":\"ds4-bench\",\n  \"model\":");
+    json_write_string(fp, cfg->model_path);
+    fprintf(fp,
+            ",\n  \"backend\":\"%s\",\n  \"quality\":%s,\n"
+            "  \"quant_bits\":%d,\n  \"prompt_tokens\":%d,\n"
+            "  \"frontier_tokens\":%d,\n  \"prefill_tokens\":%d,\n"
+            "  \"ctx\":%d,\n  \"vocab\":%d,\n"
+            "  \"argmax_id\":%d,\n  \"argmax_logit\":%.9g,\n  \"logits\":[",
+            ds4_backend_name(cfg->backend),
+            cfg->quality ? "true" : "false",
+            ds4_engine_routed_quant_bits(engine),
+            frontier,
+            frontier,
+            frontier - previous,
+            cfg->ctx_alloc,
+            vocab,
+            argmax,
+            logits[argmax]);
+    for (int i = 0; i < vocab; i++) {
+        if (i) fputc(',', fp);
+        if ((i % 8) == 0) fputs("\n    ", fp);
+        if (isfinite(logits[i])) fprintf(fp, "%.9g", logits[i]);
+        else fputs("null", fp);
+    }
+    fputs("\n  ]\n}\n", fp);
+    if (fclose(fp) != 0) {
+        fprintf(stderr, "ds4-bench: failed to close %s\n", path);
+        free(logits);
+        return 1;
+    }
+    free(logits);
+    return 0;
+}
+
 static int next_frontier(const bench_config *c, int cur) {
     if (cur >= c->ctx_max) return c->ctx_max;
     int next;
@@ -271,8 +521,15 @@ static int next_frontier(const bench_config *c, int cur) {
     return next;
 }
 
-static void log_context_memory(ds4_backend backend, int ctx_size) {
-    ds4_context_memory m = ds4_context_memory_estimate(backend, ctx_size);
+static void log_context_memory(ds4_backend backend,
+                               int         ctx_size,
+                               uint32_t    prefill_chunk,
+                               bool        ssd_streaming) {
+    ds4_context_memory m =
+        ds4_context_memory_estimate_with_prefill_mode(backend,
+                                                      ctx_size,
+                                                      prefill_chunk,
+                                                      ssd_streaming);
     fprintf(stderr,
             "ds4-bench: context buffers %.2f MiB (ctx=%d, backend=%s, prefill_chunk=%u, raw_kv_rows=%u, compressed_kv_rows=%u)\n",
             (double)m.total_bytes / (1024.0 * 1024.0),
@@ -283,19 +540,163 @@ static void log_context_memory(ds4_backend backend, int ctx_size) {
             m.comp_cap);
 }
 
+static int wait_distributed_route(ds4_session *session) {
+    char err[256] = {0};
+    char last[256] = {0};
+    unsigned ticks = 0;
+    const struct timespec delay = {0, 250000000L};
+
+    for (;;) {
+        int ready = ds4_session_distributed_route_ready(session, err, sizeof(err));
+        if (ready > 0) {
+            if (ticks) fprintf(stderr, "ds4-bench: distributed route ready\n");
+            return 0;
+        }
+        if (ready < 0) {
+            fprintf(stderr,
+                    "ds4-bench: distributed route readiness failed: %s\n",
+                    err[0] ? err : "unknown error");
+            return 1;
+        }
+        const char *why = err[0] ? err : "route incomplete";
+        if (strcmp(last, why) != 0 || (ticks % 20u) == 0) {
+            fprintf(stderr, "ds4-bench: waiting for distributed route: %s\n", why);
+            snprintf(last, sizeof(last), "%s", why);
+        }
+        nanosleep(&delay, NULL);
+        ticks++;
+    }
+}
+
+static void maybe_warn_distributed_step_shape(const bench_config *cfg, ds4_session *session) {
+    if (!cfg || !session || cfg->dist.role != DS4_DISTRIBUTED_COORDINATOR) return;
+    uint32_t chunk = cfg->dist.prefill_chunk;
+    if (chunk == 0) {
+        const int cap = ds4_session_prefill_cap(session);
+        if (cap > 0) chunk = (uint32_t)cap;
+    }
+    if (chunk == 0) return;
+    if (cfg->step_mul == 1.0 &&
+        cfg->step_incr > 0 &&
+        (uint32_t)cfg->step_incr < chunk &&
+        cfg->ctx_start < cfg->ctx_max)
+    {
+        fprintf(stderr,
+                "ds4-bench: note: --step-incr=%d is smaller than distributed prefill chunk %u; "
+                "suffix rows will not show multi-chunk pipeline overlap\n",
+                cfg->step_incr,
+                chunk);
+    }
+}
+
+static void close_engine(ds4_engine *engine, ds4_tp *tp) {
+    if (tp) ds4_tp_send_stop(tp);
+    ds4_engine_close(engine);
+    ds4_tp_free(tp);
+}
+
 int main(int argc, char **argv) {
     bench_config cfg = parse_options(argc, argv);
-    log_context_memory(cfg.backend, cfg.ctx_alloc);
+
+    /* Hint the packer at the largest ctx this bench run will exercise
+     * so per-layer KV bytes are priced for the real session size, not
+     * a stale 4096 default. Single-tier and CPU paths ignore this. */
+    int placement_ctx_hint = cfg.ctx_max;
+    if (cfg.ctx_alloc > placement_ctx_hint) placement_ctx_hint = cfg.ctx_alloc;
+
+    ds4_gpu_config gpu_cfg = {0};
+    bool skip_cuda = false;
+    const bool have_gpu_config = cfg.gpu_vram_arg || cfg.gpu_devices_arg;
+    if (have_gpu_config) {
+        char gpu_err[256];
+        if (parse_gpu_vram_arg(cfg.gpu_vram_arg, cfg.gpu_devices_arg,
+                               &gpu_cfg, &skip_cuda,
+                               gpu_err, sizeof(gpu_err)) != 0) {
+            fprintf(stderr, "ds4-bench: %s\n", gpu_err);
+            return 2;
+        }
+        cfg.backend = skip_cuda ? DS4_BACKEND_CPU : DS4_BACKEND_CUDA;
+    }
 
     ds4_engine_options opt = {
         .model_path = cfg.model_path,
         .backend = cfg.backend,
         .n_threads = cfg.threads,
+        .context_size = cfg.ctx_alloc,
+        .prefill_chunk = cfg.prefill_chunk,
+        .ssd_streaming_cache_experts = cfg.ssd_streaming_cache_experts,
+        .ssd_streaming_cache_bytes = cfg.ssd_streaming_cache_bytes,
+        .ssd_streaming_full_layers = cfg.ssd_streaming_full_layers,
+        .ssd_streaming_preload_experts = cfg.ssd_streaming_preload_experts,
+        .simulate_used_memory_bytes = cfg.simulate_used_memory_bytes,
+        .power_percent = cfg.power_percent,
         .warm_weights = cfg.warm_weights,
         .quality = cfg.quality,
+        .cuda_tensor_parallel = cfg.cuda_tensor_parallel,
+        .ssd_streaming = cfg.ssd_streaming,
+        .ssd_streaming_cold = cfg.ssd_streaming_cold,
+        .ssd_streaming_full_layers_set = cfg.ssd_streaming_full_layers_set,
+        .expert_profile_path = cfg.expert_profile_path,
+        .distributed = cfg.dist,
+        .tp = cfg.tp,
     };
+    char dist_err[256];
+    if (ds4_dist_prepare_engine_options(&cfg.dist, &opt, dist_err, sizeof(dist_err)) != 0) {
+        fprintf(stderr, "ds4-bench: %s\n", dist_err);
+        return 2;
+    }
+    char tp_err[256];
+    if (!ds4_tp_validate_engine_options(&opt, tp_err, sizeof(tp_err))) {
+        fprintf(stderr, "ds4-bench: %s\n", tp_err);
+        return 2;
+    }
     ds4_engine *engine = NULL;
-    if (ds4_engine_open(&engine, &opt) != 0) return 1;
+    if (have_gpu_config && !skip_cuda) {
+        const bool was_auto =
+            (cfg.gpu_vram_arg && !strcmp(cfg.gpu_vram_arg, "auto")) ||
+            (!cfg.gpu_vram_arg && cfg.gpu_devices_arg);
+        char layout[256];
+        if (format_gpu_layout_line(&gpu_cfg, was_auto,
+                                   layout, sizeof(layout)) > 0) {
+            fprintf(stdout, "%s\n", layout);
+            fflush(stdout);
+        }
+        if (ds4_engine_create_with_gpu_config(
+                &engine, &opt, &gpu_cfg) != 0) return 1;
+    } else if (ds4_engine_open(&engine, &opt) != 0) {
+        return 1;
+    }
+    ds4_tp *tp_leader = NULL;
+    if (cfg.tp.role == DS4_TP_LEADER) {
+        ds4_tp_identity tp_id = {
+            .gguf_bytes = ds4_engine_model_bytes(engine),
+            .model_id = (uint32_t)ds4_engine_model_id(engine),
+            .n_layer = (uint32_t)ds4_engine_layer_count(engine),
+            .n_embd = (uint32_t)ds4_engine_embd_dim(engine),
+            .n_vocab = (uint32_t)ds4_engine_vocab_size(engine),
+            .quant_bits = (uint32_t)ds4_engine_routed_quant_bits(engine),
+            .ctx_size = (uint32_t)cfg.ctx_alloc,
+        };
+        ds4_engine_tp_gate_schedule(engine,
+                                    &tp_id.gate_slot_start,
+                                    &tp_id.gate_slot_step,
+                                    &tp_id.gates_per_token,
+                                    tp_id.gate_slot_mask);
+        if (!ds4_tp_create(&tp_leader,
+                           &cfg.tp,
+                           &tp_id,
+                           tp_err,
+                           sizeof(tp_err)) ||
+            !ds4_engine_tp_bind(engine, tp_leader, tp_err, sizeof(tp_err))) {
+            fprintf(stderr, "ds4-bench: %s\n", tp_err);
+            close_engine(engine, tp_leader);
+            return 1;
+        }
+    }
+    log_context_memory(opt.backend,
+                       cfg.ctx_alloc,
+                       ds4_engine_prefill_chunk(engine),
+                       cfg.ssd_streaming);
 
     char *text = read_file(cfg.prompt_path ? cfg.prompt_path : cfg.chat_prompt_path);
     ds4_tokens prompt = {0};
@@ -306,13 +707,16 @@ int main(int argc, char **argv) {
     }
     free(text);
 
-    if (prompt.len < cfg.ctx_max) {
+    const int prompt_tokens_needed = cfg.ctx_max +
+        (cfg.teacher_forced_decode ? cfg.gen_tokens : 0);
+    if (prompt.len < prompt_tokens_needed) {
         fprintf(stderr,
-                "ds4-bench: prompt has %d tokens, need at least --ctx-max=%d\n",
+                "ds4-bench: prompt has %d tokens, need at least %d%s\n",
                 prompt.len,
-                cfg.ctx_max);
+                prompt_tokens_needed,
+                cfg.teacher_forced_decode ? " for teacher-forced decode" : "");
         ds4_tokens_free(&prompt);
-        ds4_engine_close(engine);
+        close_engine(engine, tp_leader);
         return 1;
     }
 
@@ -320,9 +724,18 @@ int main(int argc, char **argv) {
     if (ds4_session_create(&session, engine, cfg.ctx_alloc) != 0) {
         fprintf(stderr, "ds4-bench: failed to create session\n");
         ds4_tokens_free(&prompt);
-        ds4_engine_close(engine);
+        close_engine(engine, tp_leader);
         return 1;
     }
+    if (cfg.dist.role == DS4_DISTRIBUTED_COORDINATOR &&
+        wait_distributed_route(session) != 0)
+    {
+        ds4_session_free(session);
+        ds4_tokens_free(&prompt);
+        close_engine(engine, tp_leader);
+        return 1;
+    }
+    maybe_warn_distributed_step_shape(&cfg, session);
 
     FILE *out = stdout;
     if (cfg.csv_path) {
@@ -331,15 +744,20 @@ int main(int argc, char **argv) {
             fprintf(stderr, "ds4-bench: failed to open %s: %s\n", cfg.csv_path, strerror(errno));
             ds4_session_free(session);
             ds4_tokens_free(&prompt);
-            ds4_engine_close(engine);
+            close_engine(engine, tp_leader);
             return 1;
         }
     }
-    fprintf(out, "ctx_tokens,prefill_tokens,prefill_tps,gen_tokens,gen_tps,kvcache_bytes\n");
+    fprintf(out, "ctx_tokens,prefill_tokens,prefill_tps,gen_tokens,gen_tps,gen_first_ms,gen_steady_tokens,gen_steady_tps,kvcache_bytes\n");
     fflush(out);
 
     const int eos = ds4_token_eos(engine);
+    const bool distributed =
+        cfg.dist.role == DS4_DISTRIBUTED_COORDINATOR ||
+        cfg.tp.role == DS4_TP_LEADER;
     ds4_session_snapshot snap = {0};
+    const uint64_t snapshot_max_bytes = bench_snapshot_max_bytes();
+    bool warned_large_snapshot = false;
     char err[256];
     int previous = 0;
     int rc = 0;
@@ -352,7 +770,18 @@ int main(int argc, char **argv) {
         };
 
         const double prefill_t0 = bench_now_sec();
-        if (ds4_session_sync(session, &prefix, err, sizeof(err)) != 0) {
+#if defined(DS4_BENCH_HAVE_CUDA_PROFILER)
+        const bool cuda_profile_prefill =
+            getenv("DS4_BENCH_CUDA_PROFILE_PREFILL") != NULL &&
+            cudaProfilerStart && cudaProfilerStop;
+        if (cuda_profile_prefill) (void)cudaProfilerStart();
+#endif
+        const int prefill_rc =
+            ds4_session_sync(session, &prefix, err, sizeof(err));
+#if defined(DS4_BENCH_HAVE_CUDA_PROFILER)
+        if (cuda_profile_prefill) (void)cudaProfilerStop();
+#endif
+        if (prefill_rc != 0) {
             fprintf(stderr, "ds4-bench: prefill to %d failed: %s\n", frontier, err);
             rc = 1;
             break;
@@ -361,49 +790,145 @@ int main(int argc, char **argv) {
         const double prefill_sec = prefill_t1 - prefill_t0;
         const int prefill_tokens = frontier - previous;
 
-        if (ds4_session_save_snapshot(session, &snap, err, sizeof(err)) != 0) {
-            fprintf(stderr, "ds4-bench: snapshot at %d failed: %s\n", frontier, err);
+        if (write_frontier_logits_json(&cfg, engine, session, frontier, previous) != 0) {
             rc = 1;
             break;
         }
 
+        const bool need_restore_after_generation =
+            cfg.gen_tokens > 0 && frontier < cfg.ctx_max;
+        bool have_snapshot = false;
+        if (need_restore_after_generation && !distributed &&
+            getenv("DS4_BENCH_DISABLE_SNAPSHOT") == NULL) {
+            const uint64_t payload_bytes = ds4_session_payload_bytes(session);
+            const bool large_snapshot_forced =
+                getenv("DS4_BENCH_FORCE_SNAPSHOT") != NULL;
+            if (payload_bytes > snapshot_max_bytes && !large_snapshot_forced) {
+                if (!warned_large_snapshot) {
+                    fprintf(stderr,
+                            "ds4-bench: session payload snapshot is %.2f GiB, above the %.2f GiB benchmark limit; "
+                            "replaying prefixes instead (set DS4_BENCH_FORCE_SNAPSHOT=1 to force snapshots)\n",
+                            bytes_to_gib(payload_bytes),
+                            bytes_to_gib(snapshot_max_bytes));
+                    warned_large_snapshot = true;
+                }
+            } else if (payload_bytes > 0) {
+                if (ds4_session_save_snapshot(session, &snap, err, sizeof(err)) != 0) {
+                    fprintf(stderr, "ds4-bench: snapshot at %d failed: %s\n", frontier, err);
+                    rc = 1;
+                    break;
+                }
+                have_snapshot = true;
+            }
+        }
+
         const double gen_t0 = bench_now_sec();
+        double gen_first_sec = 0.0;
+        double gen_steady_sec = 0.0;
+        int gen_done = 0;
+        int *gen_token_buf = cfg.show_output && cfg.gen_tokens > 0
+            ? malloc((size_t)cfg.gen_tokens * sizeof(gen_token_buf[0]))
+            : NULL;
+        int gen_token_count = 0;
+        int cuda_profile_start = -1;
+        int cuda_profile_tokens = 0;
+        const char *cuda_profile_range =
+            getenv("DS4_BENCH_CUDA_PROFILE_RANGE");
+        if (cuda_profile_range &&
+            (sscanf(cuda_profile_range, "%d:%d",
+                    &cuda_profile_start, &cuda_profile_tokens) != 2 ||
+             cuda_profile_start < 0 || cuda_profile_tokens <= 0)) {
+            fprintf(stderr,
+                    "ds4-bench: invalid DS4_BENCH_CUDA_PROFILE_RANGE=%s "
+                    "(expected START_TOKEN:TOKEN_COUNT)\n",
+                    cuda_profile_range);
+            cuda_profile_start = -1;
+            cuda_profile_tokens = 0;
+        }
         for (int i = 0; i < cfg.gen_tokens; i++) {
             if (ds4_session_pos(session) + 1 >= ds4_session_ctx(session)) {
                 fprintf(stderr, "ds4-bench: generation would exceed allocated context at frontier %d\n", frontier);
                 rc = 1;
                 break;
             }
-            const int token = ds4_session_argmax_excluding(session, eos);
+            const int token = cfg.teacher_forced_decode
+                ? prompt.v[frontier + i]
+                : ds4_session_argmax_excluding(session, eos);
             if (token < 0) {
                 fprintf(stderr, "ds4-bench: failed to choose non-EOS token at frontier %d\n", frontier);
                 rc = 1;
                 break;
             }
+            const double token_t0 = bench_now_sec();
+#if defined(DS4_BENCH_HAVE_CUDA_PROFILER)
+            if (i == cuda_profile_start && cudaProfilerStart) {
+                (void)cudaProfilerStart();
+            }
+#endif
             if (ds4_session_eval(session, token, err, sizeof(err)) != 0) {
                 fprintf(stderr, "ds4-bench: decode at frontier %d failed: %s\n", frontier, err);
                 rc = 1;
                 break;
             }
+#if defined(DS4_BENCH_HAVE_CUDA_PROFILER)
+            if (cuda_profile_start >= 0 &&
+                i + 1 == cuda_profile_start + cuda_profile_tokens &&
+                cudaProfilerStop) {
+                (void)cudaProfilerStop();
+            }
+#endif
+            const double token_t1 = bench_now_sec();
+            if (i == 0) gen_first_sec = token_t1 - token_t0;
+            else gen_steady_sec += token_t1 - token_t0;
+            if (gen_token_buf) gen_token_buf[gen_token_count++] = token;
+            gen_done++;
         }
         const double gen_t1 = bench_now_sec();
+        if (cfg.show_output && gen_token_buf && gen_token_count > 0) {
+            fprintf(stderr, "ds4-bench: gen[ctx=%d] decoded text: \"", frontier);
+            for (int i = 0; i < gen_token_count; i++) {
+                size_t tlen = 0;
+                char *txt = ds4_token_text(engine, gen_token_buf[i], &tlen);
+                if (txt) {
+                    fwrite(txt, 1, tlen, stderr);
+                    free(txt);
+                }
+            }
+            fprintf(stderr, "\"\n");
+            fflush(stderr);
+        }
+        free(gen_token_buf);
         if (rc != 0) break;
 
-        if (ds4_session_load_snapshot(session, &snap, err, sizeof(err)) != 0) {
-            fprintf(stderr, "ds4-bench: restore at %d failed: %s\n", frontier, err);
-            rc = 1;
-            break;
+        if (!need_restore_after_generation) {
+            /* Nothing later depends on the frontier state. */
+        } else if (distributed || !have_snapshot) {
+            if (ds4_session_sync(session, &prefix, err, sizeof(err)) != 0) {
+                fprintf(stderr, "ds4-bench: replay restore at %d failed: %s\n", frontier, err);
+                rc = 1;
+                break;
+            }
+        } else {
+            if (ds4_session_load_snapshot(session, &snap, err, sizeof(err)) != 0) {
+                fprintf(stderr, "ds4-bench: restore at %d failed: %s\n", frontier, err);
+                rc = 1;
+                break;
+            }
         }
 
         const double gen_sec = gen_t1 - gen_t0;
+        const int gen_steady_tokens = gen_done > 1 ? gen_done - 1 : 0;
         fprintf(out,
-                "%d,%d,%.2f,%d,%.2f,%llu\n",
+                "%d,%d,%.2f,%d,%.2f,%.3f,%d,%.2f,%llu\n",
                 frontier,
                 prefill_tokens,
                 prefill_sec > 0.0 ? (double)prefill_tokens / prefill_sec : 0.0,
-                cfg.gen_tokens,
-                gen_sec > 0.0 ? (double)cfg.gen_tokens / gen_sec : 0.0,
-                (unsigned long long)snap.len);
+                gen_done,
+                gen_sec > 0.0 ? (double)gen_done / gen_sec : 0.0,
+                gen_first_sec * 1000.0,
+                gen_steady_tokens,
+                gen_steady_sec > 0.0 ? (double)gen_steady_tokens / gen_steady_sec : 0.0,
+                (unsigned long long)(have_snapshot ? snap.len : 0));
         fflush(out);
 
         previous = frontier;
@@ -414,6 +939,6 @@ int main(int argc, char **argv) {
     ds4_session_snapshot_free(&snap);
     ds4_session_free(session);
     ds4_tokens_free(&prompt);
-    ds4_engine_close(engine);
+    close_engine(engine, tp_leader);
     return rc;
 }

ds4_distributed.h

@@ -0,0 +1,125 @@
+#ifndef DS4_DISTRIBUTED_H
+#define DS4_DISTRIBUTED_H
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+
+#include "ds4.h"
+
+/* Distributed inference is an engine backend, not a separate frontend API.
+ * Programs parse distributed options here, then keep using the normal
+ * ds4_session_* calls. Only worker/coordinator serving modes call
+ * ds4_dist_run() directly.
+ */
+typedef ds4_distributed_options ds4_dist_options;
+typedef struct ds4_dist_session ds4_dist_session;
+
+/* Options used by standalone `./ds4 --role coordinator -p ...` generation.
+ * Interactive tools and the server go through the normal ds4_session API.
+ */
+typedef struct {
+    const char *prompt;
+    const char *system;
+    const char *dump_logits_path;
+    const char *dump_logprobs_path;
+    int dump_logprobs_top_k;
+    int n_predict;
+    int ctx_size;
+    float temperature;
+    float top_p;
+    float min_p;
+    uint64_t seed;
+    ds4_think_mode think_mode;
+} ds4_dist_generation_options;
+
+typedef enum {
+    DS4_DIST_CLI_ERROR = -1,
+    DS4_DIST_CLI_NOT_MATCHED = 0,
+    DS4_DIST_CLI_MATCHED = 1,
+} ds4_dist_cli_parse_result;
+
+/* Shared option parsing. */
+bool ds4_dist_enabled(const ds4_dist_options *opt);
+ds4_dist_options *ds4_dist_options_create(void);
+void ds4_dist_options_free(ds4_dist_options *opt);
+void ds4_dist_usage(FILE *fp);
+ds4_dist_cli_parse_result ds4_dist_parse_cli_arg(
+        const char *arg,
+        int *index,
+        int argc,
+        char **argv,
+        ds4_dist_options *opt,
+        char *err,
+        size_t errlen);
+
+/* Applies distributed layer-loading choices to the engine options before the
+ * model is loaded.
+ */
+int ds4_dist_prepare_engine_options(
+        const ds4_dist_options *opt,
+        ds4_engine_options *engine,
+        char *err,
+        size_t errlen);
+
+/* Coordinator session backend used by ds4.c. These mirror the normal session
+ * operations; callers outside the engine should not need to call them directly.
+ */
+int ds4_dist_session_create(
+        ds4_dist_session **out,
+        ds4_engine *engine,
+        const ds4_dist_options *opt,
+        ds4_session *owner,
+        int ctx_size,
+        char *err,
+        size_t errlen);
+void ds4_dist_session_free(ds4_dist_session *d);
+
+/* Returns 1 when the coordinator has full layer coverage, 0 when workers are
+ * still missing, and -1 for configuration or internal errors.
+ */
+int ds4_dist_session_route_ready(ds4_dist_session *d, char *err, size_t errlen);
+
+/* Synchronize the distributed KV state to the requested prompt timeline. */
+int ds4_dist_session_sync(
+        ds4_dist_session *d,
+        ds4_session *owner,
+        const ds4_tokens *checkpoint,
+        const ds4_tokens *prompt,
+        float *logits,
+        char *err,
+        size_t errlen);
+
+/* Evaluate one additional token across the current distributed route. */
+int ds4_dist_session_eval(
+        ds4_dist_session *d,
+        ds4_session *owner,
+        const ds4_tokens *checkpoint,
+        int token,
+        float *logits,
+        char *err,
+        size_t errlen);
+
+/* Save/load use the normal DSV4 payload format. The coordinator gathers or
+ * pushes remote layer shards internally so saved files are topology-neutral.
+ */
+int ds4_dist_session_save_payload(
+        ds4_dist_session *d,
+        ds4_session *owner,
+        FILE *fp,
+        char *err,
+        size_t errlen);
+int ds4_dist_session_load_payload(
+        ds4_dist_session *d,
+        ds4_session *owner,
+        FILE *fp,
+        uint64_t payload_bytes,
+        char *err,
+        size_t errlen);
+
+/* Standalone distributed mode. Workers stay in this loop; coordinator one-shot
+ * mode uses it for `./ds4 --role coordinator -p ...`.
+ */
+int ds4_dist_run(ds4_engine *engine, const ds4_dist_options *opt, const ds4_dist_generation_options *gen);
+
+#endif

ds4_glm53_vision_gpu.cuh

@@ -0,0 +1,606 @@
+/* GLM-5.3 vision operations shared by CUDA and ROCm. */
+
+#ifndef DS4_GLM53_VISION_TYPES_DEFINED
+#define DS4_GLM53_VISION_TYPES_DEFINED
+#define DS4_GLM53_VISION_LAYERS 24u
+typedef struct {
+    uint64_t norm1;
+    uint64_t qkv_weight;
+    uint64_t qkv_bias;
+    uint64_t q_norm;
+    uint64_t k_norm;
+    uint64_t attn_proj_weight;
+    uint64_t attn_proj_bias;
+    uint64_t norm2;
+    uint64_t gate_weight;
+    uint64_t gate_bias;
+    uint64_t up_weight;
+    uint64_t up_bias;
+    uint64_t down_weight;
+    uint64_t down_bias;
+} ds4_glm53_vision_layer_weights;
+
+typedef struct {
+    uint64_t patch_weight;
+    uint64_t patch_bias;
+    uint64_t post_norm;
+    uint64_t downsample_weight;
+    uint64_t downsample_bias;
+    uint64_t merger_proj;
+    uint64_t merger_norm;
+    uint64_t merger_norm_bias;
+    uint64_t merger_gate;
+    uint64_t merger_up;
+    uint64_t merger_down;
+    ds4_glm53_vision_layer_weights layer[DS4_GLM53_VISION_LAYERS];
+} ds4_glm53_vision_weights;
+#endif
+
+#ifndef DS4_GLM53_VISION_STREAM
+#define DS4_GLM53_VISION_STREAM 0
+#endif
+
+__device__ __forceinline__ static float glm53_vision_bf16(
+        const uint16_t *p) {
+    return __uint_as_float((uint32_t)(*p) << 16);
+}
+
+__device__ __forceinline__ static float glm53_vision_erf_approx(float x) {
+    const float sign = x < 0.0f ? -1.0f : 1.0f;
+    const float a = fabsf(x);
+    const float t = 1.0f / (1.0f + 0.3275911f * a);
+    const float p = (((((1.061405429f * t - 1.453152027f) * t) +
+                       1.421413741f) * t - 0.284496736f) * t +
+                       0.254829592f) * t;
+    return sign * (1.0f - p * expf(-a * a));
+}
+
+__global__ static void glm53_vision_bias_kernel(
+        float          *x,
+        const uint16_t *bias,
+        const float    *residual,
+        uint64_t        count,
+        uint32_t        width) {
+    const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= count) return;
+    float v = x[i] + glm53_vision_bf16(bias + i % width);
+    if (residual) v += residual[i];
+    x[i] = v;
+}
+
+__global__ static void glm53_vision_rms_kernel(
+        float          *out,
+        const float    *x,
+        const uint16_t *weight,
+        uint32_t        width,
+        float           eps) {
+    __shared__ float partial[256];
+    const uint32_t row = blockIdx.x;
+    const uint32_t tid = threadIdx.x;
+    const float *xr = x + (uint64_t)row * width;
+    float *yr = out + (uint64_t)row * width;
+    float sum = 0.0f;
+    for (uint32_t d = tid; d < width; d += blockDim.x) {
+        sum = fmaf(xr[d], xr[d], sum);
+    }
+    partial[tid] = sum;
+    __syncthreads();
+    for (uint32_t stride = blockDim.x / 2u; stride != 0u; stride >>= 1u) {
+        if (tid < stride) partial[tid] += partial[tid + stride];
+        __syncthreads();
+    }
+    const float inv = rsqrtf(partial[0] / (float)width + eps);
+    for (uint32_t d = tid; d < width; d += blockDim.x) {
+        yr[d] = xr[d] * inv * glm53_vision_bf16(weight + d);
+    }
+}
+
+__global__ static void glm53_vision_qkv_rope_kernel(
+        float          *q,
+        float          *k,
+        float          *v,
+        const float    *qkv,
+        const uint16_t *bias,
+        const uint16_t *q_weight,
+        const uint16_t *k_weight,
+        uint32_t        rows,
+        uint32_t        grid_w,
+        float           eps) {
+    __shared__ float qsum[32];
+    __shared__ float ksum[32];
+    const uint32_t row = blockIdx.x;
+    const uint32_t head = blockIdx.y;
+    const uint32_t lane = threadIdx.x;
+    if (row >= rows || head >= 16u || lane >= 32u) return;
+    const uint64_t qkv_base = (uint64_t)row * 3072u +
+                              (uint64_t)head * 64u;
+    const uint64_t out_base = (uint64_t)row * 1024u +
+                              (uint64_t)head * 64u;
+    float q0 = qkv[qkv_base + lane] +
+               glm53_vision_bf16(bias + (uint64_t)head * 64u + lane);
+    float q1 = qkv[qkv_base + lane + 32u] +
+               glm53_vision_bf16(bias + (uint64_t)head * 64u + lane + 32u);
+    float k0 = qkv[qkv_base + 1024u + lane] +
+               glm53_vision_bf16(bias + 1024u +
+                                 (uint64_t)head * 64u + lane);
+    float k1 = qkv[qkv_base + 1024u + lane + 32u] +
+               glm53_vision_bf16(bias + 1024u +
+                                 (uint64_t)head * 64u + lane + 32u);
+    qsum[lane] = fmaf(q0, q0, q1 * q1);
+    ksum[lane] = fmaf(k0, k0, k1 * k1);
+    __syncthreads();
+    for (uint32_t stride = 16u; stride != 0u; stride >>= 1u) {
+        if (lane < stride) {
+            qsum[lane] += qsum[lane + stride];
+            ksum[lane] += ksum[lane + stride];
+        }
+        __syncthreads();
+    }
+    const float qinv = rsqrtf(qsum[0] / 64.0f + eps);
+    const float kinv = rsqrtf(ksum[0] / 64.0f + eps);
+    q0 *= qinv * glm53_vision_bf16(q_weight + lane);
+    q1 *= qinv * glm53_vision_bf16(q_weight + lane + 32u);
+    k0 *= kinv * glm53_vision_bf16(k_weight + lane);
+    k1 *= kinv * glm53_vision_bf16(k_weight + lane + 32u);
+
+    const uint32_t merge_w = grid_w / 2u;
+    const uint32_t group_index = row / 4u;
+    const uint32_t within = row & 3u;
+    const uint32_t py = (group_index / merge_w) * 2u + within / 2u;
+    const uint32_t px = (group_index % merge_w) * 2u + within % 2u;
+    const uint32_t freq_index = lane & 15u;
+    const uint32_t pos = lane < 16u ? py : px;
+    const float inv_freq = powf(10000.0f, -(float)freq_index / 16.0f);
+    const float angle = (float)pos * inv_freq;
+    const float cs = cosf(angle);
+    const float sn = sinf(angle);
+    q[out_base + lane] = q0 * cs - q1 * sn;
+    q[out_base + lane + 32u] = q1 * cs + q0 * sn;
+    k[out_base + lane] = k0 * cs - k1 * sn;
+    k[out_base + lane + 32u] = k1 * cs + k0 * sn;
+    v[out_base + lane] = qkv[qkv_base + 2048u + lane] +
+                         glm53_vision_bf16(bias + 2048u +
+                                           (uint64_t)head * 64u + lane);
+    v[out_base + lane + 32u] = qkv[qkv_base + 2048u + lane + 32u] +
+                               glm53_vision_bf16(
+                                   bias + 2048u +
+                                   (uint64_t)head * 64u + lane + 32u);
+}
+
+__global__ static void glm53_vision_attention_kernel(
+        float       *out,
+        const float *q,
+        const float *k,
+        const float *v,
+        uint32_t     rows) {
+    __shared__ float dot[32];
+    const uint32_t row = blockIdx.x;
+    const uint32_t head = blockIdx.y;
+    const uint32_t lane = threadIdx.x;
+    if (row >= rows || head >= 16u || lane >= 32u) return;
+    const uint64_t base = (uint64_t)row * 1024u +
+                          (uint64_t)head * 64u;
+    const float q0 = q[base + lane];
+    const float q1 = q[base + lane + 32u];
+    float acc0 = 0.0f;
+    float acc1 = 0.0f;
+    float max_score = -INFINITY;
+    float denom = 0.0f;
+    for (uint32_t key_row = 0; key_row < rows; key_row++) {
+        const uint64_t kb = (uint64_t)key_row * 1024u +
+                            (uint64_t)head * 64u;
+        dot[lane] = q0 * k[kb + lane] + q1 * k[kb + lane + 32u];
+        __syncthreads();
+        for (uint32_t stride = 16u; stride != 0u; stride >>= 1u) {
+            if (lane < stride) dot[lane] += dot[lane + stride];
+            __syncthreads();
+        }
+        const float score = dot[0] * 0.125f;
+        const float next_max = fmaxf(max_score, score);
+        const float old_scale = key_row == 0u ? 0.0f :
+                                expf(max_score - next_max);
+        const float new_scale = expf(score - next_max);
+        denom = denom * old_scale + new_scale;
+        acc0 = acc0 * old_scale + new_scale * v[kb + lane];
+        acc1 = acc1 * old_scale + new_scale * v[kb + lane + 32u];
+        max_score = next_max;
+        __syncthreads();
+    }
+    out[base + lane] = acc0 / denom;
+    out[base + lane + 32u] = acc1 / denom;
+}
+
+__global__ static void glm53_vision_swiglu_bias_kernel(
+        float          *out,
+        const float    *gate,
+        const uint16_t *gate_bias,
+        const float    *up,
+        const uint16_t *up_bias,
+        uint64_t        count,
+        uint32_t        width) {
+    const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= count) return;
+    const uint32_t d = (uint32_t)(i % width);
+    const float g = fminf(gate[i] + glm53_vision_bf16(gate_bias + d), 10.0f);
+    const float u = fminf(fmaxf(up[i] + glm53_vision_bf16(up_bias + d),
+                                -10.0f), 10.0f);
+    out[i] = (g / (1.0f + expf(-g))) * u;
+}
+
+__global__ static void glm53_vision_reorder_kernel(
+        float       *out,
+        const float *x,
+        uint64_t     count) {
+    const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= count) return;
+    const uint32_t d = (uint32_t)(i % 4096u);
+    const uint64_t row = i / 4096u;
+    const uint32_t channel = d / 4u;
+    const uint32_t within = d & 3u;
+    out[i] = x[(row * 4u + within) * 1024u + channel];
+}
+
+__global__ static void glm53_vision_layernorm_gelu_kernel(
+        float          *out,
+        const float    *x,
+        const uint16_t *weight,
+        const uint16_t *bias,
+        uint32_t        width,
+        float           eps) {
+    __shared__ float partial[256];
+    const uint32_t row = blockIdx.x;
+    const uint32_t tid = threadIdx.x;
+    const float *xr = x + (uint64_t)row * width;
+    float *yr = out + (uint64_t)row * width;
+    float sum = 0.0f;
+    for (uint32_t d = tid; d < width; d += blockDim.x) sum += xr[d];
+    partial[tid] = sum;
+    __syncthreads();
+    for (uint32_t stride = blockDim.x / 2u; stride != 0u; stride >>= 1u) {
+        if (tid < stride) partial[tid] += partial[tid + stride];
+        __syncthreads();
+    }
+    const float mean = partial[0] / (float)width;
+    float var = 0.0f;
+    for (uint32_t d = tid; d < width; d += blockDim.x) {
+        const float centered = xr[d] - mean;
+        var = fmaf(centered, centered, var);
+    }
+    partial[tid] = var;
+    __syncthreads();
+    for (uint32_t stride = blockDim.x / 2u; stride != 0u; stride >>= 1u) {
+        if (tid < stride) partial[tid] += partial[tid + stride];
+        __syncthreads();
+    }
+    const float inv = rsqrtf(partial[0] / (float)width + eps);
+    const float inv_sqrt2 = 0.7071067811865475f;
+    for (uint32_t d = tid; d < width; d += blockDim.x) {
+        float value = (xr[d] - mean) * inv * glm53_vision_bf16(weight + d) +
+                      glm53_vision_bf16(bias + d);
+        yr[d] = 0.5f * value *
+                (1.0f + glm53_vision_erf_approx(value * inv_sqrt2));
+    }
+}
+
+__global__ static void glm53_vision_scatter_hc_kernel(
+        float       *hc,
+        const float *image,
+        uint32_t     dst_row,
+        uint32_t     image_row,
+        uint32_t     rows,
+        uint32_t     total_rows,
+        uint32_t     width,
+        uint32_t     n_hc) {
+    const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    const uint64_t count = (uint64_t)rows * n_hc * width;
+    if (i >= count) return;
+    const uint32_t d = (uint32_t)(i % width);
+    const uint64_t linear_row = i / width;
+    const uint32_t image_delta = (uint32_t)(linear_row / n_hc);
+    const uint32_t hc_index = (uint32_t)(linear_row % n_hc);
+    if (dst_row + image_delta >= total_rows) return;
+    const uint64_t dst = ((uint64_t)(dst_row + image_delta) * n_hc +
+                          hc_index) * width + d;
+    const uint64_t src = (uint64_t)(image_row + image_delta) * width + d;
+    hc[dst] = image[src];
+}
+
+static const uint16_t *glm53_vision_weight(
+        const void *model_map,
+        uint64_t    model_size,
+        uint64_t    offset,
+        uint64_t    elements,
+        const char *label) {
+    if (!model_map || elements > UINT64_MAX / sizeof(uint16_t) ||
+        offset > model_size) return NULL;
+    const uint64_t bytes = elements * sizeof(uint16_t);
+    if (bytes > model_size - offset) return NULL;
+#if defined(__HIP_PLATFORM_AMD__)
+    return (const uint16_t *)cuda_model_range_ptr(
+            model_map, offset, bytes, label);
+#else
+    return (const uint16_t *)cuda_resolve_weight_ptr(
+            model_map, offset, bytes, 0, label);
+#endif
+}
+
+static int glm53_vision_launch_ok(const char *label) {
+    return cuda_ok(cudaGetLastError(), label);
+}
+
+extern "C" int ds4_gpu_glm53_vision_encode(
+        float                          *out,
+        const float                    *patches,
+        uint32_t                        grid_h,
+        uint32_t                        grid_w,
+        const void                     *model_map,
+        uint64_t                        model_size,
+        const ds4_glm53_vision_weights *weights) {
+    if (!out || !patches || !model_map || !weights || grid_h == 0u ||
+        grid_w == 0u || (grid_h & 1u) != 0u || (grid_w & 1u) != 0u ||
+        grid_h > UINT32_MAX / grid_w) return 0;
+    const uint32_t rows = grid_h * grid_w;
+    const uint32_t merged_rows = rows / 4u;
+    const uint64_t row1024 = (uint64_t)rows * 1024u;
+    const uint64_t row3072 = (uint64_t)rows * 3072u;
+    const uint64_t row4096 = (uint64_t)rows * 4096u;
+    const uint64_t merged10240 = (uint64_t)merged_rows * 10240u;
+    if (row4096 > SIZE_MAX / sizeof(float) ||
+        merged10240 > SIZE_MAX / sizeof(float) ||
+        merged10240 > UINT32_MAX) return 0;
+
+    ds4_gpu_tensor *patch = NULL, *a = NULL, *b = NULL, *qkv = NULL;
+    ds4_gpu_tensor *q = NULL, *k = NULL, *v = NULL, *attn = NULL;
+    ds4_gpu_tensor *gate = NULL, *up = NULL, *mid = NULL;
+    ds4_gpu_tensor *cur = NULL, *tmp = NULL;
+    const uint16_t *bias = NULL;
+    const uint16_t *merger_norm = NULL, *merger_bias = NULL;
+    const uint64_t merged4096 = (uint64_t)merged_rows * 4096u;
+    int ok = 0;
+#define VISION_ALLOC(name_, count_) do { \
+        name_ = ds4_gpu_tensor_alloc((count_) * sizeof(float)); \
+        if (!(name_)) goto cleanup; \
+    } while (0)
+    VISION_ALLOC(patch, (uint64_t)rows * 1176u);
+    VISION_ALLOC(a, row1024);
+    VISION_ALLOC(b, row1024);
+    VISION_ALLOC(qkv, row3072);
+    VISION_ALLOC(q, row1024);
+    VISION_ALLOC(k, row1024);
+    VISION_ALLOC(v, row1024);
+    VISION_ALLOC(attn, row1024);
+    VISION_ALLOC(gate, row4096);
+    VISION_ALLOC(up, row4096);
+    VISION_ALLOC(mid, row4096);
+#undef VISION_ALLOC
+    if (!ds4_gpu_tensor_write(patch, 0, patches,
+            (uint64_t)rows * 1176u * sizeof(float)) ||
+        !ds4_gpu_begin_commands()) goto cleanup;
+    ok = ds4_gpu_glm53_matmul_bf16(
+            a, model_map, model_size, weights->patch_weight,
+            1176u, 1024u, patch, rows);
+    if (ok) {
+        bias = glm53_vision_weight(model_map, model_size,
+                weights->patch_bias, 1024u, "vision patch bias");
+        if (!bias) ok = 0;
+    }
+    if (ok) {
+        glm53_vision_bias_kernel<<<
+            (unsigned)((row1024 + 255u) / 256u), 256u, 0,
+            DS4_GLM53_VISION_STREAM>>>(
+                (float *)a->ptr, bias, NULL, row1024, 1024u);
+        ok = glm53_vision_launch_ok("GLM-5.3 vision patch bias");
+    }
+
+    cur = a;
+    tmp = b;
+    for (uint32_t il = 0; ok && il < DS4_GLM53_VISION_LAYERS; il++) {
+        const ds4_glm53_vision_layer_weights *w = &weights->layer[il];
+        const uint16_t *norm = glm53_vision_weight(
+                model_map, model_size, w->norm1, 1024u, "vision norm1");
+        if (!norm) { ok = 0; break; }
+        glm53_vision_rms_kernel<<<rows, 256u, 0,
+            DS4_GLM53_VISION_STREAM>>>(
+                (float *)tmp->ptr, (const float *)cur->ptr,
+                norm, 1024u, 1.0e-5f);
+        ok = glm53_vision_launch_ok("GLM-5.3 vision norm1");
+        if (ok) ok = ds4_gpu_glm53_matmul_bf16(
+                qkv, model_map, model_size, w->qkv_weight,
+                1024u, 3072u, tmp, rows);
+        const uint16_t *qkv_bias = NULL, *q_norm = NULL, *k_norm = NULL;
+        if (ok) {
+            qkv_bias = glm53_vision_weight(model_map, model_size,
+                    w->qkv_bias, 3072u, "vision QKV bias");
+            q_norm = glm53_vision_weight(model_map, model_size,
+                    w->q_norm, 64u, "vision Q norm");
+            k_norm = glm53_vision_weight(model_map, model_size,
+                    w->k_norm, 64u, "vision K norm");
+            if (!qkv_bias || !q_norm || !k_norm) ok = 0;
+        }
+        if (ok) {
+            glm53_vision_qkv_rope_kernel<<<dim3(rows, 16u, 1u), 32u, 0,
+                DS4_GLM53_VISION_STREAM>>>(
+                    (float *)q->ptr, (float *)k->ptr, (float *)v->ptr,
+                    (const float *)qkv->ptr, qkv_bias, q_norm, k_norm,
+                    rows, grid_w, 1.0e-5f);
+            ok = glm53_vision_launch_ok("GLM-5.3 vision QKV");
+        }
+        if (ok) {
+            glm53_vision_attention_kernel<<<dim3(rows, 16u, 1u), 32u, 0,
+                DS4_GLM53_VISION_STREAM>>>(
+                    (float *)attn->ptr, (const float *)q->ptr,
+                    (const float *)k->ptr, (const float *)v->ptr, rows);
+            ok = glm53_vision_launch_ok("GLM-5.3 vision attention");
+        }
+        if (ok) ok = ds4_gpu_glm53_matmul_bf16(
+                tmp, model_map, model_size, w->attn_proj_weight,
+                1024u, 1024u, attn, rows);
+        if (ok) {
+            bias = glm53_vision_weight(model_map, model_size,
+                    w->attn_proj_bias, 1024u, "vision attention bias");
+            if (!bias) ok = 0;
+        }
+        if (ok) {
+            glm53_vision_bias_kernel<<<
+                (unsigned)((row1024 + 255u) / 256u), 256u, 0,
+                DS4_GLM53_VISION_STREAM>>>(
+                    (float *)tmp->ptr, bias, (const float *)cur->ptr,
+                    row1024, 1024u);
+            ok = glm53_vision_launch_ok("GLM-5.3 vision attention residual");
+        }
+        ds4_gpu_tensor *swap = cur; cur = tmp; tmp = swap;
+
+        if (ok) {
+            norm = glm53_vision_weight(model_map, model_size,
+                    w->norm2, 1024u, "vision norm2");
+            if (!norm) ok = 0;
+        }
+        if (ok) {
+            glm53_vision_rms_kernel<<<rows, 256u, 0,
+                DS4_GLM53_VISION_STREAM>>>(
+                    (float *)tmp->ptr, (const float *)cur->ptr,
+                    norm, 1024u, 1.0e-5f);
+            ok = glm53_vision_launch_ok("GLM-5.3 vision norm2");
+        }
+        if (ok) ok = ds4_gpu_glm53_matmul_bf16(
+                gate, model_map, model_size, w->gate_weight,
+                1024u, 4096u, tmp, rows);
+        if (ok) ok = ds4_gpu_glm53_matmul_bf16(
+                up, model_map, model_size, w->up_weight,
+                1024u, 4096u, tmp, rows);
+        const uint16_t *gate_bias = NULL, *up_bias = NULL;
+        if (ok) {
+            gate_bias = glm53_vision_weight(model_map, model_size,
+                    w->gate_bias, 4096u, "vision gate bias");
+            up_bias = glm53_vision_weight(model_map, model_size,
+                    w->up_bias, 4096u, "vision up bias");
+            if (!gate_bias || !up_bias) ok = 0;
+        }
+        if (ok) {
+            glm53_vision_swiglu_bias_kernel<<<
+                (unsigned)((row4096 + 255u) / 256u), 256u, 0,
+                DS4_GLM53_VISION_STREAM>>>(
+                    (float *)mid->ptr, (const float *)gate->ptr, gate_bias,
+                    (const float *)up->ptr, up_bias, row4096, 4096u);
+            ok = glm53_vision_launch_ok("GLM-5.3 vision SwiGLU");
+        }
+        if (ok) ok = ds4_gpu_glm53_matmul_bf16(
+                tmp, model_map, model_size, w->down_weight,
+                4096u, 1024u, mid, rows);
+        if (ok) {
+            bias = glm53_vision_weight(model_map, model_size,
+                    w->down_bias, 1024u, "vision down bias");
+            if (!bias) ok = 0;
+        }
+        if (ok) {
+            glm53_vision_bias_kernel<<<
+                (unsigned)((row1024 + 255u) / 256u), 256u, 0,
+                DS4_GLM53_VISION_STREAM>>>(
+                    (float *)tmp->ptr, bias, (const float *)cur->ptr,
+                    row1024, 1024u);
+            ok = glm53_vision_launch_ok("GLM-5.3 vision down residual");
+        }
+        swap = cur; cur = tmp; tmp = swap;
+    }
+    if (ok) {
+        const uint16_t *norm = glm53_vision_weight(model_map, model_size,
+                weights->post_norm, 1024u, "vision post norm");
+        if (!norm) ok = 0;
+        else {
+            glm53_vision_rms_kernel<<<rows, 256u, 0,
+                DS4_GLM53_VISION_STREAM>>>(
+                    (float *)tmp->ptr, (const float *)cur->ptr,
+                    norm, 1024u, 1.0e-5f);
+            ok = glm53_vision_launch_ok("GLM-5.3 vision post norm");
+        }
+    }
+    if (ok) {
+        glm53_vision_reorder_kernel<<<
+            (unsigned)((merged4096 + 255u) / 256u), 256u, 0,
+            DS4_GLM53_VISION_STREAM>>>(
+                (float *)cur->ptr, (const float *)tmp->ptr, merged4096);
+        ok = glm53_vision_launch_ok("GLM-5.3 vision downsample reorder");
+    }
+    if (ok) ok = ds4_gpu_glm53_matmul_bf16(
+            tmp, model_map, model_size, weights->downsample_weight,
+            4096u, 4096u, cur, merged_rows);
+    if (ok) {
+        bias = glm53_vision_weight(model_map, model_size,
+                weights->downsample_bias, 4096u, "vision downsample bias");
+        if (!bias) ok = 0;
+    }
+    if (ok) {
+        glm53_vision_bias_kernel<<<
+            (unsigned)((merged4096 + 255u) / 256u), 256u, 0,
+            DS4_GLM53_VISION_STREAM>>>(
+                (float *)tmp->ptr, bias, NULL, merged4096, 4096u);
+        ok = glm53_vision_launch_ok("GLM-5.3 vision downsample bias");
+    }
+    if (ok) ok = ds4_gpu_glm53_matmul_bf16(
+            cur, model_map, model_size, weights->merger_proj,
+            4096u, 4096u, tmp, merged_rows);
+    if (ok) {
+        merger_norm = glm53_vision_weight(model_map, model_size,
+                weights->merger_norm, 4096u, "vision merger norm");
+        merger_bias = glm53_vision_weight(model_map, model_size,
+                weights->merger_norm_bias, 4096u, "vision merger norm bias");
+        if (!merger_norm || !merger_bias) ok = 0;
+    }
+    if (ok) {
+        glm53_vision_layernorm_gelu_kernel<<<merged_rows, 256u, 0,
+            DS4_GLM53_VISION_STREAM>>>(
+                (float *)tmp->ptr, (const float *)cur->ptr,
+                merger_norm, merger_bias, 4096u, 1.0e-5f);
+        ok = glm53_vision_launch_ok("GLM-5.3 vision merger norm");
+    }
+    if (ok) ok = ds4_gpu_glm53_matmul_bf16(
+            gate, model_map, model_size, weights->merger_gate,
+            4096u, 10240u, tmp, merged_rows);
+    if (ok) ok = ds4_gpu_glm53_matmul_bf16(
+            up, model_map, model_size, weights->merger_up,
+            4096u, 10240u, tmp, merged_rows);
+    if (ok) ok = ds4_gpu_swiglu_tensor(
+            mid, gate, up, (uint32_t)merged10240, 10.0f, 1.0f);
+    if (ok) ok = ds4_gpu_glm53_matmul_bf16(
+            cur, model_map, model_size, weights->merger_down,
+            10240u, 4096u, mid, merged_rows);
+    if (ds4_gpu_end_commands() == 0) ok = 0;
+    if (ok) ok = ds4_gpu_tensor_read(
+            cur, 0, out, merged4096 * sizeof(float));
+
+cleanup:
+    ds4_gpu_tensor_free(mid);
+    ds4_gpu_tensor_free(up);
+    ds4_gpu_tensor_free(gate);
+    ds4_gpu_tensor_free(attn);
+    ds4_gpu_tensor_free(v);
+    ds4_gpu_tensor_free(k);
+    ds4_gpu_tensor_free(q);
+    ds4_gpu_tensor_free(qkv);
+    ds4_gpu_tensor_free(b);
+    ds4_gpu_tensor_free(a);
+    ds4_gpu_tensor_free(patch);
+    return ok;
+}
+
+extern "C" int ds4_gpu_glm53_scatter_image_hc(
+        ds4_gpu_tensor       *hc,
+        const ds4_gpu_tensor *image,
+        uint32_t              dst_row,
+        uint32_t              image_row,
+        uint32_t              rows,
+        uint32_t              total_rows,
+        uint32_t              n_embd,
+        uint32_t              n_hc) {
+    if (!hc || !image || rows == 0u || n_embd == 0u || n_hc == 0u ||
+        dst_row > total_rows || rows > total_rows - dst_row) return 0;
+    const uint64_t count = (uint64_t)rows * n_hc * n_embd;
+    glm53_vision_scatter_hc_kernel<<<
+        (unsigned)((count + 255u) / 256u), 256u, 0,
+        DS4_GLM53_VISION_STREAM>>>(
+            (float *)hc->ptr, (const float *)image->ptr,
+            dst_row, image_row, rows, total_rows, n_embd, n_hc);
+    return glm53_vision_launch_ok("GLM-5.3 vision scatter");
+}
+
+#undef DS4_GLM53_VISION_STREAM

ds4_gpu_args.c

@@ -0,0 +1,230 @@
+/* ds4_gpu_args.c — shared parser for --gpu-vram and --gpu-devices.
+ *
+ * See ds4_gpu_args.h. The actual CUDA probing for the "auto" branch
+ * lives in ds4_cuda.cu so this .c file never includes <cuda_runtime.h>
+ * directly (mirroring the rest of the tree's separation between .c and
+ * .cu compilation units).
+ */
+#include "ds4_gpu_args.h"
+
+#include <ctype.h>
+#include <errno.h>
+#include <stdarg.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#define DS4_GPU_ARGS_DEFAULT_SAFETY_MARGIN ((size_t)512 * 1024 * 1024)
+
+/* Helper: copy at most outlen-1 bytes of msg into out and NUL-terminate. */
+static void errbuf_set(char *out, size_t outlen, const char *msg) {
+    if (!out || outlen == 0) return;
+    size_t n = strlen(msg);
+    if (n >= outlen) n = outlen - 1;
+    memcpy(out, msg, n);
+    out[n] = '\0';
+}
+
+static void errbuf_setf(char *out, size_t outlen, const char *fmt, ...) {
+    if (!out || outlen == 0) return;
+    va_list ap;
+    va_start(ap, fmt);
+    vsnprintf(out, outlen, fmt, ap);
+    va_end(ap);
+}
+
+/* Maximum per-entry values. Generous, but small enough to keep all
+ * downstream arithmetic safely inside size_t/int. */
+#define DS4_GPU_ARGS_MAX_VRAM_GB        16384   /* 16 TiB per device */
+#define DS4_GPU_ARGS_MAX_DEVICE_INDEX   1023    /* CUDA visible-device ID cap */
+
+/* Parse a comma-separated list of nonnegative integers into out_count
+ * entries of out[]. Returns 0 on success, nonzero on parse error.
+ * Empty entries (",,") or trailing/leading commas error. Per-entry
+ * upper bound enforced via max_value (use LONG_MAX for "no cap"). */
+static int parse_csv_int_list(const char *s,
+                              long       *out,
+                              int        *out_count,
+                              int         max_count,
+                              long        max_value,
+                              char       *errbuf,
+                              size_t      errbuflen,
+                              const char *what) {
+    *out_count = 0;
+    if (!s || !*s) {
+        errbuf_setf(errbuf, errbuflen, "%s: empty value", what);
+        return 1;
+    }
+    const char *p = s;
+    while (*p) {
+        /* Disallow leading whitespace inside the list. */
+        while (*p == ' ' || *p == '\t') p++;
+        char *end = NULL;
+        errno = 0;
+        long v = strtol(p, &end, 10);
+        if (end == p) {
+            errbuf_setf(errbuf, errbuflen, "%s: not a number: '%s'", what, p);
+            return 1;
+        }
+        if (errno != 0 || v < 0 || v > max_value) {
+            errbuf_setf(errbuf, errbuflen,
+                        "%s: out-of-range value: '%s' (allowed 0..%ld)",
+                        what, p, max_value);
+            return 1;
+        }
+        if (*out_count >= max_count) {
+            errbuf_setf(errbuf, errbuflen, "%s: too many entries (max %d)", what, max_count);
+            return 1;
+        }
+        out[*out_count] = v;
+        (*out_count)++;
+        while (*end == ' ' || *end == '\t') end++;
+        if (*end == ',') {
+            end++;
+            if (*end == '\0' || *end == ',') {
+                errbuf_setf(errbuf, errbuflen, "%s: empty entry near '%s'", what, end);
+                return 1;
+            }
+        } else if (*end != '\0') {
+            errbuf_setf(errbuf, errbuflen, "%s: junk after value: '%s'", what, end);
+            return 1;
+        }
+        p = end;
+    }
+    if (*out_count == 0) {
+        errbuf_setf(errbuf, errbuflen, "%s: no values parsed", what);
+        return 1;
+    }
+    return 0;
+}
+
+int parse_gpu_vram_arg(const char     *vram_arg,
+                       const char     *devices_arg,
+                       ds4_gpu_config *out,
+                       bool           *out_skip_cuda,
+                       char           *errbuf,
+                       size_t          errbuflen) {
+    if (!out || !out_skip_cuda) {
+        errbuf_set(errbuf, errbuflen, "internal: NULL out parameter");
+        return 1;
+    }
+    *out_skip_cuda = false;
+    /* Caller zeroed *out — just confirm. */
+    out->n_gpus = 0;
+    out->safety_margin_bytes = DS4_GPU_ARGS_DEFAULT_SAFETY_MARGIN;
+
+    /* Parse devices_arg first so we have the filter when handling "auto". */
+    long dev_list[DS4_MAX_GPUS] = {0};
+    int  dev_count = 0;
+    if (devices_arg) {
+        if (parse_csv_int_list(devices_arg, dev_list, &dev_count,
+                                DS4_MAX_GPUS,
+                                DS4_GPU_ARGS_MAX_DEVICE_INDEX,
+                                errbuf, errbuflen,
+                                "--gpu-devices") != 0) {
+            return 1;
+        }
+    }
+
+    /* Case 1: vram_arg unset. */
+    if (vram_arg == NULL) {
+        if (devices_arg == NULL) {
+            errbuf_set(errbuf, errbuflen, "internal: parse_gpu_vram_arg called with both args NULL");
+            return 1;
+        }
+        /* --gpu-devices given without --gpu-vram: implicit auto with filter. */
+        vram_arg = "auto";
+    }
+
+    /* Case 2: --gpu-vram 0 — short-circuit. */
+    if (!strcmp(vram_arg, "0")) {
+        if (devices_arg) {
+            errbuf_set(errbuf, errbuflen,
+                       "--gpu-vram 0 conflicts with --gpu-devices");
+            return 1;
+        }
+        *out_skip_cuda = true;
+        return 0;
+    }
+
+    /* Case 3: --gpu-vram auto. */
+    if (!strcmp(vram_arg, "auto")) {
+#if !defined(DS4_NO_GPU) && !defined(__APPLE__)
+        int filter_buf[DS4_MAX_GPUS];
+        const int *filter_ptr = NULL;
+        if (dev_count > 0) {
+            for (int i = 0; i < dev_count; i++) filter_buf[i] = (int)dev_list[i];
+            filter_ptr = filter_buf;
+        }
+        if (ds4_gpu_args_probe_auto_cuda(filter_ptr, dev_count, out,
+                                          out->safety_margin_bytes,
+                                          errbuf, errbuflen) != 0) {
+            return 1;
+        }
+        return 0;
+#else
+        (void)dev_list; (void)dev_count;
+        errbuf_set(errbuf, errbuflen,
+                   "--gpu-vram auto requires a CUDA build");
+        return 1;
+#endif
+    }
+
+    /* Case 4: explicit GB budgets — "N[,N...]" */
+    long vram_list[DS4_MAX_GPUS] = {0};
+    int  vram_count = 0;
+    if (parse_csv_int_list(vram_arg, vram_list, &vram_count,
+                            DS4_MAX_GPUS,
+                            DS4_GPU_ARGS_MAX_VRAM_GB,
+                            errbuf, errbuflen,
+                            "--gpu-vram") != 0) {
+        return 1;
+    }
+    if (dev_count > 0 && dev_count != vram_count) {
+        errbuf_setf(errbuf, errbuflen,
+                    "--gpu-devices count (%d) does not match --gpu-vram count (%d)",
+                    dev_count, vram_count);
+        return 1;
+    }
+    out->n_gpus = vram_count;
+    for (int i = 0; i < vram_count; i++) {
+        int dev_id = (dev_count > 0) ? (int)dev_list[i] : i;
+        out->device_indices[i] = dev_id;
+        out->vram_bytes[i] = (size_t)vram_list[i] * (size_t)1024 * 1024 * 1024;
+    }
+    return 0;
+}
+
+int format_gpu_layout_line(const ds4_gpu_config *cfg,
+                           bool                  was_auto,
+                           char                 *out,
+                           size_t                outlen) {
+    if (!out || outlen == 0 || !cfg) return -1;
+    /* "ds4: GPU config: N devices [d0,d1,...] requested, budgets G0,G1,... GB; auto=<bool>" */
+    int written = snprintf(out, outlen,
+                           "ds4: GPU config: %d device%s [",
+                           cfg->n_gpus, cfg->n_gpus == 1 ? "" : "s");
+    if (written < 0 || (size_t)written >= outlen) return written;
+    for (int i = 0; i < cfg->n_gpus; i++) {
+        int n = snprintf(out + written, outlen - (size_t)written,
+                         "%s%d", i ? "," : "", cfg->device_indices[i]);
+        if (n < 0 || (size_t)(written + n) >= outlen) return written + n;
+        written += n;
+    }
+    int n = snprintf(out + written, outlen - (size_t)written,
+                     "] requested, budgets ");
+    if (n < 0 || (size_t)(written + n) >= outlen) return written + n;
+    written += n;
+    for (int i = 0; i < cfg->n_gpus; i++) {
+        double gb = (double)cfg->vram_bytes[i] / (1024.0 * 1024.0 * 1024.0);
+        int m = snprintf(out + written, outlen - (size_t)written,
+                         "%s%.0f", i ? "," : "", gb);
+        if (m < 0 || (size_t)(written + m) >= outlen) return written + m;
+        written += m;
+    }
+    int t = snprintf(out + written, outlen - (size_t)written,
+                     " GB; auto=%s", was_auto ? "true" : "false");
+    if (t < 0) return t;
+    written += t;
+    return written;
+}

ds4_gpu_args.h

@@ -0,0 +1,84 @@
+/* ds4_gpu_args.h — shared CLI argument parser for --gpu-vram and
+ * --gpu-devices flags. Used by ds4_cli, ds4_server, ds4_agent, and
+ * ds4_bench. See mgpu-cli-wiring task (wave 2) for design.
+ *
+ * The parser does not link CUDA headers directly. On CUDA builds the
+ * "auto" path is gated and dispatches to ds4_gpu_args_probe_auto_cuda,
+ * which is defined in ds4_cuda.cu. On Metal/CPU builds, "auto" returns
+ * an error.
+ */
+#ifndef DS4_GPU_ARGS_H
+#define DS4_GPU_ARGS_H
+
+#include <stdbool.h>
+#include <stddef.h>
+
+#include "ds4_gpu_mgpu.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/* Single canonical parser. Populates *out and *out_skip_cuda on
+ * success; returns 0 on success, nonzero on parse error (errbuf
+ * populated). Caller zeroes out and out_skip_cuda before calling.
+ *
+ *   vram_arg may be:
+ *     NULL       -- flag unset
+ *     "0"        -- skip CUDA entirely; sets *out_skip_cuda = true
+ *     "auto"     -- query free VRAM via CUDA (CUDA build only)
+ *     "N[,N...]" -- explicit GB budgets per device
+ *
+ *   devices_arg may be:
+ *     NULL       -- 0..n-1 of visible devices (or 0..vram_count-1)
+ *     "N[,N...]" -- explicit CUDA device indices
+ *
+ * Constraints:
+ *   - If both args are given and counts differ, error.
+ *   - Negative or non-integer values error.
+ *   - Counts > DS4_MAX_GPUS error.
+ */
+int parse_gpu_vram_arg(
+    const char     *vram_arg,
+    const char     *devices_arg,
+    ds4_gpu_config *out,
+    bool           *out_skip_cuda,
+    char           *errbuf,
+    size_t          errbuflen);
+
+/* Format the layout echo line written to stdout right before
+ * ds4_engine_create_with_gpu_config is called.
+ *
+ * Example output:
+ *   ds4: GPU config: 2 devices [0,1] requested, budgets 40,12 GB; auto=true
+ *
+ * Returns the number of bytes written (excluding the trailing NUL), or
+ * -1 if out is NULL/outlen is 0.
+ */
+int format_gpu_layout_line(const ds4_gpu_config *cfg,
+                           bool                  was_auto,
+                           char                 *out,
+                           size_t                outlen);
+
+/* CUDA-only auto probe. Defined in ds4_cuda.cu under the same gate.
+ * Returns 0 on success, nonzero on failure (errbuf populated).
+ * device_filter may be NULL (use all visible devices) or a list of
+ * length filter_len (use only those device indices).
+ *
+ * Declared here for the parser's use; not called by application code
+ * directly. The Mac/CPU builds skip this branch via the
+ * #if !defined(DS4_NO_GPU) && !defined(__APPLE__) gate in the parser. */
+#if !defined(DS4_NO_GPU) && !defined(__APPLE__)
+int ds4_gpu_args_probe_auto_cuda(const int      *device_filter,
+                                  int             filter_len,
+                                  ds4_gpu_config *out,
+                                  size_t          safety_margin_bytes,
+                                  char           *errbuf,
+                                  size_t          errbuflen);
+#endif
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* DS4_GPU_ARGS_H */

ds4_gpu_mgpu.h

@@ -0,0 +1,237 @@
+/* ds4_gpu_mgpu.h — multi-GPU plumbing types and APIs (v0).
+ *
+ * This header carries the new multi-GPU additions for the multi-GPU plumbing PP work
+ * (device-aware CUDA). It is included from ds4_cuda.cu and from
+ * downstream tasks that need access to g_gpu[], g_n_gpus, g_gpu_peer_ok[],
+ * the ds4_gpu_config struct, and the new tensor APIs.
+ *
+ * Why not in ds4_gpu.h? The legacy ds4_gpu.h is included from C-only
+ * callers (ds4.c, ds4_cli.c, etc.) and from the Metal build, but is NOT
+ * included from ds4_cuda.cu historically. That asymmetry hid pre-existing
+ * signature mismatches between the legacy header and ds4_cuda.cu. We keep
+ * the legacy header opaque and put the new shared types here, so this
+ * file is the single source of truth for both ds4_cuda.cu and downstream
+ * multi-GPU tasks without disturbing the legacy contract.
+ *
+ * The struct definitions reference CUDA-specific handle types via void *
+ * placeholders so the header is safe to include from C builds, Metal
+ * builds, and the CUDA build (where ds4_cuda.cu casts the void * back
+ * to cudaStream_t / cublasHandle_t / cudaEvent_t internally).
+ */
+#ifndef DS4_GPU_MGPU_H
+#define DS4_GPU_MGPU_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#define DS4_MAX_GPUS 16
+
+/* Complete definition of the previously-opaque ds4_gpu_tensor, plus a
+ * typedef so the new API prototypes below can use the bare name
+ * `ds4_gpu_tensor *` in both C and C++ without forcing callers to
+ * include ds4_gpu.h first. Callers that include this header can
+ * stack-allocate or struct-embed tensors and pass to
+ * ds4_gpu_tensor_alloc_on. */
+struct ds4_gpu_tensor {
+    void    *ptr;
+    uint64_t bytes;
+    int      owner;
+    int      device_id;   /* -1 means legacy/untagged → treat as device 0 */
+};
+#ifndef DS4_GPU_TENSOR_DEFINED
+#define DS4_GPU_TENSOR_DEFINED
+typedef struct ds4_gpu_tensor ds4_gpu_tensor;
+#endif
+
+#ifndef DS4_GPU_ATTENTION_DECODE_ROW_DEFINED
+#define DS4_GPU_ATTENTION_DECODE_ROW_DEFINED
+#define DS4_GPU_ATTENTION_DECODE_BATCH_MAX 32u
+typedef struct {
+    uint64_t raw_kv;
+    uint64_t comp_kv;
+    uint64_t topk;
+    uint32_t pos;
+    uint32_t n_raw;
+    uint32_t raw_cap;
+    uint32_t raw_start;
+    uint32_t n_comp;
+    uint32_t top_k;
+    uint32_t window;
+    uint32_t ratio;
+    uint32_t indexed;
+} ds4_gpu_attention_decode_row;
+#endif
+
+/* Tagged so headers (notably ds4.h) can forward-declare `struct
+ * ds4_gpu_config` without dragging in this entire header. */
+typedef struct ds4_gpu_config {
+    int    device_indices[DS4_MAX_GPUS];   /* CUDA device IDs to use */
+    /* Explicit per-device budget in bytes. The engine does NOT auto-fill
+     * missing budgets - a value of 0 means "zero bytes of budget for
+     * this slot" and (combined with reserves) will push placement to
+     * CPU spill. Auto-detection (e.g. mapping --gpu-vram auto to
+     * cudaMemGetInfo) is the caller's job; see CLI flag wiring for the
+     * canonical CLI path. The engine emits a clear stderr and refuses
+     * if n_gpus > 0 and every vram_bytes[] is 0 (almost certainly a
+     * caller bug from zero-initializing the struct). */
+    size_t vram_bytes[DS4_MAX_GPUS];
+    int    n_gpus;
+    size_t safety_margin_bytes;            /* per-device reserve */
+} ds4_gpu_config;
+
+typedef struct {
+    int    device_id;
+    void  *stream;             /* cudaStream_t under CUDA */
+    void  *cublas;             /* cublasHandle_t under CUDA */
+    int    cublas_ready;
+    void  *scratch;
+    size_t scratch_bytes;
+    size_t budget_bytes;
+    size_t used_bytes;
+    void  *boundary_event;     /* cudaEvent_t under CUDA */
+} ds4_gpu_ctx;
+
+extern ds4_gpu_ctx g_gpu[DS4_MAX_GPUS];
+extern int         g_n_gpus;
+extern int         g_gpu_peer_ok[DS4_MAX_GPUS][DS4_MAX_GPUS];
+
+/* Primary multi-device init. The existing ds4_gpu_init (declared in
+ * ds4_gpu.h) is a thin shim that builds a single-device config for
+ * device 0 and calls this. */
+int ds4_gpu_init_multi(const ds4_gpu_config *cfg);
+
+/* Caller-supplied struct alloc on a specific device. Returns 0 on
+ * success, nonzero on error. Pair with ds4_gpu_tensor_free_in_place. */
+int  ds4_gpu_tensor_alloc_on(ds4_gpu_tensor *t, int device_id, uint64_t bytes);
+void ds4_gpu_tensor_free_in_place(ds4_gpu_tensor *t);
+
+/* Heap-allocated tensor on a specific logical tier; mirrors the legacy
+ * ds4_gpu_tensor_alloc ABI (returns ds4_gpu_tensor *) but with a tier
+ * parameter. Returns NULL on failure. Used by the multi-tier graph
+ * allocations in ds4.c. Single-tier callers can continue using the
+ * legacy ds4_gpu_tensor_alloc(bytes) which is equivalent to
+ * ds4_gpu_tensor_alloc_ptr_on(0, bytes). */
+ds4_gpu_tensor *ds4_gpu_tensor_alloc_ptr_on(int tier, uint64_t bytes);
+int ds4_gpu_tensor_copy_async(ds4_gpu_tensor *dst, const ds4_gpu_tensor *src, uint64_t bytes);
+void ds4_gpu_enable_q8_dequant_gemm(void);
+
+/* Heap-allocated MANAGED-memory tensor on a specific logical tier. Used
+ * for the KV-cache pool. Stamps tier on the tensor so subsequent
+ * ds4_gpu_tensor_free runs under the correct device. Note managed memory
+ * pages on first-touch and is not strictly device-bound; the stamped
+ * tier records the home tier for accounting + free. */
+ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed_on(int tier, uint64_t bytes);
+
+/* Cross-device tensor copy. Same-device → cudaMemcpyAsync; peer-capable
+ * cross-device → cudaMemcpyPeerAsync with event sync; non-peer → pinned
+ * host bounce (per src→dst pair). Honors DS4_FORCE_HOST_BOUNCE=1. */
+int ds4_gpu_tensor_copy_xdev(ds4_gpu_tensor *dst,
+                              const ds4_gpu_tensor *src,
+                              uint64_t bytes);
+int ds4_gpu_tensor_copy_xdev_default(ds4_gpu_tensor *dst,
+                                      const ds4_gpu_tensor *src,
+                                      uint64_t bytes);
+
+/* Grouped default-stream handoff whose copies execute on the destination
+ * device. The source default stream records readiness; the destination waits,
+ * performs all three copies, then naturally orders its following kernels.
+ * This allows source-device work submitted afterward to overlap the handoff. */
+int ds4_gpu_tensor_copy_xdev3_default_dst(
+        ds4_gpu_tensor       *dst0,
+        const ds4_gpu_tensor *src0,
+        uint64_t              bytes0,
+        ds4_gpu_tensor       *dst1,
+        const ds4_gpu_tensor *src1,
+        uint64_t              bytes1,
+        ds4_gpu_tensor       *dst2,
+        const ds4_gpu_tensor *src2,
+        uint64_t              bytes2);
+
+/* Cross-device copy of three buffers between the same source/destination
+ * tiers, recording a single readiness event for the destination stream. This
+ * is meant for tiny grouped activation handoffs where three independent
+ * ds4_gpu_tensor_copy_xdev calls would spend more time in event plumbing than
+ * in the copies themselves. Falls back internally when a grouped peer copy is
+ * not applicable. */
+int ds4_gpu_tensor_copy_xdev3(ds4_gpu_tensor       *dst0,
+                              const ds4_gpu_tensor *src0,
+                              uint64_t              bytes0,
+                              ds4_gpu_tensor       *dst1,
+                              const ds4_gpu_tensor *src1,
+                              uint64_t              bytes1,
+                              ds4_gpu_tensor       *dst2,
+                              const ds4_gpu_tensor *src2,
+                              uint64_t              bytes2);
+
+/* Cross-device copy that also orders against prior work on the destination
+ * stream before overwriting dst. This is needed for pipelined prefill slots
+ * where the next producer reuses a destination buffer that the destination
+ * stage just consumed. */
+int ds4_gpu_tensor_copy_xdev_ordered(ds4_gpu_tensor *dst,
+                                      const ds4_gpu_tensor *src,
+                                      uint64_t bytes);
+
+/* Order a peer read without copying: records readiness on src's stream and
+ * makes dst_tier's stream wait for it. */
+int ds4_gpu_tensor_wait_xdev(const ds4_gpu_tensor *src, int dst_tier);
+int ds4_gpu_tensor_wait_xdev_default(const ds4_gpu_tensor *src, int dst_tier);
+
+/* Cross-device float add used by CUDA tensor-parallel reductions.
+ * See ds4_gpu.h for the full contract. */
+int ds4_gpu_add_xdev_tensor(ds4_gpu_tensor       *out,
+                             const ds4_gpu_tensor *local,
+                             const ds4_gpu_tensor *remote,
+                             ds4_gpu_tensor       *remote_tmp,
+                             uint32_t              n);
+
+/* Returns the device_id recorded on the tensor; -1 if untagged. */
+int ds4_gpu_tensor_device(const ds4_gpu_tensor *t);
+
+/* Set the current CUDA device by LOGICAL tier index (0..g_n_gpus-1).
+ *
+ * This is the canonical shim for per-layer device routing in the
+ * multi-tier execution path. The caller passes a logical tier index;
+ * the shim internally indexes g_gpu[tier].device_id and calls
+ * cudaSetDevice. Returns 0 on success, nonzero on error or if the
+ * tier index is out of range.
+ *
+ * Wave-2 multi-GPU placement scaffolding adds this shim but does not
+ * exercise the multi-tier execution path; multi-GPU execution
+ * (follow-up) is its first caller. */
+int ds4_gpu_set_current_device(int logical_tier);
+int ds4_gpu_set_current_device_fenced(int logical_tier);
+
+/* Register the mmap'd host model pointer for selective-cache lookups
+ * WITHOUT triggering any device-side copy. This bypasses the
+ * DS4_CUDA_COPY_MODEL environment-variable branch that
+ * ds4_gpu_set_model_map normally honors, which is essential for
+ * multi-tier startup: we want only per-device selective tensor caches,
+ * never the whole-model copy. Returns 1 on success, 0 on error. */
+int ds4_gpu_register_model_map_no_copy(const void *model_map, uint64_t model_size);
+
+/* Strict per-device selective-cache lookup (no fallback).
+ *
+ * Returns 1 only if a covering entry exists whose device_id matches the
+ * caller-supplied expected_device (a PHYSICAL CUDA device id). Otherwise
+ * returns 0 — no host-pointer fallback, no different-device match. The
+ * caller is expected to have cudaSetDevice'd to expected_device before
+ * invoking; the returned pointer is valid to consume from that device's
+ * kernel.
+ *
+ * Used by multi-tier kernel-dispatch resolvers in
+ * multi-GPU execution (multi-GPU execution). Single-tier callers should
+ * keep using ds4_gpu_lookup_cache for back-compat behavior. */
+int ds4_gpu_lookup_cache_strict(uint64_t source_offset,
+                                 uint64_t bytes,
+                                 int      expected_device,
+                                 void   **out_device_ptr);
+
+#ifdef __cplusplus
+} /* extern "C" */
+#endif
+
+#endif /* DS4_GPU_MGPU_H */

ds4_help.c

@@ -0,0 +1,595 @@
+#include "ds4_help.h"
+
+#include <stdbool.h>
+#include <stdio.h>
+#include <string.h>
+#include <unistd.h>
+
+typedef struct {
+    const char *off;
+    const char *cyan;
+    const char *title;
+    const char *yellow;
+    const char *grey;
+    const char *red;
+    const char *white;
+    const char *bright;
+} help_colors;
+
+static help_colors help_make_colors(FILE *fp) {
+    bool color = isatty(fileno(fp));
+    help_colors c = {0};
+    if (!color) return c;
+    c.off = "\x1b[0m";
+    c.cyan = "\x1b[38;5;81m";
+    c.title = "\x1b[1;38;5;250m";
+    c.yellow = "\x1b[38;5;179m";
+    c.grey = "\x1b[38;5;240m";
+    c.red = "\x1b[38;5;203m";
+    c.white = "\x1b[38;5;252m";
+    c.bright = "\x1b[1;38;5;231m";
+    return c;
+}
+
+static void title(FILE *fp, const help_colors *c, const char *s) {
+    fprintf(fp, "%s%s%s\n", c->title ? c->title : "", s, c->off ? c->off : "");
+}
+
+static void title_red(FILE *fp, const help_colors *c, const char *s) {
+    fprintf(fp, "%s%s%s\n", c->red ? c->red : "", s, c->off ? c->off : "");
+}
+
+static bool option_name_has_switch(const char *name) {
+    bool word_start = true;
+    while (*name) {
+        if (word_start && (*name == '-' || *name == '/')) return true;
+        word_start = (*name == ' ');
+        name++;
+    }
+    return false;
+}
+
+static void print_colored_option_name(FILE *fp, const help_colors *c, const char *name) {
+    bool has_switch = option_name_has_switch(name);
+    bool word_start = true;
+    while (*name) {
+        const char *start = name;
+        while (*name && *name != ' ') name++;
+        bool is_option = !has_switch || *start == '-' || *start == '/' ||
+                         (word_start && has_switch && *start != '[');
+        const char *color = is_option ? c->cyan : c->bright;
+        if (color) fputs(color, fp);
+        fwrite(start, 1, (size_t)(name - start), fp);
+        if (color && c->off) fputs(c->off, fp);
+        if (*name == ' ') {
+            fputc(*name++, fp);
+            word_start = false;
+        }
+    }
+}
+
+static void opt(FILE *fp, const help_colors *c, const char *name, const char *desc) {
+    if (c->cyan) {
+        fputs("  ", fp);
+        print_colored_option_name(fp, c, name);
+        fprintf(fp, " %s|%s ", c->grey ? c->grey : "", c->grey ? c->off : "");
+        fprintf(fp, "%s%s%s\n", c->white ? c->white : "", desc,
+                c->white ? c->off : "");
+        return;
+    }
+
+    const int col = 30;
+    int n = (int)strlen(name);
+    if (n > col) {
+        fprintf(fp, "  %s\n      %s\n", name, desc);
+    } else {
+        fprintf(fp, "  %-30s %s\n", name, desc);
+    }
+}
+
+static void para(FILE *fp, const help_colors *c, const char *s) {
+    fprintf(fp, "%s%s%s\n",
+            c->yellow ? c->yellow : "", s, c->yellow ? c->off : "");
+}
+
+static bool streq(const char *a, const char *b) {
+    return a && b && strcmp(a, b) == 0;
+}
+
+static bool topic_is(const char *topic, const char *name) {
+    return topic && strcmp(topic, name) == 0;
+}
+
+static const char *tool_name(ds4_help_tool tool) {
+    switch (tool) {
+    case DS4_HELP_DS4: return "ds4";
+    case DS4_HELP_SERVER: return "ds4-server";
+    case DS4_HELP_AGENT: return "ds4-agent";
+    case DS4_HELP_BENCH: return "ds4-bench";
+    case DS4_HELP_EVAL: return "ds4-eval";
+    }
+    return "ds4";
+}
+
+static const char *tool_usage(ds4_help_tool tool) {
+    switch (tool) {
+    case DS4_HELP_DS4:
+        return "Usage: ds4 [(-p PROMPT | --prompt-file FILE)] [options]";
+    case DS4_HELP_SERVER:
+        return "Usage: ds4-server [options]";
+    case DS4_HELP_AGENT:
+        return "Usage: ds4-agent [options]";
+    case DS4_HELP_BENCH:
+        return "Usage: ds4-bench (--prompt-file FILE | --chat-prompt-file FILE) [options]";
+    case DS4_HELP_EVAL:
+        return "Usage: ds4-eval [options]";
+    }
+    return "Usage: ds4 [options]";
+}
+
+static const char *tool_summary(ds4_help_tool tool) {
+    switch (tool) {
+    case DS4_HELP_DS4:
+        return "Chat with a local DwarfStar model, run one-shot prompts, inspect models, or coordinate distributed inference.";
+    case DS4_HELP_SERVER:
+        return "Serve one loaded DwarfStar model through OpenAI, Responses, Anthropic, and completion-compatible HTTP APIs.";
+    case DS4_HELP_AGENT:
+        return "Run the native terminal coding agent with live tools, session save/restore, and a responsive prompt while the model works.";
+    case DS4_HELP_BENCH:
+        return "Measure prefill, decode, context growth, and KV-cache size across repeatable context frontiers.";
+    case DS4_HELP_EVAL:
+        return "Run the built-in reasoning, math, science, and security evaluation harness with a live terminal UI.";
+    }
+    return "";
+}
+
+static void print_model_runtime(FILE *fp, const help_colors *c,
+                                ds4_help_tool tool, bool full) {
+    title(fp, c, "Model And Runtime");
+    opt(fp, c, "-m, --model FILE", "GGUF model path. Default: ds4flash.gguf");
+    if (tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT || tool == DS4_HELP_SERVER) {
+        opt(fp, c, "--vision FILE", "GLM 5.3 vision encoder GGUF.");
+    }
+#ifdef DS4_ROCM_BUILD
+    opt(fp, c, "--metal | --rocm | --cpu", "Select the backend explicitly.");
+    opt(fp, c, "--backend NAME", "Backend name: metal, rocm, or cpu.");
+#else
+    opt(fp, c, "--metal | --cuda | --cpu", "Select the backend explicitly.");
+    opt(fp, c, "--backend NAME", "Backend name: metal, cuda, or cpu.");
+    opt(fp, c, "--gpu-vram N[,N,...]|auto", "CUDA VRAM budgets per device, in GiB, or auto-detect free VRAM.");
+    opt(fp, c, "--gpu-devices N[,N,...]", "CUDA device indices used by multi-GPU placement.");
+    if (tool != DS4_HELP_EVAL) {
+        opt(fp, c, "--cuda-tensor-parallel", "Enable the paired DeepSeek tensor/expert path on an even multi-GPU CUDA placement.");
+    }
+#endif
+    if (tool != DS4_HELP_BENCH) {
+        opt(fp, c, "-c, --ctx N", "Allocated context tokens.");
+    }
+    if (tool == DS4_HELP_SERVER) {
+        opt(fp, c, "-n, --tokens N", "Default max output tokens when clients omit a limit.");
+    }
+    opt(fp, c, "-t, --threads N", "CPU helper threads for host-side/reference work.");
+    opt(fp, c, "--power N", "GPU duty-cycle target, 1..100. Default: 100");
+    opt(fp, c, "--ssd-streaming", "Metal/CUDA/ROCm: opt in to SSD-backed model streaming instead of full residency.");
+    opt(fp, c, "--ssd-streaming-cold", "SSD streaming: skip default popularity-based expert-cache preload.");
+    opt(fp, c, "--ssd-streaming-cache-experts N|NGB", "SSD streaming: N is an exact dynamic expert count; NGB is a routed memory budget that also reserves two full prefill layers. Auto: 80% working set minus non-routed weights; GLM Metal caps lower.");
+    opt(fp, c, "--ssd-streaming-full-layers N", "GLM Metal streaming: keep the first N routed layers fully resident. Default: auto from NGB expert budget; use 0 to disable.");
+    opt(fp, c, "--ssd-streaming-preload-experts N", "SSD streaming: upfront popularity preload count. DeepSeek auto-seeds by default; GLM demand-fills unless N is explicit.");
+    opt(fp, c, "--simulate-used-memory NGB", "Diagnostic: lock N GiB before model load to simulate a smaller-memory machine.");
+    opt(fp, c, "--prefill-chunk N", "Graph prefill chunk size. Default: CUDA TP 2048; PRO long prompts 8192; others 4096.");
+    if (full) {
+        if (tool == DS4_HELP_EVAL) {
+            opt(fp, c, "--mtp-model FILE", "External MTP or DSpark support GGUF.");
+        }
+        if (tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT || tool == DS4_HELP_SERVER) {
+            opt(fp, c, "--mtp", "Enable model-embedded MTP speculation.");
+            opt(fp, c, "--mtp-model FILE", "External MTP or DSpark support GGUF.");
+            opt(fp, c, "--mtp-draft N", "Maximum autoregressive MTP draft tokens. Default: 1");
+            opt(fp, c, "--mtp-margin F", "Verifier confidence margin for fast MTP acceptance. Default: 3");
+            opt(fp, c, "--mtp-timing", "Enable embedded MTP and print acceptance/timing counters.");
+            opt(fp, c, "--dspark", "Enable DSpark using the support GGUF passed with --mtp-model.");
+            opt(fp, c, "--dspark-confidence F", "Enable DSpark with confidence pruning threshold 0..1. Greedy/opportunistic default: Metal 0.6, CUDA/ROCm 0.7; exact sampling: 0.8");
+            opt(fp, c, "--mtp-exact-sampling", "Preserve the ordinary temperature distribution instead of accepting target-matching greedy drafts directly.");
+            opt(fp, c, "--dspark-strict", "Load DSpark support but keep target-only decode.");
+        }
+        opt(fp, c, "--quality", "Prefer exact kernels where faster approximate paths exist.");
+        opt(fp, c, "--warm-weights", "Touch mapped tensor pages at startup to reduce first-use stalls.");
+        if (tool == DS4_HELP_DS4 || tool == DS4_HELP_BENCH) {
+            opt(fp, c, "--expert-profile FILE", "Metal-only: write routed expert locality/cache simulation JSON.");
+        }
+    }
+    fputc('\n', fp);
+}
+
+static void print_sampling(FILE *fp, const help_colors *c, bool full) {
+    title(fp, c, "Prompt And Sampling");
+    opt(fp, c, "-n, --tokens N", "Maximum generated tokens.");
+    opt(fp, c, "--temp F", "Sampling temperature. 0 is greedy/deterministic.");
+    opt(fp, c, "--top-p F", "Nucleus sampling probability.");
+    opt(fp, c, "--min-p F", "Keep tokens scoring at least F times the top token.");
+    opt(fp, c, "--seed N", "Sampling seed for reproducible non-greedy runs.");
+    para(fp, c, "GLM CLI and agent runs default to temperature 1.0, top-p 0.95, and min-p 0 unless those options are set explicitly.");
+    opt(fp, c, "--think", "Use normal thinking mode.");
+    opt(fp, c, "--think-max", "Use Think Max when context is large enough.");
+    opt(fp, c, "--nothink", "Disable thinking and ask for direct replies.");
+    if (full) {
+        opt(fp, c, "-sys, --system TEXT", "System prompt. Empty string disables the default where supported.");
+        opt(fp, c, "-p, --prompt TEXT", "One-shot prompt text.");
+        opt(fp, c, "--prompt-file FILE", "Read one-shot prompt text from FILE.");
+        opt(fp, c, "--raw-prompt", "Tokenize the one-shot prompt without chat markers.");
+    }
+    fputc('\n', fp);
+}
+
+static void print_steering(FILE *fp, const help_colors *c) {
+    title(fp, c, "Directional Steering");
+    opt(fp, c, "--dir-steering-file FILE", "Load one f32 direction vector per layer.");
+    opt(fp, c, "--dir-steering-ffn F", "Apply steering after FFN outputs. Default with file: 1");
+    opt(fp, c, "--dir-steering-attn F", "Apply steering after attention outputs. Default: 0");
+    fputc('\n', fp);
+}
+
+static void print_distributed(FILE *fp, const help_colors *c) {
+    title(fp, c, "Distributed Inference");
+    fputc('\n', fp);
+    para(fp, c, "Distributed mode runs one logical session across several machines by assigning contiguous model layer ranges to workers. Workers own their layer slice and KV-cache shard; the coordinator owns the prompt, sampling loop, and client/API flow. Start workers first, then start the coordinator. The coordinator waits for a complete route and streams hidden states through the workers.");
+    fputc('\n', fp);
+    opt(fp, c, "--role ROLE", "Distributed role: coordinator or worker.");
+    opt(fp, c, "--layers A:B", "Inclusive layer slice, e.g. 0:20 or 21:output.");
+    opt(fp, c, "--listen HOST PORT", "Coordinator listen address; workers may use it for their data listener.");
+    opt(fp, c, "--coordinator HOST PORT", "Coordinator address for --role worker.");
+    opt(fp, c, "--dist-prefill-chunk N", "Coordinator prefill pipeline chunk size. Default: session cap.");
+    opt(fp, c, "--dist-prefill-window N", "Max prefill chunks in flight. Default: workers+2, capped at 8.");
+    opt(fp, c, "--dist-activation-bits N", "Hidden-state transport width: 32, 16, or 8. Default: 32");
+    opt(fp, c, "--dist-replay-check", "Diagnostic: reset and replay prompt, then compare logits.");
+    opt(fp, c, "--debug", "Print coordinator route/debug logs.");
+    fputc('\n', fp);
+    title(fp, c, "Tensor Parallelism");
+    fputc('\n', fp);
+    para(fp, c, "Tensor parallelism uses the same coordinator/worker addresses as distributed mode, but always runs one 50/50 worker. Add --tensor-parallel, omit --layers, start the worker, then start the coordinator.");
+    fputc('\n', fp);
+    opt(fp, c, "--tensor-parallel", "Switch --role/--listen/--coordinator to two-machine tensor parallelism.");
+    opt(fp, c, "--transport auto|rdma|tcp", "Tensor gate transport. Default: auto");
+    opt(fp, c, "--rdma-device NAME", "Select a verbs device when auto-detection is ambiguous.");
+    opt(fp, c, "--rdma-gid-index N", "Select the local verbs GID index.");
+    opt(fp, c, "--tensor-parallel-token-prefill", "GLM diagnostic: prefill one token at a time for exact arithmetic.");
+    opt(fp, c, "--debug-hash N", "Cross-check hidden state every N tokens.");
+    fputc('\n', fp);
+}
+
+static void print_cli_diagnostics(FILE *fp, const help_colors *c);
+
+static void print_cli_specific(FILE *fp, const help_colors *c, bool full) {
+    title(fp, c, "CLI Modes");
+    opt(fp, c, "ds4", "Start the interactive prompt.");
+    opt(fp, c, "ds4 -p TEXT", "Run one prompt and exit.");
+    opt(fp, c, "ds4 --prompt-file FILE", "Run a long prompt from a file and exit.");
+    fputc('\n', fp);
+    if (full) {
+        print_cli_diagnostics(fp, c);
+    }
+}
+
+static void print_cli_diagnostics(FILE *fp, const help_colors *c) {
+    title(fp, c, "Diagnostics And Data Collection");
+    opt(fp, c, "--inspect", "Load the model and print a summary only.");
+    opt(fp, c, "--dump-tokens", "Print the exact CLI prompt token stream, then exit. Use --raw for literal text.");
+    opt(fp, c, "--dump-logits FILE", "Write full next-token logits as JSON.");
+    opt(fp, c, "--dump-logprobs FILE", "Write greedy continuation top-logprobs as JSON.");
+    opt(fp, c, "--logprobs-top-k N", "Alternatives stored by --dump-logprobs. Default: 20");
+    opt(fp, c, "--decode-consistency N", "Compare N-token decode logits with a fresh full prefill.");
+    opt(fp, c, "--expert-profile FILE", "Metal-only: write routed expert locality/cache simulation JSON.");
+    opt(fp, c, "--perplexity-file FILE", "Score raw text with teacher-forced NLL.");
+    opt(fp, c, "--imatrix-dataset FILE", "Rendered prompt dataset for imatrix collection.");
+    opt(fp, c, "--imatrix-out FILE", "Write llama-compatible routed-MoE imatrix .dat.");
+    opt(fp, c, "--imatrix-max-prompts N", "Stop imatrix collection after N prompts.");
+    opt(fp, c, "--imatrix-max-tokens N", "Stop imatrix collection after N prompt tokens.");
+    opt(fp, c, "--head-test", "Run the output HC/logits head after the native slice.");
+    opt(fp, c, "--first-token-test", "Run exact CPU whole-model pass for the first prompt token.");
+    opt(fp, c, "--metal-graph-test", "Compare first GPU-resident graph stages with CPU.");
+    opt(fp, c, "--metal-graph-full-test", "Run the GPU-resident self-token graph across all layers.");
+    opt(fp, c, "--metal-graph-prompt-test", "Compare CPU and GPU graph logits for the full prompt.");
+    fputc('\n', fp);
+}
+
+static void print_cli_commands(FILE *fp, const help_colors *c) {
+    title_red(fp, c, "Interactive Commands");
+    opt(fp, c, "/help", "Show interactive commands.");
+    opt(fp, c, "/think, /think-max, /nothink", "Switch thinking mode.");
+    opt(fp, c, "/ctx N", "Restart the interactive session with a new context size.");
+    opt(fp, c, "/power N", "Set GPU duty cycle percentage, 1..100.");
+    opt(fp, c, "/read FILE", "Submit a text file, PNG, or JPEG as the next user message.");
+    opt(fp, c, "/quit, /exit", "Leave the prompt.");
+    opt(fp, c, "Ctrl+C", "Stop current generation and return to ds4>.");
+    fputc('\n', fp);
+}
+
+static void print_agent_specific(FILE *fp, const help_colors *c) {
+    title(fp, c, "Agent Options");
+    opt(fp, c, "-p, --prompt TEXT", "Submit an initial prompt after startup.");
+    opt(fp, c, "--prompt-file FILE", "Read the initial prompt from FILE.");
+    opt(fp, c, "--non-interactive", "Run without TUI. With an initial prompt: one turn; otherwise: repeated stdin prompts.");
+    opt(fp, c, "--raw-prompt", "Non-interactive initial prompt only: omit agent chat/tool text.");
+    opt(fp, c, "--edit-upto", "Enable anchored [upto] edits and automatic marker insertion.");
+    opt(fp, c, "-sys, --system TEXT", "Extra system prompt. Empty disables extra text.");
+    opt(fp, c, "--trace FILE", "Write prompt, token, and DSML debug trace.");
+    opt(fp, c, "--chdir DIR", "Change working directory before loading runtime assets.");
+    fputc('\n', fp);
+}
+
+static void print_agent_sessions(FILE *fp, const help_colors *c) {
+    title(fp, c, "Agent Runtime Commands");
+    opt(fp, c, "/save", "Save the current session in ~/.ds4/kvcache.");
+    opt(fp, c, "/compact", "Compact the current session context now.");
+    opt(fp, c, "/list", "List saved sessions, sorted by recent update time.");
+    opt(fp, c, "/switch ID", "Load a saved session and show recent history.");
+    opt(fp, c, "/del ID", "Delete a saved session.");
+    opt(fp, c, "/strip ID", "Remove KV payload; the text history can be rebuilt later.");
+    opt(fp, c, "/history [N]", "Show N recent user turns from the current session.");
+    opt(fp, c, "/power N", "Set GPU duty cycle percentage, 1..100.");
+    opt(fp, c, "/new", "Start a fresh session from the system prompt.");
+    opt(fp, c, "/quit, /exit", "Exit.");
+    fputc('\n', fp);
+}
+
+static void print_server_api(FILE *fp, const help_colors *c) {
+    title(fp, c, "HTTP API");
+    opt(fp, c, "--host HOST", "Bind address. Default: 127.0.0.1");
+    opt(fp, c, "--port N", "Bind port. Default: 8000");
+    opt(fp, c, "--cors", "Add Access-Control-Allow-* headers for browser JS clients.");
+    opt(fp, c, "--trace FILE", "Write prompts, cache decisions, output, and tool calls.");
+    opt(fp, c, "--batched-session N", "Keep N resident sessions and batch decode-ready requests.");
+    opt(fp, c, "--mixed-prefill-quantum N", "Prefill chunk while generations are active. Default: 128; GLM-5.3 minimum: 1024");
+    para(fp, c, "Endpoints: /v1/chat/completions, /v1/responses, /v1/completions, and /v1/messages.");
+    para(fp, c, "Model endpoint aliases include deepseek-v4-flash and deepseek-v4-pro; both serve the loaded GGUF.");
+    fputc('\n', fp);
+}
+
+static void print_server_thinking(FILE *fp, const help_colors *c) {
+    title(fp, c, "Server Thinking Defaults");
+    para(fp, c, "DeepSeek-compatible chat requests default to high-effort thinking.");
+    para(fp, c, "reasoning_effort=max or output_config.effort=max requests Think Max.");
+    para(fp, c, "Think Max requires --ctx >= 393216; smaller contexts use high.");
+    para(fp, c, "thinking={type:disabled}, think=false, or model=deepseek-chat selects non-thinking mode.");
+    para(fp, c, "In thinking mode, client sampling knobs are ignored like the official API.");
+    fputc('\n', fp);
+}
+
+static void print_kv_cache(FILE *fp, const help_colors *c) {
+    title(fp, c, "Disk KV Cache");
+    opt(fp, c, "--kv-disk-dir DIR", "Enable disk KV checkpoints in DIR.");
+    opt(fp, c, "--kv-disk-space-mb N", "Disk budget. Default when enabled: 4096");
+    opt(fp, c, "--kv-cache-min-tokens N", "Do not save/load checkpoints shorter than N. Default: 512");
+    opt(fp, c, "--kv-cache-cold-max-tokens N", "Save cold first prompts up to N tokens. 0 disables. Default: 30000");
+    opt(fp, c, "--kv-cache-continued-interval-tokens N", "Save aligned continued frontiers. 0 disables. Default: 10000");
+    opt(fp, c, "--kv-cache-boundary-trim-tokens N", "Trim tail tokens for cold boundary saves. Default: 32");
+    opt(fp, c, "--kv-cache-boundary-align-tokens N", "Align cold boundary saves to this multiple. Default: 2048");
+    opt(fp, c, "--kv-cache-reject-different-quant", "Reject checkpoints written with different routed-expert quantization.");
+    opt(fp, c, "--disable-exact-dsml-tool-replay", "Disable exact sampled DSML tool replay map.");
+    opt(fp, c, "--tool-memory-max-ids N", "Exact tool-call IDs kept in RAM. Default: 100000");
+    fputc('\n', fp);
+}
+
+static void print_bench_specific(FILE *fp, const help_colors *c) {
+    title(fp, c, "Benchmark Input");
+    opt(fp, c, "--prompt-file FILE", "Raw benchmark text; token sequence is sliced at each frontier.");
+    opt(fp, c, "--chat-prompt-file FILE", "Render FILE as one no-thinking chat user message.");
+    opt(fp, c, "-sys, --system TEXT", "System prompt used only with --chat-prompt-file.");
+    fputc('\n', fp);
+    title(fp, c, "Benchmark Sweep");
+    opt(fp, c, "--ctx-start N", "First measured frontier. Default: 2048");
+    opt(fp, c, "--ctx-max N", "Last measured frontier. Default: 32768");
+    opt(fp, c, "--ctx-alloc N", "Allocated context. Default: ctx-max + gen-tokens + 1");
+    opt(fp, c, "--step-mul F", "Multiplicative step. Default: 1");
+    opt(fp, c, "--step-incr N", "Linear step when --step-mul is 1. Default: 2048");
+    opt(fp, c, "--gen-tokens N", "Greedy decode tokens per frontier. 0 for pure prefill. Default: 128");
+    opt(fp, c, "--teacher-forced-decode", "Decode the following prompt tokens instead of each predicted argmax.");
+    opt(fp, c, "--csv FILE", "Write CSV there instead of stdout.");
+    opt(fp, c, "--dump-frontier-logits-dir DIR", "Write one full-logit JSON file per frontier.");
+    fputc('\n', fp);
+}
+
+static void print_eval_specific(FILE *fp, const help_colors *c) {
+    title(fp, c, "Evaluation");
+    opt(fp, c, "-n, --tokens N", "Max generated tokens per question. Default: 16000");
+    opt(fp, c, "--questions N", "Run only the first N embedded questions.");
+    opt(fp, c, "--case-sequence LIST", "Run 1-based case numbers in this comma-separated order.");
+    opt(fp, c, "--trace FILE", "Write questions, outputs, and grading decisions.");
+    opt(fp, c, "--regrade-trace FILE", "Regrade a prior trace without loading the model.");
+    opt(fp, c, "--soft-limit-reply-budget N", "Soft close thinking near the end of reply budget. Default: 1024");
+    opt(fp, c, "--hard-limit-reply-budget N", "Force </think> with N tokens left. Default: 512");
+    opt(fp, c, "--soft-limit-think-close-rank N", "Soft-close when </think> is in top N tokens. Default: 3");
+    opt(fp, c, "--pause-ms N", "Pause after each result in the TTY UI. Default: 350");
+    opt(fp, c, "--plain", "Disable split-screen ANSI UI.");
+    opt(fp, c, "--self-test-extractors", "Run answer-extractor self-tests and exit.");
+    fputc('\n', fp);
+}
+
+static bool tool_has_topic(ds4_help_tool tool, const char *topic) {
+    if (!topic) return true;
+    if (streq(topic, "all")) return true;
+    if (streq(topic, "runtime") || streq(topic, "distributed")) return true;
+    if (streq(topic, "sampling"))
+        return tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT || tool == DS4_HELP_EVAL;
+    if (streq(topic, "steering"))
+        return tool == DS4_HELP_DS4 || tool == DS4_HELP_SERVER || tool == DS4_HELP_AGENT;
+    switch (tool) {
+    case DS4_HELP_DS4:
+        return streq(topic, "diagnostics") || streq(topic, "commands");
+    case DS4_HELP_SERVER:
+        return streq(topic, "api") || streq(topic, "kv-cache") || streq(topic, "thinking");
+    case DS4_HELP_AGENT:
+        return streq(topic, "sessions") || streq(topic, "commands") || streq(topic, "tools");
+    case DS4_HELP_BENCH:
+        return streq(topic, "benchmark");
+    case DS4_HELP_EVAL:
+        return streq(topic, "evaluation");
+    }
+    return false;
+}
+
+static void more_line(FILE *fp, const help_colors *c, const char *label, const char *topic) {
+    static const char *colors[] = {
+        "\x1b[38;5;81m", "\x1b[38;5;114m", "\x1b[38;5;179m",
+        "\x1b[38;5;141m", "\x1b[38;5;147m"
+    };
+    static size_t idx;
+    const char *on = c->cyan ? colors[idx++ % (sizeof(colors) / sizeof(colors[0]))] : "";
+    if (streq(label, "Interactive commands:") && c->red) on = c->red;
+    const char *off = c->off ? c->off : "";
+    fprintf(fp, "    %s%-26s%s --help %s\n", on, label, off, topic);
+}
+
+static void print_more_info(FILE *fp, const help_colors *c, ds4_help_tool tool) {
+    title(fp, c, "More Info");
+    more_line(fp, c, "Runtime full info:", "runtime");
+    if (tool_has_topic(tool, "sampling"))
+        more_line(fp, c, "Sampling full info:", "sampling");
+    more_line(fp, c, "Distributed inference:", "distributed");
+    if (tool_has_topic(tool, "steering"))
+        more_line(fp, c, "Steering full info:", "steering");
+    if (tool == DS4_HELP_DS4) {
+        more_line(fp, c, "Interactive commands:", "commands");
+        more_line(fp, c, "Diagnostics:", "diagnostics");
+    } else if (tool == DS4_HELP_SERVER) {
+        more_line(fp, c, "HTTP API:", "api");
+        more_line(fp, c, "Disk KV cache:", "kv-cache");
+        more_line(fp, c, "Thinking behavior:", "thinking");
+    } else if (tool == DS4_HELP_AGENT) {
+        more_line(fp, c, "Agent sessions:", "sessions");
+        more_line(fp, c, "Agent commands:", "commands");
+        more_line(fp, c, "Agent tool system:", "tools");
+    } else if (tool == DS4_HELP_BENCH) {
+        more_line(fp, c, "Benchmark sweep:", "benchmark");
+    } else if (tool == DS4_HELP_EVAL) {
+        more_line(fp, c, "Evaluation options:", "evaluation");
+    }
+    fputc('\n', fp);
+}
+
+static void print_examples(FILE *fp, const help_colors *c, ds4_help_tool tool, const char *topic) {
+    title(fp, c, "Examples");
+    if (topic_is(topic, "distributed")) {
+        opt(fp, c, "worker", "./ds4 --role worker --layers 21:output --coordinator 192.168.0.181 9000 -m ds4flash.gguf");
+        opt(fp, c, "coordinator", "./ds4 --role coordinator --layers 0:20 --listen 0.0.0.0 9000 -p \"Hello\" -m ds4flash.gguf");
+    } else if (topic_is(topic, "runtime")) {
+        if (tool == DS4_HELP_SERVER) {
+            opt(fp, c, "Metal API", "./ds4-server -m ds4flash.gguf --metal --ctx 100000");
+            opt(fp, c, "quiet API", "./ds4-server --power 60 --host 127.0.0.1 --port 8000");
+        } else if (tool == DS4_HELP_AGENT) {
+            opt(fp, c, "agent", "./ds4-agent -m ds4flash.gguf --ctx 100000");
+            opt(fp, c, "quiet agent", "./ds4-agent --power 50");
+        } else if (tool == DS4_HELP_BENCH) {
+            opt(fp, c, "bench", "./ds4-bench --prompt-file long.txt --ctx-max 32768");
+            opt(fp, c, "quiet bench", "./ds4-bench --prompt-file long.txt --power 70");
+        } else if (tool == DS4_HELP_EVAL) {
+            opt(fp, c, "eval", "./ds4-eval --questions 10 --ctx 100000");
+            opt(fp, c, "CPU debug", "./ds4-eval --cpu --questions 1 --tokens 32");
+        } else {
+            opt(fp, c, "Metal", "./ds4 -m ds4flash.gguf --metal -c 100000");
+            opt(fp, c, "quiet thermals", "./ds4 -p \"Summarize README\" --power 50");
+        }
+    } else if (topic_is(topic, "steering")) {
+        opt(fp, c, "steer FFN", "./ds4 -p \"Write tersely\" --dir-steering-file dir.bin --dir-steering-ffn 0.8");
+    } else if (tool == DS4_HELP_SERVER || topic_is(topic, "api") || topic_is(topic, "kv-cache")) {
+        opt(fp, c, "local API", "./ds4-server --ctx 100000 --kv-disk-dir ~/.ds4/server-kv --kv-disk-space-mb 8192");
+        opt(fp, c, "curl", "curl http://127.0.0.1:8000/v1/models");
+    } else if (tool == DS4_HELP_AGENT || topic_is(topic, "sessions") || topic_is(topic, "tools")) {
+        opt(fp, c, "interactive", "./ds4-agent");
+        opt(fp, c, "one shot", "./ds4-agent --non-interactive -p \"Create /tmp/hello.c\"");
+    } else if (tool == DS4_HELP_BENCH || topic_is(topic, "benchmark")) {
+        opt(fp, c, "csv", "./ds4-bench --prompt-file long.txt --ctx-max 32768 --csv speed.csv");
+        opt(fp, c, "prefill only", "./ds4-bench --prompt-file long.txt --gen-tokens 0");
+    } else if (tool == DS4_HELP_EVAL || topic_is(topic, "evaluation")) {
+        opt(fp, c, "first 10", "./ds4-eval --questions 10 --trace eval.trace");
+        opt(fp, c, "plain", "./ds4-eval --plain --nothink --tokens 512");
+    } else {
+        opt(fp, c, "chat", "./ds4");
+        opt(fp, c, "one shot", "./ds4 -p \"Explain mmap in C\"");
+        opt(fp, c, "long prompt", "./ds4 --think-max --prompt-file prompt.txt --ctx 393216");
+    }
+    fputc('\n', fp);
+}
+
+static void print_topic(FILE *fp, const help_colors *c, ds4_help_tool tool, const char *topic) {
+    if (streq(topic, "all")) {
+        print_model_runtime(fp, c, tool, true);
+        if (tool_has_topic(tool, "sampling")) print_sampling(fp, c, true);
+        if (tool_has_topic(tool, "steering")) print_steering(fp, c);
+        print_distributed(fp, c);
+        if (tool == DS4_HELP_DS4) {
+            print_cli_specific(fp, c, true);
+            print_cli_commands(fp, c);
+        } else if (tool == DS4_HELP_SERVER) {
+            print_server_api(fp, c);
+            print_server_thinking(fp, c);
+            print_kv_cache(fp, c);
+        } else if (tool == DS4_HELP_AGENT) {
+            print_agent_specific(fp, c);
+            print_agent_sessions(fp, c);
+        } else if (tool == DS4_HELP_BENCH) {
+            print_bench_specific(fp, c);
+        } else if (tool == DS4_HELP_EVAL) {
+            print_eval_specific(fp, c);
+        }
+        return;
+    }
+
+    if (streq(topic, "runtime")) print_model_runtime(fp, c, tool, true);
+    else if (streq(topic, "sampling")) print_sampling(fp, c, true);
+    else if (streq(topic, "steering")) print_steering(fp, c);
+    else if (streq(topic, "distributed")) print_distributed(fp, c);
+    else if (tool == DS4_HELP_DS4 && streq(topic, "diagnostics")) print_cli_diagnostics(fp, c);
+    else if (tool == DS4_HELP_DS4 && streq(topic, "commands")) print_cli_commands(fp, c);
+    else if (tool == DS4_HELP_SERVER && streq(topic, "api")) print_server_api(fp, c);
+    else if (tool == DS4_HELP_SERVER && streq(topic, "kv-cache")) print_kv_cache(fp, c);
+    else if (tool == DS4_HELP_SERVER && streq(topic, "thinking")) print_server_thinking(fp, c);
+    else if (tool == DS4_HELP_AGENT && streq(topic, "sessions")) print_agent_sessions(fp, c);
+    else if (tool == DS4_HELP_AGENT && streq(topic, "commands")) print_agent_sessions(fp, c);
+    else if (tool == DS4_HELP_AGENT && streq(topic, "tools")) {
+        title(fp, c, "Agent Tool System");
+        para(fp, c, "The agent can read, search, write, edit, run bash, and browse through Chrome-backed web tools.");
+        para(fp, c, "DeepSeek-family models emit DSML tool calls; GLM models use native <tool_call> syntax. Both are rendered live in the terminal.");
+        para(fp, c, "Edit uses exact old/new replacement. --edit-upto enables anchored replacements between a unique head and tail.");
+        fputc('\n', fp);
+    } else if (tool == DS4_HELP_BENCH && streq(topic, "benchmark")) print_bench_specific(fp, c);
+    else if (tool == DS4_HELP_EVAL && streq(topic, "evaluation")) print_eval_specific(fp, c);
+}
+
+static void print_default(FILE *fp, const help_colors *c, ds4_help_tool tool) {
+    print_model_runtime(fp, c, tool, false);
+
+    if (tool == DS4_HELP_DS4) {
+        print_cli_specific(fp, c, true);
+        print_sampling(fp, c, false);
+    } else if (tool == DS4_HELP_SERVER) {
+        print_server_api(fp, c);
+        print_kv_cache(fp, c);
+    } else if (tool == DS4_HELP_AGENT) {
+        print_agent_specific(fp, c);
+        print_agent_sessions(fp, c);
+    } else if (tool == DS4_HELP_BENCH) {
+        print_bench_specific(fp, c);
+    } else if (tool == DS4_HELP_EVAL) {
+        print_eval_specific(fp, c);
+    }
+}
+
+void ds4_help_print(FILE *fp, ds4_help_tool tool, const char *topic) {
+    help_colors c = help_make_colors(fp);
+    if (topic && !tool_has_topic(tool, topic)) {
+        fprintf(fp, "%s: unknown help topic '%s'\n\n", tool_name(tool), topic);
+        topic = NULL;
+    }
+
+    fprintf(fp, "%s%s%s\n", c.bright ? c.bright : "", tool_name(tool), c.off ? c.off : "");
+    fprintf(fp, "%s\n\n", tool_summary(tool));
+    fprintf(fp, "%s\n\n", tool_usage(tool));
+
+    if (topic) print_topic(fp, &c, tool, topic);
+    else {
+        print_default(fp, &c, tool);
+        print_more_info(fp, &c, tool);
+    }
+    print_examples(fp, &c, tool, topic);
+}

ds4_help.h

@@ -0,0 +1,16 @@
+#ifndef DS4_HELP_H
+#define DS4_HELP_H
+
+#include <stdio.h>
+
+typedef enum {
+    DS4_HELP_DS4,
+    DS4_HELP_SERVER,
+    DS4_HELP_AGENT,
+    DS4_HELP_BENCH,
+    DS4_HELP_EVAL,
+} ds4_help_tool;
+
+void ds4_help_print(FILE *fp, ds4_help_tool tool, const char *topic);
+
+#endif

ds4_image.c

@@ -0,0 +1,570 @@
+#include "ds4_image.h"
+
+#include <errno.h>
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#define JPEG_IMPLEMENTATION
+#define PNG_IMPLEMENTATION
+#include "third_party/iris/jpeg.h"
+#include "third_party/iris/png.h"
+
+typedef struct {
+    uint32_t state[8];
+    uint64_t bytes;
+    uint8_t block[64];
+    size_t used;
+} ds4_sha256;
+
+static const uint32_t ds4_sha256_k[64] = {
+    0x428a2f98u, 0x71374491u, 0xb5c0fbcfu, 0xe9b5dba5u,
+    0x3956c25bu, 0x59f111f1u, 0x923f82a4u, 0xab1c5ed5u,
+    0xd807aa98u, 0x12835b01u, 0x243185beu, 0x550c7dc3u,
+    0x72be5d74u, 0x80deb1feu, 0x9bdc06a7u, 0xc19bf174u,
+    0xe49b69c1u, 0xefbe4786u, 0x0fc19dc6u, 0x240ca1ccu,
+    0x2de92c6fu, 0x4a7484aau, 0x5cb0a9dcu, 0x76f988dau,
+    0x983e5152u, 0xa831c66du, 0xb00327c8u, 0xbf597fc7u,
+    0xc6e00bf3u, 0xd5a79147u, 0x06ca6351u, 0x14292967u,
+    0x27b70a85u, 0x2e1b2138u, 0x4d2c6dfcu, 0x53380d13u,
+    0x650a7354u, 0x766a0abbu, 0x81c2c92eu, 0x92722c85u,
+    0xa2bfe8a1u, 0xa81a664bu, 0xc24b8b70u, 0xc76c51a3u,
+    0xd192e819u, 0xd6990624u, 0xf40e3585u, 0x106aa070u,
+    0x19a4c116u, 0x1e376c08u, 0x2748774cu, 0x34b0bcb5u,
+    0x391c0cb3u, 0x4ed8aa4au, 0x5b9cca4fu, 0x682e6ff3u,
+    0x748f82eeu, 0x78a5636fu, 0x84c87814u, 0x8cc70208u,
+    0x90befffau, 0xa4506cebu, 0xbef9a3f7u, 0xc67178f2u,
+};
+
+static uint32_t ds4_rotr32(uint32_t value, unsigned bits) {
+    return (value >> bits) | (value << (32u - bits));
+}
+
+static void ds4_sha256_block(ds4_sha256 *sha, const uint8_t block[64]) {
+    uint32_t w[64];
+    for (unsigned i = 0; i < 16; i++) {
+        w[i] = ((uint32_t)block[i * 4] << 24) |
+               ((uint32_t)block[i * 4 + 1] << 16) |
+               ((uint32_t)block[i * 4 + 2] << 8) |
+               (uint32_t)block[i * 4 + 3];
+    }
+    for (unsigned i = 16; i < 64; i++) {
+        uint32_t s0 = ds4_rotr32(w[i - 15], 7) ^ ds4_rotr32(w[i - 15], 18) ^
+                      (w[i - 15] >> 3);
+        uint32_t s1 = ds4_rotr32(w[i - 2], 17) ^ ds4_rotr32(w[i - 2], 19) ^
+                      (w[i - 2] >> 10);
+        w[i] = w[i - 16] + s0 + w[i - 7] + s1;
+    }
+
+    uint32_t a = sha->state[0], b = sha->state[1];
+    uint32_t c = sha->state[2], d = sha->state[3];
+    uint32_t e = sha->state[4], f = sha->state[5];
+    uint32_t g = sha->state[6], h = sha->state[7];
+    for (unsigned i = 0; i < 64; i++) {
+        uint32_t s1 = ds4_rotr32(e, 6) ^ ds4_rotr32(e, 11) ^ ds4_rotr32(e, 25);
+        uint32_t ch = (e & f) ^ (~e & g);
+        uint32_t t1 = h + s1 + ch + ds4_sha256_k[i] + w[i];
+        uint32_t s0 = ds4_rotr32(a, 2) ^ ds4_rotr32(a, 13) ^ ds4_rotr32(a, 22);
+        uint32_t maj = (a & b) ^ (a & c) ^ (b & c);
+        uint32_t t2 = s0 + maj;
+        h = g; g = f; f = e; e = d + t1;
+        d = c; c = b; b = a; a = t1 + t2;
+    }
+    sha->state[0] += a; sha->state[1] += b;
+    sha->state[2] += c; sha->state[3] += d;
+    sha->state[4] += e; sha->state[5] += f;
+    sha->state[6] += g; sha->state[7] += h;
+}
+
+static void ds4_sha256_init(ds4_sha256 *sha) {
+    static const uint32_t initial[8] = {
+        0x6a09e667u, 0xbb67ae85u, 0x3c6ef372u, 0xa54ff53au,
+        0x510e527fu, 0x9b05688cu, 0x1f83d9abu, 0x5be0cd19u,
+    };
+    memcpy(sha->state, initial, sizeof(initial));
+    sha->bytes = 0;
+    sha->used = 0;
+}
+
+static void ds4_sha256_update(ds4_sha256 *sha, const void *data, size_t len) {
+    const uint8_t *p = data;
+    sha->bytes += len;
+    while (len) {
+        size_t take = 64 - sha->used;
+        if (take > len) take = len;
+        memcpy(sha->block + sha->used, p, take);
+        sha->used += take;
+        p += take;
+        len -= take;
+        if (sha->used == 64) {
+            ds4_sha256_block(sha, sha->block);
+            sha->used = 0;
+        }
+    }
+}
+
+static void ds4_sha256_final(ds4_sha256 *sha, uint8_t out[32]) {
+    uint64_t bits = sha->bytes * 8u;
+    sha->block[sha->used++] = 0x80;
+    if (sha->used > 56) {
+        memset(sha->block + sha->used, 0, 64 - sha->used);
+        ds4_sha256_block(sha, sha->block);
+        sha->used = 0;
+    }
+    memset(sha->block + sha->used, 0, 56 - sha->used);
+    for (unsigned i = 0; i < 8; i++) {
+        sha->block[63 - i] = (uint8_t)(bits >> (i * 8));
+    }
+    ds4_sha256_block(sha, sha->block);
+    for (unsigned i = 0; i < 8; i++) {
+        out[i * 4] = (uint8_t)(sha->state[i] >> 24);
+        out[i * 4 + 1] = (uint8_t)(sha->state[i] >> 16);
+        out[i * 4 + 2] = (uint8_t)(sha->state[i] >> 8);
+        out[i * 4 + 3] = (uint8_t)sha->state[i];
+    }
+}
+
+static void ds4_image_error(char *error, size_t cap, const char *message) {
+    if (!error || cap == 0) return;
+    snprintf(error, cap, "%s", message);
+}
+
+static uint16_t ds4_exif_u16(const uint8_t *p, int little) {
+    return little ? (uint16_t)(p[0] | (p[1] << 8))
+                  : (uint16_t)((p[0] << 8) | p[1]);
+}
+
+static uint32_t ds4_exif_u32(const uint8_t *p, int little) {
+    if (little) {
+        return (uint32_t)p[0] | ((uint32_t)p[1] << 8) |
+               ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
+    }
+    return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) |
+           ((uint32_t)p[2] << 8) | (uint32_t)p[3];
+}
+
+static int ds4_jpeg_orientation(const uint8_t *data, size_t len) {
+    size_t pos = 2;
+    while (pos + 4 <= len && data[pos] == 0xff) {
+        uint8_t marker = data[pos + 1];
+        pos += 2;
+        if (marker == 0xd8 || (marker >= 0xd0 && marker <= 0xd9)) continue;
+        if (pos + 2 > len) break;
+        size_t segment = ((size_t)data[pos] << 8) | data[pos + 1];
+        if (segment < 2 || segment > len - pos) break;
+        if (marker == 0xe1 && segment >= 14 &&
+            memcmp(data + pos + 2, "Exif\0\0", 6) == 0) {
+            const uint8_t *tiff = data + pos + 8;
+            size_t tiff_len = segment - 8;
+            if (tiff_len < 8) break;
+            int little = tiff[0] == 'I' && tiff[1] == 'I';
+            int big = tiff[0] == 'M' && tiff[1] == 'M';
+            if ((!little && !big) || ds4_exif_u16(tiff + 2, little) != 42) break;
+            uint32_t ifd = ds4_exif_u32(tiff + 4, little);
+            if (ifd > tiff_len || tiff_len - ifd < 2) break;
+            uint16_t count = ds4_exif_u16(tiff + ifd, little);
+            size_t entries = (size_t)count * 12;
+            if (entries > tiff_len - ifd - 2) break;
+            for (uint16_t i = 0; i < count; i++) {
+                const uint8_t *entry = tiff + ifd + 2 + (size_t)i * 12;
+                if (ds4_exif_u16(entry, little) == 0x0112 &&
+                    ds4_exif_u16(entry + 2, little) == 3 &&
+                    ds4_exif_u32(entry + 4, little) == 1) {
+                    int orientation = ds4_exif_u16(entry + 8, little);
+                    return orientation >= 1 && orientation <= 8 ? orientation : 1;
+                }
+            }
+        }
+        if (marker == 0xda) break;
+        pos += segment;
+    }
+    return 1;
+}
+
+static int ds4_oriented_rgb(
+        ds4_image *out,
+        const uint8_t *pixels,
+        uint32_t width,
+        uint32_t height,
+        uint32_t channels,
+        int orientation) {
+    uint32_t out_width = orientation >= 5 ? height : width;
+    uint32_t out_height = orientation >= 5 ? width : height;
+    size_t bytes = (size_t)out_width * out_height * 3;
+    uint8_t *rgb = malloc(bytes);
+    if (!rgb) return 0;
+
+    for (uint32_t y = 0; y < height; y++) {
+        for (uint32_t x = 0; x < width; x++) {
+            uint32_t ox = x, oy = y;
+            switch (orientation) {
+                case 2: ox = width - 1 - x; break;
+                case 3: ox = width - 1 - x; oy = height - 1 - y; break;
+                case 4: oy = height - 1 - y; break;
+                case 5: ox = y; oy = x; break;
+                case 6: ox = height - 1 - y; oy = x; break;
+                case 7: ox = height - 1 - y; oy = width - 1 - x; break;
+                case 8: ox = y; oy = width - 1 - x; break;
+                default: break;
+            }
+            const uint8_t *src = pixels + ((size_t)y * width + x) * channels;
+            uint8_t *dst = rgb + ((size_t)oy * out_width + ox) * 3;
+            if (channels < 3) {
+                dst[0] = dst[1] = dst[2] = src[0];
+            } else {
+                dst[0] = src[0]; dst[1] = src[1]; dst[2] = src[2];
+            }
+        }
+    }
+
+    out->width = out_width;
+    out->height = out_height;
+    out->rgb = rgb;
+    ds4_sha256 sha;
+    ds4_sha256_init(&sha);
+    ds4_sha256_update(&sha, &out_width, sizeof(out_width));
+    ds4_sha256_update(&sha, &out_height, sizeof(out_height));
+    ds4_sha256_update(&sha, rgb, bytes);
+    ds4_sha256_final(&sha, out->fingerprint);
+    return 1;
+}
+
+int ds4_image_decode_memory(
+        ds4_image *out,
+        const uint8_t *encoded,
+        size_t encoded_len,
+        char *error,
+        size_t error_cap) {
+    if (!out) return 0;
+    memset(out, 0, sizeof(*out));
+    if (!encoded || encoded_len == 0 || encoded_len > DS4_IMAGE_MAX_ENCODED_BYTES) {
+        ds4_image_error(error, error_cap, "image is empty or exceeds the 64 MiB encoded limit");
+        return 0;
+    }
+
+    if (encoded_len >= 8 &&
+        memcmp(encoded, "\x89PNG\r\n\x1a\n", 8) == 0) {
+        png_image *decoded = png_load_mem(encoded, encoded_len);
+        if (!decoded) {
+            ds4_image_error(error, error_cap, "invalid or unsupported PNG image");
+            return 0;
+        }
+        int ok = ds4_oriented_rgb(out, decoded->data,
+                                  (uint32_t)decoded->width,
+                                  (uint32_t)decoded->height,
+                                  (uint32_t)decoded->channels, 1);
+        png_free(decoded);
+        if (!ok) ds4_image_error(error, error_cap, "unable to allocate decoded PNG pixels");
+        return ok;
+    }
+    if (encoded_len >= 2 && encoded[0] == 0xff && encoded[1] == 0xd8) {
+        jpeg_image *decoded = jpeg_load_mem(encoded, encoded_len);
+        if (!decoded) {
+            ds4_image_error(error, error_cap, "invalid or unsupported JPEG image");
+            return 0;
+        }
+        int orientation = ds4_jpeg_orientation(encoded, encoded_len);
+        int ok = ds4_oriented_rgb(out, decoded->data,
+                                  (uint32_t)decoded->width,
+                                  (uint32_t)decoded->height,
+                                  (uint32_t)decoded->channels, orientation);
+        jpeg_free(decoded);
+        if (!ok) ds4_image_error(error, error_cap, "unable to allocate decoded JPEG pixels");
+        return ok;
+    }
+
+    ds4_image_error(error, error_cap, "image must be JPEG or PNG");
+    return 0;
+}
+
+int ds4_image_decode_file(
+        ds4_image *out,
+        const char *path,
+        char *error,
+        size_t error_cap) {
+    if (!out || !path || !path[0]) {
+        ds4_image_error(error, error_cap, "image path is empty");
+        return 0;
+    }
+    FILE *fp = fopen(path, "rb");
+    if (!fp) {
+        if (error && error_cap) snprintf(error, error_cap, "%s: %s", path, strerror(errno));
+        return 0;
+    }
+    if (fseek(fp, 0, SEEK_END) != 0) goto io_error;
+    long end = ftell(fp);
+    if (end <= 0 || (unsigned long)end > DS4_IMAGE_MAX_ENCODED_BYTES) {
+        fclose(fp);
+        ds4_image_error(error, error_cap, "image is empty or exceeds the 64 MiB encoded limit");
+        return 0;
+    }
+    if (fseek(fp, 0, SEEK_SET) != 0) goto io_error;
+    size_t len = (size_t)end;
+    uint8_t *data = malloc(len);
+    if (!data) {
+        fclose(fp);
+        ds4_image_error(error, error_cap, "unable to allocate encoded image buffer");
+        return 0;
+    }
+    if (fread(data, 1, len, fp) != len) {
+        free(data);
+        goto io_error;
+    }
+    fclose(fp);
+    int ok = ds4_image_decode_memory(out, data, len, error, error_cap);
+    free(data);
+    return ok;
+
+io_error:
+    if (error && error_cap) snprintf(error, error_cap, "%s: %s", path, strerror(errno));
+    fclose(fp);
+    return 0;
+}
+
+void ds4_image_free(ds4_image *image) {
+    if (!image) return;
+    free(image->rgb);
+    memset(image, 0, sizeof(*image));
+}
+
+static uint32_t ds4_align_u32(uint32_t value, uint32_t factor) {
+    return (value + factor - 1) / factor * factor;
+}
+
+static int ds4_glm53_smart_resize(
+        uint32_t height,
+        uint32_t width,
+        uint32_t min_tokens,
+        uint32_t max_tokens,
+        uint32_t *target_height,
+        uint32_t *target_width) {
+    const uint32_t temporal = 2, factor = 28;
+    uint64_t pixels_per_token = (uint64_t)temporal * factor * factor;
+    uint64_t min_pixels = (uint64_t)min_tokens * pixels_per_token;
+    uint64_t max_pixels = (uint64_t)max_tokens * pixels_per_token;
+    uint32_t aligned_height = ds4_align_u32(height, factor);
+    uint32_t aligned_width = ds4_align_u32(width, factor);
+    uint64_t budget = (uint64_t)temporal * aligned_height * aligned_width;
+
+    if (budget < min_pixels) {
+        double scale = sqrt((double)min_pixels /
+                            ((double)temporal * height * width));
+        aligned_height = ds4_align_u32((uint32_t)ceil(height * scale), factor);
+        aligned_width = ds4_align_u32((uint32_t)ceil(width * scale), factor);
+        budget = (uint64_t)temporal * aligned_height * aligned_width;
+    }
+    if (budget > max_pixels) {
+        if (max_pixels < (uint64_t)temporal * factor * factor) return 0;
+        uint32_t low = 1, high = height;
+        aligned_height = aligned_width = factor;
+        while (low <= high) {
+            uint32_t content_height = low + (high - low) / 2;
+            uint32_t content_width = (uint32_t)floor(
+                    (double)width * content_height / height);
+            if (content_width < 1) content_width = 1;
+            uint32_t candidate_height = ds4_align_u32(content_height, factor);
+            uint32_t candidate_width = ds4_align_u32(content_width, factor);
+            uint64_t candidate = (uint64_t)temporal * candidate_height * candidate_width;
+            if (candidate <= max_pixels) {
+                aligned_height = candidate_height;
+                aligned_width = candidate_width;
+                low = content_height + 1;
+            } else {
+                high = content_height - 1;
+            }
+        }
+    }
+    *target_height = aligned_height;
+    *target_width = aligned_width;
+    return 1;
+}
+
+static double ds4_cubic(double x, double a) {
+    x = fabs(x);
+    if (x < 1.0) return ((a + 2.0) * x - (a + 3.0)) * x * x + 1.0;
+    if (x < 2.0) return ((a * x - 5.0 * a) * x + 8.0 * a) * x - 4.0 * a;
+    return 0.0;
+}
+
+static void ds4_resize_rgb_bicubic(
+        const uint8_t *src,
+        uint32_t src_width,
+        uint32_t src_height,
+        float *dst,
+        uint32_t dst_width,
+        uint32_t dst_height,
+        uint32_t dst_stride) {
+    double scale_x = (double)src_width / dst_width;
+    double scale_y = (double)src_height / dst_height;
+    double filter_x = scale_x >= 1.0 ? 1.0 / scale_x : 1.0;
+    double filter_y = scale_y >= 1.0 ? 1.0 / scale_y : 1.0;
+    double support_x = scale_x >= 1.0 ? 2.0 * scale_x : 2.0;
+    double support_y = scale_y >= 1.0 ? 2.0 * scale_y : 2.0;
+
+    for (uint32_t dy = 0; dy < dst_height; dy++) {
+        double center_y = scale_y * ((double)dy + 0.5);
+        int y0 = (int)(center_y - support_y + 0.5);
+        int y1 = (int)(center_y + support_y + 0.5);
+        if (y0 < 0) y0 = 0;
+        if (y1 > (int)src_height) y1 = (int)src_height;
+        for (uint32_t dx = 0; dx < dst_width; dx++) {
+            double center_x = scale_x * ((double)dx + 0.5);
+            int x0 = (int)(center_x - support_x + 0.5);
+            int x1 = (int)(center_x + support_x + 0.5);
+            if (x0 < 0) x0 = 0;
+            if (x1 > (int)src_width) x1 = (int)src_width;
+            double sum[3] = {0, 0, 0};
+            double weight_sum = 0;
+            for (int iy = y0; iy < y1; iy++) {
+                double wy = ds4_cubic(((double)iy + 0.5 - center_y) * filter_y,
+                                      -0.5);
+                for (int ix = x0; ix < x1; ix++) {
+                    double wx = ds4_cubic(((double)ix + 0.5 - center_x) * filter_x,
+                                          -0.5);
+                    double weight = wx * wy;
+                    const uint8_t *pixel = src +
+                        ((size_t)iy * src_width + (uint32_t)ix) * 3;
+                    sum[0] += pixel[0] * weight;
+                    sum[1] += pixel[1] * weight;
+                    sum[2] += pixel[2] * weight;
+                    weight_sum += weight;
+                }
+            }
+            float *pixel = dst + ((size_t)dy * dst_stride + dx) * 3;
+            for (unsigned c = 0; c < 3; c++) {
+                double value = round(sum[c] / weight_sum);
+                if (value < 0.0) value = 0.0;
+                if (value > 255.0) value = 255.0;
+                pixel[c] = (float)value;
+            }
+        }
+    }
+}
+
+int ds4_image_preprocess_glm53(
+        ds4_image_patches *out,
+        const ds4_image *image,
+        uint32_t min_image_tokens,
+        uint32_t max_image_tokens,
+        char *error,
+        size_t error_cap) {
+    static const float mean[3] = {0.48145466f, 0.4578275f, 0.40821073f};
+    static const float stddev[3] = {0.26862954f, 0.26130258f, 0.27577711f};
+    if (!out) return 0;
+    memset(out, 0, sizeof(*out));
+    if (!image || !image->rgb || image->width == 0 || image->height == 0 ||
+        min_image_tokens == 0 || max_image_tokens < min_image_tokens ||
+        max_image_tokens > 8000) {
+        ds4_image_error(error, error_cap, "invalid GLM 5.3 image preprocessing parameters");
+        return 0;
+    }
+
+    uint32_t target_height, target_width;
+    if (!ds4_glm53_smart_resize(image->height, image->width,
+                                min_image_tokens, max_image_tokens,
+                                &target_height, &target_width)) {
+        ds4_image_error(error, error_cap, "image token budget is too small");
+        return 0;
+    }
+    double scale = fmin((double)target_height / image->height,
+                        (double)target_width / image->width);
+    uint64_t min_pixels = (uint64_t)2 * 28 * 28 * min_image_tokens;
+    if ((uint64_t)2 * image->height * image->width >= min_pixels && scale > 1.0) {
+        scale = 1.0;
+    }
+    uint32_t content_height = (uint32_t)floor(image->height * scale);
+    uint32_t content_width = (uint32_t)floor(image->width * scale);
+    if (content_height < 1) content_height = 1;
+    if (content_width < 1) content_width = 1;
+    if (content_height > target_height) content_height = target_height;
+    if (content_width > target_width) content_width = target_width;
+
+    size_t canvas_values = (size_t)target_height * target_width * 3;
+    float *canvas = calloc(canvas_values, sizeof(float));
+    if (!canvas) {
+        ds4_image_error(error, error_cap, "unable to allocate resized image");
+        return 0;
+    }
+    if (content_width == image->width && content_height == image->height) {
+        for (uint32_t y = 0; y < content_height; y++) {
+            for (uint32_t x = 0; x < content_width; x++) {
+                const uint8_t *src = image->rgb + ((size_t)y * image->width + x) * 3;
+                float *dst = canvas + ((size_t)y * target_width + x) * 3;
+                dst[0] = src[0]; dst[1] = src[1]; dst[2] = src[2];
+            }
+        }
+    } else {
+        ds4_resize_rgb_bicubic(image->rgb, image->width, image->height,
+                               canvas, content_width, content_height, target_width);
+    }
+
+    for (uint32_t y = 0; y < target_height; y++) {
+        for (uint32_t x = 0; x < target_width; x++) {
+            float *pixel = canvas + ((size_t)y * target_width + x) * 3;
+            for (unsigned c = 0; c < 3; c++) {
+                float value = (x < content_width && y < content_height) ? pixel[c] : 0.0f;
+                pixel[c] = (value / 255.0f - mean[c]) / stddev[c];
+            }
+        }
+    }
+
+    uint32_t grid_height = target_height / 14;
+    uint32_t grid_width = target_width / 14;
+    uint32_t patch_count = grid_height * grid_width;
+    size_t patch_values = (size_t)patch_count * 3 * 2 * 14 * 14;
+    float *patches = malloc(patch_values * sizeof(float));
+    if (!patches) {
+        free(canvas);
+        ds4_image_error(error, error_cap, "unable to allocate vision patches");
+        return 0;
+    }
+
+    size_t index = 0;
+    for (uint32_t block_y = 0; block_y < grid_height / 2; block_y++) {
+        for (uint32_t block_x = 0; block_x < grid_width / 2; block_x++) {
+            for (uint32_t merge_y = 0; merge_y < 2; merge_y++) {
+                for (uint32_t merge_x = 0; merge_x < 2; merge_x++) {
+                    uint32_t patch_y = block_y * 2 + merge_y;
+                    uint32_t patch_x = block_x * 2 + merge_x;
+                    for (uint32_t channel = 0; channel < 3; channel++) {
+                        for (uint32_t temporal = 0; temporal < 2; temporal++) {
+                            (void)temporal;
+                            for (uint32_t y = 0; y < 14; y++) {
+                                for (uint32_t x = 0; x < 14; x++) {
+                                    const float *pixel = canvas +
+                                        ((size_t)(patch_y * 14 + y) * target_width +
+                                         patch_x * 14 + x) * 3;
+                                    patches[index++] = pixel[channel];
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+    free(canvas);
+    if (index != patch_values) {
+        free(patches);
+        ds4_image_error(error, error_cap, "internal vision patch layout mismatch");
+        return 0;
+    }
+
+    out->content_width = content_width;
+    out->content_height = content_height;
+    out->padded_width = target_width;
+    out->padded_height = target_height;
+    out->grid_width = grid_width;
+    out->grid_height = grid_height;
+    out->patch_count = patch_count;
+    out->image_token_count = patch_count / 4;
+    out->patches = patches;
+    return 1;
+}
+
+void ds4_image_patches_free(ds4_image_patches *patches) {
+    if (!patches) return;
+    free(patches->patches);
+    memset(patches, 0, sizeof(*patches));
+}

ds4_image.h

@@ -0,0 +1,55 @@
+#ifndef DS4_IMAGE_H
+#define DS4_IMAGE_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+#define DS4_IMAGE_MAX_ENCODED_BYTES (64u * 1024u * 1024u)
+#define DS4_IMAGE_MAX_DIMENSION 16384u
+#define DS4_IMAGE_MAX_PIXELS (64u * 1024u * 1024u)
+
+typedef struct {
+    uint32_t width;
+    uint32_t height;
+    uint8_t *rgb;
+    uint8_t fingerprint[32];
+} ds4_image;
+
+typedef struct {
+    uint32_t content_width;
+    uint32_t content_height;
+    uint32_t padded_width;
+    uint32_t padded_height;
+    uint32_t grid_width;
+    uint32_t grid_height;
+    uint32_t patch_count;
+    uint32_t image_token_count;
+    float *patches;
+} ds4_image_patches;
+
+int ds4_image_decode_memory(
+        ds4_image     *out,
+        const uint8_t *encoded,
+        size_t         encoded_len,
+        char          *error,
+        size_t         error_cap);
+
+int ds4_image_decode_file(
+        ds4_image *out,
+        const char *path,
+        char *error,
+        size_t error_cap);
+
+void ds4_image_free(ds4_image *image);
+
+int ds4_image_preprocess_glm53(
+        ds4_image_patches *out,
+        const ds4_image   *image,
+        uint32_t           min_image_tokens,
+        uint32_t           max_image_tokens,
+        char              *error,
+        size_t             error_cap);
+
+void ds4_image_patches_free(ds4_image_patches *patches);
+
+#endif

ds4_kvstore.h

@@ -0,0 +1,217 @@
+#ifndef DS4_KVSTORE_H
+#define DS4_KVSTORE_H
+
+#include "ds4.h"
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+
+#define DS4_KVSTORE_FIXED_HEADER 48u
+#define DS4_KVSTORE_DEFAULT_MB 4096
+#define DS4_KVSTORE_HIT_HALF_LIFE_SECONDS (6ull * 60ull * 60ull)
+
+#define DS4_KVSTORE_EXT_TOOL_MAP          (1u << 0)
+#define DS4_KVSTORE_EXT_RESPONSES_VISIBLE (1u << 1)
+#define DS4_KVSTORE_EXT_THINKING_VISIBLE  (1u << 2)
+#define DS4_KVSTORE_EXT_SESSION_TITLE     (1u << 3)
+
+typedef enum {
+    DS4_KVSTORE_REASON_UNKNOWN   = 0,
+    DS4_KVSTORE_REASON_COLD      = 1,
+    DS4_KVSTORE_REASON_CONTINUED = 2,
+    DS4_KVSTORE_REASON_EVICT     = 3,
+    DS4_KVSTORE_REASON_SHUTDOWN  = 4,
+    DS4_KVSTORE_REASON_AGENT_SYSTEM  = 5,
+    DS4_KVSTORE_REASON_AGENT_SESSION = 6,
+} ds4_kvstore_reason;
+
+typedef enum {
+    DS4_KVSTORE_LOG_DEFAULT,
+    DS4_KVSTORE_LOG_KVCACHE,
+    DS4_KVSTORE_LOG_WARNING,
+} ds4_kvstore_log_type;
+
+typedef struct {
+    /* The file name is the rendered byte prefix, not the token sequence. The
+     * payload still carries the exact tokens and graph state; the hash only
+     * answers "does this checkpoint represent the bytes at the front of the
+     * incoming prompt?" */
+    char sha[41];
+    char *path;
+    uint8_t quant_bits;
+    /* Stored in header byte 7.  Flash is 0 for backward compatibility with
+     * older cache files where this reserved byte was always written as zero. */
+    uint8_t model_id;
+    uint8_t reason;
+    uint32_t tokens;
+    uint32_t hits;
+    uint32_t ctx_size;
+    uint8_t ext_flags;
+    uint64_t created_at;
+    uint64_t last_used;
+    uint64_t payload_bytes;
+    uint64_t text_bytes;
+    uint64_t file_size;
+} ds4_kvstore_entry;
+
+typedef struct {
+    int min_tokens;
+    int cold_max_tokens;
+    int continued_interval_tokens;
+    int boundary_trim_tokens;
+    int boundary_align_tokens;
+} ds4_kvstore_options;
+
+typedef struct {
+    bool enabled;
+    char *dir;
+    uint64_t budget_bytes;
+    bool reject_different_quant;
+    ds4_kvstore_options opt;
+    int continued_last_store_tokens;
+    ds4_kvstore_entry *entry;
+    int len;
+    int cap;
+    const char *log_name;
+    void *log_ud;
+    void (*log)(void *ud, ds4_kvstore_log_type type, const char *msg);
+} ds4_kvstore;
+
+typedef struct {
+    const char *text;
+    size_t text_len;
+    uint8_t model_id;
+    uint8_t quant_bits;
+    uint32_t ctx_size;
+    bool reject_different_quant;
+} ds4_kvstore_eviction_context;
+
+typedef struct {
+    void *ud;
+    uint8_t ext_flag;
+    bool (*serialized_size)(void *ud, const char *text, uint64_t *bytes_out);
+    bool (*write)(void *ud, FILE *fp, const char *text, uint64_t *written_bytes);
+    int (*load)(void *ud, FILE *fp, const void *wanted);
+    const void *load_wanted;
+} ds4_kvstore_trailer_hooks;
+
+typedef struct {
+    int tokens;
+    uint32_t text_bytes;
+    uint8_t quant_bits;
+    uint8_t ext_flags;
+    double load_ms;
+    char *path;
+} ds4_kvstore_load_result;
+
+ds4_kvstore_options ds4_kvstore_default_options(void);
+uint8_t ds4_kvstore_reason_code(const char *reason);
+const char *ds4_kvstore_key_kind(uint8_t ext_flags);
+
+bool ds4_kvstore_open(ds4_kvstore *kc, const char *dir, uint64_t budget_mb,
+                      bool reject_different_quant, ds4_kvstore_options opt,
+                      const char *log_name,
+                      void (*log)(void *ud, ds4_kvstore_log_type type, const char *msg),
+                      void *log_ud);
+void ds4_kvstore_close(ds4_kvstore *kc);
+void ds4_kvstore_clear(ds4_kvstore *kc);
+void ds4_kvstore_entry_free(ds4_kvstore_entry *e);
+
+char *ds4_kvstore_render_tokens_text(ds4_engine *engine,
+                                     const ds4_tokens *tokens,
+                                     size_t *out_len);
+bool ds4_kvstore_byte_prefix_match(const char *text, size_t text_len,
+                                   const char *prefix, size_t prefix_len);
+void ds4_kvstore_tokens_copy_prefix(ds4_tokens *dst, const ds4_tokens *src, int n);
+void ds4_kvstore_build_prompt_from_exact_prefix_and_text_suffix(
+        ds4_engine *engine,
+        const ds4_tokens *exact_prefix,
+        const char *suffix_text,
+        ds4_tokens *out);
+
+int ds4_kvstore_store_len(const ds4_kvstore *kc, int tokens);
+int ds4_kvstore_chat_anchor_pos(const ds4_kvstore *kc,
+                                const ds4_tokens *prompt,
+                                int user_token_id,
+                                int assistant_token_id);
+int ds4_kvstore_continued_store_target(const ds4_kvstore *kc, int live_tokens);
+void ds4_kvstore_note_store(ds4_kvstore *kc, int tokens);
+int ds4_kvstore_suppress_continued_store(ds4_kvstore *kc, int tokens);
+void ds4_kvstore_restore_suppressed_continued(ds4_kvstore *kc,
+                                              int old_tokens,
+                                              int suppressed_tokens);
+
+bool ds4_kvstore_file_size_fits(const ds4_kvstore *kc,
+                                uint64_t text_bytes,
+                                uint64_t payload_bytes,
+                                uint64_t trailer_bytes,
+                                uint64_t *file_bytes_out,
+                                uint64_t *required_bytes_out);
+double ds4_kvstore_entry_eviction_score(const ds4_kvstore_entry *e,
+                                        const ds4_tokens *live,
+                                        uint64_t now,
+                                        const ds4_kvstore_eviction_context *incoming);
+void ds4_kvstore_evict(ds4_kvstore *kc, const ds4_tokens *live,
+                       uint64_t extra_bytes,
+                       const ds4_kvstore_eviction_context *incoming);
+int ds4_kvstore_find_text_prefix(ds4_kvstore *kc, const char *prompt_text,
+                                 int model_id, int quant_bits, int ctx_size);
+
+bool ds4_kvstore_store_live_prefix_text(ds4_kvstore *kc,
+                                        ds4_engine *engine,
+                                        ds4_session *session,
+                                        const ds4_tokens *tokens,
+                                        int store_len,
+                                        const char *reason,
+                                        const char *cache_text_override,
+                                        uint8_t cache_text_ext,
+                                        const char *cache_text_key,
+                                        const ds4_kvstore_trailer_hooks *hooks,
+                                        char *err,
+                                        size_t err_len);
+bool ds4_kvstore_store_live_prefix(ds4_kvstore *kc,
+                                   ds4_engine *engine,
+                                   ds4_session *session,
+                                   const ds4_tokens *tokens,
+                                   int store_len,
+                                   const char *reason,
+                                   const ds4_kvstore_trailer_hooks *hooks,
+                                   char *err,
+                                   size_t err_len);
+bool ds4_kvstore_maybe_store_continued(ds4_kvstore *kc,
+                                       ds4_engine *engine,
+                                       ds4_session *session,
+                                       const ds4_kvstore_trailer_hooks *hooks,
+                                       char *err,
+                                       size_t err_len);
+int ds4_kvstore_try_load_text(ds4_kvstore *kc,
+                              ds4_engine *engine,
+                              ds4_session *session,
+                              const char *prompt_text,
+                              ds4_tokens *effective_prompt,
+                              ds4_kvstore_load_result *result,
+                              const ds4_kvstore_trailer_hooks *hooks,
+                              bool responses_protocol);
+void ds4_kvstore_load_result_free(ds4_kvstore_load_result *result);
+
+bool ds4_kvstore_read_header(FILE *fp, ds4_kvstore_entry *e,
+                             uint32_t *text_bytes);
+bool ds4_kvstore_read_entry_file(const char *path, const char sha[41],
+                                 ds4_kvstore_entry *out);
+void ds4_kvstore_fill_header(uint8_t h[DS4_KVSTORE_FIXED_HEADER],
+                             uint8_t model_id, uint8_t quant_bits,
+                             uint8_t reason, uint8_t ext_flags,
+                             uint32_t tokens, uint32_t hits, uint32_t ctx_size,
+                             uint64_t created_at, uint64_t last_used,
+                             uint64_t payload_bytes);
+bool ds4_kvstore_touch_file(const char *path, uint32_t hits);
+bool ds4_kvstore_sha_hex_name(const char *name, char sha[41]);
+void ds4_kvstore_sha1_bytes_hex(const void *ptr, size_t len, char out[41]);
+char *ds4_kvstore_path_join(const char *dir, const char *name);
+char *ds4_kvstore_path_for_sha(ds4_kvstore *kc, const char sha[41]);
+void ds4_kvstore_le_put32(uint8_t *p, uint32_t v);
+uint32_t ds4_kvstore_le_get32(const uint8_t *p);
+
+#endif

ds4_layer_pack.c

@@ -0,0 +1,149 @@
+/* Monotonic-contiguous multi-tier layer placement packer.
+ *
+ * Pure C99: no CUDA, no platform-specific code. Used by both the CUDA and
+ * Metal/CPU builds. See ds4_layer_pack.h for the API contract and the
+ * design doc docs/superpowers/specs/2026-05-26-multi-gpu-pp-v0-design.md
+ * for the rationale. */
+
+#include "ds4_layer_pack.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+int ds4_compute_layer_placement(const size_t *entry_bytes,
+                                int n_entries,
+                                const ds4_layer_pack_config *cfg,
+                                int *device_for_entry) {
+    if (!entry_bytes || !cfg || !device_for_entry) return 1;
+    if (n_entries < 0) return 2;
+    if (cfg->n_gpus < 0 || cfg->n_gpus > DS4_LAYER_PACK_MAX_GPUS) return 3;
+
+    /* Work over a local copy of budgets so the caller's struct is untouched. */
+    size_t budget[DS4_LAYER_PACK_MAX_GPUS];
+    for (int d = 0; d < cfg->n_gpus; d++) {
+        budget[d] = cfg->gpu_budget_bytes[d];
+    }
+
+    int d = 0;
+    for (int e = 0; e < n_entries; e++) {
+        /* Advance until the current device can hold this entry, or we run
+         * out of devices. Strict greater-than: exact fits stay on d. */
+        while (d < cfg->n_gpus && entry_bytes[e] > budget[d]) {
+            d++;
+        }
+        if (d < cfg->n_gpus) {
+            device_for_entry[e] = d;
+            budget[d] -= entry_bytes[e];
+        } else {
+            device_for_entry[e] = DS4_LAYER_PACK_CPU;
+            /* d stays at cfg->n_gpus so every subsequent entry also lands
+             * on the CPU tier. */
+        }
+    }
+    return 0;
+}
+
+/* Internal helper: append " + embedding" or " + output head" to a buffer if
+ * the matching pseudo-layer lives on this tier. */
+static void append_pseudo_layer_tags(char *buf, size_t buflen,
+                                      const int *device_for_entry,
+                                      int tier,
+                                      int n_entries,
+                                      int n_layers) {
+    /* entry 0 -> embedding, entry n_layers+1 -> output head */
+    if (n_entries > 0 && device_for_entry[0] == tier) {
+        strncat(buf, " + embedding", buflen - strlen(buf) - 1);
+    }
+    if (n_entries > 1 && n_layers + 1 < n_entries &&
+        device_for_entry[n_layers + 1] == tier) {
+        strncat(buf, " + output head", buflen - strlen(buf) - 1);
+    }
+}
+
+/* Emit "L-R" for a contiguous transformer-layer range assigned to a tier.
+ * Returns 1 if anything was printed (i.e. tier owned at least one layer). */
+static int print_layer_range_for_tier(FILE *out, const int *device_for_entry,
+                                      int tier, int n_entries, int n_layers) {
+    /* Transformer layers live at indices 1..n_layers in device_for_entry.
+     * Their human-facing numbers are 0..n_layers-1 (i.e. entry index - 1). */
+    int first = -1;
+    int last  = -1;
+    for (int i = 1; i <= n_layers && i < n_entries; i++) {
+        if (device_for_entry[i] == tier) {
+            if (first < 0) first = i - 1;  /* human-facing index */
+            last = i - 1;
+        }
+    }
+    if (first < 0) {
+        fprintf(out, "(no transformer layers)");
+        return 0;
+    }
+    if (first == last) {
+        fprintf(out, "layer %d", first);
+    } else {
+        fprintf(out, "layers %d-%d", first, last);
+    }
+    return 1;
+}
+
+void ds4_layer_pack_print(FILE *out,
+                          const int *device_for_entry,
+                          int n_entries,
+                          int n_layers,
+                          const size_t *entry_bytes,
+                          const size_t *gpu_used_bytes,
+                          const size_t *gpu_budget_bytes,
+                          int n_gpus) {
+    if (!out || !device_for_entry) return;
+    (void)entry_bytes; /* not needed for the textual layout */
+
+    fprintf(out, "multi-GPU layout:\n");
+    for (int d = 0; d < n_gpus; d++) {
+        fprintf(out, "  GPU%d: ", d);
+        int have_layers = print_layer_range_for_tier(out, device_for_entry,
+                                                     d, n_entries, n_layers);
+        /* Pseudo-layer tags. */
+        char tags[64];
+        tags[0] = '\0';
+        append_pseudo_layer_tags(tags, sizeof(tags), device_for_entry, d,
+                                  n_entries, n_layers);
+        if (tags[0]) {
+            if (!have_layers) {
+                /* Erase the "(no transformer layers)" stub and just write
+                 * the tag with leading "+ " stripped. */
+                /* The stub was already written; appending " + embedding"
+                 * is still informative. */
+            }
+            fprintf(out, "%s", tags);
+        }
+        /* Usage line. */
+        if (gpu_used_bytes && gpu_budget_bytes) {
+            const double used_gb   = (double)gpu_used_bytes[d]   / 1073741824.0;
+            const double budget_gb = (double)gpu_budget_bytes[d] / 1073741824.0;
+            /* Two-space pad before the parens matches the design doc. */
+            fprintf(out, "  (%.1f / %.1f GB)", used_gb, budget_gb);
+        }
+        fputc('\n', out);
+    }
+    /* CPU tier: only print if at least one entry is on CPU. */
+    int cpu_present = 0;
+    for (int i = 0; i < n_entries; i++) {
+        if (device_for_entry[i] == DS4_LAYER_PACK_CPU) { cpu_present = 1; break; }
+    }
+    if (cpu_present) {
+        fprintf(out, "  CPU : ");
+        int have_layers = print_layer_range_for_tier(out, device_for_entry,
+                                                      DS4_LAYER_PACK_CPU,
+                                                      n_entries, n_layers);
+        char tags[64];
+        tags[0] = '\0';
+        append_pseudo_layer_tags(tags, sizeof(tags), device_for_entry,
+                                  DS4_LAYER_PACK_CPU, n_entries, n_layers);
+        if (tags[0]) {
+            (void)have_layers;
+            fprintf(out, "%s", tags);
+        }
+        fputc('\n', out);
+    }
+}

ds4_layer_pack.h

@@ -0,0 +1,77 @@
+#ifndef DS4_LAYER_PACK_H
+#define DS4_LAYER_PACK_H
+
+#include <stddef.h>
+#include <stdio.h>
+
+#define DS4_LAYER_PACK_MAX_GPUS 16
+#define DS4_LAYER_PACK_CPU      (-1)
+
+typedef struct {
+    /* Per-device budget in BYTES (already net of per-device overhead such as
+     * cuBLAS workspace, global scratch, logits buffer, MTP, steering, and the
+     * configured safety margin). The packer treats these as raw capacity. */
+    size_t gpu_budget_bytes[DS4_LAYER_PACK_MAX_GPUS];
+    int    n_gpus;
+} ds4_layer_pack_config;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/* Compute a monotonic-contiguous layer placement.
+ *
+ * Inputs:
+ *   entry_bytes[]     — per-entry byte footprint, in FORWARD ORDER:
+ *                         entry 0           = embedding pseudo-layer
+ *                         entry 1..n_layers = transformer layers
+ *                         entry n_layers+1  = output-head pseudo-layer
+ *   n_entries         — total number of entries (typically n_layers + 2)
+ *   cfg               — per-device budgets
+ *
+ * Output:
+ *   device_for_entry[] — filled with target device for each entry; CPU is
+ *                        DS4_LAYER_PACK_CPU. Caller-owned, must be
+ *                        large enough for n_entries int slots.
+ *
+ * Returns:
+ *   0 on success.
+ *   Nonzero on configuration errors (null pointer, n_entries < 0,
+ *   n_gpus out of range).
+ *
+ * Per the design doc, an entry that exceeds every GPU budget spills to CPU;
+ * by the monotonicity rule every entry after it also goes to CPU. There is
+ * no error path for "entry too large" — the CPU tier is always available. */
+int ds4_compute_layer_placement(const size_t *entry_bytes,
+                                int n_entries,
+                                const ds4_layer_pack_config *cfg,
+                                int *device_for_entry);
+
+/* Print a human-readable layout summary matching the design doc's format:
+ *
+ *   multi-GPU layout:
+ *     GPU0: layers 0-21 + embedding   (38.4 / 40.0 GB)
+ *     GPU1: layers 22-31              (11.7 / 12.0 GB)
+ *     CPU : layers 32-42 + output head
+ *
+ * Entry-naming convention:
+ *   entry 0          -> "embedding"
+ *   entry i in 1..n_layers -> transformer layer (numbered i-1)
+ *   entry n_layers+1 -> "output head"
+ *
+ * gpu_used_bytes[] and gpu_budget_bytes[] may be NULL — in that case the
+ * "(used / budget)" line is omitted for the GPU lines. */
+void ds4_layer_pack_print(FILE *out,
+                          const int *device_for_entry,
+                          int n_entries,
+                          int n_layers,
+                          const size_t *entry_bytes,
+                          const size_t *gpu_used_bytes,
+                          const size_t *gpu_budget_bytes,
+                          int n_gpus);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* DS4_LAYER_PACK_H */

ds4_rocm.cu

@@ -0,0 +1,232 @@
+#ifdef __HIP_PLATFORM_AMD__
+#include "ds4_rocm.h"
+#include <hipblaslt/hipblaslt.h>
+
+#define FULL_WARP_MASK 0xFFFFFFFFFFFFFFFFULL
+#define MASK_T uint64_t
+#define DS4_GPU_BACKEND_NAME "ROCm"
+#define DS4_GPU_LOG_PREFIX "ds4: ROCm "
+#define DS4_GPU_BLAS_NAME "hipBLAS"
+#else
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <mma.h>
+#include <cublas_v2.h>
+#include <cub/block/block_radix_sort.cuh>
+
+#define FULL_WARP_MASK 0xFFFFFFFFu
+#define MASK_T uint32_t
+#define DS4_GPU_BACKEND_NAME "CUDA"
+#define DS4_GPU_LOG_PREFIX "ds4: CUDA "
+#define DS4_GPU_BLAS_NAME "cuBLAS"
+#endif
+
+#include <stdint.h>
+#include <ctype.h>
+#include <errno.h>
+#include <limits.h>
+#include <math.h>
+#include <fcntl.h>
+#include <pthread.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/stat.h>
+#include <time.h>
+#include <unistd.h>
+#include <algorithm>
+#include <unordered_map>
+#include <vector>
+
+#include "ds4_gpu.h"
+
+#ifndef M_PI
+#define M_PI 3.14159265358979323846
+#endif
+
+#define CUDA_QK_K 256
+#define DS4_ROCM_UNUSED __attribute__((unused))
+
+enum {
+    /* attention_decode_mixed_kernel stores raw-window scores plus visible
+     * compressed scores in shared memory.  The host routes larger unmasked
+     * decode calls to the online attention kernel so this fixed buffer never
+     * becomes an out-of-bounds write at long context. */
+    DS4_ROCM_ATTENTION_SCORE_CAP = 8192u,
+    DS4_ROCM_ATTENTION_RAW_SCORE_CAP = 256u
+};
+
+struct ds4_gpu_tensor {
+    void *ptr;
+    uint64_t bytes;
+    int owner;
+};
+
+typedef struct {
+    uint8_t scales[CUDA_QK_K / 16];
+    uint8_t qs[CUDA_QK_K / 4];
+    uint16_t d;
+    uint16_t dmin;
+} cuda_block_q2_K;
+
+typedef struct {
+    uint16_t d;
+    uint16_t dmin;
+    uint8_t scales[12];
+    uint8_t qs[CUDA_QK_K / 2];
+} cuda_block_q4_K;
+
+typedef struct {
+    float d;
+    int8_t qs[CUDA_QK_K];
+    int16_t bsums[CUDA_QK_K / 16];
+} cuda_block_q8_K;
+
+typedef struct {
+    uint16_t d;
+    uint16_t qs[CUDA_QK_K / 8];
+} cuda_block_iq2_xxs;
+
+typedef struct {
+    uint8_t e;
+    uint8_t qs[16];
+} cuda_block_mxfp4;
+
+static_assert(sizeof(cuda_block_mxfp4) == 17, "cuda_block_mxfp4 must match the GGUF MXFP4 block layout");
+
+/* Twice the MXFP4 values so each 32-value sub-block can use signed-int8
+ * dp4a; the factor of 1/2 is folded into the sub-block scale. */
+__device__ __constant__ static const int8_t cuda_mxfp4_values_x2[16] = {
+     0,  1,  2,  3,  4,  6,  8,  12,
+     0, -1, -2, -3, -4, -6, -8, -12,
+};
+
+#include "ds4_iq2_tables_cuda.inc"
+
+#include "rocm/ds4_rocm_runtime.cuh"
+
+#include "rocm/ds4_rocm_common.cuh"
+
+#include "rocm/ds4_rocm_q8.cuh"
+
+#include "rocm/ds4_rocm_norm_rope.cuh"
+
+#include "rocm/ds4_rocm_fp8_kv.cuh"
+
+#include "rocm/ds4_rocm_attention.cuh"
+
+#include "rocm/ds4_rocm_hc.cuh"
+
+#include "rocm/ds4_rocm_output.cuh"
+
+#include "rocm/ds4_rocm_indexer.cuh"
+
+#include "rocm/ds4_rocm_embedding_launch.cuh"
+
+#include "rocm/ds4_rocm_matmul.cuh"
+
+#include "rocm/ds4_rocm_fp8_kv_launch.cuh"
+
+#include "rocm/ds4_rocm_compressor.cuh"
+
+#include "rocm/ds4_rocm_attention_launch.cuh"
+
+#include "rocm/ds4_rocm_shared_expert.cuh"
+
+#include "rocm/ds4_rocm_misc_launch.cuh"
+#include "rocm/ds4_rocm_router.cuh"
+
+#include "rocm/ds4_rocm_moe.cuh"
+
+#include "rocm/ds4_rocm_moe_launch.cuh"
+
+#include "rocm/ds4_rocm_glm.cuh"
+
+#include "rocm/ds4_rocm_hc_output_launch.cuh"
+
+#include "rocm/ds4_rocm_current_api_compat.cuh"
+
+#include "ds4_glm53_vision_gpu.cuh"
+
+/* Tensor-parallel gates are Metal-only; stubs keep shared graph code
+ * linkable (TP option validation rejects non-Metal backends). */
+extern "C" int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate) {
+    (void)layer; (void)gate;
+    fprintf(stderr, DS4_GPU_LOG_PREFIX "tensor parallelism is Metal-only\n");
+    return 0;
+}
+
+extern "C" void ds4_gpu_tp_set_batch_exchange(ds4_gpu_tp_batch_exchange_fn fn) {
+    (void)fn;
+}
+
+extern "C" void ds4_gpu_tp_suspend_expert_sharding(int suspend) {
+    (void)suspend;
+}
+
+extern "C" void ds4_gpu_tp_keepalive_pause(int paused) {
+    (void)paused;
+}
+
+extern "C" void ds4_gpu_tp_set_attn_head_split(int enabled) {
+    (void)enabled;
+}
+
+extern "C" void ds4_gpu_model_residency_skip(int skip) {
+    (void)skip;
+}
+
+extern "C" void ds4_gpu_tp_set_big_exchange(ds4_gpu_tp_big_exchange_fn fn) {
+    (void)fn;
+}
+
+extern "C" int ds4_gpu_tp_big_gate_encode(uint32_t layer, uint32_t rows,
+                                          const ds4_gpu_tensor *out_t,
+                                          ds4_gpu_tensor *in_t,
+                                          uint64_t bytes) {
+    (void)layer; (void)rows; (void)out_t; (void)in_t; (void)bytes;
+    return 0;
+}
+
+extern "C" int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows) {
+    (void)layer; (void)rows;
+    fprintf(stderr, DS4_GPU_LOG_PREFIX "tensor parallelism is Metal-only\n");
+    return 0;
+}
+
+extern "C" int ds4_gpu_matmul_q8_0_kslice_tensor(
+        ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
+        uint64_t weight_offset, uint64_t full_in_dim, uint64_t k_off,
+        uint64_t k_cnt, uint64_t out_dim, const ds4_gpu_tensor *x,
+        uint64_t x_elem_off) {
+    (void)out; (void)model_map; (void)model_size; (void)weight_offset;
+    (void)full_in_dim; (void)k_off; (void)k_cnt; (void)out_dim; (void)x;
+    (void)x_elem_off;
+    fprintf(stderr, DS4_GPU_LOG_PREFIX "tensor parallelism is Metal-only\n");
+    return 0;
+}
+
+extern "C" int ds4_gpu_attention_output_q8_tp_tensor(
+        ds4_gpu_tensor *out, ds4_gpu_tensor *low, const void *model_map,
+        uint64_t model_size, uint64_t out_a_offset, uint64_t out_b_offset,
+        uint64_t group_dim, uint64_t rank, uint32_t n_groups_total,
+        uint32_t group0, uint32_t group_cnt, uint64_t out_dim,
+        const ds4_gpu_tensor *heads) {
+    (void)out; (void)low; (void)model_map; (void)model_size;
+    (void)out_a_offset; (void)out_b_offset; (void)group_dim; (void)rank;
+    (void)n_groups_total; (void)group0; (void)group_cnt; (void)out_dim;
+    (void)heads;
+    fprintf(stderr, DS4_GPU_LOG_PREFIX "tensor parallelism is Metal-only\n");
+    return 0;
+}
+
+extern "C" int ds4_gpu_hc_expand_add_tensor(
+        ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out,
+        const ds4_gpu_tensor *block_add, const ds4_gpu_tensor *residual_hc,
+        const ds4_gpu_tensor *post, const ds4_gpu_tensor *comb,
+        uint32_t n_embd, uint32_t n_hc) {
+    (void)out_hc; (void)block_out; (void)block_add; (void)residual_hc;
+    (void)post; (void)comb; (void)n_embd; (void)n_hc;
+    fprintf(stderr, DS4_GPU_LOG_PREFIX "tensor parallelism is Metal-only\n");
+    return 0;
+}

ds4_rocm.h

@@ -0,0 +1,142 @@
+#pragma once
+
+#include <hip/hip_runtime.h>
+#include <hipblas/hipblas.h>
+#include <hip/hip_fp16.h>
+#include <hipcub/hipcub.hpp>
+#include <rocwmma/rocwmma-version.hpp>
+#include <rocwmma/rocwmma.hpp>
+
+#define cudaError_t hipError_t
+#define cudaStream_t hipStream_t
+#define cudaEvent_t hipEvent_t
+#define cudaDeviceProp hipDeviceProp_t
+#define cudaMemLocation hipMemLocation
+
+#define cudaSuccess hipSuccess
+#define cudaErrorNotSupported hipErrorNotSupported
+#define cudaErrorInvalidValue hipErrorInvalidValue
+#define cudaGetLastError hipGetLastError
+#define cudaGetErrorString hipGetErrorString
+
+#define cudaGetDevice hipGetDevice
+#define cudaSetDevice hipSetDevice
+#define cudaDeviceSynchronize hipDeviceSynchronize
+#define cudaDeviceGetAttribute hipDeviceGetAttribute
+#define cudaGetDeviceProperties hipGetDeviceProperties
+#define cudaDevAttrPageableMemoryAccess hipDeviceAttributePageableMemoryAccess
+#define cudaDevAttrMaxSharedMemoryPerBlockOptin hipDeviceAttributeSharedMemPerBlockOptin
+#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize
+#define cudaFuncSetAttribute(func, attr, value) hipFuncSetAttribute((const void *)(func), (attr), (value))
+#define cudaMemLocationTypeDevice hipMemLocationTypeDevice
+
+#define cudaMalloc hipMalloc
+#define cudaMallocHost hipHostMalloc
+#define cudaMallocManaged hipMallocManaged
+#define cudaFree hipFree
+#define cudaFreeHost hipFreeHost
+#define cudaMemset hipMemset
+#define cudaMemcpy hipMemcpy
+#define cudaMemcpyAsync hipMemcpyAsync
+#define cudaMemcpyHostToDevice hipMemcpyHostToDevice
+#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost
+#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice
+#define cudaMemGetInfo hipMemGetInfo
+#define cudaMemsetAsync hipMemsetAsync
+
+#define cudaHostRegister hipHostRegister
+#define cudaHostUnregister hipHostUnregister
+#define cudaHostGetDevicePointer hipHostGetDevicePointer
+#define cudaHostRegisterMapped hipHostRegisterMapped
+#define cudaHostRegisterReadOnly hipHostRegisterReadOnly
+
+#define cudaMemAdvise(p1, p2, p3, p4) hipMemAdvise(p1, p2, p3, p4.id)
+#define cudaMemPrefetchAsync(devPtr, count, location, flags, stream) hipMemPrefetchAsync(devPtr, count, location.id, stream)
+#define cudaMemAdviseSetReadMostly hipMemAdviseSetReadMostly
+#define cudaMemAdviseSetPreferredLocation hipMemAdviseSetPreferredLocation
+
+#define cudaStreamCreateWithFlags hipStreamCreateWithFlags
+#define cudaStreamSynchronize hipStreamSynchronize
+#define cudaStreamDestroy hipStreamDestroy
+#define cudaStreamNonBlocking hipStreamNonBlocking
+
+#define cudaEventCreate hipEventCreate
+#define cudaEventCreateWithFlags hipEventCreateWithFlags
+#define cudaEventDestroy hipEventDestroy
+#define cudaEventRecord hipEventRecord
+#define cudaEventSynchronize hipEventSynchronize
+#define cudaEventElapsedTime hipEventElapsedTime
+#define cudaEventDisableTiming hipEventDisableTiming
+
+#define cublasHandle_t hipblasHandle_t
+#define cublasStatus_t hipblasStatus_t
+#define cublasMath_t hipblasMath_t
+
+#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS
+#define CUBLAS_OP_N HIPBLAS_OP_N
+#define CUBLAS_OP_T HIPBLAS_OP_T
+#define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT
+#define CUBLAS_DEFAULT_MATH HIPBLAS_DEFAULT_MATH
+#define CUBLAS_COMPUTE_32F HIPBLAS_COMPUTE_32F
+#define CUBLAS_TF32_TENSOR_OP_MATH HIPBLAS_TF32_TENSOR_OP_MATH
+#define CUDA_R_16F HIPBLAS_R_16F
+#define CUDA_R_32F HIPBLAS_R_32F
+
+#define cublasCreate hipblasCreate
+#define cublasDestroy hipblasDestroy
+#define cublasSetMathMode hipblasSetMathMode
+#define cublasSgemm hipblasSgemm
+#define cublasSgemmStridedBatched hipblasSgemmStridedBatched
+#define cublasGemmEx hipblasGemmEx
+#define cublasGemmStridedBatchedEx hipblasGemmStridedBatchedEx
+
+namespace cub = hipcub;
+
+static __device__ __forceinline__ int32_t __vcmpne4(uint32_t a, uint32_t b) {
+    // For each byte: 0xFF if a != b, 0x00 if a == b
+    uint32_t diff = a ^ b;
+    // Spread any set bit in each byte to fill the whole byte
+    diff |= (diff >> 1); diff |= (diff >> 2); diff |= (diff >> 4);
+    diff &= 0x01010101u;
+    diff *= 0xFFu; // 0x01 -> 0xFF per byte
+    return (int32_t)diff;
+}
+
+static __device__ __forceinline__ int32_t __vsub4(int32_t a, int32_t b) {
+    // Per-byte subtraction (wrapping, not saturating)
+    uint32_t ua = (uint32_t)a, ub = (uint32_t)b;
+    // Trick: subtract bytes in parallel avoiding cross-byte borrows
+    uint32_t diff = ((ua | 0x80808080u) - (ub & 0x7F7F7F7Fu)) ^ ((ua ^ ~ub) & 0x80808080u);
+    return (int32_t)diff;
+}
+
+// __dp4a: dot product of 4 signed int8s packed in an int32.
+// gfx11-class AMD GPUs expose this as a single v_dot4_i32_i8 instruction;
+// using the clang builtin avoids expanding every Q8/Q8_K dot into scalar byte
+// multiplies in the ROCm compatibility layer.
+static __device__ __forceinline__ int32_t __dp4a(int32_t a, int32_t b, int32_t c) {
+#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+    union ds4_i8x4_bits { int32_t i; char4 v; } av, bv;
+    av.i = a;
+    bv.i = b;
+    return amd_mixed_dot(av.v, bv.v, c, false);
+#else
+    const int8_t *a_bytes = reinterpret_cast<const int8_t*>(&a);
+    const int8_t *b_bytes = reinterpret_cast<const int8_t*>(&b);
+    return c + (int32_t)a_bytes[0] * b_bytes[0]
+             + (int32_t)a_bytes[1] * b_bytes[1]
+             + (int32_t)a_bytes[2] * b_bytes[2]
+             + (int32_t)a_bytes[3] * b_bytes[3];
+#endif
+}
+
+// Precise transcendentals for the MoE router top-k scores, immune to -fapprox-func.
+// These functions are to be used on paths where small error can be translated to
+// some macro effect - like expert selection kernels
+extern "C" __device__ __attribute__((pure))  float __ocml_exp_f32(float);
+extern "C" __device__ __attribute__((pure))  float __ocml_log1p_f32(float);
+extern "C" __device__ __attribute__((const)) float __ocml_sqrt_f32(float);
+
+static __device__ __forceinline__ float ds4_precise_expf(float x)   { return __ocml_exp_f32(x); }
+static __device__ __forceinline__ float ds4_precise_log1pf(float x) { return __ocml_log1p_f32(x); }
+static __device__ __forceinline__ float ds4_precise_sqrtf(float x)  { return __ocml_sqrt_f32(x); }

ds4_rocm_compat.cu

@@ -0,0 +1,412 @@
+#include <hip/hip_runtime.h>
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include "ds4_gpu_mgpu.h"
+#include "ds4_gpu.h"
+#include "ds4_gpu_args.h"
+
+ds4_gpu_ctx g_gpu[DS4_MAX_GPUS] = {};
+int g_n_gpus = 1;
+int g_gpu_peer_ok[DS4_MAX_GPUS][DS4_MAX_GPUS] = {{1}};
+
+extern "C" int ds4_gpu_matmul_q4_K_tensor(
+        ds4_gpu_tensor       *out,
+        const void           *model_map,
+        uint64_t              model_size,
+        uint64_t              weight_offset,
+        uint64_t              in_dim,
+        uint64_t              out_dim,
+        const ds4_gpu_tensor *x,
+        uint64_t              n_rows);
+
+static int rocm_tier_valid(int tier) {
+    return tier == 0 && g_n_gpus == 1;
+}
+
+/* Decode-island graph capture is CUDA-only for now; ROCm decodes eagerly. */
+extern "C" int ds4_gpu_decode_graphs_supported(void) { return 0; }
+extern "C" int ds4_gpu_decode_graph_begin(const ds4_decode_graph_key *key) { (void)key; return -1; }
+extern "C" int ds4_gpu_decode_graph_end(const ds4_decode_graph_key *key) { (void)key; return -1; }
+extern "C" void ds4_gpu_decode_graph_abort(const ds4_decode_graph_key *key) { (void)key; }
+extern "C" void ds4_gpu_decode_graphs_invalidate(void) {}
+
+extern "C" int ds4_gpu_init_multi(const ds4_gpu_config *cfg) {
+    if (!cfg || cfg->n_gpus != 1) {
+        fprintf(stderr, "ds4: ROCm supports one GPU per process\n");
+        return 0;
+    }
+    g_gpu[0].device_id = cfg->device_indices[0];
+    if (hipSetDevice(g_gpu[0].device_id) != hipSuccess) return 0;
+    return ds4_gpu_init();
+}
+
+extern "C" int ds4_gpu_set_current_device(int tier) {
+    if (!rocm_tier_valid(tier)) return 1;
+    return hipSetDevice(g_gpu[0].device_id) == hipSuccess ? 0 : 1;
+}
+
+extern "C" int ds4_gpu_set_current_device_fenced(int tier) {
+    return ds4_gpu_set_current_device(tier);
+}
+
+extern "C" int ds4_gpu_tensor_alloc_on(ds4_gpu_tensor *t, int tier,
+                                         uint64_t bytes) {
+    if (!t) return 1;
+    if (!rocm_tier_valid(tier)) return 2;
+    if (bytes == 0) bytes = 1;
+    if (hipMalloc(&t->ptr, (size_t)bytes) != hipSuccess) return 3;
+    t->bytes = bytes;
+    t->owner = 1;
+    t->device_id = 0;
+    return 0;
+}
+
+extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_ptr_on(int tier,
+                                                         uint64_t bytes) {
+    if (!rocm_tier_valid(tier)) return NULL;
+    return ds4_gpu_tensor_alloc(bytes);
+}
+
+extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed_on(int tier,
+                                                             uint64_t bytes) {
+    if (!rocm_tier_valid(tier)) return NULL;
+    return ds4_gpu_tensor_alloc_managed(bytes);
+}
+
+extern "C" void ds4_gpu_tensor_free_in_place(ds4_gpu_tensor *t) {
+    if (!t) return;
+    if (t->owner && t->ptr) (void)hipFree(t->ptr);
+    memset(t, 0, sizeof(*t));
+}
+
+extern "C" int ds4_gpu_tensor_device(const ds4_gpu_tensor *t) {
+    return t ? 0 : -1;
+}
+
+extern "C" int ds4_gpu_tensor_copy_async(ds4_gpu_tensor *dst,
+                                           const ds4_gpu_tensor *src,
+                                           uint64_t bytes) {
+    if (!dst || !src || bytes > dst->bytes || bytes > src->bytes) return 0;
+    if (bytes == 0) return 1;
+    return hipMemcpyAsync(dst->ptr, src->ptr, (size_t)bytes,
+                          hipMemcpyDeviceToDevice, 0) == hipSuccess;
+}
+
+extern "C" int ds4_gpu_tensor_copy_xdev(ds4_gpu_tensor *dst,
+                                          const ds4_gpu_tensor *src,
+                                          uint64_t bytes) {
+    return ds4_gpu_tensor_copy(dst, 0, src, 0, bytes);
+}
+
+extern "C" int ds4_gpu_tensor_copy_xdev_default(ds4_gpu_tensor *dst,
+                                                  const ds4_gpu_tensor *src,
+                                                  uint64_t bytes) {
+    return ds4_gpu_tensor_copy_xdev(dst, src, bytes);
+}
+
+extern "C" int ds4_gpu_tensor_copy_xdev_ordered(ds4_gpu_tensor *dst,
+                                                  const ds4_gpu_tensor *src,
+                                                  uint64_t bytes) {
+    return ds4_gpu_tensor_copy_xdev(dst, src, bytes);
+}
+
+extern "C" int ds4_gpu_tensor_copy_xdev3(
+        ds4_gpu_tensor *dst0, const ds4_gpu_tensor *src0, uint64_t bytes0,
+        ds4_gpu_tensor *dst1, const ds4_gpu_tensor *src1, uint64_t bytes1,
+        ds4_gpu_tensor *dst2, const ds4_gpu_tensor *src2, uint64_t bytes2) {
+    return (bytes0 == 0 || ds4_gpu_tensor_copy_xdev(dst0, src0, bytes0)) &&
+           (bytes1 == 0 || ds4_gpu_tensor_copy_xdev(dst1, src1, bytes1)) &&
+           (bytes2 == 0 || ds4_gpu_tensor_copy_xdev(dst2, src2, bytes2));
+}
+
+extern "C" int ds4_gpu_tensor_copy_xdev3_default_dst(
+        ds4_gpu_tensor *dst0, const ds4_gpu_tensor *src0, uint64_t bytes0,
+        ds4_gpu_tensor *dst1, const ds4_gpu_tensor *src1, uint64_t bytes1,
+        ds4_gpu_tensor *dst2, const ds4_gpu_tensor *src2, uint64_t bytes2) {
+    return ds4_gpu_tensor_copy_xdev3(dst0, src0, bytes0, dst1, src1, bytes1,
+                                     dst2, src2, bytes2);
+}
+
+extern "C" int ds4_gpu_tensor_wait_xdev(const ds4_gpu_tensor *src,
+                                          int dst_tier) {
+    return src && rocm_tier_valid(dst_tier);
+}
+
+extern "C" int ds4_gpu_tensor_wait_xdev_default(const ds4_gpu_tensor *src,
+                                                  int dst_tier) {
+    return ds4_gpu_tensor_wait_xdev(src, dst_tier);
+}
+
+extern "C" uint64_t ds4_gpu_tier_free_vram(int tier) {
+    size_t free_bytes = 0;
+    size_t total_bytes = 0;
+    if (!rocm_tier_valid(tier) ||
+        hipMemGetInfo(&free_bytes, &total_bytes) != hipSuccess) {
+        return 0;
+    }
+    return (uint64_t)free_bytes;
+}
+
+extern "C" int ds4_gpu_args_probe_auto_cuda(
+        const int *device_filter, int filter_len, ds4_gpu_config *out,
+        size_t safety_margin_bytes, char *errbuf, size_t errbuflen) {
+    if (!out) {
+        if (errbuf && errbuflen) snprintf(errbuf, errbuflen, "internal: NULL out");
+        return 1;
+    }
+    int visible = 0;
+    hipError_t rc = hipGetDeviceCount(&visible);
+    if (rc != hipSuccess || visible <= 0) {
+        if (errbuf && errbuflen) {
+            snprintf(errbuf, errbuflen, "hipGetDeviceCount failed: %s",
+                     rc == hipSuccess ? "no devices" : hipGetErrorString(rc));
+        }
+        return 1;
+    }
+    if (filter_len > 1 || (!device_filter && visible > 1)) {
+        if (errbuf && errbuflen) {
+            snprintf(errbuf, errbuflen,
+                     "ROCm supports one GPU per process; select one device");
+        }
+        return 1;
+    }
+    const int device = device_filter && filter_len == 1 ? device_filter[0] : 0;
+    if (device < 0 || device >= visible || hipSetDevice(device) != hipSuccess) {
+        if (errbuf && errbuflen) {
+            snprintf(errbuf, errbuflen, "invalid ROCm device %d", device);
+        }
+        return 1;
+    }
+    size_t free_bytes = 0;
+    size_t total_bytes = 0;
+    rc = hipMemGetInfo(&free_bytes, &total_bytes);
+    if (rc != hipSuccess) {
+        if (errbuf && errbuflen) {
+            snprintf(errbuf, errbuflen, "hipMemGetInfo failed: %s",
+                     hipGetErrorString(rc));
+        }
+        return 1;
+    }
+    const size_t reserve_floor = (size_t)2ull * 1024ull * 1024ull * 1024ull;
+    const size_t reserve_pct = free_bytes / 20u;
+    const size_t reserve = reserve_floor > reserve_pct ? reserve_floor : reserve_pct;
+    memset(out, 0, sizeof(*out));
+    out->device_indices[0] = device;
+    out->vram_bytes[0] = free_bytes > reserve ? free_bytes - reserve : 0;
+    out->n_gpus = 1;
+    out->safety_margin_bytes = safety_margin_bytes;
+    return 0;
+}
+
+extern "C" void ds4_gpu_enable_q8_dequant_gemm(void) {
+}
+
+static int g_rocm_q8_cache_suppressed = 0;
+
+extern "C" int ds4_gpu_q8_cache_suppressed(void) {
+    return g_rocm_q8_cache_suppressed;
+}
+
+extern "C" void ds4_gpu_set_q8_cache_suppressed(int suppressed) {
+    g_rocm_q8_cache_suppressed = suppressed != 0;
+}
+
+extern "C" int ds4_gpu_set_decode_fast_attention(int enabled) {
+    (void)enabled;
+    return 0;
+}
+
+extern "C" int ds4_gpu_set_decode_score_vec4(int enabled) {
+    (void)enabled;
+    return 0;
+}
+
+extern "C" int ds4_gpu_matmul_q8_0_decode_rows_exact_tensor(
+        ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
+        uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim,
+        const ds4_gpu_tensor *x, uint32_t n_rows) {
+    return ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size,
+                                      weight_offset, in_dim, out_dim, x,
+                                      n_rows);
+}
+
+extern "C" int ds4_gpu_matmul_q8_0_pair_decode_rows_exact_tensor(
+        ds4_gpu_tensor *out0, ds4_gpu_tensor *out1, const void *model_map,
+        uint64_t model_size, uint64_t weight0_offset,
+        uint64_t weight1_offset, uint64_t in_dim, uint64_t out0_dim,
+        uint64_t out1_dim, const ds4_gpu_tensor *x, uint32_t n_rows) {
+    return ds4_gpu_matmul_q8_0_pair_tensor(
+            out0, out1, model_map, model_size, weight0_offset, weight1_offset,
+            in_dim, out0_dim, out1_dim, x, n_rows);
+}
+
+extern "C" int ds4_gpu_matmul_f16_router_rows_exact_tensor(
+        ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
+        uint64_t weight_offset, const ds4_gpu_tensor *x, uint32_t n_rows) {
+    return ds4_gpu_matmul_f16_tensor(out, model_map, model_size, weight_offset,
+                                     4096u, 256u, x, n_rows);
+}
+
+extern "C" int ds4_gpu_dsv4_qkv_rms_norm_rows_kv_rope_tensor(
+        ds4_gpu_tensor *q_out, const ds4_gpu_tensor *q,
+        const void *model_map, uint64_t model_size,
+        uint64_t q_weight_offset, uint32_t q_n,
+        ds4_gpu_tensor *kv_out, const ds4_gpu_tensor *kv,
+        uint64_t kv_weight_offset, uint32_t kv_n, uint32_t rows,
+        uint32_t kv_n_head, uint32_t kv_head_dim, uint32_t n_rot,
+        uint32_t pos0, uint32_t n_ctx_orig, bool inverse,
+        float freq_base, float freq_scale, float ext_factor,
+        float attn_factor, float beta_fast, float beta_slow, float eps) {
+    return ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
+                   q_out, q, model_map, model_size, q_weight_offset, q_n,
+                   kv_out, kv, kv_weight_offset, kv_n, rows, eps) != 0 &&
+           ds4_gpu_rope_tail_tensor(
+                   kv_out, rows, kv_n_head, kv_head_dim, n_rot, pos0,
+                   n_ctx_orig, inverse, freq_base, freq_scale, ext_factor,
+                   attn_factor, beta_fast, beta_slow) != 0;
+}
+
+extern "C" int ds4_gpu_embed_token_quant_tensor(
+        ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
+        uint64_t weight_offset, uint32_t weight_type, uint32_t n_vocab,
+        uint32_t token, uint32_t n_embd) {
+    if (weight_type != 8u) return 0;
+    return ds4_gpu_embed_token_q8_0_tensor(out, model_map, model_size,
+                                           weight_offset, n_vocab, token,
+                                           n_embd);
+}
+
+extern "C" int ds4_gpu_embed_tokens_quant_tensor(
+        ds4_gpu_tensor *out, const ds4_gpu_tensor *tokens,
+        const void *model_map, uint64_t model_size, uint64_t weight_offset,
+        uint32_t weight_type, uint32_t n_vocab, uint32_t n_tokens,
+        uint32_t n_embd) {
+    if (weight_type != 8u) return 0;
+    return ds4_gpu_embed_tokens_q8_0_tensor(out, tokens, model_map, model_size,
+                                            weight_offset, n_vocab, n_tokens,
+                                            n_embd);
+}
+
+extern "C" int ds4_gpu_matmul_quant_tensor(
+        ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
+        uint64_t weight_offset, uint32_t weight_type, uint64_t in_dim,
+        uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
+    if (weight_type == 8u) {
+        return ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size,
+                                          weight_offset, in_dim, out_dim, x,
+                                          n_tok);
+    }
+    if (weight_type == 1u) {
+        return ds4_gpu_matmul_f16_tensor(out, model_map, model_size,
+                                         weight_offset, in_dim, out_dim, x,
+                                         n_tok);
+    }
+    if (weight_type == 12u) {
+        return ds4_gpu_matmul_q4_K_tensor(out, model_map, model_size,
+                                          weight_offset, in_dim, out_dim, x,
+                                          n_tok);
+    }
+    return 0;
+}
+
+extern "C" int ds4_gpu_matmul_quant_decode_mpp_model_view_tensor(
+        ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
+        uint64_t weight_offset, uint32_t weight_type, uint64_t in_dim,
+        uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
+    if (weight_type == 8u) {
+        return ds4_gpu_matmul_q8_0_decode_mpp_model_view_tensor(
+                out, model_map, model_size, weight_offset, in_dim, out_dim,
+                x, n_tok);
+    }
+    return ds4_gpu_matmul_quant_tensor(out, model_map, model_size,
+                                       weight_offset, weight_type, in_dim,
+                                       out_dim, x, n_tok);
+}
+
+extern "C" int ds4_gpu_glm_k_b_project_typed_tensor(
+        ds4_gpu_tensor *out, const ds4_gpu_tensor *kv_norm,
+        const void *model_map, uint64_t model_size, uint64_t weight_offset,
+        uint32_t weight_type, uint32_t n_tokens, uint32_t kv_lora_dim,
+        uint32_t qk_nope, uint32_t n_head) {
+    if (weight_type != 8u) return 0;
+    return ds4_gpu_glm_k_b_project_tensor(out, kv_norm, model_map, model_size,
+                                           weight_offset, n_tokens,
+                                           kv_lora_dim, qk_nope, n_head);
+}
+
+extern "C" int ds4_gpu_glm_qk_lowrank_typed_tensor(
+        ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *q,
+        const void *model_map, uint64_t model_size, uint64_t weight_offset,
+        uint32_t weight_type, uint32_t n_head, uint32_t kv_lora_dim,
+        uint32_t qk_nope, uint32_t qk_dim) {
+    if (weight_type != 8u) return 0;
+    return ds4_gpu_glm_qk_lowrank_q8_0_tensor(
+            qk_low, q, model_map, model_size, weight_offset, n_head,
+            kv_lora_dim, qk_nope, qk_dim);
+}
+
+extern "C" int ds4_gpu_glm_qk_lowrank_typed_batch_tensor(
+        ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *q,
+        const void *model_map, uint64_t model_size, uint64_t weight_offset,
+        uint32_t weight_type, uint32_t n_tokens, uint32_t n_head,
+        uint32_t kv_lora_dim, uint32_t qk_nope, uint32_t qk_dim) {
+    if (weight_type != 8u) return 0;
+    return ds4_gpu_glm_qk_lowrank_q8_0_batch_tensor(
+            qk_low, q, model_map, model_size, weight_offset, n_tokens, n_head,
+            kv_lora_dim, qk_nope, qk_dim);
+}
+
+extern "C" int ds4_gpu_glm_value_project_typed_batch_heads_tensor(
+        ds4_gpu_tensor *heads, const ds4_gpu_tensor *lora,
+        const void *model_map, uint64_t model_size, uint64_t weight_offset,
+        uint32_t weight_type, uint32_t n_tokens, uint32_t n_head,
+        uint32_t kv_lora_dim, uint32_t value_dim) {
+    if (weight_type != 8u) return 0;
+    return ds4_gpu_glm_value_project_q8_0_batch_heads_tensor(
+            heads, lora, model_map, model_size, weight_offset, n_tokens,
+            n_head, kv_lora_dim, value_dim);
+}
+
+extern "C" int ds4_gpu_glm_attention_indexed_decode_typed_tensor(
+        ds4_gpu_tensor *heads, const ds4_gpu_tensor *q,
+        const ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *kv_lora_cache,
+        const ds4_gpu_tensor *k_rope_cache, const void *model_map,
+        uint64_t model_size, uint64_t value_weight_offset,
+        uint32_t value_weight_type, const ds4_gpu_tensor *selected,
+        uint32_t n_selected, uint32_t cache_cap, bool cache_f16,
+        uint32_t n_head, uint32_t kv_lora_dim, uint32_t qk_nope,
+        uint32_t qk_rope, uint32_t value_dim, uint32_t n_ctx_orig,
+        float freq_base, float freq_scale, float ext_factor,
+        float attn_factor, float beta_fast, float beta_slow) {
+    if (value_weight_type != 8u) return 0;
+    return ds4_gpu_glm_attention_indexed_decode_tensor(
+            heads, q, qk_low, kv_lora_cache, k_rope_cache, model_map,
+            model_size, value_weight_offset, selected, n_selected, cache_cap,
+            cache_f16, n_head, kv_lora_dim, qk_nope, qk_rope, value_dim,
+            n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor,
+            beta_fast, beta_slow);
+}
+
+extern "C" int ds4_gpu_glm_attention_indexed_batch_typed_tensor(
+        ds4_gpu_tensor *heads, const ds4_gpu_tensor *q,
+        const ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *kv_lora_cache,
+        const ds4_gpu_tensor *k_rope_cache, const void *model_map,
+        uint64_t model_size, uint64_t value_weight_offset,
+        uint32_t value_weight_type, const ds4_gpu_tensor *selected,
+        uint32_t n_tokens, uint32_t n_selected, uint32_t cache_cap,
+        bool cache_f16, uint32_t n_head, uint32_t kv_lora_dim,
+        uint32_t qk_nope, uint32_t qk_rope, uint32_t value_dim,
+        uint32_t n_ctx_orig, float freq_base, float freq_scale,
+        float ext_factor, float attn_factor, float beta_fast,
+        float beta_slow) {
+    if (value_weight_type != 8u) return 0;
+    return ds4_gpu_glm_attention_indexed_batch_tensor(
+            heads, q, qk_low, kv_lora_cache, k_rope_cache, model_map,
+            model_size, value_weight_offset, selected, n_tokens, n_selected,
+            cache_cap, cache_f16, n_head, kv_lora_dim, qk_nope, qk_rope,
+            value_dim, n_ctx_orig, freq_base, freq_scale, ext_factor,
+            attn_factor, beta_fast, beta_slow);
+}

ds4_rocm_unavailable.cu

@@ -0,0 +1,39 @@
+/* Shared graph code references CUDA-only fused and multi-device hooks even
+ * when ROCm validation guarantees that those paths cannot be selected.  Keep
+ * the unavailable surface in one place.  These C-linkage stubs intentionally
+ * accept any argument list because they never inspect their arguments. */
+
+#include <stdint.h>
+
+#define ROCM_UNAVAILABLE_INT(name) extern "C" int name(...) { return 0; }
+
+ROCM_UNAVAILABLE_INT(ds4_gpu_add_xdev_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_attention_decode_rows_rope_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_attention_output_low_q4_K_slice_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_attention_output_low_q8_rows_exact_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_attention_output_q4_K_batch_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_attention_prefill_raw_heads_range_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_attention_prefill_static_mixed_heads_range_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_device_cache_support_tensors)
+ROCM_UNAVAILABLE_INT(ds4_gpu_device_cache_tensors)
+ROCM_UNAVAILABLE_INT(ds4_gpu_indexer_top1_value_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_kv_fp8_store_raw_decode_rows_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_matmul_q8_0_kslice_hc_expand_add_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_matmul_q8_0_kslice_rows_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_matmul_q8_0_top1_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_matmul_quant_kslice_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_moe_handoff_pack_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_register_model_map_no_copy)
+ROCM_UNAVAILABLE_INT(ds4_gpu_register_support_map)
+ROCM_UNAVAILABLE_INT(ds4_gpu_rope_tail_decode_rows_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_routed_moe_batch_owned_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_routed_moe_one_owned_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_routed_moe_owned_packed_combine_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_routed_moe_owned_slots_combine_rows_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_routed_moe_owned_slots_combine_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_shared_down_hc_expand_add_q8_0_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_shared_down_hc_expand_owned_q8_0_tensor)
+ROCM_UNAVAILABLE_INT(ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor)
+
+extern "C" uint64_t ds4_gpu_tp_big_gate_kick(...) { return 0; }
+ROCM_UNAVAILABLE_INT(ds4_gpu_tp_big_gate_wait)

ds4_ssd.c

@@ -0,0 +1,210 @@
+#include "ds4_ssd.h"
+
+#include <ctype.h>
+#include <errno.h>
+#include <inttypes.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/mman.h>
+#include <unistd.h>
+
+#ifndef MAP_ANONYMOUS
+#define MAP_ANONYMOUS MAP_ANON
+#endif
+
+static const uint64_t DS4_GIB = 1024ull * 1024ull * 1024ull;
+
+bool ds4_parse_gib_arg(const char *s, uint64_t *bytes) {
+    if (bytes) *bytes = 0;
+    if (!s || !s[0] || !bytes) return false;
+
+    size_t len = strlen(s);
+    if (len > 2 &&
+        (s[len - 2] == 'g' || s[len - 2] == 'G') &&
+        (s[len - 1] == 'b' || s[len - 1] == 'B')) {
+        len -= 2;
+    }
+    if (len == 0) return false;
+    for (size_t i = 0; i < len; i++) {
+        if (!isdigit((unsigned char)s[i])) return false;
+    }
+
+    char numbuf[32];
+    if (len >= sizeof(numbuf)) return false;
+    memcpy(numbuf, s, len);
+    numbuf[len] = '\0';
+
+    errno = 0;
+    unsigned long long v = strtoull(numbuf, NULL, 10);
+    if (errno != 0 || v == 0 || v > UINT64_MAX / DS4_GIB) return false;
+
+    *bytes = (uint64_t)v * DS4_GIB;
+    return true;
+}
+
+bool ds4_parse_streaming_cache_experts_arg(const char *s,
+                                           uint32_t   *experts,
+                                           uint64_t   *bytes) {
+    if (experts) *experts = 0;
+    if (bytes) *bytes = 0;
+    if (!s || !s[0] || !experts || !bytes) return false;
+
+    const size_t len = strlen(s);
+    if (len > 2 &&
+        (s[len - 2] == 'g' || s[len - 2] == 'G') &&
+        (s[len - 1] == 'b' || s[len - 1] == 'B')) {
+        return ds4_parse_gib_arg(s, bytes);
+    }
+
+    for (size_t i = 0; i < len; i++) {
+        if (!isdigit((unsigned char)s[i])) return false;
+    }
+
+    errno = 0;
+    unsigned long v = strtoul(s, NULL, 10);
+    if (errno != 0 || v == 0 || v > UINT32_MAX) return false;
+
+    *experts = (uint32_t)v;
+    return true;
+}
+
+uint32_t ds4_ssd_cache_experts_for_byte_budget(uint64_t bytes,
+                                               uint64_t per_expert_bytes) {
+    if (bytes == 0 || per_expert_bytes == 0) return 0;
+    const uint64_t experts = bytes / per_expert_bytes;
+    if (experts == 0 || experts > UINT32_MAX) return 0;
+    return (uint32_t)experts;
+}
+
+static uint64_t ds4_ssd_auto_cache_percent(void) {
+    const char *env = getenv("DS4_SSD_AUTO_CACHE_PCT");
+    if (env && env[0]) {
+        errno = 0;
+        char *end = NULL;
+        unsigned long v = strtoul(env, &end, 10);
+        if (end != env && *end == '\0' && errno == 0 &&
+            v >= 50 && v <= 95) {
+            return (uint64_t)v;
+        }
+        fprintf(stderr,
+                "ds4: invalid DS4_SSD_AUTO_CACHE_PCT=%s (want 50..95); "
+                "using default\n",
+                env);
+    }
+    /*
+     * Decode on the ROCm streaming path is SSD-bandwidth bound: every routed
+     * expert miss is a random NVMe read, so a larger resident expert cache is
+     * the biggest decode-throughput lever.  BUT the expert cache lives in the
+     * same physical RAM as the OS on a unified-memory APU, and transient
+     * spikes (pinned read/upload staging, the prefill headroom, cache growth)
+     * ride on top of the steady-state plan.  Pushing the split too high runs
+     * the machine out of RAM and trips the Linux OOM killer.  80% was measured
+     * safe here; opt into more only deliberately via DS4_SSD_AUTO_CACHE_PCT.
+     */
+    return 80;
+}
+
+bool ds4_ssd_auto_cache_plan(uint64_t            recommended_bytes,
+                             uint64_t            non_routed_bytes,
+                             uint64_t            per_expert_bytes,
+                             uint64_t            max_model_experts,
+                             ds4_ssd_cache_plan *out) {
+    if (!out) return false;
+    memset(out, 0, sizeof(*out));
+    if (recommended_bytes == 0 || per_expert_bytes == 0) return false;
+
+    const uint64_t pct = ds4_ssd_auto_cache_percent();
+    out->model_target_bytes =
+        recommended_bytes > UINT64_MAX / pct ?
+            UINT64_MAX : (recommended_bytes * pct) / 100ull;
+    if (out->model_target_bytes > non_routed_bytes) {
+        out->cache_bytes = out->model_target_bytes - non_routed_bytes;
+    }
+
+    uint64_t cache_experts = out->cache_bytes / per_expert_bytes;
+    if (cache_experts == 0) cache_experts = 1;
+    if (max_model_experts != 0 && cache_experts > max_model_experts) {
+        cache_experts = max_model_experts;
+    }
+    if (cache_experts > UINT32_MAX) cache_experts = UINT32_MAX;
+
+    out->cache_experts = (uint32_t)cache_experts;
+    out->effective_cache_bytes = cache_experts * per_expert_bytes;
+    return out->cache_experts != 0;
+}
+
+bool ds4_ssd_memory_lock_acquire(ds4_ssd_memory_lock *lock,
+                                 uint64_t             bytes) {
+    if (!lock) return false;
+    lock->ptr = NULL;
+    lock->bytes = 0;
+    if (bytes == 0) return true;
+    if (bytes > (uint64_t)SIZE_MAX) {
+        fprintf(stderr,
+                "ds4: --simulate-used-memory is too large for this process\n");
+        return false;
+    }
+
+    void *ptr = mmap(NULL,
+                     (size_t)bytes,
+                     PROT_READ | PROT_WRITE,
+                     MAP_PRIVATE | MAP_ANONYMOUS,
+                     -1,
+                     0);
+    if (ptr == MAP_FAILED) {
+        fprintf(stderr,
+                "ds4: --simulate-used-memory mmap %.2f GiB failed: %s\n",
+                (double)bytes / (double)DS4_GIB,
+                strerror(errno));
+        return false;
+    }
+
+    const long page_long = sysconf(_SC_PAGESIZE);
+    const uint64_t page = page_long > 0 ? (uint64_t)page_long : 4096ull;
+    const uint64_t chunk_bytes = 256ull * 1024ull * 1024ull;
+    volatile unsigned char *p = (volatile unsigned char *)ptr;
+
+    /*
+     * Touch and lock in bounded chunks.  A single very large mlock() is harder
+     * to diagnose when it fails and can create long uninterruptible VM work on
+     * macOS; chunking mirrors the standalone diagnostic utility.
+     */
+    uint64_t locked = 0;
+    for (uint64_t off = 0; off < bytes; off += chunk_bytes) {
+        uint64_t len = bytes - off;
+        if (len > chunk_bytes) len = chunk_bytes;
+
+        for (uint64_t pos = off; pos < off + len; pos += page) {
+            p[pos] = (unsigned char)(pos / page);
+        }
+        if (len != 0) p[off + len - 1u] = 1;
+
+        if (mlock((void *)(p + off), (size_t)len) != 0) {
+            fprintf(stderr,
+                    "ds4: --simulate-used-memory mlock failed after %.2f/%.2f GiB: %s\n",
+                    (double)locked / (double)DS4_GIB,
+                    (double)bytes / (double)DS4_GIB,
+                    strerror(errno));
+            if (locked != 0) munlock(ptr, (size_t)locked);
+            munmap(ptr, (size_t)bytes);
+            return false;
+        }
+        locked += len;
+    }
+
+    lock->ptr = ptr;
+    lock->bytes = bytes;
+    fprintf(stderr,
+            "ds4: simulated used memory: locked %.2f GiB before model load\n",
+            (double)bytes / (double)DS4_GIB);
+    return true;
+}
+
+void ds4_ssd_memory_lock_release(ds4_ssd_memory_lock *lock) {
+    if (!lock || !lock->ptr || lock->bytes == 0) return;
+    munlock(lock->ptr, (size_t)lock->bytes);
+    munmap(lock->ptr, (size_t)lock->bytes);
+    lock->ptr = NULL;
+    lock->bytes = 0;
+}

ds4_ssd.h

@@ -0,0 +1,36 @@
+#ifndef DS4_SSD_H
+#define DS4_SSD_H
+
+#include <stdbool.h>
+#include <stdint.h>
+
+typedef struct {
+    void *ptr;
+    uint64_t bytes;
+} ds4_ssd_memory_lock;
+
+typedef struct {
+    uint64_t model_target_bytes;
+    uint64_t cache_bytes;
+    uint64_t effective_cache_bytes;
+    uint32_t cache_experts;
+} ds4_ssd_cache_plan;
+
+bool ds4_parse_gib_arg(const char *s, uint64_t *bytes);
+bool ds4_parse_streaming_cache_experts_arg(const char *s,
+                                           uint32_t   *experts,
+                                           uint64_t   *bytes);
+
+uint32_t ds4_ssd_cache_experts_for_byte_budget(uint64_t bytes,
+                                               uint64_t per_expert_bytes);
+bool ds4_ssd_auto_cache_plan(uint64_t            recommended_bytes,
+                             uint64_t            non_routed_bytes,
+                             uint64_t            per_expert_bytes,
+                             uint64_t            max_model_experts,
+                             ds4_ssd_cache_plan *out);
+
+bool ds4_ssd_memory_lock_acquire(ds4_ssd_memory_lock *lock,
+                                 uint64_t             bytes);
+void ds4_ssd_memory_lock_release(ds4_ssd_memory_lock *lock);
+
+#endif

ds4_tp.h

@@ -0,0 +1,244 @@
+#ifndef DS4_TP_H
+#define DS4_TP_H
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+
+#include "ds4.h"
+
+/* Tensor-parallel transport and lockstep protocol.
+ *
+ * Two ranks run the same logical model, each with one contiguous half of the
+ * routed experts resident. Rank 0 (leader) is a normal frontend session that
+ * mirrors every ds4_session_sync()/ds4_session_eval() call to rank 1 (worker)
+ * over a TCP control socket, so both engines execute the identical graph
+ * sequence.
+ * Inside each decoded token, partial block outputs are exchanged through a
+ * registered memory slab: two-sided RDMA SEND/RECV when RDMA over
+ * Thunderbolt is available, or a full-duplex TCP exchange as fallback.
+ *
+ * Layering: ds4.c calls the session-mirroring and slab entry points;
+ * ds4_metal.m only ever sees ds4_tp_gate_exchange() through a callback
+ * registered with the GPU gate machinery.  Nothing here touches tensors.
+ */
+
+typedef struct ds4_tp ds4_tp;
+
+enum {
+    DS4_TP_GATE_ATTN = 0,
+    DS4_TP_GATE_FFN = 1,
+    DS4_TP_GATES_PER_LAYER = 2,
+    /* Max rows in a verify-block batch gate (speculative blocks are <=5). */
+    DS4_TP_BATCH_MAX_ROWS = 8,
+};
+
+/* Engine identity exchanged in the hello so a mismatched pair aborts before
+ * any inference runs. */
+typedef struct {
+    uint64_t gguf_bytes;
+    uint32_t model_id;
+    uint32_t n_layer;
+    uint32_t n_embd;
+    uint32_t n_vocab;
+    uint32_t quant_bits;
+    uint32_t ctx_size;
+    /* Decode gate schedule, used to place RDMA recvs into the right slab
+     * slot. A nonempty mask lists the exact slots in firing order; otherwise
+     * slot(seq) = start + ((seq-1) % per_token) * step.
+     * per_token 0 falls back to the identity mapping over all slots
+     * (DS4: every layer fires ATTN then FFN). GLM fires one FFN gate per
+     * sparse layer in streaming mode. Hybrid GLM-5.3 uses the mask because
+     * DSA layers also fire ATTN while KDA layers do not. Exchanged in the
+     * hello; both sides must agree. */
+    uint32_t gate_slot_start;
+    uint32_t gate_slot_step;
+    uint32_t gates_per_token;
+    uint64_t gate_slot_mask[DS4_TP_GATE_MASK_WORDS];
+} ds4_tp_identity;
+
+bool ds4_tp_enabled(const ds4_tp_options *opt);
+
+typedef enum {
+    DS4_TP_CLI_ERROR = -1,
+    DS4_TP_CLI_NOT_MATCHED = 0,
+    DS4_TP_CLI_MATCHED = 1,
+} ds4_tp_cli_parse_result;
+
+/* CLI parsing, same contract as ds4_dist_parse_cli_arg(): returns 1 when the
+ * argument was consumed, 0 when not matched, -1 on error (err filled). */
+int ds4_tp_parse_cli_arg(
+        const char *arg,
+        int *index,
+        int argc,
+        char **argv,
+        ds4_tp_options *opt,
+        char *err,
+        size_t errlen);
+int ds4_tp_adopt_distributed_options(
+        ds4_tp_options *tp,
+        ds4_distributed_options *dist,
+        char *err,
+        size_t errlen);
+void ds4_tp_usage(FILE *fp);
+
+/* Validates option combinations that TP cannot run with (SSD streaming,
+ * distributed mode, MTP drafting, CPU backend). */
+int ds4_tp_validate_engine_options(
+        const ds4_engine_options *opt,
+        char *err,
+        size_t errlen);
+
+/* Connection bring-up.  The leader listens and accepts one worker; the
+ * worker dials with retry.  Both then exchange and validate identities.
+ * Blocking; call after the engine is loaded (identity needs the shape). */
+int ds4_tp_create(
+        ds4_tp **out,
+        const ds4_tp_options *opt,
+        const ds4_tp_identity *id,
+        char *err,
+        size_t errlen);
+void ds4_tp_free(ds4_tp *tp);
+
+int ds4_tp_rank(const ds4_tp *tp);
+bool ds4_tp_is_rdma(const ds4_tp *tp);
+uint32_t ds4_tp_peer_ctx(const ds4_tp *tp);
+bool ds4_tp_failed(const ds4_tp *tp);
+void ds4_tp_mark_failed(ds4_tp *tp);
+
+/* Gate slab.  The engine allocates one shared GPU-visible block and hands
+ * its base VA here; ds4_tp registers it with the NIC (RDMA) and exchanges
+ * remote keys.  Layout, all offsets from base, S = n_layer * 2 slots:
+ *
+ *   out vectors   S * vec_bytes   written by local GPU kernels
+ *   in  vectors   S * vec_bytes   RDMA/TCP-written with the peer partials
+ *   in  seq flags S * 8           written strictly after each in vector
+ *   token slot    16              {seq u64, token i32, pad} leader->worker
+ *   (gpu flags, then batch out/in: n_layer * BATCH_MAX_ROWS * vec_bytes
+ *    each, row partials for the speculative verify-block gates)
+ *
+ * vec_bytes = n_embd * 4 (f32 partials, never quantized on the wire). */
+uint64_t ds4_tp_slab_bytes(uint32_t n_layer, uint32_t n_embd);
+uint64_t ds4_tp_slab_out_offset(const ds4_tp *tp, uint32_t layer, uint32_t gate);
+uint64_t ds4_tp_slab_in_offset(const ds4_tp *tp, uint32_t layer, uint32_t gate);
+uint64_t ds4_tp_slab_batch_out_offset(const ds4_tp *tp, uint32_t layer);
+uint64_t ds4_tp_slab_batch_in_offset(const ds4_tp *tp, uint32_t layer);
+uint64_t ds4_tp_slab_gpu_flags_offset(const ds4_tp *tp);
+int ds4_tp_attach_slab(ds4_tp *tp, void *base, char *err, size_t errlen);
+
+/* Exchange one gate: send out[layer][gate] to the peer's in[layer][gate]
+ * and wait until the peer's partial for `seq` has fully landed locally.
+ * Called from the GPU gate service thread.  Returns 0 on failure. */
+int ds4_tp_gate_exchange(ds4_tp *tp, uint32_t layer, uint32_t gate, uint64_t seq);
+
+/* Verify-block batch gate: exchange `rows` row partials for one layer in one
+ * bulk RDMA transfer, with a symmetric TCP transfer as fallback. Called from
+ * the GPU gate service thread. */
+int ds4_tp_batch_gate_exchange(ds4_tp *tp, uint32_t layer, uint32_t rows,
+                               uint64_t seq);
+
+/* Prefill batch gate: arbitrary-size symmetric payload exchange over bulk
+ * RDMA, with interleaved 2MB TCP rounds as fallback (see ds4_tp.c). */
+int ds4_tp_big_gate_exchange(ds4_tp *tp, uint32_t layer, uint64_t seq,
+                             const void *out, void *in, uint64_t bytes);
+
+/* Lockstep mirroring (leader side) and worker loop primitives. */
+typedef struct {
+    uint64_t session_id;
+    int32_t token;
+    uint32_t reserved;
+} ds4_tp_batch_item;
+
+int ds4_tp_send_session_create(ds4_tp *tp, uint64_t session_id, int ctx_size);
+int ds4_tp_send_session_destroy(ds4_tp *tp, uint64_t session_id);
+int ds4_tp_send_sync(ds4_tp *tp, uint64_t session_id,
+                     const int *tokens, uint32_t n_tokens);
+int ds4_tp_send_sync_multimodal(ds4_tp *tp, uint64_t session_id,
+                                const int *tokens, uint32_t n_tokens,
+                                const ds4_vision_span *images,
+                                uint32_t image_count);
+int ds4_tp_send_eval(ds4_tp *tp, uint64_t session_id,
+                     uint64_t seq, int token);
+int ds4_tp_send_rewind(ds4_tp *tp, uint64_t session_id, int pos);
+int ds4_tp_send_invalidate(ds4_tp *tp, uint64_t session_id);
+int ds4_tp_send_eval_batch(ds4_tp *tp, const ds4_tp_batch_item *items,
+                           uint32_t count);
+int ds4_tp_send_mixed_batch(ds4_tp *tp, uint64_t prefill_session_id,
+                            const int *prompt, uint32_t prompt_count,
+                            const ds4_tp_batch_item *items,
+                            uint32_t count);
+int ds4_tp_send_command_ack(ds4_tp *tp, uint64_t session_id, int status);
+int ds4_tp_wait_command_ack(ds4_tp *tp, uint64_t session_id,
+                            const char *operation, char *err, size_t errlen);
+int ds4_tp_send_stop(ds4_tp *tp);
+
+/* Worker: blocks for the next mirrored command.  Frame types below; for
+ * DS4_TP_FRAME_SYNC the token array is returned in *tokens / *n_tokens
+ * (malloc'd, caller frees), for DS4_TP_FRAME_EVAL seq/token are filled. */
+typedef enum {
+    DS4_TP_FRAME_ERROR = -1,
+    DS4_TP_FRAME_SYNC = 1,
+    DS4_TP_FRAME_EVAL = 2,
+    DS4_TP_FRAME_REWIND = 3,
+    DS4_TP_FRAME_INVALIDATE = 4,
+    DS4_TP_FRAME_STOP = 5,
+    DS4_TP_FRAME_HASH = 6,
+    DS4_TP_FRAME_RDMA_INFO = 7,
+    DS4_TP_FRAME_SYNC_ACK = 8,
+    DS4_TP_FRAME_RDMA_READY = 9,
+    DS4_TP_FRAME_LOGITS = 10,
+    DS4_TP_FRAME_VERIFY = 11,
+    DS4_TP_FRAME_VERIFY_COMMIT = 12,
+    DS4_TP_FRAME_SESSION_CREATE = 13,
+    DS4_TP_FRAME_SESSION_DESTROY = 14,
+    DS4_TP_FRAME_EVAL_BATCH = 15,
+    DS4_TP_FRAME_MIXED_BATCH = 16,
+    DS4_TP_FRAME_COMMAND_ACK = 17,
+    DS4_TP_FRAME_SYNC_MULTIMODAL = 18,
+} ds4_tp_frame_type;
+
+typedef struct {
+    ds4_tp_frame_type type;
+    uint64_t session_id;
+    uint64_t seq;
+    int value;
+    int *tokens;
+    uint32_t n_tokens;
+    ds4_tp_batch_item *items;
+    uint32_t n_items;
+    ds4_vision_span *images;
+    uint32_t n_images;
+} ds4_tp_command;
+
+int ds4_tp_recv_command(
+        ds4_tp *tp,
+        ds4_tp_command *command,
+        char *err,
+        size_t errlen);
+void ds4_tp_command_free(ds4_tp_command *command);
+
+/* Debug lockstep check: both sides send their hidden-state hash for a token
+ * and compare.  Returns 0 on transport failure, -1 on hash mismatch. */
+int ds4_tp_hash_check(ds4_tp *tp, uint64_t seq, uint64_t hash, char *err, size_t errlen);
+
+/* Vocab-split output head: the worker ships its logits half to the leader
+ * after every eval (and after a sync) on the control socket. */
+int ds4_tp_send_logits_half(ds4_tp *tp, const float *half, uint32_t count);
+int ds4_tp_recv_logits_half(ds4_tp *tp, float *half, uint32_t count);
+
+/* Speculative verify mirroring.  The leader announces a draft block right
+ * before both ranks run the expert-split batch verify; the worker then blocks
+ * on the commit frame, which carries the leader's decision: full_accept keeps
+ * the pushed rows, otherwise both sides roll back and replay replay_n tokens
+ * through the gated single-token decode in lockstep. */
+int ds4_tp_send_verify(ds4_tp *tp, uint64_t session_id,
+                       const int *drafts, uint32_t n);
+int ds4_tp_send_verify_commit(ds4_tp *tp, int32_t full_accept, int32_t replay_n);
+int ds4_tp_recv_verify_commit(ds4_tp *tp, int32_t *full_accept, int32_t *replay_n);
+
+/* Standalone worker mode entry. Loads nothing itself: the engine is already
+ * open. */
+int ds4_tp_worker_run(ds4_engine *engine, const ds4_tp_options *opt);
+
+#endif

ds4_web.h

@@ -0,0 +1,33 @@
+#ifndef DS4_WEB_H
+#define DS4_WEB_H
+
+#include <stddef.h>
+#include <stdbool.h>
+
+typedef int (*ds4_web_confirm_fn)(void *privdata, const char *message,
+                                  char *err, size_t err_len);
+typedef void (*ds4_web_log_fn)(void *privdata, const char *message);
+typedef bool (*ds4_web_cancel_fn)(void *privdata);
+
+typedef struct {
+    const char *home_dir;
+    int port;
+    ds4_web_confirm_fn confirm;
+    void *confirm_privdata;
+    ds4_web_log_fn log;
+    void *log_privdata;
+    ds4_web_cancel_fn cancel;
+    void *cancel_privdata;
+} ds4_web_config;
+
+typedef struct ds4_web ds4_web;
+
+ds4_web *ds4_web_create(const ds4_web_config *cfg);
+void ds4_web_free(ds4_web *web);
+
+char *ds4_web_google_search(ds4_web *web, const char *query,
+                            char *err, size_t err_len);
+char *ds4_web_visit_page(ds4_web *web, const char *url,
+                         char *err, size_t err_len);
+
+#endif

gguf-tools/.gitignore

@@ -1,3 +1,14 @@
 deepseek4-quantize
+gguf-requantize-dense
 quality-testing/score_official
-quality-testing/data/
+quality-testing/score_llama
+quality-testing/data/*
+!quality-testing/data/README.md
+!quality-testing/data/glm52-openrouter-100/
+!quality-testing/data/glm52-openrouter-100/**
+!quality-testing/data/pro/
+!quality-testing/data/pro/**
+!quality-testing/data/flash/
+!quality-testing/data/flash/**
+!quality-testing/data/glm53-flash-openrouter-zai-fp8-100/
+!quality-testing/data/glm53-flash-openrouter-zai-fp8-100/**

gguf-tools/Makefile

@@ -1,15 +1,29 @@
 CC ?= cc
+CXX ?= c++
 UNAME_S := $(shell uname -s)
+LLAMA_CPP_DIR ?= ../../deepseek-v4-quants/llama.cpp
+LLAMA_CPP_BUILD ?= $(LLAMA_CPP_DIR)/build-remote
+LLAMA_CPP_BIN ?= $(LLAMA_CPP_BUILD)/bin
+LLAMA_CPP_CXXFLAGS ?= -O3 -Wall -Wextra -std=c++17 \
+	-I$(LLAMA_CPP_DIR)/include -I$(LLAMA_CPP_DIR)/ggml/include
 
 ifeq ($(UNAME_S),Darwin)
 NATIVE_CPU_FLAG ?= -mcpu=native
+QUANTS_SHARED := libds4quants.dylib
+SHARED_FLAGS := -dynamiclib
 QUALITY_LDLIBS := -lm -pthread -framework Foundation -framework Metal
-QUALITY_TARGETS := ds4.o ds4_metal.o
-QUALITY_OBJS := ../ds4.o ../ds4_metal.o
+QUALITY_TARGETS := ds4.o ds4_metal.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o ds4_image.o rax.o ds4_gpu_args.o
+QUALITY_OBJS := ../ds4.o ../ds4_metal.o ../ds4_distributed.o ../ds4_tp.o ../ds4_ssd.o ../ds4_layer_pack.o ../ds4_image.o ../rax.o ../ds4_gpu_args.o
 QUALITY_LINK := $(CC) $(CFLAGS) -I.. -o
+LLAMA_CPP_LDLIBS ?= -L$(LLAMA_CPP_BIN) -Wl,-rpath,$(abspath $(LLAMA_CPP_BIN)) \
+	-lllama -lggml -lggml-base -lggml-cpu -lggml-blas -lggml-metal \
+	-framework Foundation -framework Metal -framework Accelerate
 else
 NATIVE_CPU_FLAG ?= -march=native
-CUDA_HOME ?= /usr/local/cuda
+QUANTS_SHARED := libds4quants.so
+SHARED_FLAGS := -shared
+CUDA_NVCC := $(shell command -v nvcc 2>/dev/null)
+CUDA_HOME ?= $(if $(wildcard /usr/local/cuda/bin/nvcc),/usr/local/cuda,$(if $(CUDA_NVCC),$(patsubst %/bin/nvcc,%,$(CUDA_NVCC)),/usr/local/cuda))
 NVCC ?= $(CUDA_HOME)/bin/nvcc
 CUDA_ARCH ?=
 ifneq ($(strip $(CUDA_ARCH)),)
@@ -17,24 +31,37 @@ NVCC_ARCH_FLAGS := -arch=$(CUDA_ARCH)
 endif
 NVCCFLAGS ?= -O3 --use_fast_math $(NVCC_ARCH_FLAGS) -Xcompiler $(NATIVE_CPU_FLAG) -Xcompiler -pthread
 QUALITY_LDLIBS ?= -lm -Xcompiler -pthread -L$(CUDA_HOME)/targets/sbsa-linux/lib -L$(CUDA_HOME)/lib64 -lcudart -lcublas
-QUALITY_TARGETS := ds4.o ds4_cuda.o
-QUALITY_OBJS := ../ds4.o ../ds4_cuda.o
+QUALITY_TARGETS := ds4.o ds4_cuda.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o ds4_image.o rax.o ds4_gpu_args.o \
+	cuda/mmq/ds4_ggml_stubs.o cuda/mmq/ds4_mmq.o cuda/mmq/ds4_mmq_d2r.o \
+	cuda/mmq/quantize.o cuda/mmq/mmid.o cuda/mmq/mmvq.o cuda/mmq/ds4_repack.o
+QUALITY_OBJS := ../ds4.o ../ds4_cuda.o ../ds4_distributed.o ../ds4_tp.o ../ds4_ssd.o ../ds4_layer_pack.o ../ds4_image.o ../rax.o ../ds4_gpu_args.o \
+	../cuda/mmq/ds4_ggml_stubs.o ../cuda/mmq/ds4_mmq.o ../cuda/mmq/ds4_mmq_d2r.o \
+	../cuda/mmq/quantize.o ../cuda/mmq/mmid.o ../cuda/mmq/mmvq.o ../cuda/mmq/ds4_repack.o
 QUALITY_LINK := $(NVCC) $(NVCCFLAGS) -I.. -o
+LLAMA_CPP_LDLIBS ?= -L$(LLAMA_CPP_BIN) -Wl,-rpath,$(abspath $(LLAMA_CPP_BIN)) \
+	-lllama -lggml -lggml-base -lggml-cpu -ldl -lpthread
 endif
 
 CFLAGS ?= -O3 -Wall -Wextra -std=c11 $(NATIVE_CPU_FLAG)
 CPPFLAGS ?= -D_GNU_SOURCE
 
-.PHONY: all clean quality-score
+.PHONY: all clean quality-score quality-llama-score
 
-all: deepseek4-quantize
+all: deepseek4-quantize $(QUANTS_SHARED)
 
 deepseek4-quantize: deepseek4-quantize.c quants.c quants.h
 	$(CC) $(CFLAGS) $(CPPFLAGS) -o $@ deepseek4-quantize.c quants.c -lm -pthread
 
+$(QUANTS_SHARED): quants.c quants.h
+	$(CC) $(CFLAGS) $(CPPFLAGS) $(SHARED_FLAGS) -fPIC -o $@ quants.c -lm -pthread
+
 quality-score:
 	$(MAKE) -C .. $(QUALITY_TARGETS)
 	$(QUALITY_LINK) quality-testing/score_official quality-testing/score_official.c $(QUALITY_OBJS) $(QUALITY_LDLIBS)
 
+quality-llama-score:
+	$(CXX) $(LLAMA_CPP_CXXFLAGS) -o quality-testing/score_llama quality-testing/score_llama.cpp $(LLAMA_CPP_LDLIBS)
+
 clean:
-	rm -f deepseek4-quantize quality-testing/score_official
+	rm -f deepseek4-quantize libds4quants.dylib libds4quants.so \
+		quality-testing/score_official quality-testing/score_llama

gguf-tools/README.md

@@ -8,7 +8,7 @@ The important pieces are:
 - `deepseek4-quantize.c`: C HF-safetensors to GGUF quantizer.
 - `quants.[ch]`: the deliberately small local quantization implementation used
   by the quantizer.  It implements the DS4 output formats we actually ship:
-  `q8_0`, `q4_K`, `q2_K`, and `iq2_xxs`.
+  `q8_0`, `q8_K`, `q4_K`, `q2_K`, and `iq2_xxs`.
 - `imatrix/`: dataset and instructions for collecting routed-MoE activation
   importance with `ds4`.
 - `quality-testing/`: prompts and scripts used to compare local GGUF variants
@@ -86,6 +86,17 @@ gguf-tools/deepseek4-quantize \
   --imatrix gguf/DeepSeek-V4-Flash-chat-v2-routed-moe-ds4.dat
 ```
 
+True Q8_K routed experts:
+
+```sh
+gguf-tools/deepseek4-quantize \
+  --hf ../deepseek-v4-quants/hf/DeepSeek-V4-Flash \
+  --template gguf/DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-imatrix.gguf \
+  --out gguf/DeepSeek-V4-Flash-Q8KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2.gguf \
+  --experts q8_K \
+  --threads 8
+```
+
 You can override tensor families:
 
 ```sh
@@ -108,6 +119,54 @@ gguf-tools/deepseek4-quantize \
 `--compare-tensor` regenerates a single tensor and byte-compares it against the
 template or `--compare-gguf`.  `--threads N` controls routed-expert workers.
 
+## Convert A DSpark Support Checkpoint
+
+The DSpark Flash checkpoint is published as Hugging Face safetensors and stores
+the draft module under `mtp.0`, `mtp.1`, and `mtp.2`.  Before writing a support
+GGUF, inspect the official index and verify that every DSpark tensor name is
+understood by the converter:
+
+```sh
+gguf-tools/deepseek4-quantize \
+  --hf ../deepseek-v4-quants/hf/DeepSeek-V4-Flash-0731 \
+  --dspark-manifest > /tmp/dspark-manifest.tsv
+```
+
+The manifest reads only `model.safetensors.index.json`; it does not require the
+large shard files to be present.  The final summary should report three DSpark
+stages and zero unknown DSpark tensors before attempting a full conversion.
+
+To build the support GGUF used by `ds4 --mtp`, run the DSpark support mode.  This
+mode writes standalone DSpark metadata plus the packed `mtp.*` tensor payloads;
+it does not require a base-model GGUF template:
+
+```sh
+gguf-tools/deepseek4-quantize \
+  --hf ../deepseek-v4-quants/hf/DeepSeek-V4-Flash-0731 \
+  --dspark-support \
+  --out DeepSeek-V4-Flash-DSpark-support-0731.gguf
+```
+
+`--dspark-support --dry-run` reads safetensors shard headers to derive exact
+GGUF shapes and types, but it does not read tensor payloads.  The DSpark metadata
+defaults match the published Flash DSpark config: block size 5, target layers
+40,41,42, Markov rank 256, and noise token 128799.  Override them with
+`--dspark-block-size`, `--dspark-target-layers`, `--dspark-markov-rank`, and
+`--dspark-noise-token-id` if converting a different checkpoint.
+
+Before a full write, regenerate one support tensor and record its checksum:
+
+```sh
+gguf-tools/deepseek4-quantize \
+  --hf ../deepseek-v4-quants/hf/DeepSeek-V4-Flash-0731 \
+  --dspark-support \
+  --compare-tensor mtp.0.main_proj.weight
+```
+
+This reads only the payloads needed for that tensor.  Add `--compare-gguf
+DeepSeek-V4-Flash-DSpark-support-0731.gguf` to byte-compare against an existing
+support GGUF.
+
 ## When No Imatrix Is Given
 
 `iq2_xxs` requires an importance vector.  If `--imatrix` is not provided and

gguf-tools/glm53_manifest.py

@@ -0,0 +1,338 @@
+#!/usr/bin/env python3
+"""Inspect and validate an official GLM-5.3 Flash safetensors snapshot."""
+
+import argparse
+import json
+import math
+import os
+import re
+import struct
+import sys
+
+
+DTYPE_BYTES = {
+    "BOOL": 1,
+    "I8": 1,
+    "U8": 1,
+    "F8_E4M3": 1,
+    "F8_E5M2": 1,
+    "F8_E8M0": 1,
+    "I16": 2,
+    "U16": 2,
+    "F16": 2,
+    "BF16": 2,
+    "I32": 4,
+    "U32": 4,
+    "F32": 4,
+    "I64": 8,
+    "U64": 8,
+    "F64": 8,
+}
+
+LAYER_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.(.+)$")
+EXPERT_RE = re.compile(r"^mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$")
+
+
+def fail(message):
+    raise ValueError(message)
+
+
+def checked_product(values, label):
+    result = 1
+    for value in values:
+        if not isinstance(value, int) or isinstance(value, bool) or value < 0:
+            fail(f"{label}: invalid shape dimension {value!r}")
+        result *= value
+        if result > (1 << 63) - 1:
+            fail(f"{label}: tensor element count overflows int64")
+    return result
+
+
+def load_index(path):
+    with open(path, "rb") as fp:
+        document = json.load(fp)
+    weight_map = document.get("weight_map")
+    if not isinstance(weight_map, dict) or not weight_map:
+        fail(f"{path}: missing or empty weight_map")
+    for name, shard in weight_map.items():
+        if not isinstance(name, str) or not name:
+            fail(f"{path}: invalid tensor name")
+        if not isinstance(shard, str) or os.path.basename(shard) != shard:
+            fail(f"{path}: invalid shard for {name}")
+    return document, weight_map
+
+
+def load_safetensors_header(path):
+    file_size = os.path.getsize(path)
+    with open(path, "rb") as fp:
+        raw_length = fp.read(8)
+        if len(raw_length) != 8:
+            fail(f"{path}: short safetensors header length")
+        header_length = struct.unpack("<Q", raw_length)[0]
+        if header_length > file_size - 8:
+            fail(f"{path}: header extends beyond file")
+        if header_length > (1 << 30):
+            fail(f"{path}: unreasonable header length {header_length}")
+        raw_header = fp.read(header_length)
+    try:
+        document = json.loads(raw_header)
+    except (UnicodeDecodeError, json.JSONDecodeError) as error:
+        fail(f"{path}: invalid safetensors header: {error}")
+
+    payload_size = file_size - 8 - header_length
+    tensors = {}
+    spans = []
+    for name, entry in document.items():
+        if name == "__metadata__":
+            continue
+        if not isinstance(entry, dict):
+            fail(f"{path}: bad metadata for {name}")
+        dtype = entry.get("dtype")
+        shape = entry.get("shape")
+        offsets = entry.get("data_offsets")
+        if dtype not in DTYPE_BYTES:
+            fail(f"{path}: unsupported dtype {dtype!r} for {name}")
+        if not isinstance(shape, list):
+            fail(f"{path}: invalid shape for {name}")
+        if (
+            not isinstance(offsets, list)
+            or len(offsets) != 2
+            or any(not isinstance(value, int) or isinstance(value, bool) for value in offsets)
+        ):
+            fail(f"{path}: invalid data offsets for {name}")
+        start, end = offsets
+        if start < 0 or end < start or end > payload_size:
+            fail(f"{path}: out-of-range data offsets for {name}")
+        nbytes = checked_product(shape, name) * DTYPE_BYTES[dtype]
+        if end - start != nbytes:
+            fail(f"{path}: byte count mismatch for {name}: {end - start} != {nbytes}")
+        tensors[name] = {
+            "dtype": dtype,
+            "shape": shape,
+            "offset": 8 + header_length + start,
+            "nbytes": nbytes,
+        }
+        spans.append((start, end, name))
+
+    cursor = 0
+    for start, end, name in sorted(spans):
+        if start != cursor:
+            fail(f"{path}: gap or overlap before {name}: {start} != {cursor}")
+        cursor = end
+    if cursor != payload_size:
+        fail(f"{path}: unclaimed payload bytes: {payload_size - cursor}")
+    return tensors
+
+
+def tensor_scope(name):
+    if name.startswith("model.visual."):
+        return "vision"
+    if name == "lm_head.weight" or name.startswith("model.language_model."):
+        return "text"
+    return "unknown"
+
+
+def tensor_role(name):
+    if name.endswith(".weight_scale_inv"):
+        return "scale"
+    if name == "model.language_model.embed_tokens.weight":
+        return "embedding"
+    if name == "model.language_model.norm.weight":
+        return "final_norm"
+    if name == "lm_head.weight":
+        return "output"
+    match = LAYER_RE.match(name)
+    if not match:
+        return "vision" if name.startswith("model.visual.") else "unknown"
+    layer = int(match.group(1))
+    tail = match.group(2)
+    prefix = "mtp_" if layer == 45 else ""
+    expert = EXPERT_RE.match(tail)
+    if expert:
+        return f"{prefix}routed_{expert.group(2)}"
+    if tail.startswith("mlp.shared_experts."):
+        return f"{prefix}shared_expert"
+    if tail.startswith("mlp.gate."):
+        return f"{prefix}router"
+    if tail.startswith("mlp."):
+        return f"{prefix}dense_ffn"
+    if tail.startswith("self_attn.indexer."):
+        return f"{prefix}dsa_indexer"
+    if tail.startswith("self_attn."):
+        if layer < 45 and layer % 4 != 3:
+            return "linear_attention"
+        return f"{prefix}dsa"
+    if tail.startswith("hc_"):
+        return "mhc"
+    if "norm" in tail:
+        return f"{prefix}norm"
+    if tail.startswith(("eh_proj.", "shared_head.")):
+        return "mtp_head"
+    return f"{prefix}other"
+
+
+def expect_equal(actual, expected, label):
+    if actual != expected:
+        fail(f"{label}: got {actual!r}, expected {expected!r}")
+
+
+def validate_glm53_index(weight_map):
+    names = set(weight_map)
+    scopes = {tensor_scope(name) for name in names}
+    if "unknown" in scopes:
+        unknown = sorted(name for name in names if tensor_scope(name) == "unknown")
+        fail(f"unknown top-level tensors: {unknown[:5]}")
+
+    layers = set()
+    hc_layers = set()
+    linear_layers = set()
+    dsa_layers = set()
+    sparse_layers = set()
+    expert_ids = {}
+    expert_parts = {}
+    for name in names:
+        match = LAYER_RE.match(name)
+        if not match:
+            continue
+        layer = int(match.group(1))
+        tail = match.group(2)
+        layers.add(layer)
+        if tail.startswith("hc_attn_"):
+            hc_layers.add(layer)
+        if tail == "self_attn.A_log":
+            linear_layers.add(layer)
+        if tail == "self_attn.indexer.wk.weight":
+            dsa_layers.add(layer)
+        expert = EXPERT_RE.match(tail)
+        if expert:
+            sparse_layers.add(layer)
+            expert_ids.setdefault(layer, set()).add(int(expert.group(1)))
+            expert_parts.setdefault((layer, int(expert.group(1))), set()).add(expert.group(2))
+
+    expected_linear = {layer for layer in range(45) if layer % 4 != 3}
+    expected_dsa = {layer for layer in range(45) if layer % 4 == 3} | {45}
+    expected_sparse = set(range(3, 46))
+    expect_equal(layers, set(range(46)), "language-model layer set")
+    expect_equal(hc_layers, set(range(45)), "mHC layer set")
+    expect_equal(linear_layers, expected_linear, "linear-attention layer set")
+    expect_equal(dsa_layers, expected_dsa, "DSA layer set")
+    expect_equal(sparse_layers, expected_sparse, "sparse FFN layer set")
+    for layer in expected_sparse:
+        expect_equal(expert_ids.get(layer), set(range(288)), f"layer {layer} expert ids")
+        for expert in range(288):
+            expect_equal(
+                expert_parts.get((layer, expert)),
+                {"gate", "up", "down"},
+                f"layer {layer} expert {expert} projections",
+            )
+
+    required = {
+        "model.language_model.embed_tokens.weight",
+        "model.language_model.norm.weight",
+        "lm_head.weight",
+        "model.language_model.layers.45.eh_proj.weight",
+        "model.language_model.layers.45.enorm.weight",
+        "model.language_model.layers.45.hnorm.weight",
+        "model.language_model.layers.45.shared_head.norm.weight",
+    }
+    missing = sorted(required - names)
+    if missing:
+        fail(f"missing required tensors: {missing}")
+
+
+def validate_fp8_scales(tensors):
+    for name, info in tensors.items():
+        if info["dtype"] != "F8_E4M3" or not name.endswith(".weight"):
+            continue
+        if len(info["shape"]) != 2:
+            fail(f"{name}: FP8 weight is not two-dimensional")
+        scale_name = name + "_scale_inv"
+        scale = tensors.get(scale_name)
+        if scale is None:
+            fail(f"{name}: missing {scale_name}")
+        expected_shape = [math.ceil(info["shape"][0] / 128), math.ceil(info["shape"][1] / 128)]
+        if scale["dtype"] != "F32" or scale["shape"] != expected_shape:
+            fail(
+                f"{scale_name}: got {scale['dtype']} {scale['shape']}, "
+                f"expected F32 {expected_shape}"
+            )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("hf_dir", help="downloaded zai-org/GLM-5.3-Flash directory")
+    parser.add_argument(
+        "--index",
+        help="index JSON path; defaults to HF_DIR/model.safetensors.index.json",
+    )
+    parser.add_argument(
+        "--allow-missing-shards",
+        action="store_true",
+        help="emit index-only rows for shards that are still downloading",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    index_path = args.index or os.path.join(args.hf_dir, "model.safetensors.index.json")
+    document, weight_map = load_index(index_path)
+    validate_glm53_index(weight_map)
+
+    shard_names = sorted(set(weight_map.values()))
+    tensors = {}
+    missing_shards = []
+    for shard_name in shard_names:
+        shard_path = os.path.join(args.hf_dir, shard_name)
+        if not os.path.isfile(shard_path):
+            missing_shards.append(shard_name)
+            continue
+        header = load_safetensors_header(shard_path)
+        for name, info in header.items():
+            assigned = weight_map.get(name)
+            if assigned != shard_name:
+                fail(f"{shard_name}: {name} is assigned to {assigned!r} by the index")
+            if name in tensors:
+                fail(f"duplicate tensor in shard headers: {name}")
+            tensors[name] = info
+
+    if missing_shards and not args.allow_missing_shards:
+        fail(f"missing {len(missing_shards)} shards; first is {missing_shards[0]}")
+
+    present_names = set(tensors)
+    expected_present = {name for name, shard in weight_map.items() if shard not in missing_shards}
+    if present_names != expected_present:
+        missing = sorted(expected_present - present_names)
+        extra = sorted(present_names - expected_present)
+        fail(f"header/index mismatch: missing={missing[:3]} extra={extra[:3]}")
+    validate_fp8_scales(tensors)
+
+    print("scope\trole\tname\tdtype\tshape\tshard\toffset\tnbytes")
+    for name in sorted(weight_map):
+        info = tensors.get(name)
+        dtype = info["dtype"] if info else "-"
+        shape = "x".join(str(value) for value in info["shape"]) if info else "-"
+        offset = str(info["offset"]) if info else "-"
+        nbytes = str(info["nbytes"]) if info else "-"
+        print(
+            f"{tensor_scope(name)}\t{tensor_role(name)}\t{name}\t{dtype}\t{shape}\t"
+            f"{weight_map[name]}\t{offset}\t{nbytes}"
+        )
+
+    declared_size = document.get("metadata", {}).get("total_size")
+    present_bytes = sum(info["nbytes"] for info in tensors.values())
+    print(
+        f"glm53-manifest: tensors={len(weight_map)} shards={len(shard_names)} "
+        f"present_shards={len(shard_names) - len(missing_shards)} "
+        f"present_bytes={present_bytes} declared_bytes={declared_size} "
+        f"missing_shards={len(missing_shards)}",
+        file=sys.stderr,
+    )
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except (OSError, ValueError, json.JSONDecodeError) as error:
+        print(f"glm53-manifest: error: {error}", file=sys.stderr)
+        sys.exit(1)

gguf-tools/glm53_validate_gguf.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""Validate a GLM-5.3 Flash GGUF against its pinned source conversion plan."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import struct
+
+from glm53_quantize import (
+    GGUF_ALIGNMENT,
+    GGUF_ARRAY,
+    GGUF_BOOL,
+    GGUF_FLOAT32,
+    GGUF_STRING,
+    GGUF_UINT32,
+    GGUF_UINT64,
+    QTYPE_NAMES,
+    SourceDB,
+    align,
+    build_plan,
+    fail,
+    iter_native_tensor_bytes,
+    read_exact,
+    read_gguf_string,
+    read_u32,
+    read_u64,
+    skip_gguf_value,
+)
+
+
+def read_selected_metadata(fp, value_type):
+    if value_type == GGUF_STRING:
+        return read_gguf_string(fp, "GGUF metadata string")
+    if value_type == GGUF_UINT32:
+        return read_u32(fp, "GGUF metadata uint32")
+    if value_type == GGUF_UINT64:
+        return read_u64(fp, "GGUF metadata uint64")
+    if value_type == GGUF_FLOAT32:
+        return struct.unpack("<f", read_exact(fp, 4, "GGUF metadata float32"))[0]
+    if value_type == GGUF_BOOL:
+        return bool(read_exact(fp, 1, "GGUF metadata bool")[0])
+    skip_gguf_value(fp, value_type)
+    return None
+
+
+def verify_native_payload(fp, data_offset, expected, db):
+    try:
+        import numpy as np
+    except ImportError as error:
+        fail(f"NumPy is required for payload verification: {error}")
+    verified = 0
+    for index, item in enumerate(expected, 1):
+        fp.seek(data_offset + item.offset)
+        item_bytes = 0
+        for source_data in iter_native_tensor_bytes(item, db, np):
+            output_data = read_exact(fp, len(source_data), item.name)
+            if output_data != source_data:
+                fail(f"{item.name}: output payload differs from the pinned source")
+            item_bytes += len(source_data)
+        if item_bytes != item.nbytes:
+            fail(f"{item.name}: verified {item_bytes} bytes, expected {item.nbytes}")
+        padding = read_exact(fp, align(item.nbytes) - item.nbytes, f"{item.name} padding")
+        if any(padding):
+            fail(f"{item.name}: nonzero alignment padding")
+        verified += item_bytes
+        if index % 100 == 0 or index == len(expected):
+            print(f"verified payloads: {index}/{len(expected)} ({verified} bytes)")
+    return verified
+
+
+def validate(path, hf_dir, artifact, source_revision, verify_payload=False):
+    db = SourceDB(hf_dir)
+    expected = build_plan(db, artifact)
+
+    with open(path, "rb") as fp:
+        if read_exact(fp, 4, "GGUF magic") != b"GGUF":
+            fail(f"{path}: not a GGUF file")
+        version = read_u32(fp, "GGUF version")
+        if version != 3:
+            fail(f"{path}: expected GGUF v3, got v{version}")
+        tensor_count = read_u64(fp, "GGUF tensor count")
+        metadata_count = read_u64(fp, "GGUF metadata count")
+        if tensor_count != len(expected):
+            fail(f"tensor count {tensor_count} != expected {len(expected)}")
+
+        selected = {}
+        wanted = {
+            "general.architecture",
+            "general.alignment",
+            "general.source.revision",
+            "glm5-next.native_fp8",
+        }
+        for _ in range(metadata_count):
+            key = read_gguf_string(fp, "GGUF metadata key")
+            value_type = read_u32(fp, "GGUF metadata type")
+            if key in wanted:
+                selected[key] = read_selected_metadata(fp, value_type)
+            else:
+                skip_gguf_value(fp, value_type)
+
+        if selected.get("general.architecture") != "glm5-next":
+            fail(f"unexpected architecture: {selected.get('general.architecture')!r}")
+        if selected.get("general.alignment") != GGUF_ALIGNMENT:
+            fail(f"unexpected alignment: {selected.get('general.alignment')!r}")
+        if selected.get("general.source.revision") != source_revision:
+            fail(f"unexpected source revision: {selected.get('general.source.revision')!r}")
+        if artifact == "fp8" and selected.get("glm5-next.native_fp8") is not True:
+            fail("native FP8 artifact is missing its format marker")
+        if artifact != "fp8" and selected.get("glm5-next.native_fp8") is not None:
+            fail("quantized artifact unexpectedly has the native FP8 format marker")
+
+        seen = set()
+        for index, item in enumerate(expected):
+            name = read_gguf_string(fp, f"tensor {index} name")
+            n_dims = read_u32(fp, f"tensor {index} rank")
+            shape = tuple(read_u64(fp, f"tensor {index} dimension") for _ in range(n_dims))
+            qtype = read_u32(fp, f"tensor {index} type")
+            offset = read_u64(fp, f"tensor {index} offset")
+            if name in seen:
+                fail(f"duplicate tensor name: {name}")
+            seen.add(name)
+            if name != item.name:
+                fail(f"tensor {index} name {name!r} != expected {item.name!r}")
+            if shape != item.shape:
+                fail(f"{name}: shape {shape} != expected {item.shape}")
+            if qtype != item.qtype:
+                fail(
+                    f"{name}: type {QTYPE_NAMES.get(qtype, qtype)} != "
+                    f"expected {QTYPE_NAMES[item.qtype]}"
+                )
+            if offset != item.offset:
+                fail(f"{name}: offset {offset} != expected {item.offset}")
+
+        data_offset = align(fp.tell())
+        expected_size = data_offset
+        if expected:
+            last = expected[-1]
+            expected_size += last.offset + align(last.nbytes)
+        actual_size = os.fstat(fp.fileno()).st_size
+        if actual_size != expected_size:
+            fail(f"file size {actual_size} != expected {expected_size}")
+
+        role_bytes = {}
+        type_bytes = {}
+        for item in expected:
+            role_bytes[item.role] = role_bytes.get(item.role, 0) + item.nbytes
+            type_name = QTYPE_NAMES[item.qtype]
+            type_bytes[type_name] = type_bytes.get(type_name, 0) + item.nbytes
+
+        verified_bytes = 0
+        if verify_payload:
+            if artifact != "fp8":
+                fail("payload verification currently applies only to the native FP8 artifact")
+            verified_bytes = verify_native_payload(fp, data_offset, expected, db)
+
+    db.close()
+
+    print(
+        f"validated {path}: {len(expected)} tensors, {actual_size} bytes, "
+        f"data offset {data_offset}"
+    )
+    print("types: " + ", ".join(f"{key}={value}" for key, value in sorted(type_bytes.items())))
+    print("roles: " + ", ".join(f"{key}={value}" for key, value in sorted(role_bytes.items())))
+    if verify_payload:
+        print(f"source payloads matched byte for byte: {verified_bytes} bytes")
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--hf", required=True, help="official GLM-5.3 Flash snapshot")
+    parser.add_argument("--gguf", required=True, help="GGUF to validate")
+    parser.add_argument("--artifact", choices=("q4", "q2", "fp8"), required=True)
+    parser.add_argument("--verify-payload", action="store_true")
+    parser.add_argument(
+        "--source-revision",
+        default="84c6a6aa9497188e15a635ba793b0f95a79b1033",
+    )
+    args = parser.parse_args()
+    validate(args.gguf, args.hf, args.artifact, args.source_revision, args.verify_payload)
+
+
+if __name__ == "__main__":
+    main()

gguf-tools/glm53_vision.py

@@ -0,0 +1,370 @@
+#!/usr/bin/env python3
+"""Create or validate the standalone GLM-5.3 Flash vision encoder GGUF."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import io
+import json
+import os
+import shutil
+import struct
+import sys
+
+from glm53_quantize import (
+    GGUF_ALIGNMENT,
+    GGUF_ARRAY,
+    GGUF_FLOAT32,
+    GGUF_STRING,
+    GGUF_UINT32,
+    GGUF_UINT64,
+    GGUF_VERSION,
+    QTYPE_BF16,
+    QTYPE_NAMES,
+    SourceDB,
+    TensorPlan,
+    align,
+    fail,
+    kv_f32,
+    kv_string,
+    kv_u32,
+    kv_u64,
+    pack_string,
+    qtype_nbytes,
+    read_exact,
+    read_gguf_string,
+    read_u32,
+    read_u64,
+    skip_gguf_value,
+    tensor_header,
+)
+
+
+DEFAULT_SOURCE_REVISION = "84c6a6aa9497188e15a635ba793b0f95a79b1033"
+ARCHITECTURE = "glm5-next-vision"
+
+
+def kv_f32_array(key, values):
+    return (
+        pack_string(key)
+        + struct.pack("<IIQ", GGUF_ARRAY, GGUF_FLOAT32, len(values))
+        + struct.pack(f"<{len(values)}f", *values)
+    )
+
+
+def load_source_config(hf_dir):
+    with open(os.path.join(hf_dir, "config.json"), encoding="utf-8") as fp:
+        config = json.load(fp)
+    with open(os.path.join(hf_dir, "processor_config.json"), encoding="utf-8") as fp:
+        processor = json.load(fp)["image_processor"]
+    vision = config["vision_config"]
+
+    required = {
+        "depth": 24,
+        "hidden_size": 1024,
+        "intermediate_size": 4096,
+        "num_heads": 16,
+        "out_hidden_size": 4096,
+        "patch_size": 14,
+        "projection_intermediate_size": 10240,
+        "spatial_merge_size": 2,
+        "temporal_patch_size": 2,
+    }
+    for key, expected in required.items():
+        if vision.get(key) != expected:
+            fail(f"unexpected vision_config.{key}: {vision.get(key)!r}")
+    if processor.get("patch_size") != vision["patch_size"]:
+        fail("processor and model patch sizes differ")
+    if processor.get("merge_size") != vision["spatial_merge_size"]:
+        fail("processor and model merge sizes differ")
+    if processor.get("temporal_patch_size") != vision["temporal_patch_size"]:
+        fail("processor and model temporal patch sizes differ")
+    return config, vision, processor
+
+
+def build_vision_plan(db):
+    plan = []
+    offset = 0
+    for name in sorted(name for name in db.tensors if name.startswith("model.visual.")):
+        info = db.info(name)
+        if info["dtype"] != "BF16":
+            fail(f"{name}: expected BF16 source, got {info['dtype']}")
+        if ".blocks." in name:
+            role = "block"
+        elif ".patch_embed." in name:
+            role = "patch_embedding"
+        elif ".merger." in name:
+            role = "merger"
+        elif ".downsample." in name:
+            role = "downsample"
+        else:
+            role = "normalization"
+        item = TensorPlan(
+            name=name,
+            shape=tuple(reversed(info["shape"])),
+            qtype=QTYPE_BF16,
+            role=role,
+            source=name,
+            raw_copy=True,
+        )
+        item.nbytes = qtype_nbytes(item.qtype, item.shape)
+        if item.nbytes != info["nbytes"]:
+            fail(f"{name}: planned {item.nbytes} bytes, source has {info['nbytes']}")
+        item.offset = offset
+        offset += align(item.nbytes)
+        plan.append(item)
+
+    if len(plan) != 347:
+        fail(f"expected 347 vision tensors, found {len(plan)}")
+    return plan
+
+
+def vision_metadata(hf_dir, source_revision):
+    config, vision, processor = load_source_config(hf_dir)
+    prefix = "glm5-next-vision"
+    records = [
+        kv_string("general.architecture", ARCHITECTURE),
+        kv_string("general.name", "GLM-5.3-Flash Vision Encoder"),
+        kv_u32("general.alignment", GGUF_ALIGNMENT),
+        kv_string("general.source.url", "https://huggingface.co/zai-org/GLM-5.3-Flash"),
+        kv_string("general.source.revision", source_revision),
+        kv_u32(f"{prefix}.block_count", vision["depth"]),
+        kv_u32(f"{prefix}.embedding_length", vision["hidden_size"]),
+        kv_u32(f"{prefix}.feed_forward_length", vision["intermediate_size"]),
+        kv_u32(f"{prefix}.attention.head_count", vision["num_heads"]),
+        kv_u32(f"{prefix}.projection_length", vision["out_hidden_size"]),
+        kv_u32(f"{prefix}.projection.feed_forward_length", vision["projection_intermediate_size"]),
+        kv_u32(f"{prefix}.patch_size", vision["patch_size"]),
+        kv_u32(f"{prefix}.temporal_patch_size", vision["temporal_patch_size"]),
+        kv_u32(f"{prefix}.spatial_merge_size", vision["spatial_merge_size"]),
+        kv_u32(f"{prefix}.channel_count", vision["in_channels"]),
+        kv_f32(f"{prefix}.attention.layer_norm_rms_epsilon", vision["rms_norm_eps"]),
+        kv_f32(f"{prefix}.swiglu_limit", vision["swiglu_limit"]),
+        kv_u32(f"{prefix}.image_token_id", config["image_token_id"]),
+        kv_u32(f"{prefix}.image_start_token_id", config["image_start_token_id"]),
+        kv_u32(f"{prefix}.image_end_token_id", config["image_end_token_id"]),
+        kv_u32(f"{prefix}.image.min_tokens", processor["min_image_tokens"]),
+        kv_u32(f"{prefix}.image.max_tokens", processor["max_image_tokens"]),
+        kv_f32_array(f"{prefix}.image.mean", processor["image_mean"]),
+        kv_f32_array(f"{prefix}.image.std", processor["image_std"]),
+    ]
+    return records
+
+
+def layout(plan, metadata):
+    header_bytes = 4 + 4 + 8 + 8
+    header_bytes += sum(len(record) for record in metadata)
+    header_bytes += sum(len(tensor_header(item)) for item in plan)
+    data_offset = align(header_bytes)
+    data_bytes = sum(align(item.nbytes) for item in plan)
+    return data_offset, data_bytes
+
+
+def print_summary(plan, metadata):
+    data_offset, data_bytes = layout(plan, metadata)
+    print(f"tensors: {len(plan)}")
+    print(f"metadata_records: {len(metadata)}")
+    print(f"metadata_bytes: {data_offset}")
+    print(f"tensor_bytes: {sum(item.nbytes for item in plan)}")
+    print(f"file_bytes: {data_offset + data_bytes}")
+    roles = {}
+    for item in plan:
+        roles[item.role] = roles.get(item.role, 0) + item.nbytes
+    for role, size in sorted(roles.items()):
+        print(f"role_bytes: {role} {size}")
+    return data_offset, data_bytes
+
+
+def create_gguf(path, plan, metadata, db, overwrite):
+    data_offset, data_bytes = print_summary(plan, metadata)
+    required = data_offset + data_bytes + 4 * (1 << 30)
+    free = shutil.disk_usage(os.path.dirname(os.path.abspath(path))).free
+    if free < required:
+        fail(f"insufficient free space: need output plus reserve {required}, have {free}")
+    if os.path.exists(path) and not overwrite:
+        fail(f"output exists: {path}; use --overwrite")
+
+    partial = path + ".partial"
+    if os.path.exists(partial):
+        if not overwrite:
+            fail(f"partial output exists: {partial}; use --overwrite")
+        os.unlink(partial)
+
+    with open(partial, "wb") as fp:
+        fp.write(b"GGUF")
+        fp.write(struct.pack("<IQQ", GGUF_VERSION, len(plan), len(metadata)))
+        for record in metadata:
+            fp.write(record)
+        for item in plan:
+            fp.write(tensor_header(item))
+        if fp.tell() > data_offset:
+            fail("GGUF header exceeds its planned data offset")
+        fp.write(bytes(data_offset - fp.tell()))
+
+        for index, item in enumerate(plan, 1):
+            if fp.tell() != data_offset + item.offset:
+                fail(f"{item.name}: output offset mismatch")
+            written = 0
+            for chunk in db.iter_read(item.source):
+                fp.write(chunk)
+                written += len(chunk)
+            if written != item.nbytes:
+                fail(f"{item.name}: copied {written} bytes, expected {item.nbytes}")
+            fp.write(bytes(align(item.nbytes) - item.nbytes))
+            if index % 25 == 0 or index == len(plan):
+                print(f"copied tensors: {index}/{len(plan)}", file=sys.stderr, flush=True)
+        fp.flush()
+        os.fsync(fp.fileno())
+    os.replace(partial, path)
+
+
+def read_metadata_value(fp, value_type):
+    if value_type == GGUF_STRING:
+        return read_gguf_string(fp, "GGUF metadata string")
+    if value_type == GGUF_UINT32:
+        return read_u32(fp, "GGUF metadata uint32")
+    if value_type == GGUF_UINT64:
+        return read_u64(fp, "GGUF metadata uint64")
+    if value_type == GGUF_FLOAT32:
+        return struct.unpack("<f", read_exact(fp, 4, "GGUF metadata float32"))[0]
+    if value_type == GGUF_ARRAY:
+        element_type = read_u32(fp, "GGUF array element type")
+        count = read_u64(fp, "GGUF array count")
+        if element_type == GGUF_FLOAT32 and count <= 1024:
+            return list(struct.unpack(f"<{count}f", read_exact(fp, count * 4, "GGUF float array")))
+        fail(f"unsupported validation array type/count: {element_type}/{count}")
+    skip_gguf_value(fp, value_type)
+    return None
+
+
+def file_sha256(path):
+    digest = hashlib.sha256()
+    with open(path, "rb") as fp:
+        while chunk := fp.read(16 << 20):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def validate_gguf(path, plan, expected_metadata, db, verify_payload, expected_sha256):
+    expected_values = {}
+    for record in expected_metadata:
+        reader = io.BytesIO(record)
+        key = read_gguf_string(reader, "expected metadata key")
+        value_type = read_u32(reader, "expected metadata type")
+        expected_values[key] = read_metadata_value(reader, value_type)
+
+    with open(path, "rb") as fp:
+        if read_exact(fp, 4, "GGUF magic") != b"GGUF":
+            fail(f"{path}: not a GGUF file")
+        version = read_u32(fp, "GGUF version")
+        if version != GGUF_VERSION:
+            fail(f"expected GGUF v{GGUF_VERSION}, got v{version}")
+        tensor_count = read_u64(fp, "GGUF tensor count")
+        metadata_count = read_u64(fp, "GGUF metadata count")
+        if tensor_count != len(plan):
+            fail(f"tensor count {tensor_count} != expected {len(plan)}")
+        if metadata_count != len(expected_metadata):
+            fail(f"metadata count {metadata_count} != expected {len(expected_metadata)}")
+
+        actual_values = {}
+        for _ in range(metadata_count):
+            key = read_gguf_string(fp, "GGUF metadata key")
+            value_type = read_u32(fp, "GGUF metadata type")
+            actual_values[key] = read_metadata_value(fp, value_type)
+        if actual_values != expected_values:
+            fail("GGUF metadata differs from the pinned conversion metadata")
+
+        for index, item in enumerate(plan):
+            name = read_gguf_string(fp, f"tensor {index} name")
+            rank = read_u32(fp, f"tensor {index} rank")
+            shape = tuple(read_u64(fp, f"tensor {index} dimension") for _ in range(rank))
+            qtype = read_u32(fp, f"tensor {index} type")
+            offset = read_u64(fp, f"tensor {index} offset")
+            if (name, shape, qtype, offset) != (item.name, item.shape, item.qtype, item.offset):
+                fail(f"tensor {index} header differs from the conversion plan: {name}")
+
+        data_offset, data_bytes = layout(plan, expected_metadata)
+        if fp.tell() > data_offset:
+            fail("parsed GGUF header exceeds the planned data offset")
+        actual_size = os.fstat(fp.fileno()).st_size
+        if actual_size != data_offset + data_bytes:
+            fail(f"file size {actual_size} != expected {data_offset + data_bytes}")
+
+        verified = 0
+        if verify_payload:
+            for index, item in enumerate(plan, 1):
+                fp.seek(data_offset + item.offset)
+                for source in db.iter_read(item.source):
+                    if read_exact(fp, len(source), item.name) != source:
+                        fail(f"{item.name}: payload differs from the official source")
+                    verified += len(source)
+                padding = read_exact(fp, align(item.nbytes) - item.nbytes, f"{item.name} padding")
+                if any(padding):
+                    fail(f"{item.name}: nonzero alignment padding")
+                if index % 25 == 0 or index == len(plan):
+                    print(f"verified payloads: {index}/{len(plan)}", file=sys.stderr, flush=True)
+
+    print(
+        f"validated {path}: {len(plan)} {QTYPE_NAMES[QTYPE_BF16]} tensors, "
+        f"{actual_size} bytes"
+    )
+    digest = file_sha256(path)
+    if expected_sha256 and digest.lower() != expected_sha256.lower():
+        fail(f"SHA-256 {digest} != expected {expected_sha256.lower()}")
+    print(f"SHA-256: {digest}")
+    if verify_payload:
+        print(f"source payloads matched byte for byte: {verified} bytes")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--hf", required=True, help="official GLM-5.3 Flash snapshot")
+    output = parser.add_mutually_exclusive_group()
+    output.add_argument("--out", help="create this vision GGUF")
+    output.add_argument("--validate", metavar="GGUF", help="validate an existing vision GGUF")
+    parser.add_argument("--source-revision", default=DEFAULT_SOURCE_REVISION)
+    parser.add_argument("--verify-payload", action="store_true")
+    parser.add_argument("--expected-sha256", help="expected hash when validating")
+    parser.add_argument("--overwrite", action="store_true")
+    parser.add_argument("--dry-run", action="store_true")
+    args = parser.parse_args()
+    if not args.dry_run and not args.out and not args.validate:
+        parser.error("one of --out, --validate, or --dry-run is required")
+    if args.verify_payload and not args.validate:
+        parser.error("--verify-payload requires --validate")
+    if args.expected_sha256 and not args.validate:
+        parser.error("--expected-sha256 requires --validate")
+    return args
+
+
+def main():
+    args = parse_args()
+    db = SourceDB(args.hf)
+    try:
+        plan = build_vision_plan(db)
+        metadata = vision_metadata(args.hf, args.source_revision)
+        if args.dry_run:
+            print_summary(plan, metadata)
+        elif args.out:
+            create_gguf(args.out, plan, metadata, db, args.overwrite)
+            print(f"glm53-vision: wrote {args.out}", file=sys.stderr)
+        else:
+            validate_gguf(
+                args.validate,
+                plan,
+                metadata,
+                db,
+                args.verify_payload,
+                args.expected_sha256,
+            )
+    finally:
+        db.close()
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except (KeyError, OSError, ValueError, json.JSONDecodeError) as error:
+        print(f"glm53-vision: error: {error}", file=sys.stderr)
+        sys.exit(1)

gguf-tools/imatrix/README.md

@@ -1,10 +1,12 @@
 # DS4 Imatrix Pipeline
 
 This directory contains the calibration dataset and instructions used to build
-activation importance matrices for DeepSeek V4 Flash GGUF quantization.
+activation importance matrices for DeepSeek V4 Flash and Pro GGUF
+quantization.
 
-The current imatrix target is the routed MoE path.  DS4 has 43 layers, 256
-routed experts per layer, and three routed expert tensors per layer:
+The current imatrix target is the routed MoE path.  Flash has 43 layers and
+256 routed experts per layer.  Pro has 61 layers and 384 routed experts per
+layer.  Both variants expose three routed expert tensors per layer:
 
 - `blk.N.ffn_gate_exps.weight`
 - `blk.N.ffn_up_exps.weight`
@@ -36,17 +38,22 @@ It contains DS4-rendered chat prompts, separated by visible
 - C/Metal source-review prompts from this repository.
 - Long-context snippets.
 - Agent/tool-call prompts using DS4's DSML syntax.
-- English and Italian prompts.
+- Language/prose rewriting, summarization, extraction, and translation prompts.
+- `ds4-eval` GPQA Diamond, SuperGPQA, and AIME2025 benchmark prompts.
 - Both thinking and non-thinking assistant prefixes.
 
-The current tracked dataset has 1952 rendered prompts and roughly 1.44M tokens
+The current tracked dataset has 4682 rendered prompts and roughly 2.91M tokens
 by the coarse bytes/4 estimate.  Check
 `gguf-tools/imatrix/dataset/manifest.json` for the exact generated-file
 summary.
 
 ## 2. Collect The Imatrix
 
-Use the DS4 runtime itself to collect routed MoE activation statistics:
+Use the DS4 runtime itself to collect routed MoE activation statistics.  The
+collector uses the loaded GGUF metadata, so the same command shape works for
+Flash and Pro.
+
+Flash example:
 
 ```sh
 ./ds4 \
@@ -56,6 +63,18 @@ Use the DS4 runtime itself to collect routed MoE activation statistics:
   --ctx 32768
 ```
 
+Pro example with a smaller calibration budget:
+
+```sh
+./ds4 \
+  -m ../deepseek-v4-quants/gguf/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct.gguf \
+  --imatrix-dataset gguf-tools/imatrix/dataset/rendered_prompts.txt \
+  --imatrix-out ../deepseek-v4-quants/imatrix/DeepSeek-V4-Pro-Instruct-routed-moe-ds4-small.dat \
+  --imatrix-max-prompts 16 \
+  --imatrix-max-tokens 32768 \
+  --ctx 32768
+```
+
 Useful smoke-test limits:
 
 ```sh
@@ -79,6 +98,8 @@ entry length = n_expert * n_columns
 ```
 
 The quantizer slices the right expert's segment when quantizing each expert.
+For Pro, that means each routed tensor entry contains 384 vectors instead of
+Flash's 256.
 
 ## 3. Generate GGUF Files With The Imatrix
 
@@ -109,6 +130,21 @@ gguf-tools/deepseek4-quantize \
   --imatrix ../deepseek-v4-quants/imatrix/DeepSeek-V4-Flash-chat-v2-routed-moe-ds4-1p5m.dat
 ```
 
+Example Pro Q2 regeneration from a Pro template:
+
+```sh
+gguf-tools/deepseek4-quantize \
+  --hf ../deepseek-v4-quants/hf/DeepSeek-V4-Pro \
+  --template ../deepseek-v4-quants/gguf/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct.gguf \
+  --out ../deepseek-v4-quants/gguf/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix-small.gguf \
+  --imatrix ../deepseek-v4-quants/imatrix/DeepSeek-V4-Pro-Instruct-routed-moe-ds4-small.dat \
+  --threads 24
+```
+
+`deepseek4-quantize` reads the routed expert count from the template GGUF.  Use
+`--n-experts` only when working with an older template that lacks
+`deepseek4.expert_count`.
+
 For Q4, the imatrix does not change the runtime tensor type: routed experts
 remain `Q4_K`.  It changes how quantization error is weighted while choosing
 scales and codes.  For Q2, it replaces the previous synthetic weight-energy

gguf-tools/imatrix/dataset/README.md

@@ -26,11 +26,12 @@ The renderer mirrors the server prompt shape:
 
 Some records include DSML tool schemas, sampled DSML tool calls, and tool-result
 turns so the imatrix sees the same special-token patterns used by agent clients.
-The corpus is provider-neutral and also includes multi-language programming
-prompts, Bash scripting, algorithm recall, English-to-European-language
-translation, English-to-Chinese translation, reverse translation to English,
-long-context source excerpts, Metal/C code review tasks, and
-inference-specific debugging tasks.
+The corpus is provider-neutral and also includes language/prose rewriting,
+summarization, copy-editing, extraction, multilingual translation, programming
+prompts, Bash scripting, algorithm recall, `ds4-eval` benchmark-reasoning
+prompts, long-context code synthesis, agent transcript replay, log diagnosis,
+prose fact recovery, delayed-constraint and small needle tasks, Metal/C code
+review tasks, and inference-specific debugging tasks.
 
 For normal imatrix collection, use `rendered_prompts.txt` so calibration covers
 both thinking and non-thinking modes. Split files are provided for ablations.