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.

.gitignore

@@ -4,15 +4,39 @@
 /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

Makefile

@@ -6,6 +6,9 @@ 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
 
 DEBUG_FLAGS ?= -g
 CFLAGS ?= -O3 -ffast-math $(DEBUG_FLAGS) $(NATIVE_CPU_FLAG) -Wall -Wextra -std=c99
@@ -15,25 +18,45 @@ 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_distributed.o ds4_ssd.o ds4_metal.o
-CPU_CORE_OBJS = ds4_cpu.o ds4_distributed.o ds4_ssd.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
+endif
 NVCCFLAGS ?= -O3 -g -lineinfo --use_fast_math $(NVCC_ARCH_FLAGS) -Xcompiler $(NATIVE_CPU_FLAG) -Xcompiler -pthread
-CORE_OBJS = ds4.o ds4_distributed.o ds4_ssd.o ds4_cuda.o
-CPU_CORE_OBJS = ds4_cpu.o ds4_distributed.o ds4_ssd.o
+# 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
 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)
@@ -41,9 +64,11 @@ DS4_LINK_LIBS ?= $(CUDA_LDLIBS)
 METAL_LDLIBS := $(LDLIBS)
 endif
 
-.PHONY: all help clean test cpu cuda cuda-spark cuda-generic cuda-regression strix-halo rocm
+.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)
+.PHONY: metal-decode-schedule-bench metal-prefill-variant-bench check-mxfp4-half-lut
+
 all: ds4 ds4-server ds4-bench ds4-eval ds4-agent
 
 help:
@@ -51,32 +76,75 @@ help:
 	@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 ds4_help.o linenoise.o $(CORE_OBJS)
-	$(CC) $(CFLAGS) -o $@ ds4_cli.o ds4_help.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 ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS)
-	$(CC) $(CFLAGS) -o $@ ds4_server.o ds4_help.o ds4_kvstore.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 ds4_help.o $(CORE_OBJS)
-	$(CC) $(CFLAGS) -o $@ ds4_bench.o ds4_help.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)
 
 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 $(CORE_OBJS)
-	$(CC) $(CFLAGS) -o $@ ds4_agent.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.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 $(CORE_OBJS) rax.o
-	$(CC) $(QUALITY_CFLAGS) -I. -o $@ gguf-tools/quality-testing/score_official.c $(CORE_OBJS) rax.o $(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)
 
-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 $(CPU_CORE_OBJS)
-	$(CC) $(CFLAGS) -o ds4 ds4_cli_cpu.o ds4_help.o linenoise.o $(CPU_CORE_OBJS) $(LDLIBS)
-	$(CC) $(CFLAGS) -o ds4-server ds4_server_cpu.o ds4_help.o ds4_kvstore.o rax.o $(CPU_CORE_OBJS) $(LDLIBS)
-	$(CC) $(CFLAGS) -o ds4-bench ds4_bench_cpu.o ds4_help.o $(CPU_CORE_OBJS) $(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 $(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"
@@ -90,12 +158,16 @@ help:
 	@echo "  make cuda CUDA_ARCH=sm_N Build CUDA with an explicit nvcc -arch value"
 	@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) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent CUDA_ARCH=
+	$(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent CUDA_ARCH=sm_121
 
 cuda-generic:
 	$(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent CUDA_ARCH=native
@@ -110,45 +182,77 @@ cuda:
 
 strix-halo:
 	$(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent \
-		CORE_OBJS="ds4.o ds4_distributed.o ds4_ssd.o ds4_rocm.o" \
-		CFLAGS="$(CFLAGS) -DDS4_ROCM_BUILD" \
+		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)"
 
 rocm: strix-halo
 
-ds4: ds4_cli.o ds4_help.o linenoise.o $(CORE_OBJS)
+# 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
+
+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 $(CORE_OBJS)
+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 $(CORE_OBJS)
+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 $(CORE_OBJS)
+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: gguf-tools/quality-testing/score_official.c ds4.h $(CORE_OBJS) rax.o
-	$(DS4_LINK) $(filter-out -ffast-math,$(QUALITY_CFLAGS)) -I. -o $@ gguf-tools/quality-testing/score_official.c $(CORE_OBJS) rax.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 $@ $<
 
-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 $(CPU_CORE_OBJS)
-	$(CC) $(CFLAGS) -o ds4 ds4_cli_cpu.o ds4_help.o linenoise.o $(CPU_CORE_OBJS) $(LDLIBS)
-	$(CC) $(CFLAGS) -o ds4-server ds4_server_cpu.o ds4_help.o ds4_kvstore.o rax.o $(CPU_CORE_OBJS) $(LDLIBS)
-	$(CC) $(CFLAGS) -o ds4-bench ds4_bench_cpu.o ds4_help.o $(CPU_CORE_OBJS) $(LDLIBS)
+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 $(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_ssd.h ds4_distributed.h ds4_gpu.h
 	$(CC) $(CFLAGS) -c -o $@ ds4.c
 
+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
 
@@ -158,9 +262,15 @@ ds4_cli.o: ds4_cli.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h linenoise.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
 
@@ -170,7 +280,7 @@ ds4_bench.o: ds4_bench.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.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_help.h ds4_kvstore.h ds4_web.h linenoise.h
+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
@@ -182,7 +292,7 @@ ds4_kvstore.o: ds4_kvstore.c ds4_kvstore.h ds4.h ds4_ssd.h
 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_help.h ds4_kvstore.h ds4_web.h linenoise.h
+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
@@ -195,11 +305,14 @@ linenoise.o: linenoise.c linenoise.h
 	$(CC) $(CFLAGS) -c -o $@ linenoise.c
 
 ds4_cpu.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_gpu.h
-	$(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4.c
+	$(CC) $(CFLAGS) -Wno-unused-function -DDS4_NO_GPU -c -o $@ ds4.c
 
 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_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
 
@@ -215,37 +328,248 @@ ds4_agent_cpu.o: ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kv
 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
+tests/test_glm53_kda.o: tests/test_glm53_kda.c ds4_gpu.h
+	$(CC) $(CFLAGS) -I. -c -o $@ tests/test_glm53_kda.c
+
+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
+
+tests/test_glm53_vision_engine: tests/test_glm53_vision_engine.o $(CORE_OBJS)
+ifeq ($(UNAME_S),Darwin)
+	$(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS)
+else
+	$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
+endif
+
+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
 
-ds4_rocm.o: ds4_rocm.cu ds4_gpu.h ds4_iq2_tables_cuda.inc $(ROCM_SRCS)
+# 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/cuda_long_context_smoke: tests/cuda_long_context_smoke.o ds4_cuda.o
+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
-	$(NVCC) $(NVCCFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(CUDA_LDLIBS)
+	$(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
-	$(NVCC) $(NVCCFLAGS) -o $@ ds4_agent_test.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS) $(CUDA_LDLIBS)
+	$(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
+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-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official tests/test_q4k_dot *.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

QA_BEFORE_RELEASES.md

@@ -12,10 +12,15 @@ run.
 
 Preferred release test hosts:
 
-- CUDA / DGX Spark: `toor@192.168.0.180`.
+- CUDA / DGX Spark: `toor@192.168.4.180` and `toor@192.168.4.181`.
 - Metal / distributed Mac testing: `mac-m5max-it` and `mac-m5max-us`.
 - ROCm: The Strix Halo system at antirez@strixhalo (Framework Desktop).
 
+`192.168.60.250` is permission-only. Never connect to it for QA, stop or start
+its server, build there, or run tests or benchmarks there without asking
+Salvatore and receiving explicit permission for that specific QA pass. Earlier
+permission does not carry over to later work.
+
 The Mac hosts have DNS entries and are reached through an internet VPN.  They
 are connected to each other over WiFi and also through a Thunderbolt 5
 point-to-point link.  The TB5 route is the preferred distributed-inference
@@ -36,6 +41,16 @@ in this system.
   `make clean && make`.
 - Build CPU-only binaries as a compile check only:
   `make clean && make cpu`.
+- Treat compiler warnings as build failures. Save each release and test build's
+  complete output and require no `warning:` or NVCC `warning #` lines. Fix the
+  source when possible; use a narrow target-specific suppression only when a
+  test deliberately compiles a partial translation unit.
+- Repeat the warning-free build gate on the release hardware:
+  `make clean && make` on Metal,
+  `make clean && make cuda-spark` on DGX Spark,
+  `make clean && make cuda-generic CUDA_HOME=/usr` on the multi-GPU CUDA host
+  only after receiving permission for `192.168.60.250`,
+  `make clean && make strix-halo` on Strix Halo.
 - Run whitespace checks before committing:
   `git diff --check`.
 - Confirm `./ds4 --help`, `./ds4-server --help`, and `./ds4-agent --help` render
@@ -45,15 +60,81 @@ in this system.
 
 - Run the default suite:
   `make test`.
+- Run `tests/test_gpu_args_cli.sh` explicitly after changing executable option
+  parsing or multi-GPU placement. Invalid values and device/budget count
+  mismatches must reach the shared GPU parser in all four binaries; an
+  `unknown option` response from a binary that advertises the flag is a
+  release blocker. On CUDA, also start `ds4-server` once with
+  `--gpu-vram auto` and the intended `--gpu-devices` list and preserve the
+  resolved layout line.
 - Run the vector checks explicitly after any tokenizer, template, KV, kernel,
   quantization, or prompt-rendering change:
-  `./ds4_test --logprob-vectors`
-  and `./ds4_test --local-golden-vectors`.
+  `DS4_TEST_MODEL=/path/to/0731.gguf
+  DS4_TEST_VECTOR_FILE=tests/test-vectors/flash-0731/official.vec
+  ./ds4_test --logprob-vectors`
+  and
+  `DS4_TEST_MODEL=/path/to/0731.gguf
+  DS4_TEST_LOCAL_GOLDEN_FILE=tests/test-vectors/flash-0731/local-golden.vec
+  ./ds4_test --local-golden-vectors`.
 - Run server tests when HTTP, SSE, prompt rendering, cache policy, or tool-call
   replay changed:
   `./ds4_test --server`.
 - Run `./ds4-eval --self-test-extractors`.
 
+### Critical Input And Server Regression Pass
+
+Run these checks after changing parsers, server generation, model loading,
+distributed snapshots, caches, DSpark, or CUDA build rules. Keep the item
+numbers in the QA report so omissions are visible.
+
+1. Send malformed OpenAI, Responses, and Anthropic requests with repeated
+   owned string or array fields under ASan. Each request must fail cleanly and
+   a following valid request must still work. Run `./ds4_test --server` too.
+2. Replay at least 4,096 assistant/tool-result pairs through the Responses and
+   Anthropic validators. Validation must remain linear-time and preserve the
+   same accepted and rejected histories as a short replay.
+3. Feed a distributed worker a snapshot header whose declared lengths exceed
+   the configured and protocol limits. It must reject the header before a
+   large allocation or payload read, without growing RSS materially.
+4. Run malformed safetensors and GGUF fixtures through the loader and
+   `gguf-tools/deepseek4-quantize` under ASan and UBSan. Truncated files,
+   impossible dimensions, and overflowing tensor sizes must be rejected.
+5. With a checkpoint-matched DSpark drafter, compare temperature-zero output
+   against a no-drafter run for 400 and 800 generated tokens. Record the first
+   output difference, acceptance, direct commits, replay fallbacks, and decode
+   speed. Byte identity is not required: DSpark commits the batched verifier
+   state, whose floating-point operation order differs from one-token decode.
+   Verifier errors, invalid text, or a material continuation-quality regression
+   remain release blockers. Use `--dspark-strict` for the target-only control.
+6. Exercise unterminated and twice-closed reasoning in streaming and
+   non-streaming OpenAI, Responses, and Anthropic requests, with and without
+   tools. Reasoning must never leak into answer content.
+7. On real Blackwell hardware, build the CUDA targets for `sm_120` or `sm_120a`
+   and for DGX Spark `sm_121`. Confirm the emitted architecture flags retain
+   the architecture-specific feature suffix and run `make cuda-regression`.
+8. Build with CUDA 12.8 or newer and require the CUDA translation units to
+   compile warning-free, including the `FLT_MAX` users.
+9. Force a conversation past the in-memory KV threshold, restore the same disk
+   checkpoint twice, and confirm the checkpoint file remains present after
+   both successful loads. Corrupt checkpoints must still be rejected.
+10. Run `make dspark-verify-depth` with matching 0731 target and drafter files.
+    Strict capture must skip layers without a compressor and compare every
+    captured compressor layer.
+11. Send the same long GLM 5.2 prompt twice to one server session. The second
+    request must report `cache_source: memory-rewind`, reuse through one token
+    before the prompt boundary, and produce the same greedy output as a fresh
+    session.
+12. Run `./ds4_test --think-tool-recovery`, then repeat through all three HTTP
+    APIs. A complete tool block inside unclosed reasoning must be recovered
+    once, preceding prose must remain reasoning, and no synthetic continuation
+    may be generated.
+13. Run `./ds4_agent_test` under ASan with agent-cache strings whose declared
+    lengths exceed the remaining file. Loading must fail without allocating
+    the declared size, and a valid cache must still load.
+14. Run the server parser tests under UBSan with `NaN`, positive infinity, and
+    negative infinity where integer JSON fields are expected. Conversion must
+    be defined and clamped, with no sanitizer report.
+
 ## 3. Official Continuation Quality Gates
 
 These tests are release-blocking after tokenizer, template, KV-cache, attention,
@@ -63,22 +144,54 @@ top-logprob slices, so do not replace them with one sampled chat answer.
 
 - Build the scorer:
   `make -C gguf-tools quality-score`.
-- Run the tracked DeepSeek V4 Flash smoke vectors:
-  `./ds4_test --logprob-vectors`.
-  This covers the official Flash API vectors in `tests/test-vectors/`, including
-  short prompts and long-prompt attention cases.
+- Match every Flash GGUF to the fixture captured from the same checkpoint.
+  The current release checkpoint is 0731 and uses
+  `tests/test-vectors/flash-0731/`; the older undated GGUF uses the preserved
+  `tests/test-vectors/flash-pre-0731/` fixture. Never report a cross-checkpoint
+  failure as a quality regression. New checkpoints require a new
+  `flash-CHECKPOINT/` directory before release QA; do not replace an older
+  fixture. Checkpoint-labelled GGUFs such as `-0731` must use the fixture with
+  the same label.
+- Run the tracked DeepSeek V4 Flash 0731 smoke vectors:
+  `DS4_TEST_MODEL=/path/to/0731.gguf
+  DS4_TEST_VECTOR_FILE=tests/test-vectors/flash-0731/official.vec
+  ./ds4_test --logprob-vectors`.
+  This covers short prompts and long-prompt attention cases. The runner defaults
+  to this fixture, but release logs should keep the path explicit.
 - Run the 100-case DeepSeek V4 Flash fixture for every released Flash GGUF:
   `gguf-tools/quality-testing/score_official /path/to/deepseek-v4-flash.gguf gguf-tools/quality-testing/data/flash/manifest.tsv /tmp/flash.tsv 4096`.
+  This manifest is also for the 0731 checkpoint. A later checkpoint needs a
+  separately named 100-case fixture and must not be scored against this one.
+- Treat the native MXFP4 Flash GGUF as a separate release artifact. Run the
+  same 100-case fixture on Metal, resident CUDA, and CUDA SSD streaming when
+  those backends are advertised; compare each result with the Metal baseline.
+  For resident multi-GPU CUDA, pass the normal placement flags to the scorer,
+  for example `--gpu-vram auto --gpu-devices 0,2,4,6,1,3,5,7
+  --cuda-tensor-parallel`.
 - Run the 100-case GLM 5.2 OpenRouter fixture for every released GLM GGUF:
   `gguf-tools/quality-testing/score_official models/GLM-5.2-UD-Q4_K_XL.gguf gguf-tools/quality-testing/data/glm52-openrouter-100/manifest.tsv /tmp/glm52-q4.tsv 4096`.
   Current Q4 XL reference band: first-token match `95/100`, API top-1 agreement
   about `0.942`, and API pair-order agreement about `0.880`.
+- Run the 100-case GLM 5.3 Flash fixture separately for both release artifacts:
+  `gguf-tools/quality-testing/score_official /path/to/GLM-5.3-Flash-Q2.gguf gguf-tools/quality-testing/data/glm53-flash-openrouter-zai-fp8-100/manifest.tsv /tmp/glm53-q2.tsv 4096`
+  and repeat with `GLM-5.3-Flash-Q4_K.gguf`. The current Q2 reference is
+  average NLL `0.458030488`, first-token match `89/100`, and average greedy
+  prefix `7.37`; Q4 is `0.299917952`, `90/100`, and `9.66`. These are fresh
+  GLM-5.3 Z.AI FP8 continuations and must not be replaced by GLM 5.2 fixtures.
+  The Q4 layout with Q8 KDA projections, embedding, and output head scored
+  `0.300804038`, `90/100`, and `9.48` on M3 Ultra. Its paired BF16-layout
+  control scored `0.300477636`, `90/100`, and `9.48`; the Q8 layout won 54 of
+  100 cases despite its `0.109%` higher aggregate NLL.
 - Run the same GLM fixture for reduced-precision GLM release files.  The Q2
   routed reference is lower quality but should stay near first-token match
   `92/100`, API top-1 agreement about `0.890`, and API pair-order agreement
   about `0.800` unless the quantization changed deliberately.
-- Run the 100-case DeepSeek V4 PRO fixture for every released PRO GGUF:
-  `gguf-tools/quality-testing/score_official /path/to/deepseek-v4-pro.gguf gguf-tools/quality-testing/data/pro/manifest.tsv /tmp/pro.tsv 4096`.
+- Match every PRO GGUF to its checkpoint. The June preview uses
+  `gguf-tools/quality-testing/data/pro/manifest.tsv`; the August 0813 release
+  uses `gguf-tools/quality-testing/data/pro-0813/manifest.tsv`. Never interpret
+  a cross-checkpoint score as a quantization result.
+- Run the 100-case DeepSeek V4 PRO 0813 fixture for the new release GGUF:
+  `gguf-tools/quality-testing/score_official /path/to/deepseek-v4-pro-0813.gguf gguf-tools/quality-testing/data/pro-0813/manifest.tsv /tmp/pro-0813.tsv 4096 --ssd-streaming`.
 - For SSD streaming, run the same official-continuation scorer once with full
   residency and once with `--ssd-streaming` for the release model.  The summary
   and API agreement should stay in the same quality band.
@@ -108,6 +221,134 @@ Use the normal Flash GGUF that 128 GB users run.
   run `ds4-bench` with `speed-bench/promessi_sposi.txt` and compare prefill,
   generation speed, and KV bytes with the last known good numbers for the same
   machine.
+- For native MXFP4 changes, run `make mxfp4-dot-test test-mxfp4-metal`, then a
+  short greedy prompt and the section 3 continuation fixture with the MXFP4
+  GGUF. The synthetic fused-MoE test and full-model quality gate must both pass.
+
+### DSpark / DeepSpec Runtime
+
+DSpark is opt-in, but it mutates the verifier, target-hidden capture, support
+model loading, and scheduler paths.  Run these whenever DSpark support,
+speculative verification, confidence/scheduler policy, target hidden capture,
+tiny routed-MoE verifier kernels, or shared `--mtp-model` support-model code changes:
+
+Use the 0731 DSpark support GGUF only with a Flash 0731 target. A support model
+from another checkpoint can have plausible acceptance statistics while
+producing a different greedy continuation.
+
+Normal DSpark runs commit accepted target-verifier state directly. The batched
+verifier and one-token decode use the same graph with different floating-point
+operation order, so `output_match=0` against the baseline is diagnostic rather
+than a failure. `--dspark-strict` remains the byte-identical target-only mode.
+
+- Default greedy acceptance fixture:
+  `DS4_DSPARK_MODEL=/Users/antirez/ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf DS4_DSPARK_SUPPORT=/Users/antirez/ds4/gguf/DeepSeek-V4-Flash-DSpark-support-0731.gguf make dspark-acceptance`.
+- 64-token guardrail:
+  `DS4_DSPARK_FIXTURE_TOKENS=64 DS4_DSPARK_MODEL=/Users/antirez/ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf DS4_DSPARK_SUPPORT=/Users/antirez/ds4/gguf/DeepSeek-V4-Flash-DSpark-support-0731.gguf make dspark-acceptance`.
+- Opportunistic sampled acceptance fixture:
+  `DS4_DSPARK_FIXTURE_TEMPERATURE=1 DS4_DSPARK_FIXTURE_TOP_P=0.95 DS4_DSPARK_FIXTURE_MIN_P=0.05 DS4_DSPARK_FIXTURE_SEED=12345 DS4_DSPARK_FIXTURE_TOKENS=32 DS4_DSPARK_MODEL=/Users/antirez/ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf DS4_DSPARK_SUPPORT=/Users/antirez/ds4/gguf/DeepSeek-V4-Flash-DSpark-support-0731.gguf make dspark-acceptance`.
+  Require proposals, a direct commit, and zero errors. Output identity with the
+  baseline is not expected: normally evaluated boundary tokens are sampled,
+  while verified DFlash/target greedy matches are committed without sampling.
+- Exact sampled acceptance fixture:
+  repeat the command above with `DS4_DSPARK_FIXTURE_EXACT_SAMPLING=1`.
+  This mode must preserve the requested target distribution. Seeded output
+  identity is not expected because speculative accept/reject decisions consume
+  additional random numbers.
+- `make test` includes a 100,000-draw distribution check for general `p/q`
+  correction and for the point-mass DFlash proposal used at runtime. Both
+  histograms must remain within the stated tolerance of the target
+  distribution.
+- Fixed-block direct partial commit:
+  `DS4_DSPARK_FIXTURE_CONFIDENCE=0 DS4_DSPARK_FIXTURE_TOKENS=8 DS4_DSPARK_FIXTURE_REQUIRE_PARTIAL=1 DS4_DSPARK_MODEL=/Users/antirez/ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf DS4_DSPARK_SUPPORT=/Users/antirez/ds4/gguf/DeepSeek-V4-Flash-DSpark-support-0731.gguf make dspark-acceptance`.
+- DSpark verifier invariant smoke:
+  `DS4_TEST_MODEL=/Users/antirez/ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf DS4_DSPARK_SUPPORT=/Users/antirez/ds4/gguf/DeepSeek-V4-Flash-DSpark-support-0731.gguf make dspark-verify-depth`.
+- If shared support-model or verifier structures changed, also run legacy MTP:
+  `make mtp-verify-depth` with `DS4_TEST_MTP` set to a one-stage MTP support
+  GGUF, or confirm the target skips only because the optional file is missing.
+- Record `c_add` `accepted_draft`, `direct_full`, `direct_partial`,
+  `replay_fallbacks`, `errors=0`, `verify_layer`, `net_saved`, and
+  `output_match` for both 32-token and 64-token runs. At least one direct commit
+  must occur. A faster run with lower proposal quality is a regression unless
+  it was an intentional scheduler change.
+- If verifier MoE kernels changed, run one diagnostic `c_add` profile with
+  `DS4_DSPARK_VERIFY_SELECTED_PROFILE=1` or the Metal MoE stage profiler and
+  record the selected-expert footprint or stage timing in the DSpark log.
+- On Metal, benchmark at least one predictable code continuation and one
+  deliberately unpredictable prose continuation at temperature 1. For the
+  128-token M5 Max hash-table prompt, ordinary sampling measured
+  41.71/44.49/44.73 t/s and opportunistic DSpark measured
+  47.15/49.08/48.19 t/s, a median gain of about 8.3%. A surreal-prose control
+  measured 44.46 t/s ordinary and 41.66 t/s opportunistic. DSpark remains
+  opt-in because prompts with little useful speculation can still be slower.
+
+### Session Microbatching And Metal TP
+
+Run these gates whenever session scheduling, batched decode, mixed
+prefill/decode, QKV projection, shared or routed experts, tensor parallelism,
+or backend fallback selection changes.
+
+- On a single Metal machine, run the full-vocabulary exact-logit oracle with
+  2, 4, 8, and 16 sessions:
+  `DS4_TEST_MODEL=/path/to/ds4flash.gguf DS4_TEST_SESSION_COUNT=N make test-metal-session-batch`.
+  Compatible resident Q8 runs must report `native_shared=1 native_qkv=1` at
+  every tested count. The 16-session run covers row counts above the old
+  artificial eight-row limit.
+- Repeat the four-session oracle with
+  `DS4_METAL_SESSION_BATCH_SHARED=0` and with
+  `DS4_METAL_SESSION_BATCH_QKV=0`. The first run must use the complete fallback;
+  the second may batch the shared expert only. Both must remain bit-exact.
+- The oracle must cover reversed row ordering, at least six decode steps, and a
+  mixed prefill/decode call. Any nonzero differing-logit count is a blocker
+  unless the model-specific section below declares a measured full-logit
+  tolerance; argmax-only agreement is insufficient. When
+  `DS4_TEST_LIVE_CONTROLS=1` is used, its interleaved serial evaluations are a
+  correctness stress and must remain excluded from the reported batch timing.
+- Benchmark 1, 2, 4, 8, and 16 simultaneous resident sessions on the same host
+  and model. Record model-step latency and aggregate decode tokens/second, not
+  only request completion speed. The current Metal path batches QKV and part of
+  the shared expert, but still runs attention, routed experts, shared down, and
+  the output head per session. Treat flat aggregate scaling as unfinished
+  implementation work, not evidence that Metal cannot benefit from batching.
+- On `mac-m5max-it` and `mac-m5max-us`, run the same oracle in physical TP mode
+  over explicit `tcp` and `rdma` transports. Set `DS4_TEST_TP_MODE=leader` on
+  the leader and `DS4_TEST_TP_MODE=worker DS4_TEST_TP_LEADER_HOST=HOST` on the
+  worker, with a unique `DS4_TEST_TP_PORT`. Run at least 2 and 4 sessions and
+  preserve both logs. GLM 5.3 uses native row batching below token 4096 in TP
+  too; its model-specific tolerance applies. Other unsupported TP shapes must
+  select their established ordered fallback.
+- For the current TB5 MacBook link, US is `10.99.0.2` on `en1`/`rdma_en1` and
+  IT is `10.99.0.1` on `en6`/`rdma_en6`; both use GID index 1. Before testing,
+  require `rdma_ctl status` to report `enabled` and `ibv_devinfo -v` to show
+  `PORT_ACTIVE` plus the corresponding `::ffff:10.99.0.x` GID. Force the
+  device and GID with `--rdma-device NAME --rdma-gid-index 1` if automatic
+  selection is ambiguous. A working TB IP ping alone is not RDMA evidence.
+- Kill the TP worker during one batch with `DS4_TEST_TP_DISCONNECT=1` on the
+  leader. The operation must fail cleanly, invalidate every affected session,
+  and return control without hanging. When remote control latency exceeds the
+  default one-second marker window, set
+  `DS4_TEST_TP_DISCONNECT_DELAY_MS=10000` and signal the worker as soon as
+  `TP_DISCONNECT_READY` appears. Repeat over TCP and RDMA by sending `SIGSTOP`
+  to the exact worker PID at that marker, then resume it after the leader
+  returns. Startup must report the default `gate-timeout=750ms`; both paused
+  runs must fail the gate and invalidate all affected sessions without a Metal
+  GPU watchdog error. Before accepting that deadline, run one normal
+  GLM-5.3 batch over each transport and one GLM-5.2 IQ2 RDMA prompt using the
+  larger 108 GiB shard. `DS4_TP_GATE_TIMEOUT_MS` is a diagnostic override, not
+  a setting required for normal inference.
+- Set `DS4_TEST_TP_IDENTITY_MISMATCH=1` on the test leader once and require both
+  ranks to reject the hello before inference. Also reflect a leader hello from
+  a test peer without changing its role; the leader must reject two peers that
+  both claim the coordinator rank. Separately point a worker at an unused port
+  with `DS4_TP_TIMEOUT_SEC=1`; it must return a connection error rather than
+  retrying indefinitely.
+- Verify unsupported combinations explicitly. GLM 5.2, GLM 5.3 after the
+  4,096-token sparse boundary, DSpark support models, quality/reference modes,
+  and CPU-router modes must use their established exact fallback or reject the
+  combination before evaluation. GLM 5.3 below the boundary, including
+  directional steering, and supported SSD-streaming configurations have native
+  batching and must pass their model-specific oracle instead of being forced
+  off.
 
 ## 5. Metal PRO Path
 
@@ -120,7 +361,300 @@ PRO support is experimental, but release builds must not break it silently.
   touching model shape, tensor lookup, routed expert mapping, template logic,
   and KV payload compatibility.
 
-## 6. SSD Streaming
+## 6. GLM 5.2 And GLM 5.3
+
+GLM has a different template, model shape, MTP block, attention layout,
+tensor-parallel gate width, and streaming policy. Flash or PRO success does not
+substitute for this matrix.
+
+- On a 512 GB Metal machine, run short greedy prompts with both the Q4 XL and
+  reduced-precision Q2 release GGUFs. Cover thinking and no-thinking templates,
+  and verify the server reports the GLM model family rather than a DeepSeek
+  alias internally.
+- Run the OpenRouter smoke vectors explicitly:
+  `DS4_TEST_MODEL=/path/to/glm.gguf
+  DS4_TEST_VECTOR_FILE=tests/test-vectors/glm-openrouter/official.vec
+  ./ds4_test --logprob-vectors`.
+  Preserve the report as a diagnostic. The hosted vectors include very
+  low-probability top-20 tails whose membership is not stable after GLM routed
+  expert quantization, so an individual `official top token missing locally`
+  assertion is not by itself a release blocker. Selected-token mismatches must
+  remain consistent with the model's 100-case first-token band, and the
+  section 3 scorer is the release gate for aggregate GLM quality.
+- Run the 100-case Q4 XL and Q2 official fixtures from section 3 and preserve
+  both `summary` and `api_summary` lines. Compare against the documented Q4 and
+  Q2 reference bands independently.
+- Run `tests/glm_long_context_smoke.sh` with the release-advertised context on
+  the 512 GB Metal host. The generated continuation must begin with `>` and
+  contain none of the known token-corruption markers.
+- Exercise integrated GLM MTP with `--mtp-timing` on a deterministic
+  prompt. Compare the greedy text to
+  a non-MTP run, require clean speculative cycles, and record acceptance and
+  timing. Also run once with MTP disabled to prove ordinary decode remains the
+  default.
+- Run the Metal session oracle with 2 and 4 GLM 5.2 sessions. It must report
+  `family=glm native_shared=0 native_qkv=0` and remain exact, including mixed
+  prefill/decode; the DeepSeek-only row-grid kernels must not activate.
+- Run resident and SSD-streaming GLM Q2 prompts with the same greedy input.
+  Compare first token and top-logprob sanity, and record the selected full-layer
+  prefix and dynamic expert-cache budget.
+- Run physical two-machine GLM TP over TCP and RDMA with short and long prompts.
+  Record prefill/decode speed, transport, rank residency, and clean shutdown.
+  The long prompt must cross the 4,096-token indexed-attention boundary. After
+  changes to split attention, score at least 100 teacher-forced tokens beyond
+  that boundary against the unsplit reference and run the 100-case Q2 fixture
+  once through physical TP, preserving its `summary` and `api_summary` lines.
+  Repeat one run with `--tensor-parallel-token-prefill` as the exact-arithmetic
+  diagnostic. Test both Q2 and Q4 routed-expert files whose types have
+  ownership-aware GLM TP kernels. Each rank must map only its owned experts and
+  match the accepted single-host graph. A routed type without ownership-aware
+  kernels must still reject clearly before evaluation rather than loading a
+  partial split or hanging.
+  The released GLM 5.2 IQ2_XXS file keeps `indexer.proj.weight` in FP32. Its
+  loader and Metal graph must accept that established layout; GLM 5.3 may use
+  its quantized or BF16 indexer projection instead. A short exact-output smoke
+  on the current two-M5 RDMA setup returned `GLM52_OK` with 108.63 GiB planned
+  per rank.
+- With explicit permission for the current QA pass, run one resident GLM Q2
+  prompt, a long-context prompt, integrated GLM MTP, and concurrent server
+  requests on the eight-GPU CUDA host. Use ordinary eight-GPU layer placement
+  for GLM; do not pass the Flash-specific
+  `--cuda-tensor-parallel` option. Multi-tier GLM prefill must
+  report progress through the tier-switching token-major path, and decode,
+  cache updates, and output-head/logit assembly must complete without CPU spill.
+  Auto-placement must reserve each layer's compact DSA/indexer cache and the
+  graph workspace before loading weights; a late graph-allocation failure is a
+  release blocker. Confirm the long-context layout stays within every device's
+  budget and uses additional tiers when the cache no longer fits on the earlier
+  ones.
+  The long-context harness can select this backend with
+  `DS4_GLM_BACKEND=cuda` and pass placement flags through
+  `DS4_GLM_EXTRA_ARGS="--gpu-vram auto --gpu-devices 0,2,4,6,1,3,5,7"`.
+- Through `ds4-server`, exercise OpenAI chat, Responses, and Anthropic requests
+  against GLM, including thinking and SSE. DeepSeek compatibility endpoint
+  aliases may resolve to the loaded model, but rendered prompts and generated
+  text must use the GLM template.
+- Compare a non-tool prompt and a complete assistant tool-call/tool-result
+  transition against the model's `chat_template.jinja` byte for byte. Include
+  reasoning effort, a client system message, a function description and JSON
+  schema, assistant reasoning, and an observation. Token-count agreement alone
+  is insufficient. The fixed one-tool fixture is 901 bytes with SHA-256
+  `f1718f3ebb1c41532bcd5eedd9ebd5c84ae930b018e31962d8743efdbc5affc3`.
+  Run `./ds4_test --server` as the model-free regression for the same exact
+  schema instructions and transition delimiters.
+
+### GLM 5.3 Flash
+
+GLM 5.3 adds mHC, recurrent KDA, a pool-4 DSA indexer, and a different MTP
+block. A GLM 5.2 pass does not cover these paths.
+
+- Run the section 3 GLM 5.3 Q2 and Q4 100-case fixtures before and after any
+  graph, quantization, attention, KDA, mHC, TP, or cache change.
+- Build and run the focused primitive test:
+  `make tests/test_glm53_kda && ./tests/test_glm53_kda`.
+  It covers BF16 projections, pool-4 state construction and expansion, grouped
+  scorer arithmetic and causal visibility, recurrent KDA prefill versus
+  sequential decode, and exact repeated causal-attention output on ROCm.
+- Treat session construction as the attention-memory admission point. Every
+  owned DSA cache and indexer pool/tail, every KDA recurrent state, and the
+  complete supported prefill workspace must allocate before a request is
+  accepted. A first prefill must not grow a per-layer cache. Check both a
+  4,096-token session and a long session in the memory report.
+- Keep GLM-5.3 in absorbed MLA form: attend densely over the shared compact
+  latent cache through token 4,096, then use the pool-4 sparse selector. Do not
+  restore the 2.75 GiB expanded per-head K/V cache as a presumed quality fix.
+  The complete Q2 fixture on the compact Metal graph scored NLL `0.458177271`,
+  first-token agreement `90/100`, and average greedy prefix `7.390`, matching
+  the accepted release band. Fresh Z.AI FP8 long-context controls also favored
+  compact attention: weighted NLL was `0.539823254` versus `0.820997888` for
+  expanded K/V over 24 synthetic cases, and `0.808160860` versus `0.820309487`
+  over 12 natural source-context cases. Extending dense attention to 16K made
+  the natural set worse at `0.825709092`, so keep the 4K crossover.
+- At 100K on an M5 Max, require the compact Q2 plan to remain near 94.09 GiB:
+  89.87 GiB model, 1.11 GiB compact history, and 3.11 GiB fixed graph buffers.
+  An 8,192-token one-shot control on the same graph reached 479.09 prefill and
+  29.89 steady decode t/s. Repeat the continuation fixture after changing the
+  compact cache type, absorbed projections, FlashAttention staging, or the 4K
+  crossover; numerical similarity to the old expanded graph is not the gate.
+- On this 128 GB M3 Max, run the resident Q2 through the generic non-NAX Metal
+  path. Repeat the 4,096-4,100 boundary, official-continuation, MTP, snapshot,
+  server-session, and continued-prefill gates used on M5. Record the different
+  M3 performance floor rather than borrowing the M5 result. Run Q4 only with a
+  bounded SSD-streaming cache; never try to make it fully resident.
+- On one M5 Max, run resident Q2 with a prompt whose actionable instruction
+  begins after token 4096. The model must recover the tail instruction and
+  complete a tool or exact-output task. A short coherent continuation at token
+  4109 is not sufficient: it previously missed a broken sparse selector.
+- Run physical Q4 50/50 TP over the explicit TB5 RDMA devices. Graduate context
+  allocation through 10K, 25K, and 50K, checking both ranks before advancing.
+  The original 50K Q4 run reported 102.20 GiB per rank. Correct accounting for
+  the dynamically allocated mHC/KDA prefill workspace adds about 0.44 GiB, so
+  the expected plan is about 102.64 GiB per rank; record the exact new value on
+  the next run. It remains below the fixed 110 GiB ceiling. Swap must not grow
+  from idle, SSH must stay responsive, and both roles must exit cleanly.
+- The 25K gate must place a read/edit/test task after a long inert archive and
+  include one harmless tool failure that the agent must recover from. The 50K
+  gate must contain at least 30K live prompt tokens and a real source repair:
+  reproduce a failing test, edit the implementation without weakening tests,
+  and pass a warning-strict build. Inspect the resulting diff manually.
+- Confirm the logs select `rdma_en1` with GID 1 on US and `rdma_en6` with GID 1
+  on IT. A TCP fallback does not satisfy this gate. Keep Q2 and Q4 in
+  `~/ds4/gguf` on both hosts after testing.
+- Run ordinary greedy decode, opportunistic MTP, and `--mtp-exact-sampling`.
+  Greedy MTP must preserve the accepted continuation and provide a measured
+  gain; the current short Q2 control improved from 34.21 to 41.97 t/s.
+- After directional-steering changes, verify that a zero `45 x 4096` GLM vector
+  is output-identical to the unsteered CLI for both FFN and attention hooks. A
+  46-row file must be rejected with the expected 737,280-byte size. Build a
+  `/path/to/glm53-direction.f32` test vector with the command in
+  `dir-steering/README.md`, then run it through ordinary decode, `--mtp`, and a
+  two-session `ds4-server` smoke. Run the native batch oracle with:
+
+  ```sh
+  DS4_TEST_MODEL=/path/to/GLM-5.3-Flash-Q2.gguf \
+  DS4_TEST_SESSION_COUNT=4 DS4_TEST_LOGIT_TOLERANCE=0.001 \
+  DS4_TEST_DIRECTIONAL_STEERING_FILE=/path/to/glm53-direction.f32 \
+  DS4_TEST_DIRECTIONAL_STEERING_FFN=1 \
+  DS4_TEST_DIRECTIONAL_STEERING_ATTN=0.25 \
+  make test-metal-session-batch
+  ```
+
+  It must cover native decode and mixed prefill/decode without changing any
+  selected token. Repeat a short physical TP run over RDMA with the same file
+  and scales on both ranks. On CUDA and ROCm release targets, require a
+  warning-free build and one short steered GLM 5.3 Q2 prompt. Finally rerun a
+  held-out target/control sweep; an effective edit that makes control answers
+  repetitive or incoherent does not pass.
+- Run the two- and four-session GLM 5.3 server oracle below token 4096 on one
+  M5 and physical TP. For both native single-M5 and TP paths, set
+  `DS4_TEST_LOGIT_TOLERANCE=0.001`: row-batched reductions may differ from the
+  serial launch order, but every selected token must match and the maximum
+  full-logit delta must remain below that bound. Also run the serial rollback
+  with zero tolerance. The current four- and eight-session Q2 maxima are both
+  `0.0001297`; the current six-step two-session RDMA maximum is `0.000175238`. Past
+  4096, require the exact ordered fallback until a sparse native batch oracle
+  proves full-vocabulary correctness.
+- Exercise every pool remainder at the dense-to-sparse boundary with prompts
+  ending at tokens 4096 through 4100. Token 4096 must remain bit-identical to
+  the serial control. The four sparse cases must retain the same greedy token,
+  contain no nonfinite logits, and pass a multi-token exact-output task. Small
+  batched-reduction logit differences are acceptable only when the official
+  continuation and long-task gates remain in band. Build each prompt against
+  `--dump-tokens`; word counts are not a valid substitute for rendered-token
+  counts. The current `.180` CUDA reference returned exactly `BOUNDARY_OK` at
+  all five frontiers.
+- Run the session snapshot test across the sparse boundary:
+  `DS4_TEST_MODEL=/path/to/GLM-5.3-Flash-Q2.gguf
+  DS4_TEST_SNAPSHOT_PROMPT=/path/to/a-4k-plus-prompt.txt
+  DS4_TEST_SNAPSHOT_CTX=8192 ./ds4_test --session-snapshot`.
+  Restored top logits before and after one continued token must match the
+  uninterrupted session within the test's `1e-6` tolerance.
+- Repeat that command with `DS4_TEST_GLM_MTP=1`. The test must replay 16
+  integrated-MTP cycles across the snapshot, including both one- and two-token
+  outcomes, and match every committed token plus the final top-eight logits.
+  The M5 Max reference produced 10 one-token cycles, 6 two-token cycles, and
+  22 committed tokens without a verifier failure. It then synced back to the
+  original long prompt, reproduced its top-eight logits within `1e-6`, and
+  completed four more MTP cycles. The current `.180` CUDA run produced 5
+  one-token and 11 two-token outcomes, committed 27 tokens, and also passed.
+- Measure continued prefill as actual appends to one live session, not only as
+  a single cold prompt. On an M5 Max, run:
+  `./ds4-bench -m /path/to/GLM-5.3-Flash-Q2.gguf --metal
+  --prompt-file /path/to/a-25k-prompt.txt --ctx-start 4096 --ctx-max 12288
+  --ctx-alloc 16384 --step-incr 2048 --gen-tokens 0 --csv /tmp/glm53.csv`.
+  The current 2K append results are 452.56, 423.74, 409.79, and 396.72 t/s at
+  4K, 6K, 8K, and 10K resident prefixes. A warmed first append below 350 t/s
+  requires investigation. The serial rollback control measured 29.68 t/s.
+  On `.180` CUDA, the current 4K, 6K, and 8K append results are 522.66, 506.20,
+  and 502.50 t/s.
+- Keep the 33987-token Q4 TP agent run as the final long-state gate. The old
+  serial sparse path took about 24 minutes to reach its first tool call. The
+  batched path processed a 34023-token initial suffix in 109.846 seconds
+  (309.74 t/s), completed the full read/edit/test task in 178.93 seconds, and
+  used explicit RDMA. A faster result must still pass the tail task, fixture
+  inspection, and official-continuation gates.
+- On one DGX Spark, run Q2 through CUDA and repeat the primitive, official
+  continuation, 4,096-4,100 boundary, continued-prefill, snapshot, MTP, server,
+  and coding-agent gates. Validate independently on `.180` and `.181`; they are
+  separate single-host runs, not CUDA TP. Q4 and Spark-to-Spark RDMA are not
+  supported in this pass. The accepted `.180` 100-case reference is average
+  NLL `0.461783551`, first-token agreement `90/100`, and average greedy prefix
+  `7.49`.
+- For the default compact CUDA graph, dump the complete first-token logits at
+  a 1,024-token frontier twice with the same context allocation. Both files
+  must be byte-identical. Run the 100-case GLM-5.3 fixture from the same linked
+  objects, then run the two-session single-GPU oracle with
+  `DS4_TEST_CUDA_SINGLE_GPU=1 DS4_TEST_SESSION_COUNT=2`. Require
+  `nonexact_logits=0`. Finally run the fused D2R kernels under CUDA memcheck;
+  an argmax-only comparison or coherent text does not replace these gates.
+
+### GLM 5.3 Vision
+
+Vision is a separate sidecar on Metal, single-GPU CUDA, and ROCm, and has its
+own release gate. Text-only GLM success does not exercise image preprocessing,
+the vision graph, multimodal prompt spans, or image-aware KV identity.
+
+- Download `glm53-vision` and verify
+  `GLM-5.3-Flash-Vision-Encoder.gguf` has SHA-256
+  `ae23e14c6979e889051b2e4a39351abcdafb161e18e606fae4d8c40095a4bf3a`.
+- Build `tests/test_glm53_vision_engine` and
+  `tests/test_glm53_vision_prompt`. Run them with the release Q2 text GGUF, the
+  vision sidecar, and a fixed PNG. The prompt test must generate a visual
+  answer, reuse an unchanged image without repeated prefill, and rebuild when
+  only the image fingerprint changes. It must also hold image-token positions
+  fixed, replace the visual embedding with zeros, and observe changed output
+  logits. Use a 4,096-token test session so a large screenshot plus the output
+  cannot hit the old 2,048-token test ceiling. After explicit session
+  invalidation, and again after restoring the zeroed embedding, require the
+  complete image-conditioned logits to match the original within `1e-6`.
+  This catches a compact-prefill path that processes placeholders but silently
+  ignores the image data, as well as incomplete multimodal state rebuilds.
+- Keep one accepted Metal embedding from a fixed image and compare CUDA and
+  ROCm output with `tests/compare_glm53_vision_embeddings.py`. Require finite
+  output, cosine similarity at least `0.995`, mean absolute error at most
+  `0.001`, and maximum absolute error at most `0.06`. This permits normal BF16
+  GEMM ordering differences but rejects a changed vision graph.
+- Run a fixed model-level vision fixture containing photographs, screenshots,
+  diagrams, readable text, spatial questions, and unrelated-image controls.
+  Compare complete answers with the official GLM-5.3-Flash vision service and
+  repeat through CLI, server, and `ds4-agent`. A valid decoder, expected image
+  token count, or plausible but ungrounded prose does not pass this gate. With
+  a local vision-enabled server running, require:
+
+  ```sh
+  python3 tests/run_glm53_vision_quality.py
+  ```
+
+  to report `6/6 passed`. For agent tests, expose only the raster fixtures;
+  source SVGs or expected-answer files beside them let the agent bypass vision
+  with text tools.
+- Run the decoder over RGB, RGBA, grayscale, and palette PNG, baseline and
+  progressive JPEG, EXIF orientation, truncated files, wrong CRCs, huge
+  dimensions, and decompression-bomb fixtures under ASan and UBSan. Invalid
+  files must fail without a sanitizer report or large allocation.
+- In `./ds4`, submit one PNG and one JPEG with `/read`, then continue each chat
+  with a text turn. Repeat once with `--mtp`; verification after image prefill
+  must complete without a GLM MTP failure.
+- In `ds4-agent`, require `view_image` to inspect a real file and use the
+  resulting multimodal observation in a later read/edit/test tool loop. Image
+  observations must enter as user-role multimodal turns; GLM loses grounding
+  when several images are packed into a tool-response role. Run the five-image
+  fixture in one turn so the prompt exceeds 4K, and require the same facts as
+  the official Z.AI control. Text-only tool observations must remain tool-role.
+- Through `ds4-server`, test OpenAI Chat data URIs, Responses `input_image`,
+  and Anthropic base64 image blocks. Include two images in one message and an
+  image in a later turn. Local paths, `file:` URLs, remote URLs, malformed
+  base64, unsupported media, more than 16 images, and bodies over 64 MiB must
+  return 4xx without reading local files or making network requests.
+- Run Q4 across `mac-m5max-us` and `mac-m5max-it` over explicit TB5 RDMA with
+  `--vision` on both ranks. The leader must encode once, both ranks must keep
+  matching multimodal KV state, and the answer must remain correct. Record
+  image encode, prefill, first-token, and decode timing separately.
+- Build CPU, CUDA, and ROCm targets warning-free after vision changes. Run the
+  encoder comparison, prompt replay test, and six-case server fixture on one
+  DGX Spark and on `strixhalo`; ROCm Q2 must use bounded SSD streaming.
+
+## 7. SSD Streaming
 
 SSD streaming is a capacity path, so test both correctness and user experience.
 
@@ -130,7 +664,7 @@ SSD streaming is a capacity path, so test both correctness and user experience.
   with boosted Q4 routed-expert layers and a prompt long enough to exercise the
   selected-address prefill path; it must not fail with "model range is not
   covered by mapped model views":
-  `./ds4 -m gguf/DeepSeek-V4-Flash-Layers37-42Q4KExperts-OtherExpertLayersIQ2XXSGateUp-Q2KDown-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-fixed.gguf --ssd-streaming --ssd-streaming-cache-experts 16GB --ctx 4096 --tokens 1 --nothink --prompt-file /tmp/ds4_600tok_prompt.txt`.
+  `./ds4 -m gguf/DeepSeek-V4-Flash-Layers37-42Q4KExperts-OtherExpertLayersIQ2XXSGateUp-Q2KDown-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-fixed-0731.gguf --ssd-streaming --ssd-streaming-cache-experts 16GB --ctx 4096 --tokens 1 --nothink --prompt-file /tmp/ds4_600tok_prompt.txt`.
 - Cold streaming measurement:
   run once with `--ssd-streaming-cold` and verify no deadlock, missing expert,
   or impossible slowdown.
@@ -139,25 +673,83 @@ SSD streaming is a capacity path, so test both correctness and user experience.
 - If streaming cache internals changed, test the same prompt twice and compare
   first-token/logprob sanity between runs.
 
-## 7. CUDA / DGX Spark
+## 8. CUDA / DGX Spark
 
 Before a release, ask the user for CUDA access if it is not already configured.
-Use the DGX Spark / GB10 host `toor@192.168.0.180`.  Do not claim CUDA is
-release-ready without this pass.
+Use either DGX Spark / GB10 host, `toor@192.168.4.180` or
+`toor@192.168.4.181`. Do not claim CUDA is release-ready without this pass.
+
+Both Sparks normally run vLLM. Before stopping it, record its process, service
+or container, model, ports, and exact launch command. Confirm all vLLM workers
+have exited before loading DwarfStar. At the end, stop every DwarfStar process,
+restore the exact vLLM service, and verify its original ports and model health.
+Do not use high-performance Hugging Face Xet mode while vLLM is resident.
 
 - Fetch or push the exact release commit to the CUDA machine.
 - Build:
   `make clean && make cuda-spark`.
+- Require both the DGX Spark build and the eight-GPU CUDA build to complete
+  without compiler warnings. The eight-GPU build is performed only after
+  receiving explicit permission to use `192.168.60.250` for this QA pass.
 - Run:
   `make cuda-regression`.
+- For native MXFP4 changes, run
+  `make test-mxfp4-cuda CUDA_ARCH=native` on the multi-GPU CUDA host only after
+  receiving explicit permission for `192.168.60.250`, and
+  `make test-mxfp4-cuda CUDA_ARCH=sm_121` on DGX Spark. Dense MMQ, routed MMQ,
+  routed MMVQ, fused gate/up, and fused down must pass. The Spark run must also
+  pass the Blackwell K-tile guard. This synthetic parity test does not replace
+  full-model continuation scoring.
+- With that permission, run the native MXFP4 GGUF resident on the multi-GPU
+  host, and run it with `--ssd-streaming` on DGX Spark. Use the same greedy prompt and continuation
+  fixture on both. Record prefill and generation speed, require finite logits,
+  and compare quality with the Metal MXFP4 result. Blackwell MMQ quantizes
+  activations to native FP4 for batched work; decode MMVQ keeps Q8 activations,
+  so quality must be checked rather than inferred from kernel-only parity.
 - Run a short CLI prompt with the Flash GGUF and record generation t/s.
 - Run a longer prompt that exercises routed experts past a few thousand tokens.
+- With explicit permission for this QA pass, run the full-vocabulary decode
+  oracle on the eight-GPU CUDA host:
+  `DS4_TEST_MODEL=/path/to/flash.gguf make test-cuda-session-batch`.
+  Preserve the per-batch timing for 2, 4, and 8 rows and require
+  `nonexact_logits=0`. Run the released Q4 file and the reduced-precision Q2
+  file: Q4 exercises grouped routed/shared stages, while unsupported Q2 native
+  MoE shapes must retain the ordered exact fallback.
+- With CUDA TP attention enabled, compatible Q4 runs must use grouped
+  attention-core, QKV, KV-store, and attention-post by default and remain
+  full-vocabulary exact against isolated decode. On the eight-L40S host, the
+  16-row decode step must remain above 110 aggregate tokens/s. Repeat once with
+  `DS4_CUDA_TP_ATTN=0` only as rollback coverage; it is not the production
+  configuration.
+- Run native mixed prefill/decode at the default frontier and at compressed
+  context:
+  `DS4_TEST_MODEL=/path/to/flash.gguf make test-cuda-mixed-batch` and
+  `DS4_TEST_CONTEXT=4096 DS4_TEST_MIXED_INITIAL=2048 DS4_TEST_MIXED_ROUNDS=8
+  DS4_TEST_MODEL=/path/to/flash.gguf make test-cuda-mixed-batch`.
+  Every round must report exact logits and `mode=native`; a serialized fallback
+  is a failure for the eight-GPU TP/EP topology. Under CUDA TP attention, the
+  native mixed step must use the same exact grouped decode stages when their
+  capability checks pass; record correctness and speedup separately. Also
+  force an 800-row prefill quantum with
+  `DS4_TEST_ALLOW_FALLBACK=1`; it must report the serialized safety fallback.
+- With explicit permission for the eight-GPU host, start `ds4-server` with 8
+  and 16 batched sessions and issue at least that many simultaneous requests
+  with mixed prompt lengths. Verify no session mix-up, deadlock, or starvation
+  and record aggregate generation throughput.
+- On DGX Spark, verify the same public batch API and server concurrency use the
+  single-GPU fallback without creating peer-only TP/EP state. The eight-GPU
+  native oracle is not a valid Spark test because its topology is intentionally
+  unavailable there.
+- For GLM 5.3, use the resident Q2 artifact only. Require the dedicated CUDA
+  primitive and continuation gates in section 6, then record prefill,
+  generation, MTP, continued-prefill, server aggregate throughput, and peak
+  memory. Do not attempt the 178 GiB Q4 artifact on one 128 GB Spark.
 - If CUDA Q4, distributed, streaming hooks, tensor span loading, or model cache
   code changed, test the specific GGUF and split mode that uses that path.
 - Verify that any CUDA-only warning fixes are also clean on macOS and do not
   change Metal behavior.
 
-## 8. ROCm / Strix Halo
+## 9. ROCm / Strix Halo
 
 Use the Strix Halo Framework Desktop via the VPN hostname `strixhalo`
 (`antirez@strixhalo`).  This host validates the ROCm backend; do not use it as
@@ -166,19 +758,54 @@ a substitute for CUDA or Metal release testing.
 - Fetch or push the exact release commit to the Strix Halo machine.
 - Build:
   `make clean && make strix-halo`.
+- Require the ROCm build to complete without compiler warnings.
+- After MXFP4 or ROCm routed-MoE changes, run `make test-mxfp4-rocm`. Require
+  zero `failures` for both `mid` and `out` at 1, 3, 32, 128, and 512 tokens,
+  followed by `MXFP4 ROCm routed MoE: PASS`.
 - Use the q2 Flash imatrix GGUF for release smoke tests:
-  `DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf`.
+  `DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf`.
 - Do not use the mixed q2-q4 or Q4 Flash GGUFs for routine Strix Halo QA yet.
   They are dangerous on this machine for now because the ROCm path can hit
   system OOM instead of failing cleanly.
 - Run a short CLI prompt:
-  `./ds4 -m gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf --ctx 4096 --nothink -p "Reply with exactly: OK"`.
+  `./ds4 -m gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf --ctx 4096 --nothink -p "Reply with exactly: OK"`.
+- For DeepSeek Flash decode, confirm the default path uses prequantized Q8
+  activations. Repeat the same greedy run with
+  `DS4_ROCM_DSV4_PREQUANT_DECODE=0` only as a diagnostic control. The default
+  must be materially faster and must still pass the continuation-quality gate.
+  GLM and `--quality` must stay on the full-FP32 activation path.
+- Test DSpark with the matched 0731 target and support files:
+  `DS4_BIN=./ds4 DS4_DSPARK_MODEL=gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf DS4_DSPARK_SUPPORT=gguf/DeepSeek-V4-Flash-DSpark-support-0731.gguf DS4_DSPARK_FIXTURE_TOKENS=64 sh tests/dspark_acceptance_fixture.sh`.
+  Require proposals, accepted draft tokens, at least one direct state commit,
+  zero verifier errors, and zero replay fallbacks. Record ordinary and DSpark
+  generation speed separately. When direct verifier-state handling changes,
+  also compare with a test-only build of its immediate replay predecessor; the
+  direct build must be faster. DSpark is not currently expected to beat
+  ordinary ROCm decode, so do not describe it as a ROCm speedup without a new
+  measurement.
+- Repeat one opportunistic DSpark run with
+  `--temp 1 --top-p 0.95 --min-p 0.05`, then repeat it with
+  `--mtp-exact-sampling`. The current 128-token code references are
+  16.26 t/s ordinary, 12.28 t/s opportunistic, and 13.52 t/s exact, with no
+  verifier errors. This is a correctness gate, not a ROCm speed claim; the
+  ROCm batched verifier is still too expensive.
 - Run one longer prompt if ROCm kernels, backend hooks, tensor loading, model
   cache, KV cache, or graph prefill code changed.
+- Run the GLM Q2 release model through ROCm SSD streaming with at least four
+  generated tokens:
+  `./ds4 --rocm -m gguf/GLM-5.2-UD-Q2_K_RoutedQ2K.gguf --ssd-streaming --ctx 4096 --nothink --tokens 4 -p "Reply with exactly: OK"`.
+  Startup must select a cache budget that passes the memory guard without an
+  override, and both compact indexed prefill and decode must complete.
+- Run one longer GLM prompt with the release-advertised Strix context after
+  changes to GLM attention, typed quantized projections, streaming expert
+  caches, or memory budgeting. Record the context, cache split, and whether
+  the continuation stays free of token-corruption markers.
+- Run the same GLM model with `--mtp-timing --temp 0`. At least one draft
+  verification cycle must complete without a `glm mtp step failed` message.
 - Record startup memory/cache messages, prefill speed, generation speed, and
   whether the backend reports `ROCm backend initialized`.
 
-## 9. Distributed Inference
+## 10. Distributed Inference
 
 Distributed code has regressed around route setup, KV snapshots, request IDs,
 and split model loading.  Test it whenever distributed, KV, session, or model
@@ -196,7 +823,7 @@ loading code changes.
 - If CUDA distributed is relevant, test across the CUDA hosts and record
   generation speed, not just "it works".
 
-## 10. Disk KV Cache
+## 11. Disk KV Cache
 
 Disk KV cache bugs are high impact for server users.
 
@@ -210,7 +837,7 @@ Disk KV cache bugs are high impact for server users.
 - Test stripped agent sessions: `/strip <id>` then `/switch <id>` should rebuild
   by prefill and render sane history.
 
-## 11. Server APIs
+## 12. Server APIs
 
 The server must keep compatibility across OpenAI, Responses, and Anthropic
 clients.
@@ -220,10 +847,24 @@ clients.
 - Test OpenAI chat completion, OpenAI Responses, and Anthropic messages.
 - Test SSE streaming with thinking enabled and disabled.
 - Test keepalive during long prefill and confirm clients do not time out.
+- In batched mode, close clients while their requests are queued, prefilling,
+  and streaming decode. Repeat across OpenAI chat, Responses, Anthropic, and
+  completions. Abandoned work must stop at the next backend-safe boundary, and
+  a valid request after each cancellation must complete normally.
+- For the repeatable chat-completions cancellation and slot-reuse gate, run
+  `python3 tests/test_server_batching.py --url http://127.0.0.1:8000 --pairs 2
+  --workers 4 --case short-sampled --max-tokens 12 --cancel-first 4`. Then run
+  at least twelve short four-request waves against the same four-slot server.
+  Every pair must remain deterministic and the server must answer `/v1/models`
+  after malformed JSON and an over-context request.
+- Only after receiving explicit permission for this QA pass, start
+  `ds4-server` on the eight-L40S CUDA TP target with the release TP options and
+  verify all 16 100k-context sessions allocate. Startup must report a
+  2048-token prefill cap; a silent fallback to 4096 is an OOM regression.
 - Test `--trace` and confirm rendered prompts, cache decisions, generated text,
   and tool-parser events are useful without leaking unrelated state.
 
-## 12. ds4-agent
+## 13. ds4-agent
 
 The agent is the most stateful component.  Test it manually, not only by build.
 
@@ -234,8 +875,10 @@ The agent is the most stateful component.  Test it manually, not only by build.
 - Queue messages while the model is busy.  Queued messages must not skip tool
   execution; after tool results, the queued user text must be provided.
 - Read/search/edit/write tools:
-  create a temp project, ask for edits, verify old/new and `[upto]` anchored
-  edits fail safely on ambiguous matches and do not require retyping whole files.
+  create a temp project and ask for edits. By default, verify that exact old/new
+  replacements work and the tool prompt does not advertise `[upto]`. In a
+  separate `--edit-upto` run, verify anchored edits fail safely on ambiguous
+  matches and do not require retyping whole files.
 - Real coding edit loop:
   delete `/tmp/mymandel`, ask ds4-agent to create a small C ASCII Mandelbrot
   program there, build and run it, then in a second user turn ask for a small
@@ -243,6 +886,19 @@ The agent is the most stateful component.  Test it manually, not only by build.
   ASCII character ramp or output dimensions.  Verify the agent edits the
   existing file instead of rewriting the whole project, and that the final
   program still builds and runs.
+- For GLM-5.3 integrated MTP, repeat a long-context edit/build/test task after
+  the prompt crosses token 4,096. Require real file and shell tool calls,
+  accepted and rejected draft cycles, and no `glm mtp: GLM 5.3 verify failed`
+  message. Continue the same session after one harmless tool error and after a
+  snapshot restore so a hidden MTP fallback or damaged recurrent state is not
+  mistaken for success.
+- With a matched DSpark support file and temperature 1, repeat a coding-tool
+  turn that crosses sampled prose, greedy DSML structure, parameter text, and
+  back to sampled prose. Require the tool to execute, the final answer to be
+  valid, and DSpark stats to show zero verifier errors and replay fallbacks.
+  Run the opportunistic default and `--mtp-exact-sampling`. The M5 Max
+  opportunistic smoke created, compiled, and ran a C program printing
+  `OPPORTUNISTIC_OK`, accepting 196 of 235 draft tokens.
 - Bash tools:
   test short output, large output truncation, non-zero exit output, long-running
   jobs, `bash_status`, and `bash_stop`.
@@ -255,7 +911,7 @@ The agent is the most stateful component.  Test it manually, not only by build.
   status bar fill to terminal width, syntax highlighting in Markdown/code blocks,
   and SSH/remote terminal flicker.
 
-## 13. Download Script And Model Files
+## 14. Download Script And Model Files
 
 - Test `download_model.sh` in a temporary directory so local weights are not
   overwritten.
@@ -264,7 +920,7 @@ The agent is the most stateful component.  Test it manually, not only by build.
 - Verify legacy removed targets fail clearly.
 - Verify README model names match the script and Hugging Face repository.
 
-## 14. Performance And Power
+## 15. Performance And Power
 
 - Run `ds4-bench` on the release machine and compare with tracked CSV baselines.
 - Test `--power 100` is not throttled.
@@ -273,18 +929,89 @@ The agent is the most stateful component.  Test it manually, not only by build.
 - Confirm context buffer size, raw KV rows, compressed KV rows, and mmap behavior
   match expectations for 32k, 100k, and any release-advertised context size.
 
-## 15. Release Sign-off
+## 16. Speed Regression
+
+Performance is a release gate. A correct result that is unexpectedly much
+slower still needs an explanation before release.
+
+Use the same commit, GGUF checksum, prompt, context frontier, generated-token
+count, power setting, and backend flags as the reference run. Let the machine
+become idle, discard the first warm-up run, then record the median of three
+runs. Do not compare different model checkpoints or quantizations. For batched
+tests, record aggregate and per-session decode speed.
+
+- A slowdown over 5% requires a clean rerun and investigation.
+- A repeatable slowdown over 10% in prefill, decode, or aggregate batched
+  decode is a release blocker unless the change and tradeoff are documented.
+- Keep the complete `ds4-bench` CSV. A single short-prompt average is not enough
+  to detect a context-dependent regression.
+- Compare startup time and peak memory as well as tokens per second when model
+  loading, caches, streaming, or temporary arenas changed.
+- Run the backend-specific batch tests in sections 4 and 8. Fast single-session
+  decode does not substitute for aggregate multi-session throughput.
+
+These are the last known good observations available when this gate was added.
+They are reference points for matching hardware and workloads, not performance
+claims across different models or contexts.
+
+| System and backend | Model and workload | Prefill | Decode |
+| --- | --- | ---: | ---: |
+| MacBook Pro M3 Max 128 GB, Metal | Flash q2, 11,709-token prompt | 250.11 t/s | 21.47 t/s |
+| MacBook Pro M3 Max 128 GB, Metal | Flash 0731 q2, opportunistic temperature-1 128-token code prompt | - | 29.16 t/s ordinary; 28.18 t/s DSpark |
+| MacBook Pro M5 Max 128 GB, Metal | Flash q2, 11,707-token prompt | 463.44 t/s | 25.90 t/s |
+| MacBook Pro M5 Max 128 GB, Metal | Flash 0731 q2, opportunistic temperature-1 128-token code prompt | - | 44.49 t/s ordinary median; 48.19 t/s DSpark median |
+| Mac Studio M3 Ultra 512 GB, Metal | Flash q2, 11,709-token prompt | 468.03 t/s | 27.39 t/s |
+| Mac Studio M3 Ultra 512 GB, Metal | Flash q4, 12,018-token prompt | 448.82 t/s | 26.62 t/s |
+| Mac Studio M3 Ultra 512 GB, Metal | GLM 5.3 Flash Q4 with Q8 KDA/head, 2,048-token prompt | 437.62 t/s | 24.74 t/s |
+| Two M5 Max 128 GB Macs, Metal TP over TB5 RDMA | GLM 5.2 IQ2_XXS, 4,096-token prefill and 256-token teacher-forced decode | about 214 t/s | about 16.7 t/s |
+| MacBook Pro M5 Max 128 GB, Metal | GLM 5.3 Flash Q2, resident short prompt | 86.68 t/s | 34.45 t/s; 41.97 t/s greedy MTP |
+| MacBook Pro M5 Max 128 GB, Metal | GLM 5.3 Flash Q2, 8,192-token compact-attention prompt | 479.09 t/s | 29.89 t/s steady |
+| Two M5 Max 128 GB Macs, Metal TP over TB5 RDMA | GLM 5.3 Flash Q2, short prompt | 29.06 t/s | 32.70 t/s |
+| MacBook Pro M5 Max 128 GB, Metal | GLM 5.3 Flash Q2, 24,988/49,948-token long prompts | 424.80 / 421.75 t/s | 29.50 / 28.10 t/s |
+| Two M5 Max 128 GB Macs, Metal TP over TB5 RDMA | GLM 5.3 Flash Q2, 10,819-token prompt | 468.97 t/s | 22.85 t/s |
+| Two M5 Max 128 GB Macs, Metal TP over TB5 RDMA | GLM 5.3 Flash Q4, 34,023-token agent prefill suffix | 309.74 t/s | full coding task: 178.93 s wall |
+| MacBook Pro M5 Max 128 GB, Metal | GLM 5.3 Flash Q2, 2/4/8 resident decode sessions | - | 54.49 / 73.96 / 86.17 aggregate rows/s |
+| DGX Spark GB10, CUDA | GLM 5.3 Flash Q2, 2,048-token prefill and 16-token decode | 531.39 t/s | 14.35 t/s at 2,048 context |
+| DGX Spark GB10, CUDA | GLM 5.3 Flash Q2, live-session 4K/6K/8K continued prefills | 522.66 / 506.20 / 502.50 t/s | - |
+| DGX Spark GB10, CUDA | GLM 5.3 Flash Q2, strict two-session 24-step oracle | - | 16.5 aggregate rows/s; byte-exact full logits |
+| DGX Spark GB10, CUDA | Flash q2, 7,047-token prompt | 343.81 t/s | 13.75 t/s |
+| DGX Spark GB10, CUDA | Flash q2 DSpark, 64-token C fixture | - | 24.48 t/s direct; 13.93 t/s replay predecessor |
+| DGX Spark GB10, CUDA | pre-0731 Flash q2, exact-sampled 128-token code prompt | - | 18.17 t/s ordinary; 18.30 t/s DSpark |
+| DGX Spark GB10, CUDA | pre-0731 Flash q2, opportunistic temperature-1 128-token code prompt | - | 18.32 t/s ordinary; 18.43 t/s DSpark |
+| Strix Halo gfx1151, ROCm | Flash IQ2 resident, short section 9 smoke | - | 17.27 t/s; FP32 rollback 9.70 t/s |
+| Strix Halo gfx1151, ROCm | Flash IQ2 resident, 4,096-token context | - | 14.82 t/s; FP32 rollback 8.76 t/s |
+| Strix Halo gfx1151, ROCm | Flash IQ2 DSpark, 64-token C fixture | - | 11.40 t/s direct; 9.77 t/s replay predecessor; 16.70 t/s ordinary |
+| Strix Halo gfx1151, ROCm | Flash 0731 IQ2, exact-sampled 128-token code prompt | - | 16.55 t/s ordinary; 12.68 t/s DSpark |
+| Strix Halo gfx1151, ROCm | Flash 0731 IQ2, temperature-1 128-token code prompt | - | 16.26 t/s ordinary; 12.28 t/s opportunistic; 13.52 t/s exact |
+| 8x L40S, CUDA TP | Flash q4, 2,048-token prefill benchmark | 1,524.84 t/s | 46.93 t/s |
+| 8x L40S, CUDA TP | Flash q4, 16-row decode oracle | - | 126.0 aggregate t/s |
+
+The 8x L40S values are retained from the last recorded run on `192.168.60.250`.
+They are historical references only: never connect to that host or interrupt
+its production server without explicit permission for the current QA pass. If
+permission is granted, the existing hard floor remains 110 aggregate t/s for
+the 16-row decode oracle.
+
+## 17. Release Sign-off
 
 Do not sign off until:
 
 - macOS Metal Flash passed.
+- GLM 5.2 Metal, official-quality, MTP, batching-fallback, and applicable TP or
+  CUDA gates passed.
 - Official continuation quality gates passed for every released model family.
 - CUDA was tested on the CUDA machine or the release notes explicitly say CUDA
   was not validated.
 - ROCm was tested on Strix Halo or the release notes explicitly say ROCm was
   not validated.
+- Metal, CUDA, ROCm, CPU-only, and test builds completed without compiler
+  warnings on every release target that was validated.
 - Disk KV cache was exercised.
 - Server API streaming was exercised.
 - Agent interruption and tool loops were exercised manually.
-- Speed is within expected variance for the same hardware and model.
+- The speed-regression gate passed on every validated backend, with any skipped
+  baseline or intentional slowdown documented.
+- Metal 2/4/8/16-session exactness and forced fallback gates passed.
+- Physical Metal TP batching and CUDA native decode/mixed batching passed when
+  those backends are part of the release.
 - Any skipped item is written down with the reason.

README.md

@@ -1,47 +1,46 @@
-# DwarfStar
+<p align="center">
+  <img src="logo.svg" alt="DwarfStar logo" width="220">
+</p>
 
 **DwarfStar** is a small native inference engine optimized first for
-**DeepSeek V4 Flash**, with support for **DeepSeek V4 PRO** on very high-memory
-machines. It is
-intentionally narrow: not a generic GGUF runner, not a wrapper around another
-runtime: it is completely self-contained. Other than running the model in a
-correct and fast way, the project goal is to provide DeepSeek specific loading,
-prompt rendering, tool calling, KV state handling (RAM and on-disk), server
-API and integrated coding agent, all ready to work with coding agents or with
-the provided CLI interface. There are also tools for GGUF and imatrix generation,
-and for quality and speed testing.
+**DeepSeek V4 Flash**. It also supports **GLM 5.2**, **GLM 5.3 Flash**, and,
+on very high-memory machines, **DeepSeek V4 PRO**. It is self-contained and
+deliberately narrow, not a general GGUF runner. Model loading, prompt rendering,
+tool calls, KV state, the HTTP server, and the coding agent are built and tested together.
+The repository also includes tools and data for GGUF, imatrix, quality, and speed.
 
-We support the following backends:
-* **Metal** is our primary target. Starting from MacBooks with 96GB of RAM (or less, using SSD streaming).
-* **NVIDIA CUDA / DGX Spark**, CUDA with special care for the DGX Spark.
-* **Strix Halo (ROCm)**, systems like the Framework Desktop and other systems based on the same GPU and unified RAM design.
+Supported backends:
+
+* **Metal**, the primary target, on Macs with 96 GB or more. Smaller machines
+  can use SSD streaming.
+* **NVIDIA CUDA**, including multi-GPU systems and DGX Spark.
+* **ROCm** on Strix Halo systems such as the Framework Desktop.
 
 This project would not exist without **llama.cpp and GGML**, make sure to read
 the acknowledgements section, a big thank you to Georgi Gerganov and all the
 other contributors.
 
-**Note that DeepSeek v4** is not our only target. Right now Flash and PRO are the
-perfect fit because of capabilities, size, KV cache efficiency. If tomorrow a
-better open weight model is released for the 128GB size, we could switch, the same
-for other important size classes like 512GB of RAM. The project is stictly
-opportunistic depending on what open weight models exist in a given moment.
-If a new model will be supported, the old one may be removed completely and
-no longer supported, unless there is some kind of overlap of abilities.
+Model support is intentionally opportunistic. The project follows the best open
+weights for useful local machine sizes, especially 128 GB laptops and 512 GB
+workstations. A model may be removed when a better replacement arrives.
+
+# So, what can I do with this software?
+
+* You can run a very capable models in your consumer hardware, a MacBook, a DGX Spark, or a Strix Halo for example. Even if you have not enough RAM, with SSD streaming, you can run it at a decent speed.
+* Using the CUDA multi-GPU support and with ds4-server micro batching of decoding and generation, you can turn a server with old-ish CUDA cards (Ada Lovelace architecture), no longer supported for new models by vLLM, into a multi-user LLM server for your company. We tested this setup with 8xL40S NVIDIA cards and multiple sessions with very good results. 120 t/s aggreated generation, 2000 t/s prefill.
+* Using two MacBook M5 Max / M3 Ultra RDMA, you can run 4 bit DeepSeek Flash, GLM 5.2, or GLM 5.3 Flash with tensor parallelism.
+* You can also use pipeline paralellism to glue together multiple systems to sum their RAM and run larger models.
 
 ## Motivations
 
-* Very capable open weight models finally exist. DeepSeek v4 Flash feels quasi-frontier. The PRO is even better. Both resist 2 bit quantization very well.
-* Very capable computers like MacBooks, the DGX Spark now exist.
-* DeepSeek v4 kv cache design makes it pratical to run very big contexts. Other vendors are using this approach.
-* This few hundred billions models are strictly better than smaller (even if dense) models, regardless of what benchmarks say.
+* Capable open-weight models now fit on high-end personal machines.
+* DeepSeek V4 Flash and PRO, GLM 5.2, tolerate aggressive routed-expert quantization.
+* Compressed KV caches and fast local SSDs make long contexts practical.
+* The idea of an inference system specialized for a few models.
 
-That said, a few important things about this project:
+# AI full disclosure
 
-* The local inference landscape contains many excellent projects, but new models are released continuously, and the attention immediately gets captured by the next model to implement. This project takes a deliberately narrow bet: one model at a time, official-vector validation (logits obtained with the official implementation), long-context tests, and enough agent integration to know if it really works. The exact model may change as the landscape evolves, but the constraint remains: local inference credible on high end personal machines or Mac Studios, starting from 96/128GB of memory.
-* This software is developed with **strong assistance from GPT 5.5** and with humans leading the ideas, testing, and debugging. We say this openly because it shaped how the project was built. If you are not happy with AI-developed code, this software is not for you. The acknowledgement below is equally important: this would not exist without `llama.cpp` and GGML, largely written by hand.
-* This implementation is based on the idea that compressed KV caches like the one of DeepSeek v4 and the fast SSD disks of modern MacBooks should change our idea that KV cache belongs to RAM. **The KV cache is actually a first-class disk citizen**. Fast SSD disks also changed the inference game from the point of view of "model needs to fit RAM": while having more RAM the the model size is still preferred, SSD streaming allows to turn the available amount of RAM from a hard cutoff (can I run this model or not?) to continuous spectrum of speed levels.
-* Our vision is that local inference should be a set of three things working well together, out of the box: A) inference engine with HTTP API + B) GGUF specially crafted to run well under a given engine and given assumptions + C) testing and validation with coding agents implementations. D) Purpose built agents for specific models and execution environments. DwarfStar only runs with the GGUF files provided. It gets tested against officially obtained logits at different context sizes. This project exists because we wanted to make one local model feel finished end to end, not just runnable. However this is beta quality code, so probably we are not still there, especially since recently we introduced large new features: distributed inference, SSD streaming, and other minor improvements.
-* The optimized graph path targets **Metal on macOS** and **CUDA on Linux**. The CPU path is only for correctness checks and model/tokenizer diagnostics. For CPU-only Linux builds, use `make cpu`; it builds the normal `./ds4` and `./ds4-server` binaries without CUDA or Metal. On macOS, **warning: current macOS versions have a bug in the virtual memory implementation that will crash the kernel** if you try to run the CPU code. Remember? Software sucks. It was not possible to fix the CPU inference to avoid crashing, since each time you have to restart the computer, which is not funny. Help us, if you have the guts.
+* This software is developed with **strong assistance from GPT 5.5, 5.6, Claude Fable** and with humans leading the ideas, testing, and debugging. We say this openly because it shaped how the project was built. If you are not happy with AI-developed code, this software is not for you. The acknowledgement below is equally important: this would not exist without `llama.cpp` and GGML, largely written by hand.
 
 ## Acknowledgements to llama.cpp and GGML
 
@@ -58,14 +57,18 @@ notice in our `LICENSE` file.
 
 ## Status
 
-The code and GGUF files are to be considered of **beta quality** because
-inference and model serving is a complicated matter and all this exists
-only for a few days. It will take months to reach a more stable form.
-However, we try to keep the project in a usable state, and we are making
-progress. If you have issues, make sure to use `--trace` to log the
-sessions, and open issues including the full trace.
+The software is currently very fast changing. Consider it beta quality.
+Before each release, a big QA run is executed, however instabilities
+are definitely possible.
 
-The `ds4-agent` is alpha quality, the project was later added.
+# How to use this project?
+
+I (Salvatore) believe that the way projects should be shipped and used changed because of AI. The main differences today are:
+
+1. With AI, users can modify the software in significant ways with low efforts, costs, and even lacking deep domain knowledge about the task they want to accomplish. For instance, a DwarfStar user with a specific hardware setup can ask a coding agent to improve the inference speed of this software for the specific hardware setup, asking the model to reach the maximum prefill and generation speed without impacting correctness, and also asking to do a deep QA pass.
+2. Similiarly, because of "1", software may be shipped in a different way than before. It must be more a working template for the biggest use cases, without trying to cover every possible setup. If DwarfStar showcases a few good implementations of tensor parallel execution, the code will work as a rail for implementing the same feature in specific conditions, for a new model, and so forth.
+
+So, while this project attempts to be usable for the featured models and the most common hardware setups, I ask you, if you have access to coding agents, to consider using coding agents as an interface to discover the project, make modifications, create personalized setups. This way you can likely do more than what we ship, and certain things that are not documented or implemented, and that you require, are potentially very easy to achieve.
 
 ## More Documentation
 
@@ -75,6 +78,8 @@ next sections.
 
 - [CONTRIBUTING.md](CONTRIBUTING.md): correctness and speed regression testing
   guide for contributors. **Read this before sending a pull request**.
+- [QA_BEFORE_RELEASES.md](QA_BEFORE_RELEASES.md): the complete release test
+  matrix, including the remote Metal, CUDA, and ROCm machines.
 - [gguf-tools/README.md](gguf-tools/README.md): offline GGUF generation,
   imatrix collection, quantization tooling, and quality checks.
 - [gguf-tools/imatrix/README.md](gguf-tools/imatrix/README.md): how the
@@ -92,11 +97,12 @@ next sections.
 
 ## Model Weights
 
-This implementation only works with the DeepSeek V4 Flash and PRO GGUFs published for
-this project. It is not a general GGUF loader, and arbitrary DeepSeek/GGUF files
-will not have the tensor layout, quantization mix, metadata, or optional MTP
-state expected by the engine. The 2 bit quantizations provided here are not
-a joke: they behave well, work under coding agents, call tools in a reliable way.
+This implementation only works with the DeepSeek V4 and GLM GGUFs listed
+below. It is not a general GGUF loader, and arbitrary GGUF files will not have
+the tensor layout, quantization mix, metadata, or optional MTP state expected by
+the engine. The 2 bit quantizations provided here are verified to be actually
+high quality: they behave well, work under coding agents, call tools in a reliable way.
+
 The 2 bit quants use a very asymmetrical quantization: only the routed MoE
 experts are quantized, up/gate at `IQ2_XXS`, down at `Q2_K`. They are the
 majority of all the model space: the other components (shared experts,
@@ -105,12 +111,18 @@ projections, routing) are left untouched to guarantee quality.
 Download one main model. **Prefer the imatrix versions.**
 
 ```sh
-./download_model.sh q2-imatrix   # 96/128 GB RAM machines, imatrix-tuned q2
-./download_model.sh q2-q4-imatrix  # 96/128 GB RAM machines, q2 with last 6 layers q4
-./download_model.sh q4-imatrix   # >= 256 GB RAM machines, imatrix-tuned q4
-./download_model.sh pro-q2-imatrix  # 512 GB RAM machines, PRO q2 imatrix quant
+./download_model.sh ds4f-q2      # 96/128 GB RAM machines
+./download_model.sh ds4f-q2-q4   # q2 with the last 6 expert layers at q4
+./download_model.sh ds4f-q4      # >= 256 GB RAM machines
+./download_model.sh ds4f-mxfp4   # native MXFP4 experts, about 156 GB
+./download_model.sh pro-q2-imatrix  # 512 GB RAM machines, PRO 0813 q2 imatrix
 ```
 
+The MXFP4 GGUF preserves DeepSeek's released MXFP4 routed-expert weights rather
+than requantizing them. It runs on Metal and CUDA; Blackwell CUDA devices use
+native FP4 matrix instructions and FP4 activations for batched expert work.
+Decode and other CUDA devices use Q8 activations.
+
 For the full PRO Q4 distributed run, download one half on each machine:
 
 ```sh
@@ -118,9 +130,11 @@ For the full PRO Q4 distributed run, download one half on each machine:
 ./download_model.sh pro-q4-layers31-output  # second half of PRO Q4 split
 ```
 
-The script downloads from `https://huggingface.co/antirez/deepseek-v4-gguf`,
-stores files under `./gguf/`, resumes partial downloads with `curl -C -`, and
-updates `./ds4flash.gguf` to point at the selected main model.
+The script stores files under `./gguf/` and updates `./ds4flash.gguf` to point
+at the selected main model. DeepSeek files come from
+`antirez/deepseek-v4-gguf`; GLM targets use the repository named in the script's
+help. Smaller files resume with `curl -C -`, while large files use the official
+Hugging Face downloader.
 The `pro-q4-layers00-30`, `pro-q4-layers31-output`, and `pro-q4-split` targets
 download distributed PRO Q4 pieces and do not update `./ds4flash.gguf`.
 Authentication is optional for public downloads, but `--token TOKEN`,
@@ -128,37 +142,56 @@ Authentication is optional for public downloads, but `--token TOKEN`,
 
 If you want to regenerate GGUF files or collect a new imatrix, see
 [gguf-tools/README.md](gguf-tools/README.md). Those tools are meant for offline
-model-building work and can take a long time on the full DeepSeek V4 Flash
-weights. Flash GGUF generation is supported by the local tools. PRO GGUF
-production currently still depends on the external `llama.cpp`-based workflow;
-native tooling can be added later.
+model-building work and can take a long time on the full DeepSeek weights.
+Flash and PRO GGUF generation are supported by the local tools. PRO conversion
+uses a compatible published PRO GGUF as its metadata, tensor-layout, and output
+type template.
 
-`./download_model.sh mtp` fetches the optional speculative decoding support
-GGUF for Flash. It can be used with q2-imatrix, q2-q4-imatrix, and q4-imatrix,
-but must be enabled explicitly with `--mtp`. The current MTP/speculative
-decoding path is still experimental: it is correctness-gated and currently
-provides at most a slight speedup, not a meaningful generation-speed win.
-
-GLM 5.2 support is currently limited to the GGUF files tested by this branch:
+GLM 5.2 support is limited to the GGUF files tested by this branch:
 
 ```sh
 ./download_model.sh glm-unsloth-q4  # Unsloth UD-Q4_K_XL, 11 shards
+./download_model.sh glm-antirez-iq2xxs  # antirez routed IQ2_XXS single-file GGUF
 ./download_model.sh glm-antirez-q2  # antirez routed Q2_K single-file GGUF
 ./download_model.sh glm-antirez-q4  # antirez routed Q4_K single-file GGUF
 ```
 
-The supported GLM layout keeps dense/model-control tensors in the existing
+GLM 5.3 Flash has its own graph, artifacts, and run instructions in the
+[GLM 5.3 Flash](#glm-53-flash) section below.
+
+The supported GLM 5.2 layout keeps dense/model-control tensors in the existing
 Q8/F32 paths and supports routed expert gate/up tensors in `Q2_K`, `Q4_K`, or
 `Q5_K`; routed expert down tensors are supported in `Q2_K`, `Q4_K`, `Q5_K`, or
 `Q6_K`. Other GLM GGUF quant layouts should be treated as unsupported until they
 are added deliberately and scored against the official 100-case fixture.
 
+These GLM 5.2 formats do not all support the same execution modes. The Q4 files
+work for normal Metal and CUDA inference. Two-Mac tensor parallelism for GLM
+5.2 currently requires an ownership-aware IQ2_XXS or Q2_K routed layout. GLM
+5.3 has its own ownership-aware Q4 path.
+
+GLM's MTP block is part of the main GGUF; it does not use the separate Flash
+MTP file. Ordinary decode remains the default. `--mtp` enables experimental
+greedy speculation. `--mtp-timing` also enables it and prints acceptance
+and timing counters:
+
+```sh
+./ds4 -m gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \
+  --mtp-timing --temp 0
+```
+
+GLM 5.2 uses the Metal, CUDA, or ROCm graph backend. GLM 5.3 is validated on
+Metal, with Q2 also validated on CUDA. Directional steering, `--power` below
+100, an explicit `--prefill-chunk`, and an external `--mtp-model` file are not
+supported for GLM yet.
+
 Then build:
 
 ```sh
 make                  # macOS Metal
 make cuda-spark       # Linux CUDA, DGX Spark / GB10
 make cuda-generic     # Linux CUDA, other local CUDA GPUs
+make strix-halo       # Linux ROCm, AMD Strix Halo
 make cpu              # CPU-only diagnostics build
 ```
 
@@ -166,37 +199,230 @@ make cpu              # CPU-only diagnostics build
 select another supported GGUF from `./gguf/`. Run `./ds4 --help` and
 `./ds4-server --help` for the full flag list.
 
+## GLM 5.3 Flash
+
+GLM 5.3 Flash uses a separate graph for its recurrent KDA layers, sparse DSA
+layers, hyper-connections, and built-in MTP block. The release GGUFs were made
+from the official FP8 checkpoint:
+
+```sh
+./download_model.sh glm53-q2  # about 90 GiB
+./download_model.sh glm53-q4  # about 178 GiB
+./download_model.sh glm53-fp8 # about 305 GiB; packaged weights only
+```
+
+The Q2 file uses imatrix-guided IQ2_XXS gate/up experts and Q2_K down experts.
+It runs resident on a 128 GB M3 Max or M5 Max, and on one DGX Spark. The Q4
+file is the higher-quality control. Run it across two 128 GB Macs, or use SSD
+streaming on one Mac. The FP8 file preserves the released text weights without
+requantization; DwarfStar cannot execute that artifact yet.
+
+On one 128 GB Mac:
+
+```sh
+./ds4 -m gguf/GLM-5.3-Flash-Q2.gguf --ctx 32768
+```
+
+Q2 is close to the practical memory limit on a 128 GB Mac. Stop other
+memory-heavy workloads before loading it resident.
+
+The model's MTP block is already inside the same GGUF. `--mtp` enables it;
+`--mtp-timing` also prints acceptance and timing counters. At non-zero
+temperature the default mode directly keeps target-matching greedy drafts.
+Add `--mtp-exact-sampling` when the output must preserve the ordinary target
+sampling distribution.
+
+```sh
+./ds4 -m gguf/GLM-5.3-Flash-Q2.gguf --mtp --temp 0 --ctx 32768
+./ds4-agent -m gguf/GLM-5.3-Flash-Q2.gguf --mtp --ctx 50000
+```
+
+For a small multi-user server on one M5 Max, four 4096-token sessions fit the
+tested Q2 layout. Native decode batching is used below token 4096; longer
+sessions automatically use the ordered fallback.
+
+```sh
+./ds4-server -m gguf/GLM-5.3-Flash-Q2.gguf \
+  --ctx 4096 --batched-session 4
+```
+
+To run Q4 on one Mac without making it resident:
+
+```sh
+./ds4 -m gguf/GLM-5.3-Flash-Q4_K.gguf --ssd-streaming --ctx 4096
+```
+
+For resident Q2 or Q4 across two 128 GB Macs, use the 50/50 setup documented
+under [Tensor Parallelism over RDMA](#tensor-parallelism-over-rdma). Put the
+same GGUF at the same path on both machines, start the worker first, and select
+`--transport rdma` on both sides. Q4 is the main reason to use this setup.
+
+On one DGX Spark, use Q2 without CUDA tensor parallelism:
+
+```sh
+make cuda-spark
+./ds4 --cuda -m gguf/GLM-5.3-Flash-Q2.gguf --ctx 16384
+```
+
+Q4 does not fit one Spark, and Spark-to-Spark RDMA tensor parallelism is not
+implemented.
+
+### Vision
+
+GLM 5.3 Flash vision uses a separate 1.1 GB encoder. The text GGUF stays the
+same, and vision is enabled only when the sidecar is passed explicitly:
+
+```sh
+./download_model.sh glm53-vision
+./ds4 -m gguf/GLM-5.3-Flash-Q2.gguf \
+  --vision gguf/GLM-5.3-Flash-Vision-Encoder.gguf
+```
+
+The published encoder SHA-256 is
+`ae23e14c6979e889051b2e4a39351abcdafb161e18e606fae4d8c40095a4bf3a`.
+
+In the interactive CLI, `/read photo.jpg` or `/read image.png` submits the
+image as a user turn. `ds4-agent` exposes the same support as its `view_image`
+tool when started with `--vision`. JPEG and PNG decoding is built in; no image
+library is required.
+
+`ds4-server` accepts ordered image blocks in OpenAI Chat, Responses, and
+Anthropic requests. HTTP images must be inline: use a PNG/JPEG data URI for
+OpenAI or base64 image source for Anthropic. File paths and remote URLs are
+rejected. A request may contain up to 16 images and the HTTP body is limited to
+64 MiB.
+
+Vision runs on Metal, single-GPU CUDA, and ROCm. On a DGX Spark, use the CUDA
+command above and add `--vision FILE`. On the 128 GB Strix Halo reference host,
+Q2 needs the same SSD-streaming options as text inference:
+
+```sh
+./ds4 --rocm --ssd-streaming --ssd-streaming-cache-experts 32GB \
+  -m gguf/GLM-5.3-Flash-Q2.gguf \
+  --vision gguf/GLM-5.3-Flash-Vision-Encoder.gguf
+```
+
+In two-Mac tensor parallel mode, pass the same `--vision` file on both the
+coordinator and worker; the coordinator encodes the image and sends the
+projected visual tokens to the worker.
+
+## DSpark Speculative Decoding
+
+DSpark is an auxiliary draft model released by DeepSeek for DeepSeek V4 Flash.
+It reads hidden states from the main model and proposes up to five future
+tokens. DwarfStar checks those proposals with the main Flash model and commits
+only the accepted prefix. The main model remains authoritative; a rejected or
+low-confidence suffix falls back to ordinary target decoding.
+
+The possible gain is faster generation: when several proposed tokens are
+accepted, one target verification pass advances the stream by several tokens.
+It does not accelerate prefill, and the draft and verification work is not
+free. Predictable continuations, especially code, tend to benefit most;
+low-yield prompts can be no faster or even slower. DSpark is therefore still
+experimental and explicitly opt-in.
+
+Accepted proposals keep the state produced by the batched target verifier
+instead of running the same tokens through one-token decode again. Both paths
+execute the same inference graph, but floating-point operations are grouped in
+a different order. A long greedy DSpark run may therefore diverge from a run
+without DSpark after an otherwise valid accepted block. This is not a reduced
+precision or approximate-model mode; use ordinary decoding, `--quality`, or
+`--dspark-strict` when byte-for-byte reproducibility with one-token decode is
+required.
+
+The DSpark checkpoint for Flash 0731 is packaged here as a separate support
+GGUF of about 5.6 GiB. It is not a standalone model. Download it once:
+
+```sh
+./download_model.sh ds4f-dspark
+```
+
+The support file can be used with the 0731 Flash `ds4f-q2`, `ds4f-q2-q4`, and
+`ds4f-q4` models listed above. It is checkpoint-specific
+and must not be paired with an older Flash model. For now **DeepSeek V4 PRO**
+is not supported. On Metal, the main model may be resident or use
+`--ssd-streaming`; the support model still adds its own weights and runtime
+state to the memory requirement. DSpark replaces the legacy one-stage MTP
+support model for that run rather than stacking with it.
+
+Run it with the normal sampling defaults:
+
+```sh
+./ds4 -m ds4flash.gguf \
+  --mtp-model gguf/DeepSeek-V4-Flash-DSpark-support-0731.gguf \
+  --dspark
+```
+
+`--mtp-model` supplies the support GGUF, while `--dspark` selects the DSpark runtime.
+The default confidence threshold is `0.6` on Metal and `0.7` on CUDA and ROCm.
+It prunes suffixes that are unlikely to repay their verification cost.
+`--dspark-confidence 0` forces fixed five-token blocks and is intended for
+diagnostics.
+
+At a non-zero temperature, ordinary `--dspark` uses opportunistic sampling.
+Tokens evaluated normally are sampled with the requested temperature, top-p,
+top-k, and min-p. DFlash then proposes a temperature-zero suffix. Every draft
+token that matches the target's temperature-zero continuation is committed
+directly, even though the requested temperature is non-zero. Sampling resumes
+at the first mismatch. This is deliberately more deterministic than ordinary
+temperature sampling. On an M5 Max it retained enough of the greedy DSpark gain
+to improve a predictable code continuation by about 8% at temperature 1. A
+single M3 Max run was slightly slower, the same test was nearly neutral on DGX
+Spark, and it was slower on Strix Halo, where verification is more expensive.
+
+Use `--mtp-exact-sampling` when the output must follow the ordinary target
+distribution. With the DFlash support model, exact mode disables direct
+temperature-zero matching: it accepts each greedy proposal with its target
+probability and, on rejection, samples from the remaining target distribution.
+It uses a stricter `0.8` confidence threshold by default. Add `--temp 0` for
+fully greedy decoding. `--quality` and `--dspark-strict` keep target-only
+decoding, which is useful for reproducibility checks.
+The same DSpark flags work with `ds4-agent` and with non-batched
+`ds4-server` requests. Session-batched serving currently uses ordinary target
+decoding.
+
 ## Speed
 
-These are single-run Metal CLI numbers with `--ctx 32768`, `--nothink`, greedy
-decoding, and `-n 256`. The short prompt is a normal small Italian story
-prompt. The long prompts exercise chunked prefill plus long-context decode.
-Q4 requires the larger-memory machine class, so M3 Max Q4 numbers are `N/A`.
+The current q2 results use `ds4-bench` with the standard *Promessi sposi*
+input, 2048-token context steps, and 128 greedy generation tokens at every
+frontier. Each prefill number is for the next 2048-token chunk. The complete
+sweeps are in [m5_max.csv](speed-bench/m5_max.csv) and
+[gb10.csv](speed-bench/gb10.csv).
+
+| Machine | Backend | Context | Prefill | Generation |
+| --- | --- | ---: | ---: | ---: |
+| MacBook Pro M5 Max, 128 GB | Metal | 2048 | 790.18 t/s | 39.35 t/s |
+| MacBook Pro M5 Max, 128 GB | Metal | 16384 | 572.53 t/s | 36.14 t/s |
+| MacBook Pro M5 Max, 128 GB | Metal | 32768 | 557.04 t/s | 34.36 t/s |
+| MacBook Pro M5 Max, 128 GB | Metal | 65536 | 398.50 t/s | 27.64 t/s |
+| DGX Spark GB10, 128 GB | CUDA | 2048 | 825.76 t/s | 18.05 t/s |
+| DGX Spark GB10, 128 GB | CUDA | 16384 | 872.44 t/s | 15.10 t/s |
+| DGX Spark GB10, 128 GB | CUDA | 32768 | 855.94 t/s | 14.43 t/s |
+| DGX Spark GB10, 128 GB | CUDA | 65536 | 822.98 t/s | 13.84 t/s |
+
+Older measurements for machines and model variants not rerun in this pass are
+kept for reference. They used the earlier CLI prompt procedure and are not
+directly comparable with the table above.
 
 | Machine | Quant | Prompt | Prefill | Generation |
 | --- | ---: | ---: | ---: | ---: |
 | MacBook Pro M3 Max, 128 GB | q2 | short | 58.52 t/s | 26.68 t/s |
 | MacBook Pro M3 Max, 128 GB | q2 | 11709 tokens | 250.11 t/s | 21.47 t/s |
-| MacBook Pro M3 Max, 128 GB | q4 | short | N/A | N/A |
-| MacBook Pro M3 Max, 128 GB | q4 | long | N/A | N/A |
-| MacBook Pro M5 Max, 128 GB | q2 | short | 87.25 t/s | 34.27 t/s |
-| MacBook Pro M5 Max, 128 GB | q2 | 11707 tokens | 463.44 t/s | 25.90 t/s |
 | Mac Studio M3 Ultra, 512 GB | q2 | short | 84.43 t/s | 36.86 t/s |
 | Mac Studio M3 Ultra, 512 GB | q2 | 11709 tokens | 468.03 t/s | 27.39 t/s |
 | Mac Studio M3 Ultra, 512 GB | q4 | short | 78.95 t/s | 35.50 t/s |
 | Mac Studio M3 Ultra, 512 GB | q4 | 12018 tokens | 448.82 t/s | 26.62 t/s |
 | Mac Studio M3 Ultra, 512 GB | PRO q2 | 32768 tokens | 138.82 t/s | 9.56 t/s |
-| DGX Spark GB10, 128 GB | q2 | 7047 tokens | 343.81 t/s | 13.75 t/s |
 
-![M3 Max t/s](speed-bench/m3_max_ts.svg)
+![M5 Max t/s](speed-bench/m5_max_ts.svg)
 ![PRO model M3 Ultra t/s](speed-bench/pro_model_m3_ultra_ts.svg)
 
 ## Running models larger than RAM
 
 The normal Metal path tries to make the model resident in GPU-addressable
 memory. This is the fastest path and should remain your default when the model
-fits. When it does not fit, DwarfStar also has a Metal-only **SSD streaming**
-capacity mode. In this mode the non-routed model weights stay resident, while
+fits. DwarfStar also has an **SSD streaming** capacity mode on Metal and for
+GLM 5.2 on ROCm. In this mode the non-routed model weights stay resident, while
 routed MoE experts are kept in an in-memory cache and loaded from the GGUF file
 on cache misses.
 
@@ -224,12 +450,12 @@ The `32GB` value is a routed-expert memory budget, not a generic byte cache.
 DwarfStar first reserves headroom for the two full routed layers used by
 overlapped streaming prefill, then converts the remaining bytes to the number of
 dynamic cached experts that fit for the current GGUF. Explicit `NGB` budgets may
-also be capped after context/KV accounting so the Metal working set stays out of
+also be capped after context/KV accounting so the backend working set stays out of
 the slow pressure zone. A plain number such as
 `--ssd-streaming-cache-experts 4000` is different: it means exactly 4000 dynamic
 expert slots, with no extra accounting. Non-routed weights, KV cache, graph
 scratch, and activations need additional memory. The automatic cache budget takes
-80% of the Metal recommended working set, subtracts non-routed weights, then
+80% of the backend's recommended working set, subtracts non-routed weights, then
 applies the same routed-prefill headroom before sizing the dynamic cache. Leave
 the hot expert preload enabled for normal use; use `--ssd-streaming-cold` and
 `--ssd-streaming-preload-experts N` only for measurements.
@@ -239,7 +465,7 @@ the hot expert preload enabled for normal use; use `--ssd-streaming-cold` and
 On 64GB MacBooks, start with the 2-bit Flash GGUF and a moderate expert cache:
 
 ```sh
-./download_model.sh q2-imatrix
+./download_model.sh ds4f-q2
 
 ./ds4 \
   -m ./ds4flash.gguf \
@@ -256,7 +482,7 @@ and occasional work when you accept slow generation. Start with `--nothink`:
 ./download_model.sh pro-q2-imatrix
 
 ./ds4 \
-  -m gguf/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf \
+  -m gguf/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix-0813.gguf \
   --ssd-streaming \
   --ctx 32768 \
   --nothink
@@ -272,25 +498,47 @@ re-enable thinking with a conservative generation limit:
 
 ```sh
 ./ds4 \
-  -m gguf/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf \
+  -m gguf/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix-0813.gguf \
   --ssd-streaming \
   --ctx 32768 \
   --think \
   --tokens 1500
 ```
 
+GLM 5.2 uses the same option. Its streaming path keeps the largest full-layer
+prefix that fits resident, then uses the remaining budget for a dynamic expert
+cache. Start with the automatic budget:
+
+```sh
+./ds4 \
+  -m gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \
+  --ssd-streaming \
+  --ctx 32768
+```
+
 The important startup line is the cache report. Start conservative, then
 increase the cache if the machine has headroom.
 
-## Distributed Inference
+On a 128GB Strix Halo, use the routed Q2_K model and a 4096-token context as the
+starting point. The automatic cache budget leaves room for the GLM graph and KV
+state:
 
-Distributed inference lets DwarfStar **run a model that is too large for one machine** by
+```sh
+./download_model.sh glm-antirez-q2
+make strix-halo
+./ds4 --rocm -m gguf/GLM-5.2-UD-Q2_K_RoutedQ2K.gguf \
+  --ssd-streaming --ctx 4096
+```
+
+## Distributed inference with pipeline parallelism
+
+Pipeline parallelism lets DwarfStar **run a model that is too large for one machine** by
 splitting transformer layers across multiple machines. The main example is the
 full 4-bit Flash quant across two 128 GB MacBooks: each process maps only its
 own layer slice, activations are sent over TCP, and the coordinator keeps normal
 CLI/API behavior.
 
-Distributed inference also allows to **speed up prefill** by
+Pipeline parallelism can also **speed up prefill** by
 using multiple GPUs at the same time to process different micro-batches at
 different layers, like in an assembly line. Only prefill can be accelerated this
 way. Generation is purely autoregressive: each token must finish across the
@@ -305,6 +553,12 @@ To build an initial mental model, here are the high level concepts:
 4. Each worker keeps its slice of the KV cache.
 5. Communication is worker-to-worker, there is no need to use the coordinator as relay, so if your coordinator is `A`, and you make a request, activations will flow in `A -> B -> C -> back to A`.
 
+The resident ROCm MXFP4 routed-expert path supports the same pipeline mode. A
+tested two-host Strix Halo split uses `--layers 0:21` on the coordinator and
+`--layers 22:output` on the worker. This is a capacity configuration for a
+model that does not fit on one 128 GB system; it does not add ROCm SSD
+streaming support for Flash.
+
 ### How it works and how to configure it
 
 The prefill path is pipelined (this is why it can go faster than in a single machine).
@@ -499,6 +753,174 @@ encryption or authentication, and is not release-stable yet; coordinator and
 workers should be built from the same commit and used on trusted machines and
 trusted networks.
 
+## Tensor Parallelism over RDMA
+
+Tensor parallelism runs a single decode across two Macs connected with a
+Thunderbolt 5 cable, splitting the heavy per-layer work between the two
+GPUs and exchanging 16-24KB partial sums at synchronization gates inside the
+graph (RDMA over Thunderbolt when available, a dedicated TCP socket
+otherwise). Unlike the pipelined distributed mode above, both
+machines work on the *same token at the same time*, so it reduces
+per-token latency instead of just fitting a bigger model.
+
+Each machine keeps one contiguous half of the routed experts resident. Dense,
+attention, shared-expert, embedding, and output weights remain replicated.
+This lets a model whose routed experts do not fit on one machine run fully
+resident across the pair; routed kernels never touch the peer's expert half.
+
+### Running GLM 5.2 or GLM 5.3 across two 128 GB MacBooks
+
+One-time setup per boot, on **both** machines:
+
+```sh
+# Let the GPU wire ~117 GB (default cap is ~75% of RAM; the resident
+# expert shard needs ~97.5 GiB plus KV/scratch).
+sudo sysctl iogpu.wired_limit_mb=120000
+
+# RDMA over Thunderbolt needs an IPv4 address directly on the cabled
+# member interface (the bridge IP does not count). Use the interface
+# that is 'active' in ifconfig, e.g. en1 on one side and en6 on the
+# other. Skip this if you are fine with the TCP fallback.
+sudo ifconfig en1 inet 10.99.0.2/30 alias     # machine A
+sudo ifconfig en6 inet 10.99.0.1/30 alias     # machine B
+```
+
+Check the verbs device before loading the model:
+
+```sh
+rdma_ctl status
+ibv_devinfo -v
+```
+
+The device must be active and expose the IPv4-mapped GID for the address above,
+for example `::ffff:10.99.0.2`. A working IP ping does not prove that RDMA is
+active.
+
+Both machines need the same tree, commit, and GGUF path. Tensor parallelism is
+always a 50/50 split with one worker, so do not pass `--layers`. Start the worker
+first; it retries while the coordinator loads. The worker must dial the address
+on the Thunderbolt member interface, not the bridge address:
+
+```sh
+MODEL=gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf
+# For GLM 5.3 Q4, use MODEL=gguf/GLM-5.3-Flash-Q4_K.gguf instead.
+
+# Machine B: worker.
+./ds4 -m "$MODEL" --tensor-parallel --role worker \
+  --coordinator 10.99.0.2 9911 --transport rdma
+
+# Machine A: coordinator.
+./ds4 -m "$MODEL" --tensor-parallel --role coordinator \
+  --listen 10.99.0.2 9911 --transport rdma -c 8192 \
+  -p "Tell me something about the sea."
+```
+
+The active verbs device and IPv4-mapped GID are selected automatically. If that
+is ambiguous, add `--rdma-device rdma_en6 --rdma-gid-index 1` on the worker and
+the matching `rdma_en1` flags on the coordinator. Use `--transport tcp` on both
+sides to force TCP. Run workers with `ds4`; the coordinator may be `ds4`,
+`ds4-agent`, or `ds4-server`.
+
+Startup takes about 9 seconds per machine: each rank pre-faults its
+~100 GiB shard from SSD and pins it through a Metal residency set.
+DeepSeek V4 Flash works the same way with its own GGUF on both machines.
+DeepSeek gate vectors are 16 KB and ride as one RDMA message. GLM 5.2's
+6144-wide 24 KB vectors are split into two ordered RDMA messages. GLM 5.3 uses
+its own KDA/DSA gate schedule, exchanged and checked during TP startup.
+
+Measured on two M5 Max 128 GB MacBooks (GLM 5.2, IQ2_XXS, 188 GiB):
+
+| | two Macs, tensor parallel | one Mac, SSD streaming |
+|---|---|---|
+| decode | ~16.8 t/s (15.4 at 4k context) | ~4.8 t/s |
+| prefill (4096 tokens) | ~94 t/s | ~3-5 t/s |
+| residency | fully memory-resident | streams experts from SSD |
+
+Notes: the coordinator mirrors every prompt sync and eval to the worker, so
+both KV caches stay in lockstep; prompt processing splits both the
+routed-expert GEMMs (by expert ownership) and the attention heads (a
+contiguous half per machine) with one bulk partial-sum exchange per
+layer per stage (`--tensor-parallel-token-prefill` selects a slower
+token-by-token prefill that exactly matches the single-machine arithmetic).
+The split graph is deterministic, but its changed floating-point reduction
+order is not generally byte-identical to single-machine execution.
+
+## Tensor Parallelism across CUDA GPUs
+
+On a single CUDA server, `--cuda-tensor-parallel` splits DeepSeek V4 Flash
+tensor and routed-expert work across an even number of GPUs. This is separate
+from the Mac-to-Mac mode above: it does not use `--role`, RDMA, or the
+distributed layer pipeline. GPU placement and memory budgets are selected with
+the normal `--gpu-devices` and `--gpu-vram` options.
+
+The device order is significant. With `N` devices, the first `N/2` logical
+tiers are contiguous layer-pipeline homes and the second `N/2` tiers are their
+tensor-parallel partners. Specify all homes first and then all partners, with
+the closest P2P pair at matching positions. For example, the tested L40S host
+uses physical pairs `(0,1)`, `(2,3)`, `(4,5)`, and `(6,7)`, expressed as
+`0,2,4,6,1,3,5,7`. Each pair stores a 50/50 split of the routed experts, and
+the vocabulary head is row-sharded across the participating output tiers.
+Those large tensors are not duplicated. Dense attention, router, and shared
+expert weights are replicated within each pair.
+
+For maximum throughput on eight 48 GB L40S cards, use the imatrix Q4 model.
+Its routed `Q4_K` layout has the native grouped multi-session kernels; the Q2
+model is the lower-memory choice (including tested four-card runs), but its
+unsupported grouped routed shapes use the exact fallback and have lower
+aggregate serving throughput. Download and build the L40S target with:
+
+```sh
+./download_model.sh ds4f-q4
+make cuda CUDA_ARCH=sm_89
+```
+
+This is the interactive-agent setup used on the eight-L40S server:
+
+```sh
+MODEL=gguf/DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-imatrix-0731.gguf
+
+./ds4-agent --cuda --cuda-tensor-parallel \
+  --gpu-vram auto \
+  --gpu-devices 0,2,4,6,1,3,5,7 \
+  --model "$MODEL" \
+  --ctx 100000
+```
+
+For serving, keep multiple KV sessions resident so decode rows can be grouped
+across requests. The tested host is configured for up to 16 resident sessions:
+
+```sh
+./ds4-server --cuda --cuda-tensor-parallel \
+  --gpu-vram auto \
+  --gpu-devices 0,2,4,6,1,3,5,7 \
+  --model "$MODEL" \
+  --ctx 100000 \
+  --batched-session 16 \
+  --host 0.0.0.0
+```
+
+The equivalent local launchers are `./run-nvidia-tp-agent.sh` and
+`./run-nvidia-tp-server.sh`. The server launcher also enables the on-disk KV
+cache and defaults to the native 0731 MXFP4 GGUF. Set `DS4_MODEL` to use the Q4
+file above instead. Reduce the session count or context size if the requested
+resident KV caches do not fit after model loading. CUDA TP, half-resident expert
+ownership, output sharding, pipelined prefill, and compatible grouped decode are
+selected by `--cuda-tensor-parallel`; no `DS4_CUDA_*` environment tuning is required.
+Without an explicit `--prefill-chunk`, this mode uses 2048-token chunks so the
+tested 16-session, 100k-context layout retains enough VRAM for resident KV
+caches. An explicit `--prefill-chunk` remains an override for other topologies.
+
+Any even card count that can hold the selected model and graph scratch is a
+valid topology. On this class of 48 GB card, the useful measured endpoints are
+Q2 on four cards (two pipeline stages) and Q4 on eight cards (four stages).
+For a four-card PIX-paired subset such as physical GPUs `0,1,4,5`, the ordered
+list is `0,4,1,5`. Two cards do not have enough memory for these Flash models.
+
+This mode currently requires DeepSeek V4 Flash and an even multi-GPU
+placement. GLM 5.2 instead uses normal layer placement across the selected
+CUDA devices. DGX Spark is a single-GPU target and must not be started with
+`--cuda-tensor-parallel`.
+
 ## Reducing heat, power usage and fan noise
 
 Long local inference runs can keep the GPU busy for extended periods. If you
@@ -512,8 +934,8 @@ and inserting small sleeps between work units: during prefill it sleeps between
 layers, and during generation it sleeps between decoded tokens. This reduces
 sustained load without changing model output.
 
-The option is available on the CLI, server, agent, eval, and benchmark tools,
-for example:
+The option is available on the CLI, server, agent, eval, and benchmark tools
+for DeepSeek models. GLM 5.2 currently accepts only `--power 100`. For example:
 
 ```sh
 ./ds4 --power 50
@@ -581,12 +1003,10 @@ exponential sweeps. Output is CSV with one row per frontier: latest prefill
 interval tokens/sec, generation tokens/sec at that frontier, and
 `kvcache_bytes`.
 
-Sessions prefill long prompts in 4096-token chunks by default. Set
-`DS4_METAL_PREFILL_CHUNK=N` to compare another chunk size, for example `2048`
-to match the strict official-vector checkpoint path, or
-`DS4_METAL_PREFILL_CHUNK=0` to prefill a prompt as one whole batch when memory
-allows. Changing the chunk changes the KV checkpoint/logit path, so compare it
-as an explicit run configuration.
+Sessions prefill long prompts in 4096-token chunks by default. Use
+`--prefill-chunk 2048`, for example, to match the strict official-vector
+checkpoint path. Changing the chunk changes the KV checkpoint/logit path, so
+compare it as an explicit run configuration.
 Chunked Metal prefill reuses the same range-capable layer-major graph for each
 chunk, preserving absolute compressor/indexer boundaries while avoiding the old
 per-layer chunk dispatch path.
@@ -704,10 +1124,10 @@ conversation. Useful commands are `/help`, `/think`, `/think-max`, `/nothink`,
 and returns to `ds4>`.
 
 The CLI defaults to thinking mode. Use `/nothink` or `--nothink` for direct
-answers. `--mtp MTP.gguf --mtp-draft 2` enables the optional MTP speculative
-path; it is useful only for greedy decoding, currently uses a confidence gate
-(`--mtp-margin`) to avoid slow partial accepts, and should be treated as an
-experimental slight-speedup path.
+answers. Models with a built-in draft block use `--mtp`; models with a separate
+support GGUF use `--mtp-model MTP.gguf`. `--mtp-draft 2` sets the maximum draft
+depth. At non-zero temperature, add `--mtp-exact-sampling` when the output must
+preserve the ordinary target sampling distribution.
 
 ## Server
 
@@ -720,14 +1140,40 @@ Start a local OpenAI/Anthropic-compatible server:
 Use `--chdir /path/to/ds4` when launching `ds4-server` from another directory,
 so relative runtime files such as `metal/*.metal` resolve from the project tree.
 
-The server keeps one mutable backend/KV checkpoint in memory,
-so stateless clients that resend a longer version of the same prompt can reuse
-the shared prefix instead of pre-filling from token zero.
+By default the server keeps one mutable backend/KV checkpoint in memory, so
+stateless clients that resend a longer version of the same prompt can reuse the
+shared prefix instead of pre-filling from token zero.
 
-Request parsing and sockets run in client threads, but inference itself is
-serialized through one graph worker. The current server does not batch multiple
-independent requests together; concurrent requests wait their turn on the single
-live graph/session.
+`--batched-session N` preallocates `N` independent resident KV sessions. Ready
+decode steps are evaluated together, while long prefills alternate in bounded
+chunks so one request does not block every decoder. Requests beyond `N` wait
+for a resident slot. If disk KV caching is enabled, an idle slot is persisted
+before reuse and can be restored when that conversation returns; an active
+request is never evicted. Choose `N` and `--ctx` so all resident KV allocations
+fit in GPU memory. Without this option, inference retains the original
+single-session behavior.
+
+While generation is active, prefill yields every 128 tokens by default.
+`--mixed-prefill-quantum N` changes that interval for testing; larger values
+reduce scheduling handoffs but can make active decoders wait longer.
+
+Native decode batching runs the same model graph, but grouping rows can change
+floating-point reduction order slightly. When a native kernel is unavailable,
+DwarfStar runs the rows in a fixed order and returns the same full logits as
+separate session evaluations. The current backend behavior is:
+
+| Backend and model | Session execution |
+| --- | --- |
+| Metal, resident DeepSeek Flash | Native shared-expert and QKV batching from two rows upward when supported; ordered fallback otherwise. |
+| Metal, GLM 5.2 | Ordered exact fallback. |
+| Metal, GLM 5.3 | Native decode batching below token 4096; ordered exact fallback at longer contexts. |
+| CUDA, DeepSeek Flash on a supported multi-GPU TP/EP layout | Native decode and mixed prefill/decode, with exact fallbacks for unsupported kernel shapes. |
+| CUDA single GPU, including DGX Spark | Ordered exact fallback. |
+
+`N` resident sessions allocate `N` KV states, so a context size that fits once
+may not fit eight times. Native batching can improve aggregate throughput; an
+ordered fallback provides concurrency and fairness, but not the same speedup.
+MTP speculative decoding is disabled while native session batching is active.
 
 Supported endpoints:
 
@@ -762,8 +1208,11 @@ controls. Tool uses are returned as Anthropic `tool_use` blocks.
 
 Default sampled API generation uses `temperature=1`, `top_p=1`, and
 `min_p=0.05`, so the default filter is relative probability rather than
-nucleus mass. In thinking mode DwarfStar uses those fixed sampling defaults and
-ignores client sampling knobs, matching DeepSeek's fixed-thinking API behavior.
+nucleus mass. In thinking mode DwarfStar applies those fixed sampling defaults
+to any knob the request omits, matching DeepSeek's fixed-thinking API behavior,
+but sampling parameters set explicitly in the request always win: a
+`temperature=0` request is greedy through the whole reasoning phase, so
+benchmark harnesses get deterministic thinking-mode output.
 
 The chat, Responses, and Anthropic endpoints support SSE streaming. In thinking
 mode, reasoning is streamed in the native API shape instead of being mixed into
@@ -966,8 +1415,8 @@ matches the local `~/bin/claude-ds4` setup:
 #!/bin/sh
 unset ANTHROPIC_API_KEY
 
-export ANTHROPIC_BASE_URL="${DS4_ANTHROPIC_BASE_URL:-http://127.0.0.1:8000}"
-export ANTHROPIC_AUTH_TOKEN="${DS4_API_KEY:-dsv4-local}"
+export ANTHROPIC_BASE_URL="http://127.0.0.1:8000"
+export ANTHROPIC_AUTH_TOKEN="dsv4-local"
 export ANTHROPIC_MODEL="deepseek-v4-flash"
 
 export ANTHROPIC_CUSTOM_MODEL_OPTION="deepseek-v4-flash"
@@ -1206,6 +1655,25 @@ make cuda CUDA_ARCH=sm_120
 make cuda CUDA_ARCH=native
 ```
 
+CUDA builds accept `--gpu-vram N[,N,...]` and `--gpu-devices N[,N,...]` in the
+CLI, server, agent, and benchmark. VRAM values are per-device GiB budgets;
+`--gpu-vram auto` uses the free memory reported by CUDA. The device list controls
+the placement order and must have the same number of entries as an explicit
+budget list. Placement reserves graph and KV memory for the requested context
+and refuses to start if model layers would spill to the CPU.
+
+Without `--cuda-tensor-parallel`, CUDA uses normal layer placement across the
+listed devices. This is also the supported multi-GPU layout for GLM 5.2:
+
+```sh
+./ds4 -m gguf/GLM-5.2-UD-Q2_K_RoutedQ2K.gguf \
+  --gpu-vram auto --gpu-devices 0,2,4,6,1,3,5,7 \
+  --ctx 32768 -p "Hello"
+```
+
+For the DeepSeek Flash tensor/expert-parallel layout, see
+"Tensor Parallelism across CUDA GPUs" above.
+
 There is also a CPU reference/debug path:
 
 ```sh
@@ -1222,8 +1690,9 @@ CUDA.
 
 ## Steering
 
-This project supports steering with single-vector activation directions; see the
-`dir-steering` directory for more information. This follows the core idea of the
+DeepSeek V4 Flash and GLM 5.3 Flash support steering with single-vector
+activation directions; see the `dir-steering` directory for examples,
+model-specific shapes, and build instructions. This follows the core idea of the
 [Refusal in Language Models Is Mediated by a Single Direction](https://arxiv.org/abs/2406.11717)
 paper. You can use it to make the model more or less verbose, less likely to
 answer programming questions if it is a chatbot for your car rental web site,
@@ -1231,6 +1700,19 @@ and so forth, much faster than fine-tuning.
 This is also useful for cybersecurity researchers who want to reduce a model's
 willingness to provide dual-use or offensive security guidance.
 
+Load a direction when starting either interactive client:
+
+```sh
+./ds4 -m model.gguf --dir-steering-file direction.f32
+./ds4-agent -m model.gguf --dir-steering-file direction.f32
+```
+
+The default FFN scale is `1`. At an interactive prompt, `/steer` shows the
+current scale, `/steer 0` disables it, and `/steer F` sets a value from `-100`
+to `100` for subsequent tokens. The existing KV cache is kept. Live changes
+are currently limited to local sessions, not distributed inference or network
+tensor parallelism.
+
 ## Test Vectors
 
 `tests/test-vectors` contains short and long-context continuation vectors
@@ -1239,10 +1721,9 @@ captured from the official DeepSeek V4 Flash API. The requests use
 `top_logprobs` slice exposed by the API. Local vectors are generated with
 `./ds4 --dump-logprobs` and compared by token bytes, so tokenizer/template or
 attention regressions show up before they become long generation failures. The
-C runner pins `DS4_METAL_PREFILL_CHUNK=2048` for this strict API-vector
-comparison.
+C runner pins a 2048-token prefill chunk for this strict API-vector comparison.
 
-All project tests are driven by the C runner, with a small `ds4-eval`
+The core local tests are driven by the C runner, with a small `ds4-eval`
 extractor self-test run first:
 
 ```sh
@@ -1251,6 +1732,29 @@ make test                  # ./ds4-eval --self-test-extractors && ./ds4_test --a
 ./ds4_test --server
 ```
 
+The batching tests are model-backed and must run on the matching GPU backend:
+
+```sh
+# Metal, with DS4_TEST_SESSION_COUNT set to 2, 4, 8, and 16.
+DS4_TEST_MODEL=/path/to/model.gguf DS4_TEST_SESSION_COUNT=4 \
+  make test-metal-session-batch
+
+# GLM 5.3 native row batching uses a bounded numerical comparison.
+DS4_TEST_MODEL=/path/to/GLM-5.3-Flash-Q2.gguf \
+  DS4_TEST_SESSION_COUNT=4 DS4_TEST_LOGIT_TOLERANCE=0.001 \
+  make test-metal-session-batch
+
+# CUDA multi-GPU Flash.
+DS4_TEST_MODEL=/path/to/model.gguf make test-cuda-session-batch
+DS4_TEST_MODEL=/path/to/model.gguf make test-cuda-mixed-batch
+```
+
+For GLM, run the same Metal session test with a GLM GGUF and run
+`tests/glm_long_context_smoke.sh /path/to/model.gguf`. The official 100-case
+quality scorers, two-Mac TCP/RDMA tests, CUDA matrix, and manual agent checks are
+release gates rather than quick local tests; follow
+[QA_BEFORE_RELEASES.md](QA_BEFORE_RELEASES.md).
+
 ## Debugging Notes
 
 When a generation looks wrong, three small tools are usually enough to get a
@@ -1263,12 +1767,19 @@ first answer:
 ./ds4-server --trace /tmp/ds4-trace.txt ...
 ```
 
-- `--dump-tokens` tokenizes the `-p` or `--prompt-file` string exactly as
-  written, recognizes DS4 protocol specials, and then exits before inference
-  starts. For example, the DSML tool close marker starts as two tokens: `</`
+- `--dump-tokens` prints the exact chat prompt token stream the CLI would use,
+  then exits before inference starts. Add `--raw` to tokenize the `-p` or
+  `--prompt-file` text literally. Protocol specials are recognized in both
+  modes; for example, the DSML tool close marker starts as two tokens: `</`
   and `｜DSML｜`.
 - `--dump-logprobs` stores a greedy continuation with the top local
   alternatives at each step, which helps separate sampling choices from
   logit/model issues.
 - `ds4-server --trace` writes the rendered prompts, cache decisions, generated
   text, and tool-parser events for a whole agent session.
+
+## Logo
+
+The DwarfStar logo was designed by hand by Salvatore Sanfilippo, made more
+graphical with AI, and manually reworked by Ben Gnomino, whose human touch made
+it rock.

STRIXHALO.md

@@ -113,6 +113,16 @@ 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:

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,23 +1,28 @@
 #!/bin/sh
 set -e
 
-DEEPSEEK_REPO="antirez/deepseek-v4-gguf"
 GLM_UNSLOTH_REPO="unsloth/GLM-5.2-GGUF"
 GLM_ANTIREZ_REPO="antirez/GLM-5.2-GGUF"
-REPO=$DEEPSEEK_REPO
-Q2_IMATRIX_FILE="DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf"
-Q4_IMATRIX_FILE="DeepSeek-V4-Flash-Q4KExperts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-imatrix.gguf"
-Q2_Q4_IMATRIX_FILE="DeepSeek-V4-Flash-Layers37-42Q4KExperts-OtherExpertLayersIQ2XXSGateUp-Q2KDown-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-fixed.gguf"
-PRO_Q2_IMATRIX_FILE="DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf"
+GLM53_REPO="antirez/glm-5.3-flash-gguf"
+REPO="antirez/deepseek-v4-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"
-MTP_FILE="DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.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"}
@@ -32,38 +37,52 @@ usage() {
 DwarfStar GGUF downloader
 
 Usage:
-  ./download_model.sh q2-imatrix [--token TOKEN]
-  ./download_model.sh q2-q4-imatrix [--token TOKEN]
-  ./download_model.sh q4-imatrix [--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 mtp [--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:
 
-  q2-imatrix
+  ds4f-q2
        2-bit routed experts, about 81 GB on disk.
        Recommended model for 96 and 128 GB RAM machines.
 
-  q2-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 q2-imatrix.
+       may struggle compared to ds4f-q2.
 
-  q4-imatrix
+  ds4f-q4
        4-bit routed experts, about 153 GB on disk.
        Recommended model for machines with 256 GB RAM or more.
 
+  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.
+
+  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.
+
   pro-q2-imatrix
-       DeepSeek V4 PRO q2 imatrix quant, as a single GGUF file. About 430 GB
-       on disk; intended for 512 GB RAM machines.
+       DeepSeek V4 PRO 0813 q2 imatrix quant, as a single GGUF file. About
+       430 GB on disk; intended for 512 GB RAM machines.
 
   pro-q4-layers00-30
        First half of the DeepSeek V4 PRO Q4 routed-expert quant, layers 0..30.
@@ -78,10 +97,6 @@ Targets:
        Downloads both PRO Q4 split files into the download directory. About
        838 GB total. This target does not update ./ds4flash.gguf.
 
-  mtp  Optional speculative decoding component, about 3.5 GB on disk.
-       It is useful with q2-imatrix, q2-q4-imatrix, and q4-imatrix, but must be
-       enabled explicitly with --mtp when running ds4 or ds4-server.
-
   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.
@@ -98,6 +113,23 @@ Targets:
        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
                  cache is used if present.
@@ -113,8 +145,8 @@ 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
@@ -135,9 +167,11 @@ FORCE_HF_DOWNLOAD=0
 FLATTEN_DOWNLOADS=0
 
 case "$MODEL" in
-    q2-imatrix) MODEL_FILE=$Q2_IMATRIX_FILE ;;
-    q2-q4-imatrix) MODEL_FILE=$Q2_Q4_IMATRIX_FILE ;;
-    q4-imatrix) MODEL_FILE=$Q4_IMATRIX_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 ;;
@@ -145,7 +179,6 @@ case "$MODEL" in
         MODEL_FILES="$PRO_Q4_LAYERS00_30_FILE $PRO_Q4_LAYERS31_OUTPUT_FILE"
         LINK_MODEL=0
         ;;
-    mtp) MODEL_FILE=$MTP_FILE; LINK_MODEL=0 ;;
     glm-unsloth-q4)
         REPO=$GLM_UNSLOTH_REPO
         MODEL_FILE=$GLM_UNSLOTH_Q4_FIRST_FILE
@@ -171,6 +204,28 @@ case "$MODEL" in
         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
@@ -224,6 +279,12 @@ find_hf_command() {
         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
 }
 
@@ -252,7 +313,7 @@ download_one_hf() {
     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 PRO target." >&2
+        echo "Move or remove that partial download before retrying this target." >&2
         exit 1
     fi
 
@@ -333,11 +394,10 @@ else
     download_one "$MODEL_FILE"
 fi
 
-if [ "$MODEL" = "mtp" ]; then
+if [ "$MODEL" = "ds4f-dspark" ]; then
     echo
-    echo "MTP is an optional component for q2-imatrix, q2-q4-imatrix, and q4-imatrix."
-    echo "Enable it explicitly, for example:"
-    echo "  ./ds4 --mtp $OUT_DIR/$MTP_FILE --mtp-draft 2"
+    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,"

ds4.h

@@ -91,15 +91,50 @@ typedef struct {
     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;
@@ -112,10 +147,23 @@ typedef struct {
     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;
@@ -123,8 +171,24 @@ typedef struct {
     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);
 
@@ -150,15 +214,78 @@ typedef struct {
 } 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,
@@ -172,6 +299,7 @@ bool ds4_engine_glm_layer_payload_bytes(ds4_engine *e,
  * 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);
@@ -208,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);
@@ -242,10 +376,21 @@ bool ds4_token_is_stop_for_think_mode(ds4_engine *e,
 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;
@@ -271,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,
@@ -282,15 +434,82 @@ 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);

ds4_bench.c

@@ -1,6 +1,8 @@
 #include "ds4.h"
 #include "ds4_distributed.h"
+#include "ds4_gpu_args.h"
 #include "ds4_help.h"
+#include "ds4_tp.h"
 
 /* Purpose-built throughput benchmark.
  *
@@ -22,6 +24,14 @@
 #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 {
@@ -31,6 +41,8 @@ typedef struct {
     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;
@@ -48,11 +60,15 @@ typedef struct {
     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) {
@@ -229,6 +245,23 @@ static bench_config parse_options(int argc, char **argv) {
         }
         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")) {
@@ -268,6 +301,12 @@ static bench_config parse_options(int argc, char **argv) {
         } 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")) {
@@ -315,6 +354,10 @@ static bench_config parse_options(int argc, char **argv) {
             }
         } 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, NULL);
@@ -347,12 +390,20 @@ 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) {
+    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);
     }
@@ -538,9 +589,35 @@ static void maybe_warn_distributed_step_shape(const bench_config *cfg, ds4_sessi
     }
 }
 
+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);
 
+    /* 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,
@@ -555,22 +632,70 @@ int main(int argc, char **argv) {
         .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;
-    log_context_memory(cfg.backend,
+    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,
-                       cfg.prefill_chunk,
+                       ds4_engine_prefill_chunk(engine),
                        cfg.ssd_streaming);
 
     char *text = read_file(cfg.prompt_path ? cfg.prompt_path : cfg.chat_prompt_path);
@@ -582,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;
     }
 
@@ -596,7 +724,7 @@ 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 &&
@@ -604,7 +732,7 @@ int main(int argc, char **argv) {
     {
         ds4_session_free(session);
         ds4_tokens_free(&prompt);
-        ds4_engine_close(engine);
+        close_engine(engine, tp_leader);
         return 1;
     }
     maybe_warn_distributed_step_shape(&cfg, session);
@@ -616,7 +744,7 @@ 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;
         }
     }
@@ -624,7 +752,9 @@ int main(int argc, char **argv) {
     fflush(out);
 
     const int eos = ds4_token_eos(engine);
-    const bool distributed = cfg.dist.role == DS4_DISTRIBUTED_COORDINATOR;
+    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;
@@ -640,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;
@@ -685,30 +826,78 @@ int main(int argc, char **argv) {
         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 (!need_restore_after_generation) {
@@ -750,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_cli.c

@@ -1,5 +1,7 @@
 #include "ds4.h"
 #include "ds4_distributed.h"
+#include "ds4_gpu_args.h"
+#include "ds4_tp.h"
 #include "ds4_help.h"
 #include "linenoise.h"
 
@@ -25,6 +27,37 @@
 #include <time.h>
 #include <unistd.h>
 
+static bool cli_env_flag_enabled(const char *name, bool defval) {
+    const char *v = getenv(name);
+    if (!v || !v[0]) return defval;
+    return strcmp(v, "0") != 0;
+}
+
+static bool cli_splitkv_spec_requested(void) {
+    if (cli_env_flag_enabled("DS4_CUDA_NO_SPLITKV_SPEC", false)) return false;
+    return cli_env_flag_enabled("DS4_CUDA_SPLITKV_SPEC", false);
+}
+
+static bool cli_greedy_fast_attention_requested(void) {
+    if (!cli_env_flag_enabled("DS4_CUDA_NO_GREEDY_SPLITKV", false) &&
+        cli_env_flag_enabled("DS4_CUDA_GREEDY_SPLITKV", false))
+    {
+        return true;
+    }
+    if (!cli_env_flag_enabled("DS4_CUDA_NO_GREEDY_VEC4", false) &&
+        cli_env_flag_enabled("DS4_CUDA_GREEDY_VEC4", false))
+    {
+        return true;
+    }
+    return false;
+}
+
+static bool cli_greedy_argmax_requested(bool speculative_requested) {
+    if (cli_greedy_fast_attention_requested()) return true;
+    if (speculative_requested) return false;
+    return cli_env_flag_enabled("DS4_CUDA_GREEDY_TOP1", true);
+}
+
 typedef struct {
     const char *prompt;
     const char *system;
@@ -62,6 +95,10 @@ typedef struct {
     cli_generation_options gen;
     char *prompt_owned;
     bool inspect;
+    /* CLI flag wiring: raw argv values for --gpu-vram and --gpu-devices.
+     * Resolved post-parse via parse_gpu_vram_arg(). */
+    const char *gpu_vram_arg;
+    const char *gpu_devices_arg;
 } cli_config;
 
 static volatile sig_atomic_t cli_interrupted;
@@ -491,6 +528,9 @@ static int run_sampled_generation(ds4_engine *engine, const cli_config *cfg, con
         ds4_session_free(session);
         return 1;
     }
+    /* Pay the one-time first-submission GPU cost before the prefill timer
+     * starts (matches the TP worker's startup warmup). */
+    ds4_session_gpu_warmup(session);
 
     char err[160];
     ds4_think_mode think_mode = cli_effective_think_mode(&cfg->gen);
@@ -535,25 +575,37 @@ static int run_sampled_generation(ds4_engine *engine, const cli_config *cfg, con
     uint64_t rng = cfg->gen.seed ? cfg->gen.seed :
         ((uint64_t)time(NULL) ^ ((uint64_t)getpid() << 32) ^ (uint64_t)clock());
     int generated = 0;
+    const bool speculative_argmax = cfg->gen.temperature <= 0.0f &&
+        ((ds4_engine_mtp_draft_tokens(engine) > 1 &&
+          getenv("DS4_MTP_SPEC_DISABLE") == NULL) ||
+         cli_splitkv_spec_requested());
+    const bool greedy_argmax = cfg->gen.temperature <= 0.0f &&
+        cli_greedy_argmax_requested(speculative_argmax);
+    bool have_greedy_next = false;
+    int greedy_next = -1;
     const double t_decode0 = cli_now_sec();
     while (generated < max_tokens && !cli_interrupt_requested()) {
-        int token = ds4_session_sample(session, cfg->gen.temperature, 0,
+        int token;
+        if (greedy_argmax && have_greedy_next) {
+            token = greedy_next;
+            have_greedy_next = false;
+        } else {
+            token = ds4_session_sample(session, cfg->gen.temperature, 0,
                                        cfg->gen.top_p, cfg->gen.min_p, &rng);
+        }
         if (ds4_token_is_stop_for_think_mode(engine, token, think_mode)) break;
 
         int toks[17];
         int ntok = 0;
-        if (cfg->gen.temperature <= 0.0f && ds4_engine_mtp_draft_tokens(engine) > 1 &&
+        if (ds4_engine_mtp_draft_tokens(engine) > 1 &&
             getenv("DS4_MTP_SPEC_DISABLE") == NULL) {
             cli_dist_busy_set(cfg, true);
-            ntok = ds4_session_eval_speculative_argmax(session,
-                                                       token,
-                                                       max_tokens - generated,
-                                                       ds4_token_eos(engine),
-                                                       toks,
-                                                       (int)(sizeof(toks) / sizeof(toks[0])),
-                                                       err,
-                                                       sizeof(err));
+            ntok = ds4_session_eval_speculative(
+                session, token, max_tokens - generated,
+                ds4_token_eos(engine), cfg->gen.temperature, 0,
+                cfg->gen.top_p, cfg->gen.min_p, &rng,
+                toks, (int)(sizeof(toks) / sizeof(toks[0])),
+                err, sizeof(err));
             cli_dist_busy_set(cfg, false);
             if (ntok < 0) {
                 fprintf(stderr, "ds4: decode failed: %s\n", err);
@@ -1168,8 +1220,14 @@ static int run_generation(ds4_engine *engine, const cli_config *cfg) {
         }
     } else {
         if (cfg->engine.distributed.role == DS4_DISTRIBUTED_COORDINATOR ||
+            cfg->engine.tp.role == DS4_TP_LEADER ||
+            getenv("DS4_CLI_FORCE_SESSION") != NULL ||
             cfg->gen.temperature > 0.0f ||
             ds4_engine_mtp_draft_tokens(engine) > 1) {
+            /* TP leaders always drive the session path: the sync/eval
+             * mirroring that keeps the worker in lockstep lives there.
+             * The env override exists so TP-vs-single-node validation
+             * compares like with like. */
             rc = run_sampled_generation(engine, cfg, &prompt);
         } else {
             token_printer printer = {
@@ -1215,7 +1273,8 @@ static void print_repl_help(void) {
     puts("  /nothink       Disable thinking mode.");
     puts("  /ctx N         Set context size for following prompts.");
     puts("  /power N       Set GPU duty cycle percentage, 1..100.");
-    puts("  /read FILE     Read a prompt from FILE and run it.");
+    puts("  /steer F       Set FFN steering for subsequent tokens; no value shows it.");
+    puts("  /read FILE     Submit a text file, PNG, or JPEG.");
     puts("  /quit, /exit   Leave the prompt.");
     puts("  Ctrl+C         Stop generation and return to the prompt.");
 }
@@ -1228,6 +1287,18 @@ static bool parse_power_percent(const char *arg, int *out) {
     return true;
 }
 
+static bool parse_steering_level(const char *arg, float *out) {
+    char *end = NULL;
+    errno = 0;
+    float v = strtof(arg, &end);
+    if (!arg[0] || *end != '\0' || errno == ERANGE || !isfinite(v) ||
+        v < -100.0f || v > 100.0f) {
+        return false;
+    }
+    *out = v;
+    return true;
+}
+
 static void history_file_path(char *buf, size_t len) {
     const char *home = getenv("HOME");
     if (!home || !home[0]) home = ".";
@@ -1237,11 +1308,35 @@ static void history_file_path(char *buf, size_t len) {
 typedef struct {
     ds4_session *session;
     ds4_tokens transcript;
+    ds4_vision_span *images;
+    size_t image_count;
+    size_t image_cap;
     int ctx_size;
     int think_prefix_pos;
     int think_prefix_tokens;
 } repl_chat;
 
+static void repl_chat_trim_images(repl_chat *chat, size_t count) {
+    while (chat->image_count > count) {
+        chat->image_count--;
+        ds4_vision_embedding_free(&chat->images[chat->image_count].embedding);
+    }
+}
+
+static ds4_vision_span *repl_chat_add_image(repl_chat *chat) {
+    if (chat->image_count == chat->image_cap) {
+        size_t cap = chat->image_cap ? chat->image_cap * 2u : 4u;
+        if (cap > SIZE_MAX / sizeof(chat->images[0])) return NULL;
+        void *next = realloc(chat->images, cap * sizeof(chat->images[0]));
+        if (!next) return NULL;
+        chat->images = next;
+        chat->image_cap = cap;
+    }
+    ds4_vision_span *span = &chat->images[chat->image_count++];
+    memset(span, 0, sizeof(*span));
+    return span;
+}
+
 static void tokens_insert(ds4_tokens *dst, int pos, const ds4_tokens *src) {
     if (!src || src->len <= 0) return;
     if (pos < 0) pos = 0;
@@ -1305,10 +1400,16 @@ static void repl_chat_apply_think_prefix(ds4_engine *engine,
                        (size_t)prefix.len * sizeof(prefix.v[0]));
     }
     if (!same) {
+        const int old_prefix_tokens = chat->think_prefix_tokens;
         tokens_remove(&chat->transcript, chat->think_prefix_pos,
                       chat->think_prefix_tokens);
         tokens_insert(&chat->transcript, chat->think_prefix_pos, &prefix);
         chat->think_prefix_tokens = prefix.len;
+        const int shift = prefix.len - old_prefix_tokens;
+        for (size_t i = 0; i < chat->image_count; i++) {
+            chat->images[i].token_start =
+                (uint32_t)((int64_t)chat->images[i].token_start + shift);
+        }
         if (chat->session) ds4_session_invalidate(chat->session);
     }
     ds4_tokens_free(&prefix);
@@ -1341,6 +1442,8 @@ static void repl_chat_free(repl_chat *chat) {
     if (!chat) return;
     ds4_session_free(chat->session);
     ds4_tokens_free(&chat->transcript);
+    repl_chat_trim_images(chat, 0);
+    free(chat->images);
     memset(chat, 0, sizeof(*chat));
 }
 
@@ -1359,7 +1462,9 @@ static bool repl_chat_assistant_turn_uses_eos(ds4_engine *engine) {
  * and assistant markers, then ds4_session_sync() decides whether this is a KV
  * continuation.  If prompt processing fails, the transcript rolls back before
  * returning to the prompt. */
-static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, const char *user_text) {
+static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat,
+                         const char *user_text,
+                         ds4_vision_embedding *image) {
     if (!chat->session) {
         fprintf(stderr, "ds4: no active interactive KV cache\n");
         return 1;
@@ -1369,7 +1474,28 @@ static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, c
                                                            chat->ctx_size);
     repl_chat_apply_think_prefix(engine, chat, think_mode);
     const int rollback_len = chat->transcript.len;
-    ds4_chat_append_message(engine, &chat->transcript, "user", user_text);
+    const size_t rollback_images = chat->image_count;
+    if (image) {
+        ds4_tokens_push(&chat->transcript, ds4_token_user(engine));
+        ds4_vision_span *span = repl_chat_add_image(chat);
+        char image_error[160] = {0};
+        if (!span || !ds4_prompt_append_vision(engine, &chat->transcript,
+                                               span, image,
+                                               image_error,
+                                               sizeof(image_error))) {
+            repl_chat_trim_images(chat, rollback_images);
+            chat->transcript.len = rollback_len;
+            fprintf(stderr, "ds4: failed to add image: %s\n",
+                    image_error[0] ? image_error : "out of memory");
+            return 1;
+        }
+        if (user_text && user_text[0]) {
+            ds4_tokenize_text(engine, "\n", &chat->transcript);
+            ds4_tokenize_text(engine, user_text, &chat->transcript);
+        }
+    } else {
+        ds4_chat_append_message(engine, &chat->transcript, "user", user_text);
+    }
     ds4_chat_append_assistant_prefix(engine, &chat->transcript, think_mode);
 
     const int old_pos = ds4_session_pos(chat->session);
@@ -1389,12 +1515,18 @@ static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, c
                                      progress.use_color ? cli_prefill_progress_cb : NULL,
                                      progress.use_color ? &progress : NULL);
     cli_dist_busy_set(cfg, true);
-    int sync_rc = ds4_session_sync(chat->session, &chat->transcript, err, sizeof(err));
+    int sync_rc = ds4_session_sync_multimodal(chat->session,
+                                              &chat->transcript,
+                                              chat->images,
+                                              chat->image_count,
+                                              err,
+                                              sizeof(err));
     cli_dist_busy_set(cfg, false);
     if (sync_rc != 0) {
         ds4_session_set_progress(chat->session, NULL, NULL);
         ds4_session_set_display_progress(chat->session, NULL, NULL);
         chat->transcript.len = rollback_len;
+        repl_chat_trim_images(chat, rollback_images);
         fprintf(stderr, "ds4: prompt processing failed: %s\n", err);
         return 1;
     }
@@ -1419,29 +1551,41 @@ static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, c
     uint64_t rng = cfg->gen.seed ? cfg->gen.seed :
         ((uint64_t)time(NULL) ^ ((uint64_t)getpid() << 32) ^ (uint64_t)clock());
     int generated = 0;
+    const bool speculative_argmax = cfg->gen.temperature <= 0.0f &&
+        ((ds4_engine_mtp_draft_tokens(engine) > 1 &&
+          getenv("DS4_MTP_SPEC_DISABLE") == NULL) ||
+         cli_splitkv_spec_requested());
+    const bool greedy_argmax = cfg->gen.temperature <= 0.0f &&
+        cli_greedy_argmax_requested(speculative_argmax);
+    bool have_greedy_next = false;
+    int greedy_next = -1;
     const double t_decode0 = cli_now_sec();
     while (generated < max_tokens && !cli_interrupt_requested()) {
-        int token = ds4_session_sample(chat->session,
+        int token;
+        if (greedy_argmax && have_greedy_next) {
+            token = greedy_next;
+            have_greedy_next = false;
+        } else {
+            token = ds4_session_sample(chat->session,
                                        cfg->gen.temperature,
                                        0,
                                        cfg->gen.top_p,
                                        cfg->gen.min_p,
                                        &rng);
+        }
         if (ds4_token_is_stop_for_think_mode(engine, token, think_mode)) break;
 
         int toks[17];
         int ntok = 0;
-        if (cfg->gen.temperature <= 0.0f && ds4_engine_mtp_draft_tokens(engine) > 1 &&
+        if (ds4_engine_mtp_draft_tokens(engine) > 1 &&
             getenv("DS4_MTP_SPEC_DISABLE") == NULL) {
             cli_dist_busy_set(cfg, true);
-            ntok = ds4_session_eval_speculative_argmax(chat->session,
-                                                       token,
-                                                       max_tokens - generated,
-                                                       ds4_token_eos(engine),
-                                                       toks,
-                                                       (int)(sizeof(toks) / sizeof(toks[0])),
-                                                       err,
-                                                       sizeof(err));
+            ntok = ds4_session_eval_speculative(
+                chat->session, token, max_tokens - generated,
+                ds4_token_eos(engine), cfg->gen.temperature, 0,
+                cfg->gen.top_p, cfg->gen.min_p, &rng,
+                toks, (int)(sizeof(toks) / sizeof(toks[0])),
+                err, sizeof(err));
             cli_dist_busy_set(cfg, false);
             if (ntok < 0) {
                 fprintf(stderr, "ds4: decode failed: %s\n", err);
@@ -1491,6 +1635,7 @@ static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, c
     const bool interrupted = cli_interrupt_requested();
     if (interrupted && generated == 0) {
         chat->transcript.len = rollback_len;
+        repl_chat_trim_images(chat, rollback_images);
         ds4_session_invalidate(chat->session);
     } else if (repl_chat_assistant_turn_uses_eos(engine)) {
         ds4_tokens_push(&chat->transcript, ds4_token_eos(engine));
@@ -1507,6 +1652,16 @@ static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, c
     return 0;
 }
 
+static bool cli_file_has_image_magic(const char *path) {
+    unsigned char magic[8] = {0};
+    FILE *fp = fopen(path, "rb");
+    if (!fp) return false;
+    size_t n = fread(magic, 1, sizeof(magic), fp);
+    fclose(fp);
+    return (n >= 8 && !memcmp(magic, "\x89PNG\r\n\x1a\n", 8)) ||
+           (n >= 2 && magic[0] == 0xff && magic[1] == 0xd8);
+}
+
 static int run_repl(ds4_engine *engine, cli_config *cfg) {
     repl_chat chat;
     if (repl_chat_init(engine, &chat, cfg) != 0) return 1;
@@ -1578,6 +1733,22 @@ static int run_repl(ds4_engine *engine, cli_config *cfg) {
                     printf("Power: %d%%.\n", power);
                 }
             }
+        } else if (!strncmp(cmd, "/steer", 6) &&
+                   (cmd[6] == '\0' || isspace((unsigned char)cmd[6]))) {
+            char *arg = trim_inplace(cmd + 6);
+            if (!arg[0]) {
+                printf("Steering FFN: %g.\n",
+                       (double)ds4_session_directional_steering_ffn(chat.session));
+            } else {
+                float scale = 0.0f;
+                if (!parse_steering_level(arg, &scale)) {
+                    fprintf(stderr, "ds4: /steer must be between -100 and 100\n");
+                } else if (ds4_session_set_directional_steering_ffn(
+                                   chat.session, scale) == 0) {
+                    cfg->engine.directional_steering_ffn = scale;
+                    printf("Steering FFN: %g.\n", (double)scale);
+                }
+            }
         } else if (!strncmp(cmd, "/ctx", 4) && (cmd[4] == '\0' || isspace((unsigned char)cmd[4]))) {
             char *arg = trim_inplace(cmd + 4);
             if (!arg[0]) {
@@ -1586,7 +1757,7 @@ static int run_repl(ds4_engine *engine, cli_config *cfg) {
                 cfg->gen.ctx_size = parse_int(arg, "/ctx");
                 log_context_memory(cfg->engine.backend,
                                    cfg->gen.ctx_size,
-                                   cfg->engine.prefill_chunk,
+                                   ds4_engine_prefill_chunk(engine),
                                    cfg->engine.ssd_streaming);
                 rc = repl_chat_set_ctx(engine, &chat, cfg->gen.ctx_size);
                 if (rc != 0) {
@@ -1605,10 +1776,24 @@ static int run_repl(ds4_engine *engine, cli_config *cfg) {
             char *path = trim_inplace(cmd + 5);
             if (!path[0]) {
                 fprintf(stderr, "ds4: /read needs a file path\n");
+            } else if (cli_file_has_image_magic(path)) {
+                char image_error[256] = {0};
+                ds4_vision_embedding image = {0};
+                if (!ds4_engine_vision_encode_file(engine, path, &image,
+                                                   image_error,
+                                                   sizeof(image_error))) {
+                    fprintf(stderr, "ds4: /read image failed: %s\n", image_error);
+                } else {
+                    fprintf(stderr,
+                            "ds4: image %ux%u, %u image tokens\n",
+                            image.width, image.height, image.token_count);
+                    rc = run_chat_turn(engine, cfg, &chat, "", &image);
+                    ds4_vision_embedding_free(&image);
+                }
             } else {
                 char *prompt = read_prompt_file(path, false);
                 if (prompt) {
-                    rc = run_chat_turn(engine, cfg, &chat, prompt);
+                    rc = run_chat_turn(engine, cfg, &chat, prompt, NULL);
                     free(prompt);
                 }
             }
@@ -1616,7 +1801,7 @@ static int run_repl(ds4_engine *engine, cli_config *cfg) {
             fprintf(stderr, "ds4: unknown command: %s\n", cmd);
             fprintf(stderr, "ds4: type /help for commands\n");
         } else {
-            rc = run_chat_turn(engine, cfg, &chat, cmd);
+            rc = run_chat_turn(engine, cfg, &chat, cmd, NULL);
         }
         linenoiseFree(line);
     }
@@ -1730,6 +1915,20 @@ static cli_config parse_options(int argc, char **argv) {
         }
         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.engine.tp,
+                                                                tp_parse_err,
+                                                                sizeof(tp_parse_err));
+        if (tp_parse == DS4_TP_CLI_ERROR) {
+            fprintf(stderr, "ds4: %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, "-p") || !strcmp(arg, "--prompt")) {
             if (c.gen.prompt) {
                 fprintf(stderr, "ds4: specify only one prompt source\n");
@@ -1749,12 +1948,31 @@ static cli_config parse_options(int argc, char **argv) {
             c.gen.raw_prompt = true;
         } else if (!strcmp(arg, "-m") || !strcmp(arg, "--model")) {
             c.engine.model_path = need_arg(&i, argc, argv, arg);
+        } else if (!strcmp(arg, "--vision")) {
+            c.engine.vision_path = need_arg(&i, argc, argv, arg);
         } else if (!strcmp(arg, "--mtp")) {
+            c.engine.glm_mtp = true;
+        } else if (!strcmp(arg, "--mtp-model")) {
             c.engine.mtp_path = need_arg(&i, argc, argv, arg);
         } else if (!strcmp(arg, "--mtp-draft")) {
             c.engine.mtp_draft_tokens = parse_int(need_arg(&i, argc, argv, arg), arg);
         } else if (!strcmp(arg, "--mtp-margin")) {
             c.engine.mtp_margin = parse_float_range(need_arg(&i, argc, argv, arg), arg, 0.0f, 1000.0f);
+        } else if (!strcmp(arg, "--mtp-timing")) {
+            c.engine.glm_mtp = true;
+            c.engine.glm_mtp_timing = true;
+        } else if (!strcmp(arg, "--dspark")) {
+            c.engine.dspark = true;
+        } else if (!strcmp(arg, "--dspark-confidence")) {
+            c.engine.dspark = true;
+            c.engine.dspark_confidence_threshold =
+                parse_float_range(need_arg(&i, argc, argv, arg), arg, 0.0f, 1.0f);
+            c.engine.dspark_confidence_threshold_set = true;
+        } else if (!strcmp(arg, "--dspark-strict")) {
+            c.engine.dspark = true;
+            c.engine.dspark_strict = true;
+        } else if (!strcmp(arg, "--mtp-exact-sampling")) {
+            c.engine.dspark_exact_sampling = true;
         } else if (!strcmp(arg, "-n") || !strcmp(arg, "--tokens")) {
             c.gen.n_predict = parse_int(need_arg(&i, argc, argv, arg), arg);
         } else if (!strcmp(arg, "-c") || !strcmp(arg, "--ctx")) {
@@ -1843,6 +2061,12 @@ static cli_config parse_options(int argc, char **argv) {
         } else if (!strcmp(arg, "--cuda")) {
             c.engine.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.engine.cuda_tensor_parallel = true;
         } else if (!strcmp(arg, "--dump-tokens")) {
             c.gen.dump_tokens = true;
         } else if (!strcmp(arg, "--dump-logits")) {
@@ -1927,11 +2151,21 @@ static cli_config parse_options(int argc, char **argv) {
         fprintf(stderr, "ds4: --perplexity-file does not use -p/--prompt-file\n");
         exit(2);
     }
+    char tp_err[256];
+    if (!ds4_tp_adopt_distributed_options(&c.engine.tp, c.dist,
+                                          tp_err, sizeof(tp_err))) {
+        fprintf(stderr, "ds4: %s\n", tp_err);
+        exit(2);
+    }
     char dist_err[256];
     if (ds4_dist_prepare_engine_options(c.dist, &c.engine, dist_err, sizeof(dist_err)) != 0) {
         fprintf(stderr, "ds4: %s\n", dist_err);
         exit(2);
     }
+    if (!ds4_tp_validate_engine_options(&c.engine, tp_err, sizeof(tp_err))) {
+        fprintf(stderr, "ds4: %s\n", tp_err);
+        exit(2);
+    }
 
     return c;
 }
@@ -1944,9 +2178,18 @@ int main(int argc, char **argv) {
             free(cfg.prompt_owned);
             return 2;
         }
-        int rc = ds4_dump_text_tokenization(cfg.engine.model_path,
+        int rc;
+        if (cfg.gen.raw_prompt || is_rendered_chat_prompt(cfg.gen.prompt)) {
+            rc = ds4_dump_text_tokenization(cfg.engine.model_path,
                                             cfg.gen.prompt,
                                             stdout);
+        } else {
+            rc = ds4_dump_chat_tokenization(cfg.engine.model_path,
+                                            cfg.gen.system,
+                                            cfg.gen.prompt,
+                                            cli_effective_think_mode(&cfg.gen),
+                                            stdout);
+        }
         ds4_dist_options_free(cfg.dist);
         free(cfg.prompt_owned);
         return rc;
@@ -1955,13 +2198,84 @@ int main(int argc, char **argv) {
     cfg.engine.first_token_test = cfg.gen.first_token_test;
     cfg.engine.metal_graph_test = cfg.gen.metal_graph_test;
     cfg.engine.context_size = cfg.gen.ctx_size;
+    cfg.engine.placement_ctx_hint = cfg.gen.ctx_size;
     ds4_engine *engine = NULL;
-    if (ds4_engine_open(&engine, &cfg.engine) != 0) {
+    if (cfg.gpu_vram_arg || cfg.gpu_devices_arg) {
+        ds4_gpu_config gpu_cfg = {0};
+        bool skip_cuda = false;
+        char errbuf[256];
+        if (parse_gpu_vram_arg(cfg.gpu_vram_arg, cfg.gpu_devices_arg,
+                               &gpu_cfg, &skip_cuda,
+                               errbuf, sizeof(errbuf)) != 0) {
+            fprintf(stderr, "ds4: %s\n", errbuf);
+            ds4_dist_options_free(cfg.dist);
+            free(cfg.prompt_owned);
+            return 2;
+        }
+        cfg.engine.backend = skip_cuda ? DS4_BACKEND_CPU : DS4_BACKEND_CUDA;
+        if (skip_cuda) {
+            if (ds4_engine_open(&engine, &cfg.engine) != 0) {
+                ds4_dist_options_free(cfg.dist);
+                free(cfg.prompt_owned);
+                return 1;
+            }
+        } else {
+            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, &cfg.engine,
+                                                   &gpu_cfg) != 0) {
+                ds4_dist_options_free(cfg.dist);
+                free(cfg.prompt_owned);
+                return 1;
+            }
+        }
+    } else if (ds4_engine_open(&engine, &cfg.engine) != 0) {
         ds4_dist_options_free(cfg.dist);
         free(cfg.prompt_owned);
         return 1;
     }
     cli_apply_model_sampling_defaults(engine, &cfg.gen);
+    if (cfg.engine.tp.role == DS4_TP_WORKER) {
+        int rc = ds4_tp_worker_run(engine, &cfg.engine.tp);
+        ds4_engine_close(engine);
+        ds4_dist_options_free(cfg.dist);
+        free(cfg.prompt_owned);
+        return rc;
+    }
+    ds4_tp *tp_leader = NULL;
+    if (cfg.engine.tp.role == DS4_TP_LEADER) {
+        char tp_err[256] = "";
+        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.gen.ctx_size,
+        };
+        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.engine.tp, &tp_id, tp_err, sizeof(tp_err)) ||
+            !ds4_engine_tp_bind(engine, tp_leader, tp_err, sizeof(tp_err))) {
+            fprintf(stderr, "ds4: %s\n", tp_err);
+            ds4_tp_free(tp_leader);
+            ds4_engine_close(engine);
+            ds4_dist_options_free(cfg.dist);
+            free(cfg.prompt_owned);
+            return 1;
+        }
+    }
     if (cfg.dist && cfg.dist->role == DS4_DISTRIBUTED_WORKER) {
         ds4_dist_generation_options dist_gen = {
             .prompt = cfg.gen.prompt,
@@ -1984,6 +2298,10 @@ int main(int argc, char **argv) {
         return rc;
     }
     if (!cfg.inspect) {
+        log_context_memory(cfg.engine.backend,
+                           cfg.gen.ctx_size,
+                           ds4_engine_prefill_chunk(engine),
+                           cfg.engine.ssd_streaming);
         cli_warn_think_max_downgraded(&cfg.gen, "--think-max");
     }
     int rc = 0;
@@ -2003,7 +2321,9 @@ int main(int argc, char **argv) {
     } else {
         rc = run_generation(engine, &cfg);
     }
+    if (tp_leader) ds4_tp_send_stop(tp_leader);
     ds4_engine_close(engine);
+    ds4_tp_free(tp_leader);
     ds4_dist_options_free(cfg.dist);
     free(cfg.prompt_owned);
     return rc;

ds4_distributed.c

@@ -7012,6 +7012,17 @@ static int dist_worker_handle_snapshot_load(
         snprintf(err, sizeof(err), "invalid distributed snapshot load header");
     }
 
+    /* Reject a token_count larger than the worker context up front, before the
+     * allocation and the read loop below. token_count is only bounded to
+     * ~UINT32_MAX/4 by the header check above, so without this a peer could
+     * drive a multi-GB malloc and a full-length socket read before the
+     * matching-state check further down (which also tests token_count > ctx_size)
+     * finally rejects the request. */
+    if (!err[0] && begin.token_count > (uint32_t)state->ctx_size) {
+        dist_discard_bytes(upstream->fd, body_bytes);
+        snprintf(err, sizeof(err), "snapshot token count exceeds worker context size");
+    }
+
     int *tokens = NULL;
     if (!err[0]) {
         tokens = malloc((size_t)begin.token_count * sizeof(tokens[0]));

ds4_eval.c

@@ -1552,7 +1552,7 @@ static eval_config parse_options(int argc, char **argv) {
 
         if (!strcmp(arg, "-m") || !strcmp(arg, "--model")) {
             c.model_path = need_arg(&i, argc, argv, arg);
-        } else if (!strcmp(arg, "--mtp")) {
+        } else if (!strcmp(arg, "--mtp-model")) {
             c.mtp_path = need_arg(&i, argc, argv, arg);
         } else if (!strcmp(arg, "-c") || !strcmp(arg, "--ctx")) {
             c.ctx_size = parse_int_arg(need_arg(&i, argc, argv, arg), arg);

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

@@ -147,12 +147,20 @@ 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.");
@@ -168,14 +176,21 @@ static void print_model_runtime(FILE *fp, const help_colors *c,
     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", "Metal graph prefill chunk size. Default: auto (PRO long prompts use 8192; others use 4096).");
+    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_BENCH) {
-            opt(fp, c, "--mtp FILE", "Optional MTP support GGUF used for draft-token probes.");
+        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.");
@@ -229,6 +244,17 @@ static void print_distributed(FILE *fp, const help_colors *c) {
     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);
@@ -247,7 +273,7 @@ static void print_cli_specific(FILE *fp, const help_colors *c, bool full) {
 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", "Tokenize the prompt exactly as written, then exit.");
+    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");
@@ -272,7 +298,7 @@ static void print_cli_commands(FILE *fp, const help_colors *c) {
     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", "Read FILE and submit it as the next user message.");
+    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);
@@ -281,8 +307,10 @@ static void print_cli_commands(FILE *fp, const help_colors *c) {
 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, "--non-interactive", "Run without TUI. With -p: one turn; without -p: repeated stdin prompts.");
-    opt(fp, c, "--raw-prompt", "Non-interactive -p only: tokenize prompt without agent chat/tool text.");
+    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.");
@@ -310,6 +338,8 @@ static void print_server_api(FILE *fp, const help_colors *c) {
     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);
@@ -353,6 +383,7 @@ static void print_bench_specific(FILE *fp, const help_colors *c) {
     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);
@@ -519,7 +550,7 @@ static void print_topic(FILE *fp, const help_colors *c, ds4_help_tool tool, cons
         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; [upto] can bridge a unique head and tail for large anchored edits.");
+        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);

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.c

@@ -1318,32 +1318,19 @@ int ds4_kvstore_try_load_text(ds4_kvstore *kc,
         const double load_ms = (kv_now_sec() - load_t0) * 1000.0;
         kc->continued_last_store_tokens = loaded;
         const char *key_kind = ds4_kvstore_key_kind(hdr.ext_flags);
-        bool consumed = false;
-        if (kc->opt.cold_max_tokens > 0 && loaded > kc->opt.cold_max_tokens) {
-            unlink(path);
-            consumed = true;
-            kv_logf(kc, DS4_KVSTORE_LOG_KVCACHE,
-                    "%s: kv cache hit text%s%s tokens=%d text=%u quant=%u key=%s load=%.1f ms consumed file=%s",
-                    kv_log_name(kc),
-                    responses_protocol ? " " : "",
-                    responses_protocol ? "RESPPROTO" : "",
-                    loaded, text_bytes, hdr.quant_bits, key_kind, load_ms, path);
-        } else {
-            ds4_kvstore_touch_file(path, hdr.hits + 1);
-            kv_logf(kc, DS4_KVSTORE_LOG_KVCACHE,
-                    "%s: kv cache hit text%s%s tokens=%d text=%u quant=%u key=%s load=%.1f ms file=%s",
-                    kv_log_name(kc),
-                    responses_protocol ? " " : "",
-                    responses_protocol ? "RESPPROTO" : "",
-                    loaded, text_bytes, hdr.quant_bits, key_kind, load_ms, path);
-        }
+        ds4_kvstore_touch_file(path, hdr.hits + 1);
+        kv_logf(kc, DS4_KVSTORE_LOG_KVCACHE,
+                "%s: kv cache hit text%s%s tokens=%d text=%u quant=%u key=%s load=%.1f ms file=%s",
+                kv_log_name(kc),
+                responses_protocol ? " " : "",
+                responses_protocol ? "RESPPROTO" : "",
+                loaded, text_bytes, hdr.quant_bits, key_kind, load_ms, path);
         if (result) {
             result->tokens = loaded;
             result->text_bytes = text_bytes;
             result->quant_bits = hdr.quant_bits;
             result->ext_flags = hdr.ext_flags;
             result->load_ms = load_ms;
-            result->consumed = consumed;
             result->path = kv_xstrdup(path);
         }
     }

ds4_kvstore.h

@@ -103,7 +103,6 @@ typedef struct {
     uint8_t quant_bits;
     uint8_t ext_flags;
     double load_ms;
-    bool consumed;
     char *path;
 } ds4_kvstore_load_result;
 

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

@@ -34,6 +34,7 @@
 #include <sys/stat.h>
 #include <time.h>
 #include <unistd.h>
+#include <algorithm>
 #include <unordered_map>
 #include <vector>
 
@@ -52,8 +53,7 @@ enum {
      * 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,
-    DS4_ROCM_TOPK_MERGE_GROUP = 8u
+    DS4_ROCM_ATTENTION_RAW_SCORE_CAP = 256u
 };
 
 struct ds4_gpu_tensor {
@@ -87,6 +87,20 @@ typedef struct {
     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"
@@ -126,6 +140,93 @@ typedef struct {
 
 #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_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

@@ -77,6 +77,34 @@ uint32_t ds4_ssd_cache_experts_for_byte_budget(uint64_t bytes,
     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,
@@ -86,9 +114,10 @@ bool ds4_ssd_auto_cache_plan(uint64_t            recommended_bytes,
     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 / 4ull ?
-            UINT64_MAX : (recommended_bytes * 4ull) / 5ull;
+        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;
     }

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

gguf-tools/.gitignore

@@ -1,4 +1,5 @@
 deepseek4-quantize
+gguf-requantize-dense
 quality-testing/score_official
 quality-testing/score_llama
 quality-testing/data/*
@@ -9,3 +10,5 @@ quality-testing/data/*
 !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

@@ -9,16 +9,21 @@ LLAMA_CPP_CXXFLAGS ?= -O3 -Wall -Wextra -std=c++17 \
 
 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 ds4_distributed.o ds4_ssd.o rax.o
-QUALITY_OBJS := ../ds4.o ../ds4_metal.o ../ds4_distributed.o ../ds4_ssd.o ../rax.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)),)
@@ -26,8 +31,12 @@ 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 ds4_distributed.o ds4_ssd.o rax.o
-QUALITY_OBJS := ../ds4.o ../ds4_cuda.o ../ds4_distributed.o ../ds4_ssd.o ../rax.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
@@ -38,11 +47,14 @@ CPPFLAGS ?= -D_GNU_SOURCE
 
 .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)
@@ -51,4 +63,5 @@ 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 quality-testing/score_llama
+	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/quality-testing/README.md

@@ -15,10 +15,15 @@ calling hosted APIs:
 
 - `data/glm52-openrouter-100`: 100 GLM 5.2 continuations collected through
   OpenRouter `z-ai/glm-5.2` with `top_logprobs=20`.
-- `data/flash`: 100 DeepSeek V4 Flash continuations collected from the official
-  DeepSeek API with `top_logprobs=20`.
-- `data/pro`: 100 DeepSeek V4 PRO continuations collected from the official
-  DeepSeek API with `top_logprobs=20`.
+- `data/glm53-flash-openrouter-zai-fp8-100`: 100 GLM 5.3 Flash continuations
+  from OpenRouter's pinned Z.AI FP8 endpoint. The endpoint does not expose
+  logprobs, so these are deterministic continuation fixtures.
+- `data/flash`: 100 DeepSeek V4 Flash 0731 continuations collected from the
+  official DeepSeek API with `top_logprobs=20`.
+- `data/pro`: 100 DeepSeek V4 PRO preview continuations collected from the
+  official DeepSeek API with `top_logprobs=20`.
+- `data/pro-0813`: 100 DeepSeek V4 PRO 0813 continuations collected from the
+  official DeepSeek API with `top_logprobs=20`.
 
 DeepSeek V4 Flash also has tracked official smoke vectors in
 `tests/test-vectors/`.  Those vectors drive `./ds4_test --logprob-vectors` and
@@ -29,13 +34,20 @@ full vocabulary logits.
 
 ## 2. Collect Official Continuations
 
+For the tracked DeepSeek V4 Flash 0731 fixture:
+
 ```sh
 export DEEPSEEK_API_KEY=...
 python3 gguf-tools/quality-testing/collect_official.py \
+  --model deepseek-v4-flash \
+  --endpoint https://api.deepseek.com/chat/completions \
   --prompts gguf-tools/quality-testing/prompts.jsonl \
   --out gguf-tools/quality-testing/data/flash \
   --count 100 \
-  --max-tokens 24
+  --max-tokens 24 \
+  --top-logprobs 20 \
+  --thinking disabled \
+  --reasoning-effort omit
 ```
 
 For GLM 5.2 through OpenRouter:
@@ -58,14 +70,33 @@ python3 gguf-tools/quality-testing/collect_official.py \
   --reasoning-effort none
 ```
 
-Use one output directory per official model.  The default model is Flash, so
-`data/flash` is the recommended path for Flash continuations.  For PRO:
+For GLM 5.3 Flash through the pinned Z.AI FP8 endpoint:
+
+```sh
+export OPENROUTER_API_KEY=...
+python3 gguf-tools/quality-testing/collect_official.py \
+  --model z-ai/glm-5.3-flash \
+  --endpoint https://openrouter.ai/api/v1/chat/completions \
+  --api-key-env OPENROUTER_API_KEY \
+  --prompts gguf-tools/quality-testing/prompts.jsonl \
+  --out gguf-tools/quality-testing/data/glm53-flash-openrouter-zai-fp8-100 \
+  --count 100 \
+  --max-tokens 128 \
+  --top-logprobs 0 \
+  --token-limit-field max_tokens \
+  --provider-order z-ai/fp8 \
+  --thinking omit \
+  --reasoning-effort low
+```
+
+Use one output directory per checkpoint. For PRO 0813 through the official
+DeepSeek API:
 
 ```sh
 python3 gguf-tools/quality-testing/collect_official.py \
   --model deepseek-v4-pro \
   --prompts gguf-tools/quality-testing/prompts.jsonl \
-  --out gguf-tools/quality-testing/data/pro \
+  --out gguf-tools/quality-testing/data/pro-0813 \
   --count 100 \
   --max-tokens 24 \
   --top-logprobs 20
@@ -90,6 +121,12 @@ make -C gguf-tools quality-score
 
 The scorer links against the DS4 runtime and uses Metal by default.
 
+Build the optional llama.cpp control scorer with:
+
+```sh
+make -C gguf-tools quality-llama-score
+```
+
 ## 4. Score GGUF Variants
 
 ```sh
@@ -106,11 +143,26 @@ gguf-tools/quality-testing/score_official \
   4096
 ```
 
-Use `data/flash/manifest.tsv` for Flash GGUFs,
-`data/glm52-openrouter-100/manifest.tsv` for GLM 5.2 GGUFs, and
-`data/pro/manifest.tsv` for PRO GGUFs.  The scorer and comparator do not care
-which model produced the manifest; the manifest path selects the continuation
-set.
+Use `data/flash/manifest.tsv` for Flash GGUFs and
+`data/glm52-openrouter-100/manifest.tsv` for GLM 5.2 GGUFs. Use
+`data/glm53-flash-openrouter-zai-fp8-100/manifest.tsv` for GLM 5.3 Flash. Use
+`data/pro/manifest.tsv` for the PRO preview checkpoint and
+`data/pro-0813/manifest.tsv` for PRO 0813. The scorer and comparator do not
+care which model produced the manifest; the manifest path selects the
+continuation set.
+
+Add `--quality` to disable DS4's speed-oriented numerical shortcuts. For an
+independent llama.cpp comparison of a DeepSeek V4 GGUF, use the same manifest
+and the token-identical DS4 prompt renderer:
+
+```sh
+gguf-tools/quality-testing/score_llama \
+  /path/to/model.gguf \
+  gguf-tools/quality-testing/data/flash/manifest.tsv \
+  /tmp/llama.tsv \
+  4096 \
+  deepseek-ds4
+```
 
 For a full-residency vs SSD-streaming comparison, score the same model twice and
 add the streaming flags to one run:

gguf-tools/quality-testing/collect_official.py

@@ -139,10 +139,11 @@ def request_one(
         "model": model,
         "messages": [{"role": "user", "content": prompt}],
         "temperature": 0,
-        "logprobs": True,
-        "top_logprobs": top_logprobs,
         "stream": False,
     }
+    if top_logprobs > 0:
+        payload["logprobs"] = True
+        payload["top_logprobs"] = top_logprobs
     payload[token_limit_field] = max_tokens
     if thinking != "omit":
         payload["thinking"] = {"type": thinking}
@@ -234,6 +235,8 @@ def main() -> int:
     ap.add_argument("--allow-provider-fallbacks", action="store_true")
     ap.add_argument("--require-parameters", action="store_true",
                     help="for OpenRouter, route only to endpoints advertising the requested parameters")
+    ap.add_argument("--resume", action="store_true",
+                    help="reuse complete cases already present in the output directory")
     args = ap.parse_args()
     if args.top_logprobs < 0 or args.top_logprobs > 20:
         raise SystemExit("--top-logprobs must be between 0 and 20")
@@ -265,6 +268,8 @@ def main() -> int:
 
     manifest = out / "manifest.tsv"
     rows = []
+    observed_models = set()
+    observed_providers = set()
     total = min(args.count, len(prompts))
     print(f"model={args.model} endpoint={args.endpoint}", file=sys.stderr)
     print(f"key_env={api_key_env} token_field={token_limit_field} thinking={thinking} reasoning={reasoning_effort}",
@@ -276,34 +281,56 @@ def main() -> int:
     for i, prompt in enumerate(prompts[: args.count]):
         case_id = f"case_{i:03d}"
         print(f"official {i + 1}/{total}: {case_id}", file=sys.stderr, flush=True)
-        response = fetch_with_retry(
-            api_key,
-            args.endpoint,
-            args.model,
-            prompt,
-            args.max_tokens,
-            args.top_logprobs,
-            thinking,
-            reasoning_effort,
-            token_limit_field,
-            provider_order,
-            args.allow_provider_fallbacks,
-            provider_require_parameters,
-        )
+        prompt_path = out / "prompts" / f"{case_id}.txt"
+        cont_path = out / "continuations" / f"{case_id}.txt"
+        resp_path = out / "responses" / f"{case_id}.json"
+        complete = prompt_path.exists() and cont_path.exists() and resp_path.exists()
+        if args.resume and complete:
+            if prompt_path.read_text(encoding="utf-8") != prompt:
+                raise RuntimeError(f"{case_id}: saved prompt differs from current prompt set")
+            content = cont_path.read_text(encoding="utf-8")
+            response = json.loads(resp_path.read_text(encoding="utf-8"))
+            if not content:
+                complete = False
+
+        if not args.resume or not complete:
+            for empty_attempt in range(3):
+                response = fetch_with_retry(
+                    api_key,
+                    args.endpoint,
+                    args.model,
+                    prompt,
+                    args.max_tokens,
+                    args.top_logprobs,
+                    thinking,
+                    reasoning_effort,
+                    token_limit_field,
+                    provider_order,
+                    args.allow_provider_fallbacks,
+                    provider_require_parameters,
+                )
+                choice = response["choices"][0]
+                content = choice.get("message", {}).get("content")
+                if isinstance(content, str) and content:
+                    break
+                print(f"warning: empty continuation for {case_id}; retrying",
+                      file=sys.stderr, flush=True)
+                time.sleep(1.0 + empty_attempt)
+            else:
+                raise RuntimeError(f"{case_id}: provider returned three empty continuations")
         choice = response["choices"][0]
-        content = choice.get("message", {}).get("content", "")
-        if not content:
-            print(f"warning: empty continuation for {case_id}", file=sys.stderr)
+        if response.get("model"):
+            observed_models.add(response["model"])
+        if response.get("provider"):
+            observed_providers.add(response["provider"])
         logprob_items = (choice.get("logprobs") or {}).get("content", []) or []
         if args.top_logprobs > 0 and not logprob_items:
             raise RuntimeError(f"{case_id}: response did not include output-token logprobs")
 
-        prompt_path = out / "prompts" / f"{case_id}.txt"
-        cont_path = out / "continuations" / f"{case_id}.txt"
-        resp_path = out / "responses" / f"{case_id}.json"
-        prompt_path.write_text(prompt, encoding="utf-8")
-        cont_path.write_text(content, encoding="utf-8")
-        resp_path.write_text(json.dumps(response, ensure_ascii=False, indent=2), encoding="utf-8")
+        if not args.resume or not complete:
+            prompt_path.write_text(prompt, encoding="utf-8")
+            cont_path.write_text(content, encoding="utf-8")
+            resp_path.write_text(json.dumps(response, ensure_ascii=False, indent=2), encoding="utf-8")
         rows.append((case_id, prompt_path, cont_path, resp_path))
         time.sleep(0.05)
 
@@ -311,6 +338,27 @@ def main() -> int:
         fp.write("# id\tprompt_file\tcontinuation_file\tresponse_file\n")
         for row in rows:
             fp.write("\t".join([row[0], str(row[1]), str(row[2]), str(row[3])]) + "\n")
+    collection = {
+        "schema": "ds4-official-continuations-v1",
+        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "model": args.model,
+        "endpoint": args.endpoint,
+        "observed_models": sorted(observed_models),
+        "observed_providers": sorted(observed_providers),
+        "temperature": 0,
+        "max_tokens": args.max_tokens,
+        "top_logprobs": args.top_logprobs,
+        "thinking": thinking,
+        "reasoning_effort": reasoning_effort,
+        "token_limit_field": token_limit_field,
+        "provider_order": provider_order,
+        "allow_provider_fallbacks": args.allow_provider_fallbacks,
+        "require_parameters": provider_require_parameters,
+        "case_count": len(rows),
+    }
+    (out / "collection.json").write_text(
+        json.dumps(collection, indent=2) + "\n", encoding="utf-8"
+    )
     print(f"wrote {manifest}", file=sys.stderr)
     return 0
 

gguf-tools/quality-testing/data/README.md

@@ -5,9 +5,13 @@ safe to commit and use in release QA.
 
 - `glm52-openrouter-100`: 100 GLM 5.2 OpenRouter continuations with API
   top-logprob slices.
-- `flash`: 100 DeepSeek V4 Flash official continuations with API top-logprob
-  slices.
-- `pro`: 100 DeepSeek V4 PRO official continuations with API top-logprob
+- `glm53-flash-openrouter-zai-fp8-100`: 100 deterministic GLM 5.3 Flash
+  continuations from OpenRouter's pinned Z.AI FP8 endpoint. That endpoint did
+  not return logprobs.
+- `flash`: 100 DeepSeek V4 Flash 0731 continuations from the official DeepSeek
+  API, with API top-logprob slices.
+- `pro`: 100 DeepSeek V4 PRO preview continuations with API top-logprob slices.
+- `pro-0813`: 100 DeepSeek V4 PRO 0813 continuations with API top-logprob
   slices.
 
 Each fixture directory contains:

gguf-tools/quality-testing/data/flash/continuations/case_000.txt

@@ -1 +1 @@
-Here is a detailed explanation of B-tree insertion, covering the standard algorithm, node splits, and the special case of splitting
+To understand B-tree insertion, you have to remember the **core invariant** of a B-tree: **every node must

gguf-tools/quality-testing/data/flash/continuations/case_001.txt

@@ -1,5 +1,6 @@
 Here's a concise design for a TCP echo server that handles slow clients:
 
-## Architecture
+## Core Architecture
 
-**Event-driven, non-block
+```
+┌────────────────────────────────

gguf-tools/quality-testing/data/flash/continuations/case_002.txt

@@ -1 +1,4 @@
-Here's a comprehensive comparison of memory-mapped (mmap) model weights versus copying all weights into private buffers on macOS
+Here's a comprehensive comparison of mmaped model weights versus copying all weights into private buffers on macOS:
+
+## Memory Management
+

gguf-tools/quality-testing/data/flash/continuations/case_003.txt

@@ -1 +1 @@
-This is an excellent question because it gets to the heart of why **RMSNorm** (Root Mean Square Normalization)
+To understand why **RMSNorm** (Root Mean Square Layer Normalization) needs both a **sum of squares** and

gguf-tools/quality-testing/data/flash/continuations/case_004.txt

@@ -1 +1 @@
-This is an excellent question that gets at a core optimization challenge in modern LLM deployment. Let's break down why **
+To understand why **KV cache checkpointing** is a game-changer for agentic sessions, we first need to understand the

gguf-tools/quality-testing/data/flash/continuations/case_005.txt

@@ -1,5 +1,6 @@
 Here's pseudocode for a binary heap push (insert) operation:
 
+## Binary Heap Push Operation
+
 ```
-function push(heap, value):
-    //
+PUSH(

gguf-tools/quality-testing/data/flash/continuations/case_006.txt

@@ -1 +1 @@
-This is a great question that gets at the core of how modern LLMs generate text. Both **temperature** and **
+To understand the tradeoffs between **top-p sampling** and **temperature-zero decoding**, it helps to first clarify what they

gguf-tools/quality-testing/data/flash/continuations/case_007.txt

@@ -1 +1 @@
-Here is a summary of how speculative decoding preserves exact greedy output (i.e., the same tokens that would be generated by
+Speculative decoding preserves exact greedy output through a **strict verification and rejection mechanism**, not by trusting the draft model's guesses

gguf-tools/quality-testing/data/flash/continuations/case_008.txt

@@ -1 +1 @@
-This is an excellent question that gets to the heart of how Large Language Models (LLMs) are served efficiently. The
+To understand why a GPU prefill path is faster than single-token decode, we need to look at the fundamental difference

gguf-tools/quality-testing/data/flash/continuations/case_009.txt

@@ -1,5 +1,3 @@
-Here are three invariants for a special-token table used in a tokenizer:
+Here are three invariants for a tokenizer special-token table:
 
----
-
-### 1. **Un
+## 1. **Uniqueness Invariant